[Mlir-commits] [mlir] [MLIR][Vector] Add support for distributing masked writes (PR #71482)
Lei Zhang
llvmlistbot at llvm.org
Tue Nov 7 12:33:39 PST 2023
================
@@ -489,10 +499,26 @@ struct WarpOpTransferWrite : public OpRewritePattern<vector::TransferWriteOp> {
if (!targetType)
return failure();
+ // 2.5 Compute the distributed type for the new mask;
+ VectorType maskType;
+ if (writeOp.getMask()) {
+ // TODO: Distribution of masked writes with non-trivial permutation maps
+ // requires the distribution of the mask to elementwise match the
+ // distribution of the permuted written vector. Currently the details
+ // of which lane is responsible for which element is captured strictly
+ // by shape information on the warp op, and thus requires materializing
+ // the permutation in IR.
+ if (!writeOp.getPermutationMap().isMinorIdentity())
+ return failure();
+ maskType =
+ getDistributedType(writeOp.getMask().getType().cast<VectorType>(),
----------------
antiagainst wrote:
Nit: IIRC we have getMaskType directly?
https://github.com/llvm/llvm-project/pull/71482
More information about the Mlir-commits
mailing list