[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