[MLIR] Vectorize tensor.extract on n-D tensor

[awarzynski]

Summary

This patch implements the vectorization of tensor.extract for arbitrary
tensors. It basically extends https://reviews.llvm.org/D133786 by adding
support for n-D tensors. This works by basically flattening the
indices.

While this patch allows for more cases of tensor.extract to be
vectorised, performance on the workloads that we tested
either regressed or didn't improve. For this reason, new functionality
is hidden behind a global command-line flag:

• -linalg-vectorize-n-d-extract.

We may want remove it once the implementation is refined and we
are happy with the performance.

Extra logic in the Linalg vectorizer is added to support the new
command line option.

Related discussion: https://github.com/iree-org/iree/issues/9198

[rsuderman]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

407

You can initially construct this using numIndices for the size and create the arith::ConstantIndexOp with value 0. This avoids having branching behavior depending on the number of indices.

408

Make const auto or const int64_t.

414

This line will no longer be needed.

428

You actually don't need this saved to a vector. You only need an accumulator for mul-add. (More details below).

437

Ditto about needing a vector. You can iterate on values.

451

You can simplify the two loops below roughly with the IR equivalent below:

int64_t index[N];
int64_t shape[N];
int64_t offset = index[N - 1];

for (int i = N - 1; i > 0; --i ){
  offset = offset * shape[i] + index[i - 1]
}

It still requires the broadcast work but avoids the extraneous mul-by-1 and extra operations that are just cleaned up. It also avoids computing all the slice sizes when we just care about the global offset.

483

If you make the change to the gatherIndices declaration you can delete this entire loop.

[awarzynski]

Simplify as per comments from @rsuderman

Thank you for taking a look, Rob! Much appreciated comments and suggestions.

[dcaballe]

Thanks for working on this and addressing the feedback! LGTM. Just some minor comments!

Also wondering if we should land this under a flag (disable by default) as we know performance will degrade until we have some gather optimizations in place.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

408

nit: gatherIndices -> baseIndices? Clearer?

417

Sorry, missing something probably trivial but, where is the transpose being generated?

419

Instead of special-casing this for numIndices == 1... would it be possible to have the logic below working for just one index (and do mostly nothing)? It looks like having one index should be a base case of that code.

441

There might be a few improvements related to the offset computation. Would it make sense to move all of this to self-contained utility function? It's very likely that we play with different approaches to compute the offsets and having a that isolated into a utility function would help.

[awarzynski]

Thanks for taking a look!

In D137660#3944636, @dcaballe wrote:

Also wondering if we should land this under a flag (disable by default) as we know performance will degrade until we have some gather optimizations in place.

I agree, but will need some pointers :) The test runs TestTransformDialectInterpreterPass, but I doubt that's the pass I'd want to update here.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

417

vector.transpose is inserted in vectorizeLinalgIndex, i.e. prior to entering this method.

Note that we only need to broadcast the trailing index as that's the only one that has not been broadcast yet. ATM, I'm broadcasting "everything" to keep the code simple and rely on subsequent optimisations to remove redundant broadcasts. I'll update the broadcasting part to make this more explicit.

I suggest that we look at removing vector.transpose in a separate patch. I admit that I've not looked at all. WDTY?

419

Good point, done!

[awarzynski]

Implement suggestions from @dcaballe

• moved offset calculation to a dedicated method calculateOffsetForGLoad (perhaps there's a better name?)
• made sure that we only add broadcast for the trainling index (other indices have already been broadcast by vectorizeLinalgIndex)
• renamed gatherIndices as baseIndices

[dcaballe]

Thanks for addressing the feedback! I'm ok with landing under a flag after addressing the feedback, as stepping stone towards the final n-D gather solution. We can build on top of this iteratively.

I agree, but will need some pointers :) The test runs TestTransformDialectInterpreterPass, but I doubt that's the pass I'd want to update here.

We can use an global llvm::opt flag as this flag is temporary that will go away. The other options is to add a bool parameter to vectorize but we would have to extend the Transform dialect, etc. to provide the right value. I don't think it's worth it given that this will go away.

