Rocm jaxlib v0.5.0 warpsize #169
Conversation
zoranjovanovic-ns
requested review from
i-chaochen,
draganmladjenovic and
pemeliya
zoranjovanovic-ns
force-pushed
the
rocm-jaxlib-v0.5.0-warpsize
branch
from
971541b to
7d58776
Compare
| analysis, /*minor_dim=*/input_shape_.back(), WarpSize()); | ||
| int64_t num_warps_per_column = WarpSize(); | ||
| num_threads_ = {num_warps_per_column, WarpSize()}; | ||
| analysis, /*minor_dim=*/input_shape_.back(), kTileSize); |
There was a problem hiding this comment.
so we need to change as tile size 32 instead of WarpSize(device_info) here? may I ask why?
There was a problem hiding this comment.
This is temporary, I believe that reduction algorithm needs modifications in order to work with warp_size==64.
Without this some tests fail.
There was a problem hiding this comment.
yes, I also did not find a good solution here. This only applies for column-wise reductions.
They work as follows: one block of 1024 threads (32x32) performs column reduction for 1 vertical stripe of N rows and 32 columns. Basically each warp loads and reduces N/32 rows (each having 32 elements) and writes its resulting reduced row to a shared memory. As a result, we have 32 rows of 32 elements written to shared memory.
After that, we do syncthreads and each warp reads 1 vertical column from shared memory and performs warp-level reduction on it. So, finally each warp just writes its 1 reduced element back to global mem. As a result we have Nx32 stripe reduced to 1x32 row.
To make it working for warp_size=64, we could have 16 warps (16*64 = 1024) processing 1 vertical stripe of N rows and 64 columns. But each warp shall then process N/16 rows and perform 4 writes to shared memory (instead of 1). As a result. we would then have 1 large shared mem array of size 64x64 to be transposed. But I don't have a clear idea how to express this in terms of Indexing maps they is in the reduction emitter.
| @@ -87,7 +88,7 @@ TransposeFusion::TransposeFusion(const HloFusionAnalysis& analysis) | |||
| permutation_(transpose_.permutation), | |||
| input_shape_( | |||
| Permute(transpose_.dimensions, InversePermutation(permutation_))), | |||
| base_block_size_(WarpSize(analysis_.device_info())) { | |||
| base_block_size_(kTileSize) { | |||
|
@zoranjovanovic-ns WDYT this PR #170 |
There is number of #ifdefs that we cannot upstream, has the same issue with reduce as this PR (reduction algorithm probably should be modified), but if it fixes more tests then we can use it as temporary solution. |
No description provided.