Type System

In order for it to be compiled into WebGPU shaders, the Javascript function passed into ti.kernel() needs to obey certain rules that are not common to Javascript code. Most importantly, in normal Javascript-scope code, types of variables are not fixed and may change at runtime. This is different from kernel-scope code, where variables are "statically-typed". This means that types are fixed and determined at compile-time. When writing kernel code, it's useful to have a mental model of the type system in place, which this document will describe.

We will begin with four of the most simple and commonly used data types: primitives, vectors, matrices, and structs. We will discuss how to declare variables of these types, and common operations on them. We will also show how to declare these types as arguments to kernels. Beyond these simple types, "variables" in kernels can also refer to data containers (i.e. fields and textures) or lambdas/functions, which will be discussed in separate doc pages.

Primitives

In taichi.js, there are only two primitive types at the moment: i32 and f32. The compiler determines the type of a variable depending on the expression of its initializer:

let i = 1   // i32
let f = 1.0 // f32 

i32 values can be casted to f32 via ti.f32(..), and f32 values can be casted to i32 via ti.i32(..):

let i = ti.i32(4.2) // i32
let f = ti.f32(42)  // f32

when a valid operation is performed between an i32 and a f32, the integer is automatically promoted:

let i = ti.i32(4.2) // i32
let f = ti.f32(42)  // f32
let f2 = f + i      // f32

Due to Javascript's lack of type annotations, kernel argument types in taichi.js are specified via an extra argument to the ti.kernel() call. The following example shows how to pass an i32 and a f32 to a kernel, which returns their sum:

let k = ti.kernel(
    { i: ti.i32, f: ti.f32 },
    (i, f) => {
        return i + f
    }
)
console.log(await k(1, 2.5)) // prints 3.5

If you don't specify the arguments, then taichi.js will assume they are f32s by default:

let k = ti.kernel(
    (f1, f2) => {
        return f1 + f2
    }
)
console.log(await k(1.5, 2.5)) // prints 4.0

Vectors

Vectors are fixed-size tuples of primitives. In taichi.js, tuples are declared via the Javascript array notation [...]. For example:

let vi = [0, 1, 2]
let vf = [0.0, 1.0, 2.0, 3.0]

Individual components of of vectors can be acessed by indexing with an integerv[i]. Notice that, at the moment, indices used to access vectors must be compile-time constant expressions.

let v = [0.0, 1.0, 2.0, 3.0]
let f = v[0] // 0.0
f = v[1] // 1.0

Alternatively, for accessing the first 4 elements, you also write .x/.y/.z/.w or .r/.g/.b/.a. This is handy if your vector represents a 3D point/direction, or a color:

let v = [0.0, 1.0, 2.0, 3.0]
let f = v.x // 0.0
f = v.g // 1.0

taichi.js also supports "swizzling":

let v = [0, 1, 2, 3]
let v1 = v.xyz // [0, 1, 2]
let v2 = v.bgr // [2, 1, 0]

Vectors of the same type can have arithmetic operations among each other, where arithmetics are applied component wise. For example

let v = [1.0, 2.0] + [3.0, 4.0] // v = [4.0, 6.0]

Vectors can also have arithmetic operations with scalars, where the scalar is broadcasted:

let v = [1.0, 2.0] * 2.0 // v = [2.0, 4.0]

Beyond arithmetics, taichi.js contains many built-in mathematical operations which you can make use of. These are listed here.

For passing vectors into kernels as arguments, declare the argument as ti.types.vector( <primitive type name>, <component count> ):

let length = ti.kernel(
    {v: ti.types.vector(ti.f32, 2)},
    (v) => {
        return ti.sqrt(v.x * v.x + v.y * v.y)
    }
)
console.log(await length([3, 4])) // prints 5

Matrices

Matrices are fixed size of tuples of vectors of the same size. In taichi.js, matrices are declared via 2D arrays of primitives. For example:

let identity = [[1.0, 0.0], 
                [0.0, 1.0]]

Matrices in taichi.js are row-major. This means that each vector in the initializer is a row-vector, and when accessing matrices using an integer index i, the ith row will be returned:

let m = [[1.0, 2.0], [3.0, 4.0]]
let v = m[0] // [1.0, 2.0]

As with vectors, arithmetic operations among matrices are done component-wise, and primitives can be broadcased to operate with matrices:

let m1 = [[1.0, 2.0], [3.0, 4.0]] + [[4.0, 3.0], [2.0, 1.0]] // [[5.0, 5.0], [5.0, 5.0]]
let m2 = [[1.0, 2.0], [3.0, 4.0]] * 2.0 // [[2.0, 4.0], [6.0, 8.0]]

Common operations among matrices and vectors are available via built-in functions. For example, for matrix-vector multiplication, you can write it using the matmul built-in. There are two syntaxes for this:

let m = [[1.0, 2.0], [3.0, 4.0]]
let v = m[0] // [1.0, 2.0]

let mv1 = m.matmul(v)
let mv2 = ti.matmul(m, v)
// mv1 == mv2 == [5, 11]

For passing vectors into kernels as arguments, declare the argument as ti.types.matrix( <primitive type name>, <row count>, <column count> ):

let trace = ti.kernel(
    {m: ti.types.matrix(ti.f32, 2, 2)},
    (m) => {
        return m[0][0] + m[1][1];
    }
)
console.log(await trace([[1.0, 2.0], [3.0, 4.0]])) // prints 5

Structs

Structs are collections of values with names. These are described with the Javascript object notation:

let particle = {
    index: 0,
    position: [1.0, 2.0, 3.0],
    velocity: [4.0, 5.0, 6.0]
}

As with Javascript objects, elements of structs can be accessed via the dot notation:

let o = {i:0, v: [1.0, 2.0]};
let v = o.v; // [1.0, 2.0];

For passing vectors into kernels as arguments, declare the argument as ti.types.struct({ <element name> : <element type> ... }):

let particleType = ti.types.struct({
    position: ti.types.vector(ti.f32, 3),
    velocity: ti.types.vector(ti.f32, 3),
})
let predictNewPosition = ti.kernel(
    {p: particleType, time: ti.f32},
    (p, time) => {
        return p.position + time * p.velocity;
    }
)
console.log(await predictNewPosition({
    position: [0, 0, 0],
    velocity: [1, 1, 1]
}, 2)) 
// prints [2, 2, 2]