Compute.js

About the author

Made with ❤️ by Zakaria Chabihi

Table of contents

1. Introduction
2. Tensors
3. Features
4. How does it work
5. Installation
6. Get started
7. Example
8. Input interpolates
9. Using the transpiler
10. Methods
11. Tensor metadata
12. API
13. License

Introduction

Compute.js is an open-source hardware accelerated Typescript library for tensor computation, it supports GPGPU (General purpose computing on GPUs) with fallback to regular javascript.

Compute.js allows you to parallelize computation on GPU and expose current thread tensor value and a bunch of helper methods, It also automaticallly computes tensor's metadata such as rank, stride, bytes length and texture matrix.

Tensors

Tensors are multidimensional vector spaces that maps in a (multi-)linear manner to vectors, scalars, and other tensors to a resulting tensor.

Features

• Zero dependency codebase
• Extermely lightweight, under 2000LOC
• Transpile shaders to WebGL and CPU code
• Compute tensor's metadata
• Parallelize output computations on GPU
• Efficient input compaction
• Fallback to CPU computation
• GPU kernel precompilation
• Smartly fallback to CPU if data size threshold is faster on CPU
• Expose current tensor value both on CPU and GPU backends
• Expose helper accessors both on CPU and GPU backends

How does it work

Compute.js parallelize the work on gpu according the expected output shape.

A thread variable is exposed on both CPU and GPU compute programs, which refers to the current output index which is within 0 and output.length - 1.

We also expose input interpolates for every input on both CPU and GPU programs so you can write your custom logic.

Installation

To get started using Compute.js first install it

$ npm install compute.js

Import Compute.js

import Compute from "compute.js"; // ES6 style
const Compute = require("compute.js"); // RequireJS style
// or simply use Compute global on browser
Compute // same as window.Compute

Get started

Construct a Compute instance

const instance = new Compute();

Define an input shape

You may add as many inputs as you want

instance.input("input", Uint8Array, 2, 5, 3]); // this creates a tensor name "input" with a shape of [3, 5, 3] using an Uint8Array

Define the output shape

instance.output(Uint8Array, [2, 5, 3]); // this defines the output tensor with a shape of [3, 5, 3] using an Uint8Array

Define the CPU fallback compute function

instance.cpu<{ input: 3 }>(({ $input }) => {
    return $input;
});

Define the GPU fallback compute function

instance.gpu(`return float($input);`);

Supply the input values

input values is an object that maps name inputs to TypeArray

const input = new Uint8Array(5 * 3 * 2).map((_, i) =>
    Math.floor(Math.random() * 100)
);
instance.input("input", input);

Execute the computation

const output = instance.run(); // output would be the same as input

Example

const input = new Uint8Array(5 * 3 * 2).map((_, i) =>
    Math.floor(Math.random() * 100)
);
const backend = new Backend();
const result = backend
    .input("input", Uint8Array, [5, 3, 2])
    .output(Float32Array, [5, 3, 2])
    .cpu<{ input: 3 }>(({ $input }) => {
    return $input;
    })
    .gpu(`return float($input);`)
    .inputs({ input })
    .run(); // would return a Float32Array with the same values from input

Input interpolates

| Interpolate | Explanation | Type | |-----|----|----| | $NAME | get value on current thread | Int or Float (Value) | | $NAME__coords | get coordinates on current thread | Vector | | $NAME__access_index | input accessor helper using array index | (index) => Value | | $NAME__access_coords | input accessor helper using coordinates | (coords) => Value | | $NAME__coords_index | coordinates to index helper | (coords) => Int | | $NAME__index_coords | index to coordinates helper | (index) => Vector |

Using the transpiler

Compute.js ships with a lightweight transpile-only code parser which returns a configure Compute instance with both CPU/GPU code and input/output shapes configured

