Atomics

In taichi.js kernels, top-level for-loops are parallelized by the compiler. As a result, kernel code are subject to many of the common pitfalls in parallel programming. Consider this kernel which attempts to sum from 1 to 100:

let sum100_buggy = ti.kernel(() => {
    let sum = 0
    for (let i of ti.range(100)){
        sum = sum + (i + 1)
    }
    return sum
})
console.log(await sum100_buggy())

If this summation logic do not live in a ti.kernel, then it is perfectly valid. However, due to parallelization, there are 100 GPU threads trying to to read from and write to the variable sum at the same time. This is known as a "race condition". There is no guarantee of what are values received by each of the threads, and the behavior of this kernel is essentially undefined.

It's not always an easy task to figure out how to make kernel code free from race conditions. However, one tool that proves handy in many situation is atomics. And atomic operation is one which can be safely applied in parallel, and the hardware will take care of serializing the operation an ensuring its correctness. In our summation example, we may use an atomicAdd:

let sum100 = ti.kernel(() => {
    let sum = 0
    for (let i of ti.range(100)){
        ti.atomicAdd(sum, i + 1)
    }
    return sum
})
console.log(await sum100()) // 5050

This is the full list of atomic operations available in taichi.js:

• ti.atomicAdd(destination, value)
• ti.atomicSub(destination, value)
• ti.atomicMax(destination, value)
• ti.atomicMin(destination, value)
• ti.atomicAnd(destination, value)
• ti.atomicOr(destination, value)
• ti.atomicXor(destination, value)

Syntactic Sugar

Instead of writing out ti.atomicAdd(destination, value) which is a bit clunky, taichi.js allows you to use the self-addition notation to express atomic semantics. That is, you can also write destination += value. Using this syntax, the previous sum100 kernel can be written like this:

let sum100 = ti.kernel(() => {
    let sum = 0
    for (let i of ti.range(100)){
        sum += i + 1
    }
    return sum
})
console.log(await sum100()) // 5050

Similarly, you can also use

• -= for atomicSub
• &= for atomicAnd
• |= for atomicOr
• ^= for atomicXor

Atomics on Field Elements

In the previous example, the atomic operation is applied to a global temporary variable. You can also use atomics on elements of a ti.field. Here is a example

let f = ti.field(ti.i32, 1)
ti.addToKernelScope({ f })
let sum100 = ti.kernel(() => {
    f[0] = 0
    for (let i of ti.range(100)){
        f[0] += i + 1
        // or ti.atomicAdd(f[0], i + 1)
    }
    return f[0]
})
console.log(await sum100()) // 5050