commit | 47ae0decbf4524022c0d595a18e1313a68895d64 | [log] [tgz] |
---|---|---|
author | Charitha Saumya <136391709+charithaintc@users.noreply.github.com> | Fri Sep 12 09:37:04 2025 -0700 |
committer | Copybara-Service <copybara-worker@google.com> | Fri Sep 12 09:41:22 2025 -0700 |
tree | 612f70484bb5377620baaccd880528ed9ae87a51 | |
parent | 60684db8f1c6e98f117df27b3726037bab09f0c0 [diff] |
[mlir][xegpu] Add support for `vector.multi_reduction` and `vector.shape_cast` SIMT distribution. (#157560) Add support for distributing the `vector.multi_reduction` operation across lanes in a warp. Currently only 2D to 1D reductions are supported. Given layouts for the source and accumulator vectors, * If the reduction dimension is distributed across lanes, the reduction is non-lane-local and the reduction is done using warp shuffles. Here we simply rewrite the `MultiDimReductionOp` to a sequence of `ReductionOp`s inside the warp op body. Actual distribution will be done by `WarpOpReduction` pattern. * If the reduction dimension is not distributed across lanes, the reduction is lane-local. In this case, we yield the source and accumulator vectors from the warp op and perform the lane-local reduction outside the warp op using a sequence of `ReductionOp`s. PR also adds support for distributing `vector.shape_cast` based on layouts. GitOrigin-RevId: 9b0d7ddb04665e76cfa90b5d69c6183b90772243
See https://mlir.llvm.org/ for more information.