From e41ff11fe403f2e99cd757c4548b28f9785ce6c8 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Mon, 7 Oct 2024 15:00:38 +0100 Subject: [PATCH 1/2] [mlir][vector] Clarify the semantics of masking maps We use the term "masking map" throughout the Linalg vectorization logic, but we don't really define what it is and how it differs from Linalg indexing maps. This PR clarifies the differnces, makes sure that the new terminology is used consistenty and improves code re-use. --- .../Linalg/Transforms/Vectorization.cpp | 44 +++++++++++++------ 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp index 09c6b2683b4388..dcbdd7264f95ae 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -250,6 +250,32 @@ struct VectorizationState { LinalgOp linalgOp, std::optional maybeMaskingMap); + /// Check whether this permutation map can be used for masking. Add the + /// moment we only make sure that there are no broadcast dimensions, but this + /// might change if indexing maps evolve. + bool isValidMaskingMap(AffineMap maskingMap) { + return maskingMap.getBroadcastDims().size() == 0; + } + + /// Turn the input indexing map into a valid masking map. + /// + /// The input indexing map may contain "zero" results, e.g.: + /// (d0, d1, d2, d3) -> (d2, d1, d0, 0) + /// Applying such maps to canonical vector shapes like this one: + /// (1, 16, 16, 4) + /// would yield an invalid vector shape like this: + /// (16, 16, 1, 0) + /// Instead, drop the broadcasting dims that make no sense for masking perm. + /// maps: + /// (d0, d1, d2, d3) -> (d2, d1, d0) + /// This way, the corresponding vector/mask type will be: + /// vector<16x16x1xty> + /// rather than this invalid Vector type: + /// vector<16x16x1x0xty> + AffineMap getMaskingMapFromIndexingMap(AffineMap &indexingMap) { + return indexingMap.dropZeroResults(); + } + // Holds the compile-time static sizes of the iteration space to vectorize. // Dynamic dimensions are represented using ShapedType::kDynamic. SmallVector iterSpaceStaticSizes; @@ -360,6 +386,10 @@ VectorizationState::initState(RewriterBase &rewriter, LinalgOp linalgOp, Value VectorizationState::getOrCreateMaskFor( RewriterBase &rewriter, Operation *opToMask, LinalgOp linalgOp, std::optional maybeMaskingMap) { + + assert(!maybeMaskingMap || + isValidMaskingMap(*maybeMaskingMap) && "Ill-formed masking map."); + // No mask is needed if the operation is not maskable. auto maskableOp = dyn_cast(opToMask); if (!maskableOp) @@ -429,20 +459,8 @@ VectorizationState::maskOperation(RewriterBase &rewriter, Operation *opToMask, LDBG("Trying to mask: " << *opToMask << "\n"); std::optional maybeMaskingMap = std::nullopt; - // The Operand indexing map may contain "zero" results, e.g.: - // (d0, d1, d2, d3) -> (d0, d1, d2, 0) - // When applied to canonical vector shapes like these: - // (1, 16, 16, 4) - // we would get: - // (1, 16, 16, 0) - // Instead, we should extract the following map permutation map for masking: - // (d0, d1, d2, d3) -> (d0, d1, d2) - // This way, the corresponding vector/mask type will be: - // vector<1x16x16xty> - // rather than: - // vector<1x16x16x0xty> if (maybeIndexingMap) - maybeMaskingMap = maybeIndexingMap->dropZeroResults(); + maybeMaskingMap = getMaskingMapFromIndexingMap(*maybeIndexingMap); // Create or retrieve mask for this operation. Value mask = From 75e32daa8426b14cd0cb95b09ba2afdd3fddf578 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Fri, 18 Oct 2024 07:59:36 +0100 Subject: [PATCH 2/2] fixup! [mlir][vector] Clarify the semantics of masking maps Fix typos --- mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp index dcbdd7264f95ae..8925fec034cb9f 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -250,7 +250,7 @@ struct VectorizationState { LinalgOp linalgOp, std::optional maybeMaskingMap); - /// Check whether this permutation map can be used for masking. Add the + /// Check whether this permutation map can be used for masking. At the /// moment we only make sure that there are no broadcast dimensions, but this /// might change if indexing maps evolve. bool isValidMaskingMap(AffineMap maskingMap) { @@ -387,8 +387,8 @@ Value VectorizationState::getOrCreateMaskFor( RewriterBase &rewriter, Operation *opToMask, LinalgOp linalgOp, std::optional maybeMaskingMap) { - assert(!maybeMaskingMap || - isValidMaskingMap(*maybeMaskingMap) && "Ill-formed masking map."); + assert((!maybeMaskingMap || isValidMaskingMap(*maybeMaskingMap)) && + "Ill-formed masking map."); // No mask is needed if the operation is not maskable. auto maskableOp = dyn_cast(opToMask);