[VectorExt] Implement masked vectorization for iree_linalg_ext.gather (iree-org#21189)

Implements masked vectorization for iree_linalg_ext.gather to
iree_vector_ext.transfer_gather. The masked vector shape can be inferred
from the dps result.

Author: Groverkss

compiler/src/iree/compiler/Codegen/Common/GenericVectorization.cpp

@@ -73,6 +73,13 @@ getVectorSizes(Operation *op, bool useConfiguredVectorSizes) {
           return;
         vectorSizes = SmallVector<int64_t>(ty.getShape());
       })
+      .Case<IREE::LinalgExt::GatherOp>([&](IREE::LinalgExt::GatherOp gatherOp) {
+        std::optional<VectorizationTileSizes> result =
+            inferSizesFromIR(gatherOp.getOutput());
+        if (result) {
+          vectorSizes = result->vectorSizes;
+        }
+      })
       .Default([&](Operation *) {});
 
   if (vectorSizes) {
@@ -168,8 +175,8 @@ void GenericVectorizationPass::runOnOperation() {
           rewriter, cast<linalg::GenericOp>(op), vectorSizes, scalableVecDims,
           vectorizeGatherAccesses);
     } else if (auto gatherOp = dyn_cast<IREE::LinalgExt::GatherOp>(op)) {
-      (void)IREE::VectorExt::vectorizeLinalgExtGatherToTransferGather(rewriter,
-                                                                      gatherOp);
+      (void)IREE::VectorExt::vectorizeLinalgExtGatherToTransferGather(
+          rewriter, gatherOp, vectorSizes);
     } else {
       FailureOr<linalg::VectorizationResult> result = linalg::vectorize(
           rewriter, op, vectorSizes, scalableVecDims, vectorizeGatherAccesses);
@@ -221,6 +228,7 @@ void GenericVectorizationPass::runOnOperation() {
   if (enableVectorMasking) {
     vector::populateVectorMaskLoweringPatternsForSideEffectingOps(
         vectorizationPatterns);
+    IREE::VectorExt::populateVectorMaskLoweringPatterns(vectorizationPatterns);
     vectorizationPatterns.add<linalg::LinalgCopyVTRForwardingPattern,
                               linalg::LinalgCopyVTWForwardingPattern>(
         funcOp.getContext(), /*benefit=*/2);

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/BUILD.bazel

@@ -34,6 +34,7 @@ iree_td_library(
     deps = [
         "@llvm-project//mlir:BuiltinDialectTdFiles",
         "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
+        "@llvm-project//mlir:MaskableOpInterfaceTdFiles",
         "@llvm-project//mlir:OpBaseTdFiles",
         "@llvm-project//mlir:SideEffectInterfacesTdFiles",
         "@llvm-project//mlir:VectorInterfacesTdFiles",

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h

@@ -10,6 +10,7 @@
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtInterfaces.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/Vector/Interfaces/MaskableOpInterface.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.cpp

@@ -655,6 +655,12 @@ void TransferGatherOp::getCanonicalizationPatterns(RewritePatternSet &results,
       ctx);
 }
 
+// MaskableOpInterface methods.
+
+Type TransferGatherOp::getExpectedMaskType() {
+  return vector::inferTransferOpMaskType(getVectorType(), getPermutationMap());
+}
+
 // clang-format off
 #define GET_OP_CLASSES
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.cpp.inc" // IWYU pragma: keep

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.h

@@ -12,6 +12,7 @@
 #include "llvm/ADT/MapVector.h"
 #include "mlir/Bytecode/BytecodeImplementation.h"
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/Vector/Interfaces/MaskableOpInterface.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.td

@@ -10,6 +10,7 @@
 include "mlir/Interfaces/VectorInterfaces.td"
 include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtBase.td"
 include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtAttrs.td"
+include "mlir/Dialect/Vector/Interfaces/MaskableOpInterface.td"
 include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtInterfaces.td"
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
@@ -118,6 +119,7 @@ def IREEVectorExt_TransferGatherOp : IREEVectorExt_PureOp<"transfer_gather", [
     DeclareOpInterfaceMethods<VectorTransferOpInterface>,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
     DeclareOpInterfaceMethods<ConditionallySpeculatable>,
+    DeclareOpInterfaceMethods<MaskableOpInterface>,
     AttrSizedOperandSegments
   ]> {
   let arguments = (ins AnyShaped:$base,

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/Transforms.h

@@ -19,12 +19,27 @@ LogicalResult vectorizeGatherLikeGenericToTransferGather(
     ArrayRef<int64_t> vectorSizes = {}, ArrayRef<bool> scalableVecDims = {},
     bool vectorizeNDExtract = false);
 
+/// Vectorizes iree_linalg_ext.gather to iree_vector_ext.transfer_gather.
+/// Currently, this pattern only works when the index_depth and batch rank of
+/// the gather is 1.
+///
+/// %gather = iree_linalg_ext.gather dimension_map=[0] ins(%source, %indices)
+///                                                    outs(%output)
+///
+///  vectorizes to:
+///
+/// %indices_vec = vector.transfer_read %indices
+/// %gather_vec = iree_vector_ext.gather %source[...][%indices_vec...]
+/// %gather = vector.transfer_write %gather_vec, %output
 LogicalResult
 vectorizeLinalgExtGatherToTransferGather(RewriterBase &rewriter,
-                                         IREE::LinalgExt::GatherOp gatherOp);
+                                         IREE::LinalgExt::GatherOp gatherOp,
+                                         ArrayRef<int64_t> vectorSizes = {});
 
 void populateVectorTransferGatherLoweringPatterns(RewritePatternSet &patterns);
 
+void populateVectorMaskLoweringPatterns(RewritePatternSet &patterns);
+
 }; // namespace mlir::iree_compiler::IREE::VectorExt
 
 #endif // IREE_COMPILER_CODEGEN_DIALECT_VECTOR_EXT_TRANSFORMS_TRANSFORMS_H_

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/VectorizeIREEVectorExtOps.cpp

@@ -6,11 +6,13 @@
 
 #include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
 #include "iree/compiler/Codegen/Dialect/VectorExt/Transforms/Passes.h"
+#include "iree/compiler/Codegen/Dialect/VectorExt/Transforms/Transforms.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
 #include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -494,9 +496,17 @@ LogicalResult vectorizeGatherLikeGenericToTransferGather(
   return success();
 }
 
+Value maskOperation(RewriterBase &rewriter, Operation *op, Value mask) {
+  Value maskedOp =
+      cast<vector::MaskOp>(mlir::vector::maskOperation(rewriter, op, mask))
+          .getResult(0);
+  return maskedOp;
+}
+
 LogicalResult
 vectorizeLinalgExtGatherToTransferGather(RewriterBase &rewriter,
-                                         IREE::LinalgExt::GatherOp gatherOp) {
+                                         IREE::LinalgExt::GatherOp gatherOp,
+                                         ArrayRef<int64_t> vectorSizes) {
 
   // TODO: need to split the innermost dim of `indices` into `indexDepth`
   // vectors so that each independent index can be passed to the
@@ -520,23 +530,29 @@ vectorizeLinalgExtGatherToTransferGather(RewriterBase &rewriter,
   ShapedType gatherTy = gatherOp.getOutputType();
   ShapedType sourceTy = gatherOp.getSourceType();
 
-  auto gatherVectorTy =
-      VectorType::get(gatherTy.getShape(), gatherTy.getElementType());
+  if (vectorSizes.empty()) {
+    vectorSizes = gatherTy.getShape();
+  }
+
+  auto gatherVectorTy = VectorType::get(vectorSizes, gatherTy.getElementType());
   // Rank-reduced to remove the innermost unit dim.
-  auto indicesVecTy =
-      VectorType::get(indicesTy.getShape().take_front(gatherOp.getBatchRank()),
-                      rewriter.getIndexType());
+  auto indicesVecTy = VectorType::get(
+      vectorSizes.take_front(gatherOp.getBatchRank()), rewriter.getIndexType());
 
-  // Read `indices` tensor via `vector.transfer_read` and cast from int to
-  // index.
   Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
-  Value indicesVec = rewriter.create<vector::TransferReadOp>(
+  auto indicesVecRead = rewriter.create<vector::TransferReadOp>(
       loc, indicesVecTy.clone(indicesTy.getElementType()),
       gatherOp.getIndices(), SmallVector<Value>(indicesTy.getRank(), zero));
+  VectorType indicesMaskType = indicesVecTy.clone(rewriter.getI1Type());
+  SmallVector<OpFoldResult> gatherDims =
+      tensor::getMixedSizes(rewriter, loc, gatherOp.getOutput());
+  Value indicesMask = rewriter.create<vector::CreateMaskOp>(
+      loc, indicesMaskType,
+      ArrayRef(gatherDims).take_front(gatherOp.getBatchRank()));
+  Value indicesVec = maskOperation(rewriter, indicesVecRead, indicesMask);
   indicesVec =
       rewriter.create<arith::IndexCastOp>(loc, indicesVecTy, indicesVec);
 
-  // Create transfer_gather op
   SmallVector<Value> baseIndices(sourceTy.getRank(), zero);
   SmallVector<bool> indexed(sourceTy.getRank(), false);
   indexed[0] = true;
@@ -547,21 +563,59 @@ vectorizeLinalgExtGatherToTransferGather(RewriterBase &rewriter,
       rewriter.getMultiDimIdentityMap(sourceTy.getRank()).getMajorSubMap(1)));
   Value padding = rewriter.create<arith::ConstantOp>(
       loc, rewriter.getZeroAttr(gatherTy.getElementType()));
-
   auto transferGatherOp = rewriter.create<IREE::VectorExt::TransferGatherOp>(
       loc, gatherVectorTy, gatherOp.getSource(), baseIndices,
       ValueRange{indicesVec}, rewriter.getBoolArrayAttr(indexed), indexedMaps,
       rewriter.getMultiDimIdentityMap(gatherTy.getRank()), padding,
       /*mask=*/Value(), inBounds);
 
-  // Write back into tensor.
-  auto emptyOp = rewriter.create<tensor::EmptyOp>(loc, gatherTy.getShape(),
-                                                  gatherTy.getElementType());
+  VectorType gatherMaskType = gatherVectorTy.clone(rewriter.getI1Type());
+  Value gatherMask =
+      rewriter.create<vector::CreateMaskOp>(loc, gatherMaskType, gatherDims);
+  Value maskedGather = maskOperation(rewriter, transferGatherOp, gatherMask);
   SmallVector<Value> writeIndices(gatherTy.getRank(), zero);
   auto writeOp = rewriter.create<vector::TransferWriteOp>(
-      loc, transferGatherOp.getResult(), emptyOp, writeIndices);
-  rewriter.replaceOp(gatherOp, writeOp);
+      loc, maskedGather, gatherOp.getOutput(), writeIndices);
+  Value maskedWrite = maskOperation(rewriter, writeOp, gatherMask);
+
+  rewriter.replaceOp(gatherOp, maskedWrite);
   return success();
 }
 
+/// Lowers vector.mask %mask { iree_vector_ext.transfer_gather }
+///  into
+/// iree_vector_ext.transfer_gather %mask
+///
+/// Ideally, the mask should have just been put on transfer_gather directly,
+/// but this is done this way to match upstream vector.transfer_read masking.
+struct MaskedTransferGatherOpPattern : public OpRewritePattern<vector::MaskOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::MaskOp maskOp,
+                                PatternRewriter &rewriter) const override {
+    auto gatherOp = dyn_cast<TransferGatherOp>(maskOp.getMaskableOp());
+    if (!gatherOp) {
+      return failure();
+    }
+    // TODO: The 'vector.mask' passthru is a vector and 'transfer_gather'
+    // expects a scalar. We could only lower one to the other for cases where
+    // the passthru is a broadcast of a scalar.
+    if (maskOp.hasPassthru()) {
+      return rewriter.notifyMatchFailure(
+          maskOp, "can't lower passthru to transfer_gather");
+    }
+    rewriter.replaceOpWithNewOp<TransferGatherOp>(
+        maskOp, gatherOp.getVectorType(), gatherOp.getBase(),
+        gatherOp.getIndices(), gatherOp.getIndexVecs(), gatherOp.getIndexed(),
+        gatherOp.getIndexedMaps(), gatherOp.getPermutationMap(),
+        gatherOp.getPadding(), maskOp.getMask(), gatherOp.getInBounds());
+    return success();
+  }
+};
+
+void populateVectorMaskLoweringPatterns(RewritePatternSet &patterns) {
+  patterns.add<MaskedTransferGatherOpPattern>(patterns.getContext());
+}
+
 } // namespace mlir::iree_compiler::IREE::VectorExt

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/test/vectorize_vector_ext_ops.mlir

