taichi如何优化行reduce操作

综合讨论区
#1

比如像layernorm这个操作，需要计算一行的mean, variance，然后进行归一化

我需要用哪些操作能让它感知到我需要对每一行的数据进行缓存么

#2

hello，一般warp内的reduction是用shfl指令实现，warp外block内的shared memory，block外的就用俩kernel加起来了，想知道你那边具体的use case是啥样啊

#3

用numpy写成伪代码如下：

def norm(x, out): 
    # X是一个2D输入(Rows, Cols) 
    mean_val = np.mean(x, 1)
    var_val = np.var(x, 1)
     out = (x - mean) / sqrt(var_val)

考虑在CUDA里实现，X会被用来计算mean val，也会用于后续输出的计算。假设我们一个warp处理一行，那么每个线程需要缓存X数据到寄存器。通过warpReduce计算sum， var。然后再用之前寄存器缓存下来的x值计算后续输出（就不需要反复从global memory读取）。

（假设一行是128个元素，一个warp处理一行，那么每个线程需要缓存4个x

这在taichi可以做到么

#4

不好意思回复晚了

最简单的基础实现可以用atomics on shared memory先凑合一下，验证一下正确性和性能，进一步的实现和优化可以用ti.simt里面的shuffle指令来搞～不过shared memory的atomic要稍等一下1.5.0，如果你不介意装nightly的话里面已经有了。之前的版本跑不对。

CUDA后端的simt方法在这

另外你可以试一下vulkan后端，vulkan的subgroup和warp是一个概念，Vulkan性能也挺好的。