WebGL による並列計算

gpu.js

var GPU = (function () {
	
	function GPU(canvas) {
		this.gl = canvas.getContext(CONTEXT_TYPE);
		if (this.gl == null) {
			throw new Error(CONTEXT_TYPE);
		}
		
		this._extensions = {};
		for (var i = 0; i < EXTENSION_NAMES.length; i++) {
			var name = EXTENSION_NAMES[i];
			var extension = this.gl.getExtension(name);
			if (extension == null) {
				throw new Error(name);
			}
			this._extensions[name] = extension;
		}
		
		this.gl.bindBuffer(
			this.gl.ARRAY_BUFFER, this.gl.createBuffer());
		this.gl.bufferData(
			this.gl.ARRAY_BUFFER, new Float32Array(VERTICES),
			this.gl.STATIC_DRAW);
		
		this.gl.bindFramebuffer(
			this.gl.FRAMEBUFFER, this.gl.createFramebuffer());
		
		this._vertexShader = this._createShader(
			this.gl.VERTEX_SHADER, VERTEX_SHADER_SOURCE);
	}
	var prototype = GPU.prototype;
	
	var CONTEXT_TYPE = 'webgl';
	var EXTENSION_NAMES = [
		'OES_texture_float',
		'WEBGL_color_buffer_float'
	];
	
	var ATTRIB_NAME = 'a_position';
	var VERTICES = [-1, -1, 1, -1, -1, 1, 1, 1];
	
	GPU.PRECISION_MEDIUM = 0;
	GPU.PRECISION_LOW  = 1;
	GPU.PRECISION_HIGH = 2;
	
	prototype._createShader = function (type, source) {
		var shader = this.gl.createShader(type);
		this.gl.shaderSource(shader, source);
		this.gl.compileShader(shader);
		
		if (this.gl.getShaderParameter(
			shader, this.gl.COMPILE_STATUS))
		{
			return shader;
		}
		try {
			throw new Error(this.gl.getShaderInfoLog(shader));
		} finally {
			this.gl.deleteShader(shader);
		}
	};
	
	prototype._createProgram = function (fragmentShader) {
		var program = this.gl.createProgram();
		this.gl.attachShader(program, this._vertexShader);
		this.gl.attachShader(program, fragmentShader);
		this.gl.linkProgram(program);
		
		if (this.gl.getProgramParameter(
			program, this.gl.LINK_STATUS))
		{
			return program;
		}
		try {
			throw new Error(this.gl.getProgramInfoLog(program));
		} finally {
			this.gl.deleteProgram(program);
		}
	};
	
	prototype._linkProgram = function (fragmentShaderSource) {
		return this._createProgram(this._createShader(
			this.gl.FRAGMENT_SHADER, fragmentShaderSource));
	};
	
	prototype._getAttribLocation = function (program) {
		return this.gl.getAttribLocation(program, ATTRIB_NAME);
	};
	prototype._vertexAttrib = function (location) {
		this.gl.enableVertexAttribArray(location);
		this.gl.vertexAttribPointer(
			location, 2, this.gl.FLOAT, false, 0, 0);
	};
	
	prototype._framebufferTexture = function (texture) {
		this.gl.framebufferTexture2D(
			this.gl.FRAMEBUFFER,
			this.gl.COLOR_ATTACHMENT0,
			this.gl.TEXTURE_2D, texture, 0);
	};
	
	prototype.program = function (args, source, precision) {
		return new GPUProgram(this, args, source, precision);
	};
	prototype.alloc = function (array, length) {
		return new GPUMemory(this, array, length);
	};
	
	var VERTEX_SHADER_SOURCE = (
		'attribute vec4 ' + ATTRIB_NAME + ';' +
		'void main(void) {' + (
			'gl_Position = ' + ATTRIB_NAME + ';'
		) + '}'
	);
	
	return GPU;
})();