(perhaps there's a better name?)

calculateGatherOffset?

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

417

SG, thanks!

[awarzynski]

Hide the new functionality behind an llvm::cl option

Added a new global command line option for enabling the vectorisaiton of
tensor.extract when the input is an n-D tensor:
--linalg-vectorize-n-d-extract. The default value is "off", which means that
this patch will not change the default behaviour. I'll update the summary of
this change shortly.

Given that this is the Thanksgiving week, I'll wait until next week before
landing this. This way our friends celebrating get an opportunity to take
another look before I merge this. I will also add some reviewers that might
have an opinion about the new global options.##

[nicolasvasilache]

Please hold off with this, I am making significant changes to vectorization.
I hope I can share something tomorrow.

[nicolasvasilache]

Here is the WIP https://reviews.llvm.org/D138688

Unfortunately I am having issues on the IREE side that I cannot yet diagnose and do not know whether this is causing potentially major regressions that need to be addressed before landing.

[awarzynski]

In D137660#3951094, @nicolasvasilache wrote:

Here is the WIP https://reviews.llvm.org/D138688

Unfortunately I am having issues on the IREE side that I cannot yet diagnose and do not know whether this is causing potentially major regressions that need to be addressed before landing.

Thanks for the heads-up and for sharing! Those look like some really nice improvements, thank you!

I definitely would prefer to have this in sooner rather than later, but also don't mind waiting and definitely don't want to make anyone's life harder :) Shall I wait another week? Would that be sufficient for you?

[nicolasvasilache]

In D137660#3951361, @awarzynski wrote:

In D137660#3951094, @nicolasvasilache wrote:

Here is the WIP https://reviews.llvm.org/D138688

Unfortunately I am having issues on the IREE side that I cannot yet diagnose and do not know whether this is causing potentially major regressions that need to be addressed before landing.

Thanks for the heads-up and for sharing! Those look like some really nice improvements, thank you!

I definitely would prefer to have this in sooner rather than later, but also don't mind waiting and definitely don't want to make anyone's life harder :) Shall I wait another week? Would that be sufficient for you?

Thanks!

This is definitely a long-standing simplification for at least 2 years, now is the right time to do it before too many things are added on top.

I estimate 1 week should be good, the main blocker atm is that I have no easy way of turning downstream IREE errors into actual test cases here.

[nicolasvasilache]

In D137660#3951561, @nicolasvasilache wrote:

In D137660#3951361, @awarzynski wrote:

In D137660#3951094, @nicolasvasilache wrote:

Here is the WIP https://reviews.llvm.org/D138688

Unfortunately I am having issues on the IREE side that I cannot yet diagnose and do not know whether this is causing potentially major regressions that need to be addressed before landing.

Thanks for the heads-up and for sharing! Those look like some really nice improvements, thank you!

I definitely would prefer to have this in sooner rather than later, but also don't mind waiting and definitely don't want to make anyone's life harder :) Shall I wait another week? Would that be sufficient for you?

Thanks!

This is definitely a long-standing simplification for at least 2 years, now is the right time to do it before too many things are added on top.

I estimate 1 week should be good, the main blocker atm is that I have no easy way of turning downstream IREE errors into actual test cases here.

We iterated with @dcaballe offline. I can remove that blocker and reshuffle things on my end.

However, the CL option needs to be reworked.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

52

This is problematic, you shouldn't expect that any of this is run in a pass pipeline.
You could add a new attribute the to VectorizeOp (like we do for padding) and propagate it from there instead.

mlir/test/Dialect/Linalg/vectorization.mlir

1

This is problematic, you shouldn't expect that any of this is run in a pass pipeline.
You could add a new attribute the to VectorizeOp (like we do for padding) and propagate it from there instead.

[pzread]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

322–323

nit: This comment can be removed.

364

nit: I think it's clearer and safer to write:

extractOp.getTensor().getType()

Also do we check if the dim is dynamic? If it has been checked in another place, we can have an assert here.

[awarzynski]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

52

This is problematic, you shouldn't expect that any of this is run in a pass pipeline.

Isn't a global command-line option orthogonal to whether this functionality is wrapped in a pass or not?

You could add a new attribute the to VectorizeOp (like we do for padding) and propagate it from there instead.

I feel that this would only be sufficient for testing/experimenting with the Transform dialect. But I'd like more than that :) I'm not that familiar with the Transform dialect and might be missing something obvious - please let me know! I'm just about send and updated that preserves this global command-line flag and adds a Transform dialect attribute. Is that what you are suggesting?

364

nit: I think it's clearer and safer to write:

Thanks for the suggestion, I'll update this in the next revision!

Also do we check if the dim is dynamic? If it has been checked in another place, we can have an assert here.

It is checked here. Should I add an assert regardless?

[awarzynski]

Address PR comments

Added an attribute to VectorizeOp so that the vectorization can be controlled from the Transfer dialect level. @nicolasvasilache, is that what you had in mind?

Thanks for all the comments so far!

[dcaballe]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

52

I suggested that we could go with a global flag as this is a temporary flag that will be removed once we stabilize the performance of the gather. I actually wanted to avoid having to extend the transform dialect just for this temporary thing but happy to have it working with the transform dialect as well.

However, it's not clear to me why we need a struct (MLIRLinalgVectorizerOptions) and all the MLIRLinalgVectorizerOptions and MLIRLinalgVectorizerOptions related code. Wouldn't just adding the llvm::cl option work?

[awarzynski]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

52

However, it's not clear to me why we need a struct (MLIRLinalgVectorizerOptions) and all the MLIRLinalgVectorizerOptions and MLIRLinalgVectorizerOptions related code. Wouldn't just adding the llvm::cl option work?

Without a struct we would have a global destructor (for the vectorizeNDExtract), and that's not an option, see https://github.com/llvm/llvm-project/blob/d620bae999465f4a418c38b6451b11d80523c6de/mlir/lib/CMakeLists.txt#L1-L2 :)

This approach is actually fairly common these days (see how things are set-up in mlir-opt - it allows you to control when particular sets of options are exposed. Otherwise, all "global" options would be available at all times.

[awarzynski]

Ping.

[dcaballe]

Thanks for adding support for the Transform dialect! LGTM. I think we should be able to land this. Does anybody else have any comments?

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

52

Thanks! You are definitely more up-to-date on this than I am.

350

const SmallVector<int64, 4> & -> const SmallVectorImpl<int64_t> &

[nicolasvasilache]

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

269 ↗

(On Diff #478309)

oof this looks like a very unwelcome dependence at this level ..

@rriddle , any comments / suggestions here ?

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

53

This leaks badly all the way up to the main binary.. I'd like to see this discussed.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

1791

thanks for adding this

1853

why not have this work exactly the same way as the code 2 lines above?

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

60

If you really want a CL option, you could add a test pass (we don't want to start maintaining a pass with heuristics and perf regressions that people would depend on).
But I think the value of that is limited these days.

[dcaballe]

I suggested a global flag to avoid having to extend the Transform dialect, as the flag will go away in a few weeks. Since it's already supported in the Transform dialect, I think we could remove the global flag and add a bool option to the vectorize method. That should work for both, the Transform dialect and any other users of the vectorizer.

[awarzynski]

I see that the global option is proving more controversial than we anticipated :)

@nicolasvasilache, we don't really need the changes in mlir-opt (@dcaballe ?), but I'd like to be able to control this vectorization through other tools ( iree-compile in my case). AFAIK, a transfer dialect attribute is not sufficient for this (perhaps I'm missing something?). Global options is what folks within LLVM have been using for this sort of things, so perhaps this is not too bad?

In D137660#3972178, @dcaballe wrote:

Since it's already supported in the Transform dialect, I think we could remove the global flag and add a bool option to the vectorize method. That should work for both, the Transform dialect and any other users of the vectorizer.

This actually a bit tricky. I'd need to find a way to pass it to tensorExtractVectorizationPrecondition, but it only accepts one argument (it's type is CustomVectorizationPrecondition). Also, this would mean adding a bool argument in quite a few places, so not so tidy. I'll try again tomorrow.

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

1853

You know this stuff much better than I do, so please point me in the right direction if I'm talking nonsense.

Basically, hooks for vectorising tensor.pad and tensor.extract are registered differently ATM:

• tensor.pad goes through the OpRewritePattern API,
• for tensor.extract we just use a CustomVectorizationHook (added here).

I wanted to avoid refactoring, hence this slight inconsistency how tensor.pad and tensor.extract are treated. Do you reckon that we should go via OpRewritePattern for tensor.extract instead? If yes, then presumably we'd want that in a separate patch?

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

53

This leaks badly all the way up to the main binary

Only if you register it through linalg::registerLinalgVectorizerCLOptions(); like I did in MLIROptMain.cpp.

I'd like to see this discussed.

Are you concerned about mlir-opt specifically? Because I don't really need this in mlir-opt. I use iree-compile instead, but that's orthogonal here.

[awarzynski]

Revert changes from mlir-opt

The changes in mlir-opt are not needed for this patch to work. I'm leaving
the llvm::cl option for now, as even extending vectorize to accept additional
bool (as per Diego's sugestion) would require adding a global variable.
Otherwise, we'd need to propagage that bool to either
tensorExtractVectorizationPrecondition or VectorizeTensorExtract, but
that's not possible (i.e. we can't really change the corresponding interfaces).

ALso rebased on top of main and replaced SmallVector with SmallVectorImpl.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

1853

I see, I mistook the fact that you are vectorizing a tensor.xx op for what you are actually doing, which is vectorizing something new within the Linalg body and injecting an orthogonal control to enable it.

The easy way forward is to pass the boolean to the pattern constructor and to the vectorize function, feel free to wrap that in an options struct (i.e. your LinalgVectorizerOptions).
If you don't want to pass that forward to vectorizeAsLinalgGeneric, feel free to hoist the hooks registration one level which will also make them more visible and documented one level higher.

Despite passing the information in 3 places, this is infinitely better than an orthogonal control injection with a global, this definitely feels like an abuse to me.
Doing a quick grep in MLIR I do not see any such use and I am definitely not looking at having Linalg introduce such a mechanism, sorry.
If you feel strongly about this please start and RFC on discourse to gather more feedback.

[awarzynski]

Removed the llvm::cl logic

I removed all the logic related to llvm::cl and replaced it entirely with:

• transfer dialect attribute (vectorize_nd_extract)
• extra bool option passed to vectorize

@nicolasvasilache, I didn't really need the global, it just happened to be an interface that I am very familiar with. I'm new to this area and it's not always obvious which approach would be the least frictionless. This approach works for me, thanks for pointing me in the right direction!

I am not too happy about the extra argument in vectorizeLinalgOpPrecondition that's very specific that what I'm trying to do for tensor.extract. But equally, this hook is currently only used for tensor.extract.

[nicolasvasilache]

thanks!

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

36

spurious leftover include

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let description = [{
     This operation produces `definiteFailure` if vectorization fails for any
     reason.
     The operation always returns the handle to the target op that is expected
     to be isolated from above.
   }];
 
   let arguments = (ins PDL_Operation:$target,
                    UnitAttr:$vectorize_padding,
+                   UnitAttr:$vectorize_nd_extract,
                    UnitAttr:$disable_multi_reduction_to_contract_patterns,
                    UnitAttr:$disable_transfer_permutation_map_lowering_patterns);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let builders = [
-    OpBuilder<(ins "Value":$target, CArg<"bool", "false">:$vectorizePadding)>
+    OpBuilder<(ins "Value":$target,
+               CArg<"bool", "false">:$vectorizePadding,
+               CArg<"bool", "false">:$vectorizeNDExtract)>,
   ];
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
         ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // LINALG_TRANSFORM_OPS

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 ///   3. Take a partial slice of the full view in step 2. and copy into it.
 ///
 /// Return the modified linalg op (the modification happens in place) as well
 /// as all the copy ops created.
 FailureOr<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
                                     const LinalgPromotionOptions &options);
 
 /// Emit a suitable vector form for a Linalg op with fully static shape.
-LogicalResult vectorize(RewriterBase &builder, LinalgOp linalgOp);
+LogicalResult vectorize(RewriterBase &builder, LinalgOp linalgOp,
+                        bool vectorizeNDExtract = false);
 
 /// Emit a suitable vector form for a Copy op with fully static shape.
 LogicalResult vectorizeCopy(RewriterBase &builder, memref::CopyOp copyOp);
 
 /// Emit a loop nest of `scf.for` with the proper body for `linalgOp`.
 FailureOr<LinalgLoops> linalgOpToLoops(PatternRewriter &rewriter,
                                        LinalgOp linalgOp);
 

 // Preconditions that ensure the corresponding transformation succeeds and can
 // be applied as a rewrite pattern.
 //===----------------------------------------------------------------------===//
 /// Promote memref.subviews feeding linalg-on-buffers operations.
 LogicalResult promoteSubviewsPrecondition(Operation *op,
                                           LinalgPromotionOptions options);
 
 /// Return success if the operation can be vectorized.
-LogicalResult vectorizeLinalgOpPrecondition(LinalgOp linalgOp);
+LogicalResult vectorizeLinalgOpPrecondition(LinalgOp linalgOp,
+                                            bool vectorizeNDExtract = false);
 
 //===----------------------------------------------------------------------===//
 // Transformations exposed as rewrite patterns.
 //===----------------------------------------------------------------------===//
 
 using TileSizeComputationFunction =
     std::function<SmallVector<Value, 4>(OpBuilder &, Operation *)>;
 

 /// These patterns are meant to apply in a complementary fashion. Benefits
 /// are used to encode a certain ordering of pattern application. To avoid
 /// scattering magic constants throughout the code base, the patterns must be
 /// added with this function. `baseBenefit` can be used to offset the benefit
 /// of all tensor::PadOp vectorization patterns by a certain value.
 void populatePadOpVectorizationPatterns(RewritePatternSet &patterns,
                                         PatternBenefit baseBenefit = 1);
 
+void populateExtractOpVectorizationPatterns(RewritePatternSet &patterns,
+                                            PatternBenefit baseBenefit = 1);
+
 /// Match and rewrite for the pattern:
 /// ```
 ///    %alloc = ...
 ///    [optional] %view = memref.view %alloc ...
 ///    %subView = subview %allocOrView ...
 ///    [optional] linalg.fill(%allocOrView, %cst) ...
 ///    ...
 ///    memref.copy(%in, %subView) ...

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

void transform::TileToScfForOp::getEffects(
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // VectorizeOp
 //===----------------------------------------------------------------------===//
 
 void transform::VectorizeOp::build(OpBuilder &builder, OperationState &result,
-                                   Value target, bool vectorizePadding) {
+                                   Value target, bool vectorizePadding,
+                                   bool vectorizeExtract) {
   result.addOperands(target);
   if (vectorizePadding) {
     result.addAttribute(VectorizeOp::getVectorizePaddingAttrName(result.name),
                         builder.getUnitAttr());
   }
+  if (vectorizeExtract) {
+    result.addAttribute(VectorizeOp::getVectorizeNdExtractAttrName(result.name),
+                        builder.getUnitAttr());
+  }
   result.addTypes(pdl::OperationType::get(builder.getContext()));
 }
 
 namespace {
 /// This is an helper only to call vectorize via a pattern inside of
 /// VectorizeOp::applyToOne.
 struct VectorizationPattern : public RewritePattern {
-  explicit VectorizationPattern(MLIRContext *context)
-      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
+  explicit VectorizationPattern(MLIRContext *context,
+                                bool vectorizeExtract = false)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context),
+        vectorizeNDExtract(vectorizeExtract) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
     if (!linalgOp)
       return rewriter.notifyMatchFailure(op, "expected Linalg Op");
-    return vectorize(rewriter, linalgOp);
+    return vectorize(rewriter, linalgOp, vectorizeNDExtract);
   }
+
+private:
+  /// Controls whether to vectorize `tensor.extract` when the input tensor is
+  /// rank >= 2.
+  bool vectorizeNDExtract = false;
 };
 } // namespace
 
 DiagnosedSilenceableFailure
 transform::VectorizeOp::applyToOne(Operation *target,
                                    SmallVectorImpl<Operation *> &results,
                                    transform::TransformState &state) {
   if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
     auto diag = this->emitOpError("requires isolated-from-above targets");
     diag.attachNote(target->getLoc()) << "non-isolated target";
     return DiagnosedSilenceableFailure::definiteFailure();
   }
 
   MLIRContext *ctx = getContext();
   RewritePatternSet patterns(ctx);
-  patterns.add<VectorizationPattern>(ctx);
+  patterns.add<VectorizationPattern>(ctx, getVectorizeNdExtract());
 
   if (!getDisableTransferPermutationMapLoweringPatterns())
     vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
 
   if (!getDisableMultiReductionToContractPatterns())
     vector::populateVectorReductionToContractPatterns(patterns);
 
   patterns.add<linalg::LinalgCopyVTRForwardingPattern,
                linalg::LinalgCopyVTWForwardingPattern>(ctx,
                                                        /*benefit=*/2);
   vector::TransferReadOp::getCanonicalizationPatterns(patterns, ctx);
   vector::TransferWriteOp::getCanonicalizationPatterns(patterns, ctx);
 
   patterns.add<CopyVectorizationPattern>(ctx);
 
   if (getVectorizePadding())
     linalg::populatePadOpVectorizationPatterns(patterns);
 
   if (failed(applyPatternsAndFoldGreedily(target, std::move(patterns))))
     return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
 
   results.push_back(target);
   return DiagnosedSilenceableFailure(success());
 }
 
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <type_traits>
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 #define DEBUG_TYPE "linalg-vectorization"
 
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X)
 
 /// Try to vectorize `convOp` as a convolution.
 static FailureOr<Operation *> vectorizeConvolution(RewriterBase &rewriter,
                                                    LinalgOp convOp);
 
 /// Return the unique instance of OpType in `block` if it is indeed unique.
 /// Return null if none or more than 1 instances exist.
 template <typename OpType>
 static OpType getSingleOpOfType(Block &block) {
   OpType res;
   block.walk([&](OpType op) {
     if (res) {
       res = nullptr;
       return WalkResult::interrupt();
     }
     res = op;
     return WalkResult::advance();
   });
   return res;
 }
 
 /// Given an indexing `map` coming from a LinalgOp indexing, restricted to a