const instance = Compute.transpile(`
f32(3) (f32(3) input) {
    i8 acc = 0;
    for (i8 a=0; a < 10; a++) {
    acc = a % 2 < 1 ? 0 : acc + a;
    }
    return float(acc);
}
`);
// or
const instance = Compute.Transpile`
    Code goes here...
`;

Using interpolates

Compute.js adopts a different sugar syntax on its transpiler, to refer to an interpolate simply remplace the $ with an @ and __ with a . such as $name_coords will become @name.coords

Supported primitives

| Primitive | Equivalent | WebGL | |-----|----|----| | f32 | Float32 | float | | i32 | Int32 | int | | i16 | Int16 | int | | i8 | Int8 | int | | bool | Boolean | bool | | vecN | VecN | vecN | | vecN | VecN | ivecN |

Supported functions

Compute.js transpiler support the following funcions: not, length, normalize, distance, pow, exp, exp2, log, log2, sqrt, floor, ceil, sin, cos, tan, asin, acos, atan, mult, div, mod, plus, minus, lt, lte, gt, gte, eq, neq, and, or, float, int in addition to vecN and ivecN constructors.

Supported expressions

Compute.js transpiler support the following expressions:

• Comparison such as <, <=, > and >=
• Equality such as == and !=
• Binary such as +, -, *, /, ^ and %
• Logical such as ! and ternary operator

Supported declarations

Compute.js transpiler support the following declarations:

• For loops
• Conditions
• Break, continue and return statements

Not supported yet

Compute.js transpiler support a subset of the GLSL syntax, yet it doesn't support:

• Comments
• Other builtin functions
• Arrays and structs

Fully featured example

f32 custom_add(f32 a, f32 b) {
    return a + b;
}
f32 custom_add_mult(f32 a, f32 b) {
    return custom_add(a, b) * custom(a, b);
}
f32(3) (f32(3) left, right i8[5,2,3]) {
    i8 acc = 0;
    f32 l = $left;
    vec3<i8> r = $right.coords;
    vec3<i8> r_s = r + ivec3(1);
    i8 r_v = $right.access_coords(r_s);
    for (i8 a=0; a < 10; a++) {
        if (a % 3 == 0) {
            continue;
        }
        if (a < 5) {
            break;
        }
        acc = a % 2 < 1 ? 0 : custom_add_mult(acc, a);
        acc += l + r;
    }
    return float(acc) + l + float(r_v);
}

Methods

Backend.prototype.input

Backend.prototype.input(name: string, type: TypedArray | TypedArrayConstructor, shape: number | number[]): Backend;

Backend.prototype.input(name: string, shape: number | number[], type: TypedArray | TypedArrayConstructor): Backend;

Backend.prototype.input(name: string, shape: number | number[]): Backend;

Backend.prototype.input(name: string, type: TypedArray | TypedArrayConstructor): Backend;

Defines a input named input with a type TypedArrayConstructor with a shape of shape.

If type is a TypedArray then its data will be uploaded otherwise it will be used to define the tensor type.

If shape is a number it will be used to infer tensor's dimensions otherwise it shape dimensions will be uploaded to the shader.

• name: input name must be a compatible GLSL variable name.
• type: a typed array or a typed array constructor.
• shape: a 1D array of dimensions shapes

return: this object.

Backend.prototype.output

Backend.prototype.output(type: TypedArray | TypedArrayConstructor, shape: number | number[]): Backend;

Backend.prototype.output(shape: number | number[], type: TypedArray | TypedArrayConstructor): Backend;

Backend.prototype.output(type: TypedArray | TypedArrayConstructor): Backend;

Backend.prototype.output(shape: number | number[]): Backend;

Defines the output a type TypedArrayConstructor with a shape of shape

If type is a TypedArray then its data will be uploaded otherwise it will be used to define the tensor type.

If shape is a number it will be used to infer tensor's dimensions otherwise it shape dimensions will be uploaded to the shader.

• type: a typed array or a typed array constructor.
• shape: a 1D array of dimensions shapes

return: this object.

Backend.prototype.function