var GPUProgram = (function () {
	
	function GPUProgram(gpu, args, source, precision) {
		this.gpu = gpu;
		this.gl = gpu.gl;
		
		this._args = {};
		this._names = [];
		for (var name in args) {
			var type = args[name];
			if (type) {
				this._args[name] = '' + type;
			} else {
				this._names[this._names.length] = name;
			}
		}
		
		this._program = gpu._linkProgram(
			this._createSource(source, precision));
		
		this._attribLocation = gpu._getAttribLocation(this._program);
		this._uniformLocationObject = {};
	}
	var prototype = GPUProgram.prototype;
	
	var MAIN_NAME = 'gpu';
	
	var DEFINE_NAME = 'gpu_define';
	var TEXTURE_NAME = 'u_texture';
	var SIZE_NAME = 'u_size';
	
	var EMPTY_ARGS = {};
	
	prototype._getUniformLocation = function (name) {
		if (name in this._uniformLocationObject) {
			return this._uniformLocationObject[name];
		}
		return this._uniformLocationObject[name] =
			this.gl.getUniformLocation(this._program, name);
	};
	prototype._uniform1f = function (name, v0) {
		this.gl.uniform1f(this._getUniformLocation(name), v0);
	};
	prototype._uniform1i = function (name, v0) {
		this.gl.uniform1i(this._getUniformLocation(name), v0);
	};
	
	prototype._uniformArgs = function (args) {
		for (var name in this._args) {
			switch (this._args[name]) {
				case 'float':
				this._uniform1f(name, args[name]);
				break;
				
				case 'int':
				this._uniform1i(name, args[name]);
				break;
			}
		}
	};
	
	prototype._bind = function (index, input) {
		var alloc = !(input instanceof GPUMemory);
		if (alloc) {
			input = this.gpu.alloc(input);
		}
		this.gl.activeTexture(this.gl.TEXTURE0 + index);
		this.gl.bindTexture(this.gl.TEXTURE_2D, input._texture);
		this._uniform1i(TEXTURE_NAME + index, index);
		this._uniform1i(SIZE_NAME + index, input._size);
		return alloc ? input : null;
	};
	prototype._unbind = function (index, input) {
		this.gl.activeTexture(this.gl.TEXTURE0 + index);
		this.gl.bindTexture(this.gl.TEXTURE_2D, null);
		if (input != null) {
			input.free();
		}
	};
	
	prototype.exec = function (length, args) {
		this.gl.useProgram(this._program);
		this.gpu._vertexAttrib(this._attribLocation);
		
		var memory = this.gpu.alloc(null, length);
		this.gl.viewport(0, 0, memory._size, memory._size);
		this._uniform1i(SIZE_NAME, memory._size);
		
		if (args == null) {
			args = EMPTY_ARGS;
		}
		this._uniformArgs(args);
		
		var i;
		var memories = [];
		for (i = this._names.length - 1; i >= 0; i--) {
			memories[i] = this._bind(i, args[this._names[i]]);
		}
		
		this.gpu._framebufferTexture(memory._texture);
		this.gl.drawArrays(this.gl.TRIANGLE_STRIP, 0, 4);
		this.gpu._framebufferTexture(null);
		
		for (i = memories.length - 1; i >= 0; i--) {
			this._unbind(i, memories[i]);
		}
		return memory;
	};
	
	prototype._createSource = function (source, precision) {
		var declare;
		switch (precision) {
			default:
			case GPU.PRECISION_MEDIUM:
			declare = 'precision mediump float;';
			break;
			
			case GPU.PRECISION_LOW:
			declare = 'precision lowp float;';
			break;
			
			case GPU.PRECISION_HIGH:
			declare = 'precision highp float;';
			break;
		}
		for (var name in this._args) {
			declare += 'uniform ' +
				this._args[name] + ' ' + name + ';';
		}
		for (var i = 0; i < this._names.length; i++) {
			var textureName = TEXTURE_NAME + i;
			var sizeName = SIZE_NAME + i;
			
			declare += '\n#define ' +
				this._names[i] + '(index) ' +
				DEFINE_NAME + '(' +
				textureName + ',' + sizeName + ',index)\n';
			
			declare += 'uniform sampler2D ' + textureName + ';';
			declare += 'uniform int ' + sizeName + ';';
		}
		if (this._names.length > 0) {
			declare += DEFINE_SOURCE;
		}
		return declare + '\n' + source + '\n' + MAIN_SOURCE;
	};
	
	var DEFINE_SOURCE = (
		'float ' + DEFINE_NAME +
		'(sampler2D texture, int size, int index) {' + (
			'if (index >= 0) {' + (
				'int i = index / 4, y = i / size;' +
				'if (y < size) {' + (
					'int channel = index - i * 4;' +
					'vec4 color = texture2D(texture, (' + (
						'vec2(i - y * size, y) + .5'
					) + ') / vec2(size));' +
					'if (channel == 0) return color.r;' +
					'if (channel == 1) return color.g;' +
					'if (channel == 2) return color.b;' +
					'if (channel == 3) return color.a;'
				) + '}'
			) + '}' +
			'return 0.0;'
		) + '}'
	);
	var MAIN_SOURCE = (
		'uniform int ' + SIZE_NAME + ';' +
		'void main(void) {' + (
			'int i = 4 * (' + (SIZE_NAME +
				'* int(gl_FragCoord.y) + int(gl_FragCoord.x)'
			) + ');' +
			'gl_FragColor = vec4(' + (
				MAIN_NAME + '(i + 0),' +
				MAIN_NAME + '(i + 1),' +
				MAIN_NAME + '(i + 2),' +
				MAIN_NAME + '(i + 3)') + ');'
		) + '}'
	);
	
	return GPUProgram;
})();

var GPUMemory = (function () {
	
	function GPUMemory(gpu, array, length) {
		this.gpu = gpu;
		this.gl = gpu.gl;
		
		if (length == null && array != null) {
			length = array.length;
		}
		this._size = 1 << Math.log2(length - 1) / 2;
		
		var pixels;
		if (array == null) {
			pixels = null;
		} else {
			var minLength = this._size * this._size * 4;
			if (array instanceof Float32Array &&
				array.length >= minLength)
			{
				pixels = array;
			} else {
				pixels = new Float32Array(minLength);
				pixels.set(array);
			}
		}
		this._texture = this.gl.createTexture();
		
		this.gl.activeTexture(this.gl.TEXTURE0);
		this.gl.bindTexture(this.gl.TEXTURE_2D, this._texture);
		
		this.gl.texParameteri(
			this.gl.TEXTURE_2D,
			this.gl.TEXTURE_MAG_FILTER, this.gl.NEAREST);
		this.gl.texParameteri(
			this.gl.TEXTURE_2D,
			this.gl.TEXTURE_MIN_FILTER, this.gl.NEAREST);
		
		this.gl.texImage2D(
			this.gl.TEXTURE_2D,
			0, this.gl.RGBA, this._size, this._size,
			0, this.gl.RGBA, this.gl.FLOAT, pixels);
		this.gl.bindTexture(this.gl.TEXTURE_2D, null);
	}
	var prototype = GPUMemory.prototype;
	
	prototype.read = function () {
		var result = new Float32Array(this._size * this._size * 4);
		
		this.gpu._framebufferTexture(this._texture);
		this.gl.readPixels(
			0, 0, this._size, this._size,
			this.gl.RGBA, this.gl.FLOAT, result);
		this.gpu._framebufferTexture(null);
		
		return result;
	};
	
	prototype.free = function () {
		this.gl.deleteTexture(this._texture);
	};
	
	return GPUMemory;
})();