[mlir][vector] Add gather lowering patterns

[kuhar]

Summary

This is for targets that do not support gather-like ops, e.g., SPIR-V.

Gather is expanded into lower-level vector ops with memory accesses
guarded with scf.if.

I also considered generating vector.maskedloads, but decided against
it to keep the memref and tensor codepath closer together. There's a
good chance that if a target doesn't support gather it does not support
masked loads either.

Issue: https://github.com/llvm/llvm-project/issues/60905

[kuhar]

Remove unnecessary includes

[kuhar]

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

3251

Could someone confirm if this index calculation is correct? The current version of the spec says:

The gather operation gathers elements from memory or ranked tensor into a n-D vector as defined by a base with indices and an additional n-D index vector (each index is a 1-D offset on the base)

So I'd think that base [indices] defines a base pointer and then all accesses are offset relative to this base pointer. However, I'm not sure if can turn these into individual loads/extract by adding the offset to the last index only.

[ThomasRaoux]

mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h

342

nit: I'm not sure emulation is the right term. This is just a lowering of vector.gather to vector.maskedload as emulation would suggest we have some assumption that this maps to a native hw functionality.
Could we pick a name suggesting lowering instead?

[kuhar]

Rebase. Rename populate function (emulation -> lowering). Do not rely on temporary array refs.

[kuhar]

mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h

342

Done, thanks for the suggestion!

[kuhar]

Also rename the test

[ThomasRaoux]

LGTM

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

3251

I'm not an expert but that looks correct to me.

[dcaballe]

LGTM! A couple of comments.

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

3259

you could do simple memref load?

mlir/test/Dialect/Vector/vector-gather-lowering.mlir

18

it would be good to check that we have the patterns to canonicalize an all-one mask.

[kuhar]

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

3259

Is it generally considered simpler or prefered for some other reason? The reason I used vector.load is that it seemed lower-level (no indexing map like in memre.load).

[kuhar]

mlir/test/Dialect/Vector/vector-gather-lowering.mlir

18

We do. See https://reviews.llvm.org/D146144.

[dcaballe]

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

3259

memref.load is mostly to load scalar values and you wouldn't have to generate any extract op so I think. I don't think one is lower level than the other. IIRC, the map in memref.load is to describe the layout.

mlir/test/Dialect/Vector/vector-gather-lowering.mlir

18

Thanks!

mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h

 ///      <-------------------- join ------------------------->
 ///
 /// Other local patterns then kick in iteratively (including DCE) and compose
 /// to combine the ExtractStridedSlice/InsertStridedSlice.
 void populateVectorUnrollPatterns(RewritePatternSet &patterns,
                                   const UnrollVectorOptions &options,
                                   PatternBenefit benefit = 1);
 
+/// Expands `vector.gather` ops into a series of conditional scalar loads
+/// (`vector.load` for memrefs or `tensor.extract` for tensors). These loads are
+/// conditional to avoid out-of-bounds memory accesses and guarded with `scf.if`
+/// ops. This lowering path is intended for targets that do not feature
+/// dedicated gather ops.
+void populateVectorGatherLoweringPatterns(RewritePatternSet &patterns,
+                                          PatternBenefit benefit = 1);
+
 //===----------------------------------------------------------------------===//
 // Finer-grained patterns exposed for more control over individual lowerings.
 //===----------------------------------------------------------------------===//
 /// Apply `splitFullAndPartialTransfer` selectively via a pattern. This pattern
 /// may take an extra filter to perform selection at a finer granularity.
 struct VectorTransferFullPartialRewriter : public RewritePattern {
   using FilterConstraintType =
       std::function<LogicalResult(VectorTransferOpInterface op)>;

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

 //===- VectorTransforms.cpp - Conversion within the Vector dialect --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements target-independent rewrites as 1->N patterns.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 
+#include <cstdint>
 #include <functional>
 #include <optional>
 #include <type_traits>
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/IR/Location.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 #include "mlir/Support/LogicalResult.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
rewriter.replaceOpWithNewOp<vector::ContractionOp>(
         op.getIteratorTypes());
     return success();
   };
 
 private:
   FilterConstraintType filter;
 };
 