Backend.prototype.function(name: string, result: string, code: string): Backend;

Defines a custom shader function with name and a return value of result and a function body of code

• name: input name must be a compatible GLSL variable name.
• result: a valid GLSL function return type
• code: a valid GLSL function body code

return: this object.

Backend.prototype.gpu

Backend.prototype.gpu(syntax: string): Backend;

Defines the CPU compute shader of the current thread

This shader must return a float if you were expecting a Float32Array output and must return an int value if otherwise.

• syntax: the compute shader syntax, interpolates that starts with a dollar sign will be replaced automatically

return: this object.

Backend.prototype.cpu

Backend.prototype.cpu<T extends {[s: string]: number} = {}>(closure: (map: IMethodsMap<T> & {thread: number}) => number): Backend;

Defines the CPU compute closure of the current thread

• map: this object supplies the current thread in addition to the current input interpolate values

return: this object.

Backend.prototype.run({runtime = "fastest", threshold = 4096 : { runtime?: "gpu" | "cpu" | "fallback" | "fastest"; threshold?: number; } = {}): TypedArray;

Execute the compute kernel and returns a TypedArray

• runtime: execution runtime
• threshold: slow GPU data size threshold

| runtime | explanation | |---------|-------| | gpu | Runs the kernel on GPU exclusively | | cpu | Runs the kernel on CPU exclusively | | fallback| Runs the kernel on GPU the on CPU if GPU fails | | fastest| if the input data size is below threshold run the kernel on CPU otherwise run it on the GPU |

return: TypedArray.

Tensor metadata

A call with Compute.tensor(type: Float32, [2, 3, 5]) would compute the following values :

1. rank: number of dimensions
2. shape: depth of each dimension
3. strides: each dimension stride
4. length: total elements in tensor
5. width: tensor 2D width
6. height: tensor 2D width
7. bytes: bytes per element, 4 for floats, 2 for 16 bit numbers and 1 for 8 bits numbers

API

type TypedArray = Int8Array | Int16Array | Int32Array | Uint8Array | Uint16Array | Uint32Array | Float32Array;
type TypedArrayConstructor = Int8ArrayConstructor | Int16ArrayConstructor | Int32ArrayConstructor | Uint8Array | Uint16ArrayConstructor | Uint32ArrayConstructor | Float32ArrayConstructor;
type TensorMetadata = {
    constructor: TypedArrayConstructor;
    typed: "int8" | "uint8" | "int16" | "uint16" | "int32" | "float" | "uint31";
    native: "float" | "int" | "uint";
    precision: "low" | "medium" | "high";
    shader: "float" | "int";
    rank: number;
    bytes: number;
    shape: number[];
    length: number;
    width: number;
    height: number;
    strides: number[];
};

type IOutput = [type: TypedArray | TypedArrayConstructor]
    | [shape: number | number[]]
    | [type: TypedArray | TypedArrayConstructor, shape: number | number[]];

type IInput = [name: string, type: TypedArray | TypedArrayConstructor]
    | [name: string, shape: number | number[]]
    | [name: string, type: TypedArray | TypedArrayConstructor, shape: number | number[]]
    | [name: string, shape: number | number[], type: TypedArray | TypedArrayConstructor];

export default class Compute {
    static Transpile(string: TemplateStringsArray): Compute;
    static transpile(code: string): Compute;
    static tensor(type: TypedArray | TypedArrayConstructor, shape: number[]): TensorMetadata;

    output(...args: IOutput): this;
    input(...args: IInput): this;

    function(name: string, result: string, code: string): this;

    cpu<T extends {
        [s: string]: number;
    } = {}>(closure: (map: IMethodsMap<T> & {
        thread: number;
    }) => number): this;

    gpu(code: string): this;

    run({ runtime, threshold, safe, }?: {
        runtime?: "gpu" | "cpu" | "fallback" | "fastest";
        threshold?: number;
        safe?: boolean;
    }): TypedArray;
}

Licence

This code is released under the MIT license.