Functions and Lambdas

In taichi.js, each ti.kernel represents a sequence of GPU operations which can be invoked in Javascript code. In other words, a ti.kernel can be thought of a parallelized GPU functions, which can be invoked on the CPU. However, in order to write complex GPU kernels, you also want to have GPU functions that are callable in GPU scope. This doc page will discuss how to use functions and lambdas in taichi.js kernels.

Syntax

The simplest way to define a function is to directly define it in the kernel you want to use it in:

  let addOne = ti.kernel((f) => {
    function impl(arg) {
      return arg + 1;
    }
    return impl(f);
  });
  console.log(await addOne(42)); // 43

It is also recommened to use a lambda instead of a function, for cleaner syntax:

  let addOne = ti.kernel((f) => {
    let impl = arg => arg + 1
    return impl(f);
  });
  console.log(await addOne(42)); // 43

If you want to share a lambda/function among multiple kernels, you can define it as a normal Javascript function, and add it to kernel scope:

let addOneImpl = (x) => x + 1
ti.addToKernelScope({ addOneImpl })
let addOne = ti.kernel((x) => addOneImpl(x))
let addTwo = ti.kernel((x) => addOneImpl(addOneImpl(x)))
console.log(await addOne(42)) // 43
console.log(await addTwo(42)) // 44

The beauty with this mechanism is that, not only it does it allow you to write shared logic for different kernels, it allows you to share that some logic in non-kernel Javascript code as well!

console.log(addOneImpl(42)) // normal JS code

Argument Passing Semantics

Functions and lambdas in taichi.js always use pass-by-reference semantics when possible, even for primitive types:

  let k = ti.kernel(() => {
    let f = (i) => {
      i = i + 1;
    };
    let h = 42;
    f(h);
    return h;
  });
  console.log(await k()); // 43

This is different from normal Javascript code

Consequences of Inlining

Functions in taichi.js are implemented via "inlining". This means that taichi.js functions don't actually translate to functions in shaders. But rather, we expand the function logic in place at every call-site of the function. This has the following consequnces

• Each function can only have one return statement. For example, the following function is not allowed

  let f = (x) => {
    if(x > 0) {
        return 1;
    }
    else {
        return -1
    }
  }
  instead, the function needs to be written as

  let f = (x) => {
    let result = -1
    if(x > 0) {
        result = 1;
    }
    return result
  }
  (Having multiple return statements is theoretically possible for inlined functions. But this requires some compiler trickery which taichi.js does not have at the moment.)
• No support for recursion.

On the other hand, having inlined functions does allow some interesting syntaxes, which are not even available in WGSL, which is WebGPU's native shading language. As an example, due to inlining, you can have a function which accepts another function as argument (a.k.a. a "higher-order" function):

let add = ti.kernel((i, j) => {
    let apply = (f, x) => f(x);
    return apply((x) => i + x, j);
});
console.log(await add(12, 13));

ti.func

If you come from the world of the python taichi library, you probably know that taichi has a @ti.func decorator, which is required for any non-kernel function to be called in kernels. As you've learned in this page, in taichi.js, this is not needed, and kernels can call any Javascript function as long as it is in kernel scope. However, taichi.js still provides a ti.func(..) function, which you can use to wrap a Javascript function.

let f = ti.field(ti.i32, 1000)
let addOne = ti.func((x) => x + 1)
ti.addToKernelScope({ f, addOne })
let k = ti.kernel(() => {
    for(let i of ti.range(1000)){
        f[i] = addOne(f[i])
    }
}) 

The ti.func() function does not do anything -- it returns whatever argument is passed to it. However, it is still beneficial to use it on the functions that will only be used in kernels, for the sake of explicitness. (In a later page, we will see a case where the ti.func() wrapper makes a significant difference, when working with minifiers)