if (!write->getResults().empty())
     return write->getResult(0);
   return Value();
 }
 
 // Custom vectorization precondition function type. This is intented to be used
 // with CustomVectorizationHook. Returns success if the corresponding custom
 // hook can vectorize the op.
 using CustomVectorizationPrecondition =
-    std::function<LogicalResult(Operation *)>;
+    std::function<LogicalResult(Operation *, bool)>;
 
 // Custom vectorization function type. Produce a vector form of Operation*
 // assuming all its vectorized operands are already in the BlockAndValueMapping.
 // Return nullptr if the Operation cannot be vectorized.
 using CustomVectorizationHook = std::function<VectorizationResult(
     Operation *, const BlockAndValueMapping &)>;
 
 /// Helper function to vectorize the terminator of a `linalgOp`. New result
vectorizeLinalgIndex(RewriterBase &rewriter, Operation *op, LinalgOp linalgOp) {
   std::swap(transposition.back(), transposition[indexOp.getDim()]);
   auto transposeOp =
       rewriter.create<vector::TransposeOp>(loc, broadCastOp, transposition);
   return VectorizationResult{VectorizationStatus::NewOp, transposeOp};
 }
 
 /// Helper function to check if the tensor.extract can be vectorized by the
 /// custom hook vectorizeTensorExtract.
