Question regarding subgroup inclusive_add and block_dim

[ffwqo]

Am I wrong to assume how subgroup inclusive add is supposed to work? Does block_dim set the local_size_x or the dispatch size?

import taichi as ti
ti.init(arch=ti.vulkan)
x = ti.field(ti.int32, shape=128)
@ti.kernel
def sum(x: ti.template()) -> ti.i32:
    total = 0
    ti.loop_config(block_dim=128)
    for I in ti.grouped(x):
        val = ti.lang.simt.subgroup.inclusive_add(x[I])
        ti.simt.subgroup.barrier()
        if(ti.simt.subgroup.elect()):
            ti.atomic_add(total, val)
    return total
@ti.kernel
def init(x: ti.template()):
    for I in ti.grouped(x):
        x[I] = 1
init(x)
print(sum(x)) #returns 2 instead of 128

[bobcao3]

The subgroup operations works in the size of "subgroups" which is a smaller unit than the block_dim or dispatch size. Usually the value is 32 on Nvidia and 64 on AMD.

However in this case your code is almost correct. Subgroup inclusive_add is an inclusive scan operation, which means each thread in the subgroup gets the inclusive sum of values from subgroup IDs smaller than the local one, but elect statement by definition here elects one thread, it doesn't guarantee it will be the thread holding the summation value.

[bobcao3]

For your case, you should use reduce_add because that does a All->One reduction and stores the sum in every thread. Then you can use elect.