[Mlir-commits] [mlir] [mlir][vector] Add support for multi-dim reduction vector distribution (PR #71193)
Lei Zhang
llvmlistbot at llvm.org
Sat Nov 4 11:11:44 PDT 2023
================
@@ -569,15 +569,17 @@ struct TestVectorDistribution
});
MLIRContext *ctx = &getContext();
auto distributionFn = [](Value val) {
- // Create a map (d0, d1) -> (d1) to distribute along the inner
- // dimension. Once we support n-d distribution we can add more
- // complex cases.
+ // Create a map (d0, d1) -> (d1, d0) to distribute starting from the inner
+ // dimensions.
VectorType vecType = dyn_cast<VectorType>(val.getType());
int64_t vecRank = vecType ? vecType.getRank() : 0;
OpBuilder builder(val.getContext());
- if (vecRank == 0)
- return AffineMap::get(val.getContext());
- return AffineMap::get(vecRank, 0, builder.getAffineDimExpr(vecRank - 1));
+ SmallVector<AffineExpr, 4> vecDims = llvm::to_vector<4>(
----------------
antiagainst wrote:
`llvm::map_to_vector`?
https://github.com/llvm/llvm-project/pull/71193
More information about the Mlir-commits
mailing list