-static LogicalResult tensorExtractVectorizationPrecondition(Operation *op) {
+static LogicalResult
+tensorExtractVectorizationPrecondition(Operation *op, bool vectorizeNDExtract) {
   tensor::ExtractOp extractOp = dyn_cast<tensor::ExtractOp>(op);
   if (!extractOp)
     return failure();
 
-  // Currently only supports extraction with an 1-D index.
-  if (extractOp.getIndices().size() != 1)
+  if (extractOp.getIndices().size() != 1 && !vectorizeNDExtract)
     return failure();
 
   if (!VectorType::isValidElementType(extractOp.getIndices()[0].getType()))
     return failure();
 
   if (llvm::any_of(extractOp->getResultTypes(), [](Type type) {
         return !VectorType::isValidElementType(type);
       })) {
     return failure();
   }
 
   return success();
 }
 
+/// Calculates the offsets (`$index_vec`) for `vector.gather` operations
+/// generated from `tensor.extract`. The offset is calculated as follows
+/// (example using scalar values):
+///
+///    offset = extractOp.indices[0]
+///    for (i = 1; i < numIndices; i++)
+///      offset = extractOp.dimSize[i] * offset + extractOp.indices[i];
+///
+/// For tensor<45 x 80 x 15 x f32> and index [1, 2, 3], this leads to:
+///  offset = ( ( 1 ) * 80 +  2 ) * 15  + 3
+static Value
+calculateGatherOffset(OpBuilder &b, tensor::ExtractOp extractOp,
+                      const BlockAndValueMapping &bvm,
+                      const SmallVectorImpl<int64_t> &targetShape) {
+  // The vector of indices for GatherOp should be shaped as the output vector
+  auto indexVecType = VectorType::get(targetShape, b.getIndexType());
+  auto loc = extractOp.getLoc();
+
+  Value offset = b.create<vector::BroadcastOp>(
+      loc, indexVecType, bvm.lookup(extractOp.getIndices()[0]));
+
+  const size_t numIndices = extractOp.getIndices().size();
+  for (size_t i = 1; i < numIndices; i++) {
+    auto dimSizeBcast = b.create<vector::BroadcastOp>(
+        loc, indexVecType,
+        b.create<arith::ConstantIndexOp>(
+            loc,
+            extractOp.getTensor().getType().cast<ShapedType>().getDimSize(i)));
+    offset = b.create<arith::MulIOp>(loc, offset, dimSizeBcast);
+
+    auto originalIndexBcast = bvm.lookup(extractOp.getIndices()[i]);
+    if (i == numIndices - 1) {
+      // We only need an additional broadcast for the trailing index. All other
+      // indices have already been broadcast by `vectorizeLinalgIndex` to match
+      // the output size.
+      originalIndexBcast = b.create<vector::BroadcastOp>(
+          loc, indexVecType, bvm.lookup(extractOp.getIndices()[i]));
+    }
+
+    offset = b.create<arith::AddIOp>(loc, originalIndexBcast, offset);
+  }
+
+  return offset;
+}
+
 /// Helper function to vectorize the tensor.extract operations. Returns
 /// VectorizationStatus::NewOp to signal the vectorization algorithm that it
 /// should map the produced operations. This function is meant to be used as a
 /// CustomVectorizationHook.
 static VectorizationResult
 vectorizeTensorExtract(RewriterBase &rewriter, Operation *op, LinalgOp linalgOp,
                        const BlockAndValueMapping &bvm) {
   tensor::ExtractOp extractOp = dyn_cast<tensor::ExtractOp>(op);
   if (!extractOp)
     return VectorizationResult{VectorizationStatus::Failure, nullptr};
   auto loc = extractOp.getLoc();
 
-  // Currently only supports extraction with an 1-D index. Checked in the
-  // tensorExtractVectorizationPrecondition.
-  assert(extractOp.getIndices().size() == 1);
-
-  auto indexVec = bvm.lookup(extractOp.getIndices()[0]);
   // Compute the static loop sizes of the extract op.
   auto targetShape = linalgOp.computeStaticLoopSizes();
 
-  SmallVector<Value> gatherIndices;
-  gatherIndices.push_back(rewriter.create<arith::ConstantIndexOp>(loc, 0));
-
+  auto resultType =
+      VectorType::get(targetShape, extractOp.getResult().getType());
   auto maskConstantOp = rewriter.create<arith::ConstantOp>(
       loc, DenseIntElementsAttr::get(
                VectorType::get(targetShape, rewriter.getI1Type()),
                /*value=*/true));
-
-  auto resultType =
-      VectorType::get(targetShape, extractOp.getResult().getType());
   auto passThruConstantOp =
       rewriter.create<arith::ConstantOp>(loc, rewriter.getZeroAttr(resultType));
 
+  // Base indices are currently set to 0. We will need to re-visit if more
+  // generic scenarios are to be supported.
+  SmallVector<Value> baseIndices(
+      extractOp.getIndices().size(),
+      rewriter.create<arith::ConstantIndexOp>(loc, 0));
+
+  Value offset = calculateGatherOffset(rewriter, extractOp, bvm, targetShape);
+
+  // Generate the gather load
   auto gatherOp = rewriter.create<vector::GatherOp>(
-      loc, resultType, extractOp.getTensor(), gatherIndices, indexVec,
+      loc, resultType, extractOp.getTensor(), baseIndices, offset,
       maskConstantOp, passThruConstantOp);
 
   return VectorizationResult{VectorizationStatus::NewOp, gatherOp};
 }
 
 /// Emit reduction operations if the shapes of the value to reduce is different
 /// that the result shape.
 // Note: this is a true builder that notifies the OpBuilder listener.
 // TODO: Consider moving as a static helper on the ReduceOp.
 static Operation *reduceIfNeeded(OpBuilder &b, LinalgOp linalgOp,
                                  Operation *op, Value reduceValue,
                                  Value initialValue,
                                  const BlockAndValueMapping &bvm) {
   Value reduceVec = bvm.lookup(reduceValue);
   Value outputVec = bvm.lookup(initialValue);
   auto reduceType = reduceVec.getType().dyn_cast<VectorType>();
   auto outputType = outputVec.getType().dyn_cast<VectorType>();
   // Reduce only if needed as the value may already have been reduce for
   // contraction vectorization.
   if (!reduceType ||
       (outputType && reduceType.getShape() == outputType.getShape()))
     return nullptr;
   SmallVector<bool> reductionMask = getReductionMask(linalgOp);
   return buildMultiDimReduce(b, op, reduceVec, outputVec, reductionMask);
 }
 
 /// Generic vectorization for a single operation `op`, given already vectorized
 /// operands carried by `bvm`. Vectorization occurs as follows:
 ///   1. Try to apply any of the `customVectorizationHooks` and return its
 ///   result on success.
 ///   2. Clone any constant in the current scope without vectorization: each
 ///   consumer of the constant will later determine the shape to which the
 ///   constant needs to be broadcast to.
 ///   3. Fail on any remaining non `ElementwiseMappable` op. It is the purpose
 ///   of the `customVectorizationHooks` to cover such cases.
 ///   4. Clone `op` in vector form to a vector of shape prescribed by the first
 ///   operand of maximal rank. Other operands have smaller rank and are
 ///   broadcast accordingly. It is assumed this broadcast is always legal,
 ///   otherwise, it means one of the `customVectorizationHooks` is incorrect.
 ///
 /// This function assumes all operands of `op` have been vectorized and are in
 /// the `bvm` mapping. As a consequence, this function is meant to be called  on
for (auto &customFunc : customVectorizationHooks) {
       return result;
     }
   }
 
   // 2. Constant ops don't get vectorized but rather broadcasted at their users.
   // Clone so that the constant is not confined to the linalgOp block .
   if (isa<arith::ConstantOp, func::ConstantOp>(op))
     return VectorizationResult{VectorizationStatus::NewOp, rewriter.clone(*op)};
 
   // 3. Only ElementwiseMappable are allowed in the generic vectorization.
   if (!OpTrait::hasElementwiseMappableTraits(op))
     return VectorizationResult{VectorizationStatus::Failure, nullptr};
 
   // 4 . Check if the operation is a reduction.
   SmallVector<std::pair<Value, Value>> reductionOperands;
   for (Value operand : op->getOperands()) {
     auto arg = operand.dyn_cast<BlockArgument>();
vectorizeAsLinalgGeneric(RewriterBase &rewriter, LinalgOp linalgOp,
   // 4a. Register CustomVectorizationHook for yieldOp.
   CustomVectorizationHook vectorizeYield =
       [&](Operation *op,
           const BlockAndValueMapping &bvm) -> VectorizationResult {
     return vectorizeLinalgYield(rewriter, op, bvm, linalgOp, newResults);
   };
   hooks.push_back(vectorizeYield);
 
-  // 4rewriter. Register CustomVectorizationHook for indexOp.
+  // 4b. Register CustomVectorizationHook for indexOp.
   CustomVectorizationHook vectorizeIndex =
       [&](Operation *op,
           const BlockAndValueMapping &bvm) -> VectorizationResult {
     return vectorizeLinalgIndex(rewriter, op, linalgOp);
   };
   hooks.push_back(vectorizeIndex);
 
   // 4c. Register CustomVectorizationHook for extractOp.
if (!reduceOp || !getCombinerOpKind(reduceOp)) {
       return failure();
     }
   }
   return success();
 }
 
 static LogicalResult vectorizeStaticLinalgOpPrecondition(
     linalg::LinalgOp op,
-    ArrayRef<CustomVectorizationPrecondition> customPreconditions) {
+    ArrayRef<CustomVectorizationPrecondition> customPreconditions,
+    bool vectorizeNDExtract) {
 
   // All types in the body should be a supported element type for VectorType.
   for (Operation &innerOp : op->getRegion(0).front()) {
     // Check if any custom hook can vectorize the inner op.
     if (llvm::any_of(
             customPreconditions,
             [&](const CustomVectorizationPrecondition &customPrecondition) {
-              return succeeded(customPrecondition(&innerOp));
+              return succeeded(
+                  customPrecondition(&innerOp, vectorizeNDExtract));
             })) {
       continue;
     }
     if (llvm::any_of(innerOp.getOperandTypes(), [](Type type) {
           return !VectorType::isValidElementType(type);
         })) {
       return failure();
     }
static LogicalResult vectorizeStaticLinalgOpPrecondition(
   }
   if (failed(reductionPreconditions(op))) {
     LDBG("precondition failed: reduction preconditions");
     return failure();
   }
   return success();
 }
 
-LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(LinalgOp linalgOp) {
+LogicalResult
+mlir::linalg::vectorizeLinalgOpPrecondition(LinalgOp linalgOp,
+                                            bool vectorizeNDExtract) {
   // All types must be static shape to go to vector.
   if (linalgOp.hasDynamicShape()) {
     LDBG("precondition failed: dynamic shape");
     return failure();
   }
 
   SmallVector<CustomVectorizationPrecondition> customPreconditions;
 
   // Register CustomVectorizationPrecondition for extractOp.
   customPreconditions.push_back(tensorExtractVectorizationPrecondition);
 
-  return vectorizeStaticLinalgOpPrecondition(linalgOp, customPreconditions);
+  return vectorizeStaticLinalgOpPrecondition(linalgOp, customPreconditions,
+                                             vectorizeNDExtract);
 }
 
-LogicalResult mlir::linalg::vectorize(RewriterBase &rewriter,
-                                      LinalgOp linalgOp) {
-  if (failed(vectorizeLinalgOpPrecondition(linalgOp)))
+LogicalResult mlir::linalg::vectorize(RewriterBase &rewriter, LinalgOp linalgOp,
+                                      bool vectorizeNDExtract) {
+  if (failed(vectorizeLinalgOpPrecondition(linalgOp, vectorizeNDExtract)))
     return failure();
 
   SmallVector<Value> results;
   // TODO: isaConvolutionOpInterface that can also infer from generic
   // features. Will require stride/dilation attributes inference.
   FailureOr<Operation *> convOr = vectorizeConvolution(rewriter, linalgOp);
   if (succeeded(convOr)) {
     llvm::append_range(results, (*convOr)->getResults());
   } else {
-    if (failed(vectorizeLinalgOpPrecondition(linalgOp)))
+    if (failed(vectorizeLinalgOpPrecondition(linalgOp, vectorizeNDExtract)))
       return failure();
     LDBG("Vectorize generic by broadcasting to a common shape: " << linalgOp);
     if (failed(vectorizeAsLinalgGeneric(rewriter, linalgOp, results)))
       return failure();
   }
 
   if (!results.empty())
     rewriter.replaceOp(linalgOp, results);

mlir/test/Dialect/Linalg/vectorization.mlir

 // RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s
 
 // -----
 
 // CHECK-LABEL: contraction_dot
 func.func @contraction_dot(%A: memref<1584xf32>, %B: memref<1584xf32>, %C: memref<f32>) {
 
 // CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584xf32>
 // CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [0] : vector<1584xf32> to f32
^bb1(%arg1: !pdl.operation):
   %2 = transform.structured.vectorize %1
 }
 
 // -----
 
 #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2)>
 #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
 #map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
-func.func @not_vectorize_nd_tensor_extract(%arg0: tensor<3x3xf32>, %arg1: tensor<4x3xi32>, %arg2: tensor<4x3xi32>, %arg3: tensor<4x7x2xf32>, %arg4: tensor<4x7x3x2xf32>) -> tensor<4x7x3x2xf32> {
+func.func @vectorize_nd_tensor_extract(%arg0: tensor<3x3xf32>, %arg1: tensor<4x3xi32>, %arg2: tensor<4x3xi32>, %arg3: tensor<4x7x2xf32>, %arg4: tensor<4x7x3x2xf32>) -> tensor<4x7x3x2xf32> {
   %2 = linalg.generic {
     indexing_maps = [#map0, #map0, #map1, #map2],
     iterator_types = ["parallel", "parallel", "parallel", "parallel"]
   } ins(%arg1, %arg2, %arg3 : tensor<4x3xi32>, tensor<4x3xi32>, tensor<4x7x2xf32>) outs(%arg4 : tensor<4x7x3x2xf32>) {
   ^bb0(%arg5: i32, %arg6: i32, %arg7: f32, %arg8: f32):
     %3 = arith.index_cast %arg5 : i32 to index
     %4 = arith.index_cast %arg6 : i32 to index
     %7 = tensor.extract %arg0[%3, %4] : tensor<3x3xf32>
     linalg.yield %7 : f32
   } -> tensor<4x7x3x2xf32>
   return %2 : tensor<4x7x3x2xf32>
 }
-// CHECK-LABEL: func.func @not_vectorize_nd_tensor_extract
-// CHECK: tensor.extract
+// CHECK-LABEL: func.func @vectorize_nd_tensor_extract
+// CHECK-SAME: %[[ARG0:.*]]: tensor<3x3xf32>
+// CHECK-SAME: %[[ARG1:arg1]]: tensor<4x3xi32>
+// CHECK-SAME: %[[ARG2:arg2]]: tensor<4x3xi32>
+// CHECK-SAME: %[[ARG3:.*]]: tensor<4x7x2xf32>
+// CHECK-SAME: %[[ARG4:.*]]: tensor<4x7x3x2xf32>
+// CHECK:    %[[C0:.*]] = arith.constant 0 : index
+// CHECK:    %[[C0_i32:.*]] = arith.constant 0 : i32
+// CHECK:    %[[CST:.*]] = arith.constant dense<3> : vector<7x2x4x3xindex>
+// CHECK:    %[[CST_1:.*]] = arith.constant dense<true> : vector<4x7x3x2xi1>
+// CHECK:    %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<4x7x3x2xf32>
+// CHECK:    %[[V0:.*]] = vector.transfer_read %[[ARG1]][%[[C0]], %[[C0]]], %[[C0_i32]] {in_bounds = [true, true]} : tensor<4x3xi32>, vector<4x3xi32>
+// CHECK:    %[[V1:.*]] = vector.transfer_read %[[ARG2]][%[[C0]], %[[C0]]], %[[C0_i32]] {in_bounds = [true, true]} : tensor<4x3xi32>, vector<4x3xi32>
+// CHECK:    %[[CAST:.*]] = arith.index_cast %[[V0]] : vector<4x3xi32> to vector<4x3xindex>
+// CHECK:    %[[B1:.*]] = vector.broadcast %[[CAST]] : vector<4x3xindex> to vector<7x2x4x3xindex>
+// CHECK:    %[[CAST_1:.*]] = arith.index_cast %[[V1]] : vector<4x3xi32> to vector<4x3xindex>
+// CHECK:    %[[B2:.*]] = vector.broadcast %[[CAST_1]] : vector<4x3xindex> to vector<7x2x4x3xindex>
+// CHECK:    %[[MULI:.*]] = arith.muli %[[B1]], %[[CST]] : vector<7x2x4x3xindex>
+// CHECK:    %[[ADDI:.*]] = arith.addi %[[B2]], %[[MULI]] : vector<7x2x4x3xindex>
+// CHECK:    %[[T:.*]] = vector.transpose %[[ADDI]], [2, 0, 3, 1] : vector<7x2x4x3xindex> to vector<4x7x3x2xindex>
+// CHECK:    %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]]] [%[[T]]], %[[CST_1]], %[[PASSTHRU]] : tensor<3x3xf32>, vector<4x7x3x2xindex>, vector<4x7x3x2xi1>, vector<4x7x3x2xf32> into vector<4x7x3x2xf32>
+// CHECK:    vector.transfer_write %[[GATHER]], %[[ARG4]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true, true]} : vector<4x7x3x2xf32>, tensor<4x7x3x2xf32>
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
   %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
-  %2 = transform.structured.vectorize %1
+  %2 = transform.structured.vectorize %1 { vectorize_nd_extract }
 }
 
 // -----
 
 func.func @vectorize_map(%arg0: memref<64xf32>,
     %arg1: memref<64xf32>, %arg2: memref<64xf32>) {
   linalg.map ins(%arg0, %arg1 : memref<64xf32>, memref<64xf32>)
              outs(%arg2 : memref<64xf32>)