Version: develop

SIMT Intrinsics

For the CUDA backend, Taichi supports warp-level and block-level intrinsics, which are needed for writing high-performance SIMT kernels. You can use them in Taichi similar to the usage in CUDA kernels. Currently, the following functions are supported:

Operation	Mapped CUDA intrinsic
`ti.simt.warp.all_nonzero`	`__all_sync`
`ti.simt.warp.any_nonzero`	`__any_sync`
`ti.simt.warp.unique`	`__uni_sync`
`ti.simt.warp.ballot`	`__ballot_sync`
`ti.simt.warp.shfl_sync_i32`	`__shfl_sync`
`ti.simt.warp.shfl_sync_f32`	`__shfl_sync`
`ti.simt.warp.shfl_up_i32`	`__shfl_up_sync`
`ti.simt.warp.shfl_up_f32`	`__shfl_up_sync`
`ti.simt.warp.shfl_down_i32`	`__shfl_down_sync`
`ti.simt.warp.shfl_down_f32`	`__shfl_down_sync`
`ti.simt.warp.shfl_xor_i32`	`__shfl_xor_sync`
`ti.simt.warp.match_any`	`__match_any_sync`
`ti.simt.warp.match_all`	`__match_all_sync`
`ti.simt.warp.active_mask`	`__activemask`
`ti.simt.warp.sync`	`__syncwarp`

See Taichi's API reference for more information on each function.

Here is an example of performing data exchange within a warp in Taichi:

a = ti.field(dtype=ti.i32, shape=32)

@ti.kernel
def foo():
    ti.loop_config(block_dim=32)
    for i in range(32):
        a[i] = ti.simt.warp.shfl_up_i32(ti.u32(0xFFFFFFFF), a[i], 1)

for i in range(32):
    a[i] = i * i

foo()

for i in range(1, 32):
    assert a[i] == (i - 1) * (i - 1)