+/// Flattens 2 or more dimensional `vector.gather` ops by unrolling the
+/// outermost dimension. For example:
+/// ```
+/// %g = vector.gather %base[%c0][%v], %mask, %pass_thru :
+///        ... into vector<2x3xf32>
+///
+/// ==>
+///
+/// %0   = arith.constant dense<0.0> : vector<2x3xf32>
+/// %g0  = vector.gather %base[%c0][%v0], %mask0, %pass_thru0 : ...
+/// %1   = vector.insert %g0, %0 [0] : vector<3xf32> into vector<2x3xf32>
+/// %g1  = vector.gather %base[%c0][%v1], %mask1, %pass_thru1 : ...
+/// %g   = vector.insert %g1, %1 [1] : vector<3xf32> into vector<2x3xf32>
+/// ```
+///
+/// When applied exhaustively, this will produce a sequence of 1-d gather ops.
+struct FlattenGather : OpRewritePattern<vector::GatherOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::GatherOp op,
+                                PatternRewriter &rewriter) const override {
+    VectorType resultTy = op.getType();
+    if (resultTy.getRank() < 2)
+      return rewriter.notifyMatchFailure(op, "already flat");
+
+    Location loc = op.getLoc();
+    Value indexVec = op.getIndexVec();
+    Value maskVec = op.getMask();
+    Value passThruVec = op.getPassThru();
+
+    Value result = rewriter.create<arith::ConstantOp>(
+        loc, resultTy, rewriter.getZeroAttr(resultTy));
+
+    Type subTy = VectorType::get(resultTy.getShape().drop_front(),
+                                 resultTy.getElementType());
+
+    for (int64_t i = 0, e = resultTy.getShape().front(); i < e; ++i) {
+      int64_t thisIdx[1] = {i};
+
+      Value indexSubVec =
+          rewriter.create<vector::ExtractOp>(loc, indexVec, thisIdx);
+      Value maskSubVec =
+          rewriter.create<vector::ExtractOp>(loc, maskVec, thisIdx);
+      Value passThruSubVec =
+          rewriter.create<vector::ExtractOp>(loc, passThruVec, thisIdx);
+      Value subGather = rewriter.create<vector::GatherOp>(
+          loc, subTy, op.getBase(), op.getIndices(), indexSubVec, maskSubVec,
+          passThruSubVec);
+      result =
+          rewriter.create<vector::InsertOp>(loc, subGather, result, thisIdx);
+    }
+
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+};
+
+/// Turns 1-d `vector.gather` into a scalarized sequence of `vector.loads` or
+/// `tensor.extract`s. To avoid out-of-bounds memory accesses, these
+/// loads/extracts are made conditional using `scf.if` ops.
+struct Gather1DToConditionalLoads : OpRewritePattern<vector::GatherOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::GatherOp op,
+                                PatternRewriter &rewriter) const override {
+    VectorType resultTy = op.getType();
+    if (resultTy.getRank() != 1)
+      return rewriter.notifyMatchFailure(op, "unsupported rank");
+
+    Location loc = op.getLoc();
+    Type elemTy = resultTy.getElementType();
+    // Vector type with a single element. Used to generate `vector.loads`.
+    VectorType elemVecTy = VectorType::get({1}, elemTy);
+
+    Value condMask = op.getMask();
+    Value base = op.getBase();
+    Value indexVec = rewriter.createOrFold<arith::IndexCastOp>(
+        loc, op.getIndexVectorType().clone(rewriter.getIndexType()),
+        op.getIndexVec());
+    auto baseOffsets = llvm::to_vector(op.getIndices());
+    Value lastBaseOffset = baseOffsets.back();
+
+    Value result = op.getPassThru();
+
+    // Emit a conditional access for each vector element.
+    for (int64_t i = 0, e = resultTy.getNumElements(); i < e; ++i) {
+      int64_t thisIdx[1] = {i};
+      Value condition =
+          rewriter.create<vector::ExtractOp>(loc, condMask, thisIdx);
+      Value index = rewriter.create<vector::ExtractOp>(loc, indexVec, thisIdx);
+      baseOffsets.back() =
+          rewriter.createOrFold<arith::AddIOp>(loc, lastBaseOffset, index);
+
+      auto loadBuilder = [&](OpBuilder &b, Location loc) {
+        Value extracted;
+        if (isa<MemRefType>(base.getType())) {
+          // `vector.load` does not support scalar result; emit a vector load
+          // and extract the single result instead.
+          Value load =
+              b.create<vector::LoadOp>(loc, elemVecTy, base, baseOffsets);
+          int64_t zeroIdx[1] = {0};
+          extracted = b.create<vector::ExtractOp>(loc, load, zeroIdx);
+        } else {
+          extracted = b.create<tensor::ExtractOp>(loc, base, baseOffsets);
+        }
+
+        Value newResult =
+            b.create<vector::InsertOp>(loc, extracted, result, thisIdx);
+        b.create<scf::YieldOp>(loc, newResult);
+      };
+      auto passThruBuilder = [result](OpBuilder &b, Location loc) {
+        b.create<scf::YieldOp>(loc, result);
+      };
+
+      result =
+          rewriter
+              .create<scf::IfOp>(loc, condition, /*thenBuilder=*/loadBuilder,
+                                 /*elseBuilder=*/passThruBuilder)
+              .getResult(0);
+    }
+
+    rewriter.replaceOp(op, result);
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::vector::populateVectorMaskMaterializationPatterns(
     RewritePatternSet &patterns, bool force32BitVectorIndices,
     PatternBenefit benefit) {
   patterns.add<VectorCreateMaskOpConversion,
                MaterializeTransferMask<vector::TransferReadOp>,
                MaterializeTransferMask<vector::TransferWriteOp>>(
patterns
           patterns.getContext(), benefit);
 }
 
 void mlir::vector::populateVectorScanLoweringPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {
   patterns.add<ScanToArithOps>(patterns.getContext(), benefit);
 }
 
+void mlir::vector::populateVectorGatherLoweringPatterns(
+    RewritePatternSet &patterns, PatternBenefit benefit) {
+  patterns.add<FlattenGather, Gather1DToConditionalLoads>(patterns.getContext(),
+                                                          benefit);
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd enum attribute definitions
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Vector/Transforms/VectorTransformsEnums.cpp.inc"

mlir/test/Dialect/Vector/vector-gather-lowering.mlir

+// RUN: mlir-opt %s --test-vector-gather-lowering | FileCheck %s
+
+// CHECK-LABEL: @gather_memref_1d
+// CHECK-SAME:    ([[BASE:%.+]]: memref<?xf32>, [[IDXVEC:%.+]]: vector<2xindex>, [[MASK:%.+]]: vector<2xi1>, [[PASS:%.+]]: vector<2xf32>)
+// CHECK-DAG:     [[M0:%.+]]    = vector.extract [[MASK]][0] : vector<2xi1>
+// CHECK-DAG:     %[[IDX0:.+]]  = vector.extract [[IDXVEC]][0] : vector<2xindex>
+// CHECK-NEXT:    [[RES0:%.+]]  = scf.if [[M0]] -> (vector<2xf32>)
+// CHECK-NEXT:      [[LD0:%.+]]   = vector.load [[BASE]][%[[IDX0]]] : memref<?xf32>, vector<1xf32>
+// CHECK-NEXT:      [[ELEM0:%.+]] = vector.extract [[LD0]][0] : vector<1xf32>
+// CHECK-NEXT:      [[INS0:%.+]]  = vector.insert [[ELEM0]], [[PASS]] [0] : f32 into vector<2xf32>
+// CHECK-NEXT:      scf.yield [[INS0]] : vector<2xf32>
+// CHECK-NEXT:    else
+// CHECK-NEXT:      scf.yield [[PASS]] : vector<2xf32>
+// CHECK-DAG:     [[M1:%.+]]    = vector.extract [[MASK]][1] : vector<2xi1>
+// CHECK-DAG:     %[[IDX1:.+]]  = vector.extract [[IDXVEC]][1] : vector<2xindex>
+// CHECK-NEXT:    [[RES1:%.+]]  = scf.if [[M1]] -> (vector<2xf32>)
+// CHECK-NEXT:      [[LD1:%.+]]   = vector.load [[BASE]][%[[IDX1]]] : memref<?xf32>, vector<1xf32>
+// CHECK-NEXT:      [[ELEM1:%.+]] = vector.extract [[LD1]][0] : vector<1xf32>
+// CHECK-NEXT:      [[INS1:%.+]]  = vector.insert [[ELEM1]], [[RES0]] [1] : f32 into vector<2xf32>
+// CHECK-NEXT:      scf.yield [[INS1]] : vector<2xf32>
+// CHECK-NEXT:    else
+// CHECK-NEXT:      scf.yield [[RES0]] : vector<2xf32>
+// CHECK:         return [[RES1]] : vector<2xf32>
+func.func @gather_memref_1d(%base: memref<?xf32>, %v: vector<2xindex>, %mask: vector<2xi1>, %pass_thru: vector<2xf32>) -> vector<2xf32> {
+  %c0 = arith.constant 0 : index
+  %0 = vector.gather %base[%c0][%v], %mask, %pass_thru : memref<?xf32>, vector<2xindex>, vector<2xi1>, vector<2xf32> into vector<2xf32>
+  return %0 : vector<2xf32>
+}
+
+// CHECK-LABEL: @gather_memref_1d_i32_index
+// CHECK-SAME:    ([[BASE:%.+]]: memref<?xf32>, [[IDXVEC:%.+]]: vector<2xi32>, [[MASK:%.+]]: vector<2xi1>, [[PASS:%.+]]: vector<2xf32>)
+// CHECK-DAG:     [[C42:%.+]]   = arith.constant 42 : index
+// CHECK-DAG:     [[IDXS:%.+]]  = arith.index_cast [[IDXVEC]] : vector<2xi32> to vector<2xindex>
+// CHECK-DAG:     [[IDX0:%.+]]  = vector.extract [[IDXS]][0] : vector<2xindex>
+// CHECK-NEXT:    %[[OFF0:.+]]  = arith.addi [[IDX0]], [[C42]] : index
+// CHECK-NEXT:    [[RES0:%.+]]  = scf.if
+// CHECK-NEXT:      [[LD0:%.+]]   = vector.load [[BASE]][%[[OFF0]]] : memref<?xf32>, vector<1xf32>
+// CHECK:         else
+// CHECK:         [[IDX1:%.+]]  = vector.extract [[IDXS]][1] : vector<2xindex>
+// CHECK:         %[[OFF1:.+]]  = arith.addi [[IDX1]], [[C42]] : index
+// CHECK:         [[RES1:%.+]]  = scf.if
+// CHECK-NEXT:      [[LD1:%.+]]   = vector.load [[BASE]][%[[OFF1]]] : memref<?xf32>, vector<1xf32>
+// CHECK:         else
+// CHECK:         return [[RES1]] : vector<2xf32>
+func.func @gather_memref_1d_i32_index(%base: memref<?xf32>, %v: vector<2xi32>, %mask: vector<2xi1>, %pass_thru: vector<2xf32>) -> vector<2xf32> {
+  %c0 = arith.constant 42 : index
+  %0 = vector.gather %base[%c0][%v], %mask, %pass_thru : memref<?xf32>, vector<2xi32>, vector<2xi1>, vector<2xf32> into vector<2xf32>
+  return %0 : vector<2xf32>
+}
+
+// CHECK-LABEL: @gather_memref_2d
+// CHECK-SAME:    ([[BASE:%.+]]: memref<?x?xf32>, [[IDXVEC:%.+]]: vector<2x3xindex>, [[MASK:%.+]]: vector<2x3xi1>, [[PASS:%.+]]: vector<2x3xf32>)
+// CHECK-DAG:     %[[C0:.+]]    = arith.constant 0 : index
+// CHECK-DAG:     %[[C1:.+]]    = arith.constant 1 : index
+// CHECK-DAG:     [[PTV0:%.+]]  = vector.extract [[PASS]][0] : vector<2x3xf32>
+// CHECK-DAG:     [[M0:%.+]]    = vector.extract [[MASK]][0, 0] : vector<2x3xi1>
+// CHECK-DAG:     [[IDX0:%.+]]  = vector.extract [[IDXVEC]][0, 0] : vector<2x3xindex>
+// CHECK-NEXT:    %[[OFF0:.+]]  = arith.addi [[IDX0]], %[[C1]] : index
+// CHECK-NEXT:    [[RES0:%.+]]  = scf.if [[M0]] -> (vector<3xf32>)
+// CHECK-NEXT:      [[LD0:%.+]]   = vector.load [[BASE]][%[[C0]], %[[OFF0]]] : memref<?x?xf32>, vector<1xf32>
+// CHECK-NEXT:      [[ELEM0:%.+]] = vector.extract [[LD0]][0] : vector<1xf32>
+// CHECK-NEXT:      [[INS0:%.+]]  = vector.insert [[ELEM0]], [[PTV0]] [0] : f32 into vector<3xf32>
+// CHECK-NEXT:      scf.yield [[INS0]] : vector<3xf32>
+// CHECK-NEXT:    else
+// CHECK-NEXT:      scf.yield [[PTV0]] : vector<3xf32>
+// CHECK-COUNT-5: scf.if
+// CHECK:         [[FINAL:%.+]] = vector.insert %{{.+}}, %{{.+}} [1] : vector<3xf32> into vector<2x3xf32>
+// CHECK-NEXT:    return [[FINAL]] : vector<2x3xf32>
+ func.func @gather_memref_2d(%base: memref<?x?xf32>, %v: vector<2x3xindex>, %mask: vector<2x3xi1>, %pass_thru: vector<2x3xf32>) -> vector<2x3xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = vector.gather %base[%c0, %c1][%v], %mask, %pass_thru : memref<?x?xf32>, vector<2x3xindex>, vector<2x3xi1>, vector<2x3xf32> into vector<2x3xf32>
+  return %0 : vector<2x3xf32>
+ }
+
+// CHECK-LABEL: @gather_tensor_1d
+// CHECK-SAME:    ([[BASE:%.+]]: tensor<?xf32>, [[IDXVEC:%.+]]: vector<2xindex>, [[MASK:%.+]]: vector<2xi1>, [[PASS:%.+]]: vector<2xf32>)
+// CHECK-DAG:     [[M0:%.+]]    = vector.extract [[MASK]][0] : vector<2xi1>
+// CHECK-DAG:     %[[IDX0:.+]]  = vector.extract [[IDXVEC]][0] : vector<2xindex>
+// CHECK-NEXT:    [[RES0:%.+]]  = scf.if [[M0]] -> (vector<2xf32>)
+// CHECK-NEXT:      [[ELEM0:%.+]] = tensor.extract [[BASE]][%[[IDX0]]] : tensor<?xf32>
+// CHECK-NEXT:      [[INS0:%.+]]  = vector.insert [[ELEM0]], [[PASS]] [0] : f32 into vector<2xf32>
+// CHECK-NEXT:      scf.yield [[INS0]] : vector<2xf32>
+// CHECK-NEXT:    else
+// CHECK-NEXT:      scf.yield [[PASS]] : vector<2xf32>
+// CHECK-DAG:     [[M1:%.+]]    = vector.extract [[MASK]][1] : vector<2xi1>
+// CHECK-DAG:     %[[IDX1:.+]]  = vector.extract [[IDXVEC]][1] : vector<2xindex>
+// CHECK-NEXT:    [[RES1:%.+]]  = scf.if [[M1]] -> (vector<2xf32>)
+// CHECK-NEXT:      [[ELEM1:%.+]] = tensor.extract [[BASE]][%[[IDX1]]] : tensor<?xf32>
+// CHECK-NEXT:      [[INS1:%.+]]  = vector.insert [[ELEM1]], [[RES0]] [1] : f32 into vector<2xf32>
+// CHECK-NEXT:      scf.yield [[INS1]] : vector<2xf32>
+// CHECK-NEXT:    else
+// CHECK-NEXT:      scf.yield [[RES0]] : vector<2xf32>
+// CHECK:         return [[RES1]] : vector<2xf32>
+func.func @gather_tensor_1d(%base: tensor<?xf32>, %v: vector<2xindex>, %mask: vector<2xi1>, %pass_thru: vector<2xf32>) -> vector<2xf32> {
+  %c0 = arith.constant 0 : index
+  %0 = vector.gather %base[%c0][%v], %mask, %pass_thru : tensor<?xf32>, vector<2xindex>, vector<2xi1>, vector<2xf32> into vector<2xf32>
+  return %0 : vector<2xf32>
+}
+
+// CHECK-LABEL: @gather_tensor_2d
+// CHECK:  scf.if
+// CHECK:    tensor.extract
+// CHECK:  else
+// CHECK:  scf.if
+// CHECK:    tensor.extract
+// CHECK:  else
+// CHECK:  scf.if
+// CHECK:    tensor.extract
+// CHECK:  else
+// CHECK:  scf.if
+// CHECK:    tensor.extract
+// CHECK:  else
+// CHECK:  scf.if
+// CHECK:    tensor.extract
+// CHECK:  else
+// CHECK:  scf.if
+// CHECK:    tensor.extract
+// CHECK:  else
+// CHECK:       [[FINAL:%.+]] = vector.insert %{{.+}}, %{{.+}} [1] : vector<3xf32> into vector<2x3xf32>
+// CHECK-NEXT:  return [[FINAL]] : vector<2x3xf32>
+ func.func @gather_tensor_2d(%base: tensor<?x?xf32>, %v: vector<2x3xindex>, %mask: vector<2x3xi1>, %pass_thru: vector<2x3xf32>) -> vector<2x3xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = vector.gather %base[%c0, %c1][%v], %mask, %pass_thru : tensor<?x?xf32>, vector<2x3xindex>, vector<2x3xi1>, vector<2x3xf32> into vector<2x3xf32>
+  return %0 : vector<2x3xf32>
+ }

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
getOperation()->walk([](Operation *op) {
       Value bcast = vector::BroadcastOp::createOrFoldBroadcastOp(
           b, op->getOperand(0), targetShape, broadcastedDims);
       op->getResult(0).replaceAllUsesWith(bcast);
       op->erase();
     });
   }
 };
 
+struct TestVectorGatherLowering
+    : public PassWrapper<TestVectorGatherLowering,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorGatherLowering)
+
+  StringRef getArgument() const final { return "test-vector-gather-lowering"; }
+  StringRef getDescription() const final {
+    return "Test patterns that lower the gather op in the vector conditional "
+           "loads";
+  }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<arith::ArithDialect, func::FuncDialect,
+                    memref::MemRefDialect, scf::SCFDialect,
+                    tensor::TensorDialect, vector::VectorDialect>();
+  }
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    populateVectorGatherLoweringPatterns(patterns);
+    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+  }
+};
+
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestVectorLowerings() {
   PassRegistration<TestVectorToVectorLowering>();
 
   PassRegistration<TestVectorContractionLowering>();
void registerTestVectorLowerings() {
 
   PassRegistration<TestVectorScanLowering>();
 
   PassRegistration<TestVectorDistribution>();
 
   PassRegistration<TestVectorExtractStridedSliceLowering>();
 
   PassRegistration<TestCreateVectorBroadcast>();
+
+  PassRegistration<TestVectorGatherLowering>();
 }
 } // namespace test
 } // namespace mlir