yzh119 edited a comment on pull request #10207:
URL: https://github.com/apache/tvm/pull/10207#issuecomment-1034575574
Some other notes:
If in the following case:
```python
@T.prim_func
def reduce(a: T.handle, b: T.handle, n: T.int32) -> None:
A = T.match_buffer(a, [1, 4, 8])
B = T.match_buffer(b, [1, 4])
for i, j, k in T.grid(1, 4, 8):
with T.block("reduce"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
B[vi, vj] = 0.
B[vi, vj] = B[vi, vj] + A[vi, vj, vk]
```
we bind `j` to `threadIdx.y` and `k` to `threadIdx.x`, different `j`'s might
be mapped to the same warp, we need different masks for different `j` to
distinguish them.
Below is an example of generated code:
```python
extern "C" __global__ void __launch_bounds__(32)
default_function_kernel0(float* __restrict__ A, float* __restrict__ B) {
float red_buf0[1];
uint mask[1];
float t0[1];
red_buf0[(0)] = A[(((((int)threadIdx.y) * 8) + ((int)threadIdx.x)))];
mask[(0)] = (__activemask() & ((uint)(255 << (((int)threadIdx.y) * 8))));
t0[(0)] = __shfl_down_sync(mask[(0)], red_buf0[(0)], 4, 32);
red_buf0[(0)] = (red_buf0[(0)] + t0[(0)]);
t0[(0)] = __shfl_down_sync(mask[(0)], red_buf0[(0)], 2, 32);
red_buf0[(0)] = (red_buf0[(0)] + t0[(0)]);
t0[(0)] = __shfl_down_sync(mask[(0)], red_buf0[(0)], 1, 32);
red_buf0[(0)] = (red_buf0[(0)] + t0[(0)]);
red_buf0[(0)] = __shfl_sync(mask[(0)], red_buf0[(0)], (((int)threadIdx.y)
* 8), 32);
B[(((int)threadIdx.y))] = red_buf0[(0)];
}
```
Another thing worth noting is, we can only allow cross warp reduction by
shuffle-down, thus warp size must be a multiple of `blockDim.x` when
`blockDim.y * blockDim.z != 1`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