@@ -146,3 +146,30 @@ func.func @linalg_ext_gather_unit_dim(%source : tensor<1024x128xi32>, %indices :
 //       CHECK:   %[[CAST:.+]] = arith.index_cast %[[READ]]
 //       CHECK:   %[[GATHER:.+]] = iree_vector_ext.transfer_gather %[[ARG0]]
 //  CHECK-SAME:     [%[[C0]], %[[C0]]][%[[CAST]]: vector<10xindex>, None]
+
+// -----
+
+func.func @linalg_ext_gather_masked(%source : tensor<?x128xi32>, %indices : tensor<?x1xi32>) -> (tensor<?x128xi32>) {
+  %c0 = arith.constant 0 : index
+  %dim = tensor.dim %indices, %c0 : tensor<?x1xi32>
+  %dim_ub = util.assume.int %dim[<umin = 1, umax = 12>] : index
+  %empty = tensor.empty(%dim_ub) : tensor<?x128xi32>
+  %result = iree_linalg_ext.gather dimension_map = [0]
+    ins(%source, %indices : tensor<?x128xi32>, tensor<?x1xi32>)
+    outs(%empty: tensor<?x128xi32>) -> tensor<?x128xi32>
+  return %result : tensor<?x128xi32>
+}
+
+// CHECK-LABEL: @linalg_ext_gather_masked
+//  CHECK-SAME:    %[[SOURCE:[a-zA-Z0-9]+]]
+//  CHECK-SAME:    %[[INDICES:[a-zA-Z0-9]+]]
+//  CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+//  CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+//  CHECK: %[[DIM:.+]] = tensor.dim %[[INDICES]], %[[C0]]
+//  CHECK: %[[DIM_UB:.+]] = util.assume.int %[[DIM]]
+//  CHECK: %[[INDICES_MASK:.+]] = vector.create_mask %[[DIM_UB]]
+//  CHECK: vector.transfer_read %[[INDICES]]
+//  CHECK-SAME: %[[INDICES_MASK]]
+//  CHECK: %[[MASK:.+]] = vector.create_mask %[[DIM_UB]], %[[C128]]
+//  CHECK: iree_vector_ext.transfer_gather %[[SOURCE]]
+//  CHECK-SAME: %[[MASK]]

compiler/src/iree/compiler/Codegen/Utils/Utils.cpp

@@ -1863,37 +1863,44 @@ std::optional<VectorizationTileSizes> inferSizesFromIR(linalg::UnPackOp op) {
   return result;
 }
 
+std::optional<VectorizationTileSizes> static inferSizesFromMixedSizes(
+    SmallVector<OpFoldResult> shape) {
+  VectorizationTileSizes result;
+  for (OpFoldResult dim : shape) {
+    LLVM_DEBUG(llvm::dbgs() << "Dim #" << dim << ": ");
+    FailureOr<int64_t> maybeDimBound =
+        ValueBoundsConstraintSet::computeConstantBound(
+            presburger::BoundType::UB, dim,
+            /*stopCondition=*/nullptr, /*closedUB=*/true);
+    if (failed(maybeDimBound)) {
+      LLVM_DEBUG(llvm::dbgs() << "failed\n");
+      return std::nullopt;
+    }
+
+    LLVM_DEBUG(llvm::dbgs() << maybeDimBound.value() << "\n");
+    result.vectorSizes.push_back(maybeDimBound.value());
+    result.destShape.push_back(maybeDimBound.value());
+  }
+  return result;
+}
+
 std::optional<VectorizationTileSizes> inferSizesFromIR(Value val) {
   if (!val.getDefiningOp())
     return std::nullopt;
 
   std::optional<VectorizationTileSizes> result;
+  LLVM_DEBUG(llvm::dbgs() << "Inferring sizes for:\n" << val << "\n");
   TypeSwitch<Operation *, void>(val.getDefiningOp())
       .Case<linalg::LinalgOp>(
           [&](auto op) { result = inferSizesFromIR(op, cast<OpResult>(val)); })
       .Case<linalg::PackOp>([&](auto op) { result = inferSizesFromIR(op); })
-      .Case<tensor::ExtractSliceOp>([&](tensor::ExtractSliceOp op) {
+      .Case<tensor::ExtractSliceOp, tensor::EmptyOp>([&](auto op) {
         // tensor::ExtractSliceOp is not vectorizable, so only `destShape` has
         // the values.
-        result = VectorizationTileSizes();
-        LLVM_DEBUG(llvm::dbgs() << "Inferring sizes for:\n" << op << "\n");
-        int64_t destRank = op.getResult().getType().getRank();
-        for (int dim = 0; dim < destRank; ++dim) {
-          LLVM_DEBUG(llvm::dbgs() << "Dim #" << dim << ": ");
-          FailureOr<int64_t> maybeDimBound =
-              ValueBoundsConstraintSet::computeConstantBound(
-                  presburger::BoundType::UB, {op, dim},
-                  /*stopCondition=*/nullptr, /*closedUB=*/true);
-          if (failed(maybeDimBound)) {
-            LLVM_DEBUG(llvm::dbgs() << "failed\n");
-            result = std::nullopt;
-            return;
-          }
-          LLVM_DEBUG(llvm::dbgs() << maybeDimBound.value() << "\n");
-          result->destShape.push_back(maybeDimBound.value());
-        }
+        result = inferSizesFromMixedSizes(op.getMixedSizes());
       })
       .Default([&](Operation *) {});
+
   return result;
 }