[mlir][Vector] Initial masking support in Linalg vectorizer

[dcaballe]

Summary

This patch introduces the initial bits to support vector masking
using the vector.mask operation. Vectorization changes should be
NFC for non-masked cases. We can't test masked cases directly until
we extend the Transform dialect to support masking.

[dcaballe]

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

1809

Hey @nicolasvasilache, do you think you could help extending the Transform dialect so that we can provide the vector sizes for masked dims?

[dcaballe]

Fixing a couple of issues when no vector sizes for masked dimensions are not provided

[nicolasvasilache]

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

1809

Happy to!

Could you temporarily hardcode some size in there and add some test IR so I can see what I should expect?

This will likely require a new transform op that is not a blanket "vectorize the world" so that we can pass the information you want at a finer granularity. This will likely need some iteration to get to a reasonably scalable usage.

Left some other review comments in the meantime.

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

73

super-nit: can we make the line spacing uniform between methods(here) and members(below)?

108

plz avoid unsigned everywhere, we know by now this is not meant for expressing >=0 but should really be used for bit-twiddling or when we really really nerd the extra bit.

108

Should this be a (better named) helper on the LinalgOp interface?
This seems to reimplement common functionality (but maybe not exactly) and rely on deep internal Linalg assumptions (e.g. all possible ways of defining the extent of an iterator have to match and you can therefore take the first one).

The name makes it hard to understand what it does and we should be doing any such manipulation in a very localized place in LinalgOp.

136

unsigned purge here and everywhere plz

153

Can we call these staticUpperBounds everywhere?
And the other ones dynamicUpperBounds ?

This seems easier to me to relate to what we're looking to do instead of vecSizesForMaskedDims and extractDynamicVectorDimValues.

157

How about early exit here ?

if (!linalgOp.hasDynamicShape()) {
  canonicalVecShape = linalgOp.getStaticLoopRanges();
  return success();
}

I don't think you need the checks and debugs after that in the static case?

165

Seems fishy, what happens in this case ?

I'd expect this to be not fail gracefully ..
Make it an assert and lift logic to the precondition to avoid this?

Edit: ah scratch that, I see that this is just after the precondition, can we make it part of the precondition?

170

Can we sprinkle a few precompute prefixes in some of these APIs to make it clear what happens at init time?

177

LLVM_DEBUG(llvm::interleaveComma(canonicalVecShape, llvm::dbgs() << ...));

247

pass RewriterBase here and everywhere possible post https://reviews.llvm.org/D137922 plz

272

Plz use updateRootInPlace once RewriterBase is piped through.

810

nit: can we spell this as:

// 3.a. Convert the indexing map for this input/output to a transfer read...
...
/// 3.a.i  For input reads we use the canonical vector shape.
if (linalgOp.isDpsInput(opOperand))
  ...
} else {
  /// 3.a.ii For output reads (iteration-carried dependence, e.g., reductions)
  ...

// 3.b. If masked, set in-bounds to true.
...
// 3.c. Not all ops support 0-d vectors,

1044

Can we make this init the state as part of the precondition?

mlir/lib/IR/AffineMap.cpp

344

Fails when called on a non-projected-permutation.

is misleading here.

It expects a projected permutation otherwise it crashes.
Failing would have returned llvm::None without crashing unless I am missing something?

347

Better name and doc please, this is much too confusing.

In fact I think you can just do something like
llvm::find(map.getResults(), AffineDimExpr::get(input)) at the client and avoid adding more APIs to AffineMap

[dcaballe]

• Addressed feedback
• New changes around vectorization initState and canonical vector shape computation

[dcaballe]

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

108

plz avoid unsigned everywhere, we know by now this is not meant for expressing >=0 but should really be used for bit-twiddling or when we really really nerd the extra bit.

Sorry,dim is a misnomer as it refers to the dimension position, not the size. It has to be unsigned as that's what it's expected by the AffineMap::getXXXYYYPosition().
Added to Pos to a few names. Hopefully that's better.

Moved it to LinalgOp interface. I can't think of a better name for the utility... It's just mapping an iteration space dimension to a dimension of an operand... Any other suggestion? I'm happy to replace with with any other existing utility but I couldn't find any.

136

As explained before, renamed to indicate that it's the position, not the size.

153

I had already renamed this locally to inputVectorSizes and change a bit the meaning. The input sizes are now taken into account to compute the canonical vector shape and if they are also provided for static shapes they should match the size of the static shapes. We are passing them all now to simplify the client API, including the transform dialect, as it's easier to provide all the vector sizes than having to filter out the static ones. Let me know if that works.

157

This is gone now. This code has changed a bit in the last version.

165

Let me know if it makes more sense in the new version, where the inputVectorSizes, if provided, should match the linalgop.getNumLoops(). Otherwise, this would be a bug.

170

Much better!

177

Ah! I didn't know this utility! It's been such a pain to always print SmallVector's...
Thanks!

247

I think we are going in the opposite direction based on the review comments?

1044

Probably better to separate the concerns. I think even part of the precondition checks are reused outside of the vectorizer. A public interface was introduced recently.

mlir/lib/IR/AffineMap.cpp

347

I had renamed this like 10 times. It's a difficult name. Hopefully it's better now :).

[nicolasvasilache]

mlir/lib/IR/AffineMap.cpp

347

Wait .. what do I see just above .. literally the same functionality modulo an assert ..

Can we just have a single

Optional<int64_t> AffineMap::getResultPosition(AffineExpr e) const {
  for (int64_t i = 0, numResults = getNumResults(); i < numResults; i++)
    if (getResult(i) == e)
      return i;
  return llvm::None;
}

and let clients do the assertions they want ?

It seems very counterproductive to have all these special case functions with slightly varying assertions and hard to grok names ..

[dcaballe]

Addressed feedback + minor fixes.
Please ignore the AffineMap utility. It will be removed after rebasing on top of D138946.

[nicolasvasilache]

I do not understand the implication of

// TODO: We mask the transfer.transfer_write here because this op is
// special-cased. A linalg.yield may produced multiple vector.transfer_write
// ops and can't be mapped using BlockAndValueMapping.
AffineMap opOperandMap = linalgOp.getMatchingIndexingMap(opOperand);
write = state.maskOperation(b, write, linalgOp, opOperandMap);

.. also I am not seeing any test changes, so it seems you are adding a lot of code that is not tested and not activated ?

mlir/include/mlir/IR/AffineMap.h

178

I think this goes away with the rebase, just flagging for removal so we don't forget.

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

68

Why not make this a ctor?

247

We cannot have code like this:

Operation *maskOpTerminator = &maskOp.getMaskRegion().front().back();
for (auto &en : llvm::enumerate(opToMask->getResults()))
  en.value().replaceAllUsesExcept(maskOp.getResult(en.index()),
                                  maskOpTerminator);

it must use a RewriterBase with updateRootInPlace

283

Hmm .. what's the contract between this createRegionMask lambda and the insertion points during
builder.create<vector::MaskOp> ?

I've seen too much ugly stuff re. insertion points leaking across function call boundaries.

Let's add an OpBuilder::InsertionGuard at the top of this function.

299–300

this must use a RewriterBase with updateRootInPlace

543

nit: produce

[dcaballe]

Addressed feedback.

[dcaballe]

.. also I am not seeing any test changes, so it seems you are adding a lot of code that is not tested and not activated ?

Changes in the overall vectorization algorithm are tested with existing vectorization tests. This patch is NFC for those. Masking is not enabled if inputVectorSizes are not provided. If they are provided, only elementwise ops without reductions and fully dynamic shapes are vectorized.
I can't add unit tests until the new operation for masked vectorization is added to the transform dialect, as we no longer have the vectorizer testing pass. However, this PR has been extensively tested in IREE, both with and without masking for even more cases than the currently supported right now.
Waiting on the new transform dialect op to land to add more tests.

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

68

Because it can fail, at least for now. This will change once we support more cases with masking. Then, we could assert and turn it into a constructor.

247

I can use updateRootInPlace when we have a rewriter here but I can't do any replacement method because opToMask is moved inside the mask region, not replaced.

283

Added guard. This is a simple lambda to create the op region. We follow the same approach for scf.if and other region ops. The only contract is that the region needs to have a vector::YieldOp, which is described in the vector.mask doc.

299–300

Added TODO until we have a rewriter.

543

I do not understand the implication of

TODO: We mask the transfer.transfer_write here because this op is
special-cased. A linalg.yield may produced multiple vector.transfer_write
// ops and can't be mapped using BlockAndValueMapping.

Good point. This is a comment for an old problem. I removed it and moved this code to buildVectorWrite.

mlir/lib/IR/AffineMap.cpp

347

Extracted this to https://reviews.llvm.org/D138946

[dcaballe]

Added testing support to Transform dialect + tests

[dcaballe]

Rebase + remove dead code (wrong rebase)

[nicolasvasilache]

Waiting on the new transform dialect op to land to add more tests.

Thanks for integrating it and adding tests, the testing part LGTM, till need to make another pass on the last version of the code.

[nicolasvasilache]

Thanks @dcaballe !

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

1120

"definite failure"

1138

can you add a TODO: applyToOne plz ?

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

222

Use the find API and the iterator you get from it to avoid multi-lookups

263

nit: uppermension :) ?

490

can you add a TODO to tighten op semantics so that we don't mix inbounds and mask since this is well defined?

826

I'll need to revisit all this in light of the broadcast separation.
From a cursory glance this looks reasonable, let's land and iterate.

972

ok as a first appox.

[dcaballe]

Thanks! I addressed comments. Landing now...

[nicolasvasilache]

@dcaballe, I was looking at this PR as I was doing some spelunking and I realize we are not testing the case of the SSA value as well as the error case when the SSA value is not a constant.

Can we please add the missing tests in a a followup ?

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
         result.reserve($_op->getNumOperands());
         llvm::transform(
           this->getOperation()->getOpOperands(),
           std::back_inserter(result),
           [](OpOperand &opOperand) { return &opOperand; });
         return result;
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Given a dimension of the iteration space of a Linalg operation, finds an
+        operand in the operation that is defined on such dimension. Returns
+        whether such operand was found or not. If found, also returns the
+        operand value and the dimension position within the operand.
+      }],
+      /*retTy=*/"LogicalResult",
+      /*methodName=*/"mapIterationSpaceDimToOperandDim",
+      /*args=*/(ins "unsigned":$dimPos,
+                    "::mlir::Value &":$operand,
+                    "unsigned &":$operandDimPos),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        // Retrieve the operand and its dimension position from the first
+        // operand with a permutation map that is defined on such dimension.
+        for (auto [i, idxMap] : llvm::enumerate($_op.getIndexingMapsArray())) {
+          if (idxMap.isProjectedPermutation()) {
+            if (auto mayOperandDim = idxMap.getResultPosition(
+                getAffineDimExpr(dimPos, idxMap.getContext()))) {
+              operand = $_op->getOperand(i);
+              operandDimPos = *mayOperandDim;
+              return success();
+            }
+          }
+        }
+
+        return failure();
+      }]
+    >,
     //===------------------------------------------------------------------===//
     // Linalg generalization hooks.
     //===------------------------------------------------------------------===//
     InterfaceMethod<
       /*desc=*/[{
         Hook to provide a custom AffineMap used to compute all the operand
         subshapes given loop bounds. This is used to answer the question: "given
         an iteration space over the codomain, what are the subshapes of the

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

def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
         ::llvm::SmallVectorImpl<::mlir::Operation *> &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
+def MaskedVectorizeOp : Op<Transform_Dialect, "structured.masked_vectorize",
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     TransformOpInterface]> {
+  let description = [{
+    Vectorize the target ops, which must be Linalg ops, with masked vectors
+    of the specified size.
+
+    The vector sizes can be either static or dynamic (SSA values). In case of
+    SSA values, the handle must be mapped to exactly one payload op with
+    exactly one index-typed result.
+
+    #### Return modes:
+
+    This operation produces a definite if the dynamic vector sizes (SSA values)
+    do not satify the constraints mentioned above. It produces a silenceable
+    failure if at least one target op is not a Linalg op or fails to vectorize.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       Variadic<PDL_Operation>:$vector_sizes,
+                       DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:
+                          $static_vector_sizes);
+  let results = (outs);
+  let assemblyFormat = [{
+      $target
+      `vector_sizes` custom<DynamicIndexList>($vector_sizes,
+                                              $static_vector_sizes)
+      attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure apply(
+        ::mlir::transform::TransformResults &transformResults,
+        ::mlir::transform::TransformState &state);
+
+    ::llvm::SmallVector<::mlir::OpFoldResult> getMixedVectorSizes();
+  }];
+}
+
 #endif // LINALG_TRANSFORM_OPS

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

 ///   2. Take a full view on the buffer.
 ///   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);
+/// Emit a suitable vector form for a Linalg op. If provided, `inputVectorSizes`
+/// are used to vectorize this operation. `inputVectorSizes` must match the rank
+/// of the iteration space of the operation and the sizes must be smaller or
+/// equal than their counterpart interation space sizes, if static.
+/// `inputVectorShapes` also allows the vectorization of operations with dynamic
+/// shapes.
+LogicalResult vectorize(RewriterBase &rewriter, LinalgOp linalgOp,
+                        ArrayRef<int64_t> inputVectorSizes = {});
 
 /// 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,
+                              ArrayRef<int64_t> inputVectorSizes = {});
 
 //===----------------------------------------------------------------------===//
 // Transformations exposed as rewrite patterns.
 //===----------------------------------------------------------------------===//
 
 using TileSizeComputationFunction =
     std::function<SmallVector<Value, 4>(OpBuilder &, Operation *)>;
 

mlir/include/mlir/Dialect/Vector/Interfaces/MaskableOpInterface.td

let methods = [
     InterfaceMethod<
       /*desc=*/"Returns true if the operation is masked by a "
                "MaskingOpInterface.",
       /*retTy=*/"bool",
       /*methodName=*/"isMasked",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return mlir::isa<mlir::vector::MaskingOpInterface>($_op->getParentOp());
+        mlir::Operation *parentOp = $_op->getParentOp();
+        return parentOp &&
+               mlir::isa<mlir::vector::MaskingOpInterface>(parentOp);
     }]>,
     InterfaceMethod<
       /*desc=*/"Returns the MaskingOpInterface masking this operation.",
       /*retTy=*/"mlir::vector::MaskingOpInterface",
       /*methodName=*/"getMaskingOp",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return mlir::cast<mlir::vector::MaskingOpInterface>(
           $_op->getParentOp());
     }]>,
     InterfaceMethod<
       /*desc=*/"Returns true if the operation can have a passthru argument when"
                " masked.",
       /*retTy=*/"bool",
       /*methodName=*/"supportsPassthru",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return false;
     }]>,
     InterfaceMethod<
-      /*desc=*/"Returns the mask type expected by this operation. It requires "
-               "the operation to be vectorized.",
-      /*retTy=*/"mlir::VectorType",
+      /*desc=*/"Returns the mask type expected by this operation. Mostly used"
+               " for verification purposes. It requires the operation to be "
+               "vectorized.",
+      /*retTy=*/"mlir::Type",
       /*methodName=*/"getExpectedMaskType",
       /*args=*/(ins),
       /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-      // Default implementation is only aimed for operations that implement the
-      // `getVectorType()` method.
-        return $_op.getVectorType().cloneWith(/*shape=*/llvm::None,
-          IntegerType::get($_op.getContext(), /*width=*/1));
-    }]>,
+      /*defaultImplementation=*/"">,
   ];
 }
 
 #endif // MLIR_DIALECT_VECTOR_INTERFACES_MASKABLEOPINTERFACE_TD

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

 #include "mlir/Dialect/Vector/Transforms/Passes.h.inc"
 
 /// Creates an instance of the `vector` dialect bufferization pass.
 std::unique_ptr<Pass> createVectorBufferizePass();
 
 /// Creates an instance of the `vector.mask` lowering pass.
 std::unique_ptr<Pass> createLowerVectorMaskPass();
 
+/// Populates instances of `MaskOpRewritePattern` to lower masked operations
+/// with `vector.mask`. Patterns should rewrite the `vector.mask` operation and
+/// not its nested `MaskableOpInterface`.
+void populateVectorMaskLoweringPatternsForSideEffectingOps(
+    RewritePatternSet &patterns);
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Vector/Transforms/Passes.h.inc"
 } // namespace vector
 } // namespace mlir
 
 #endif // MLIR_DIALECT_VECTOR_TRANSFORMS_PASSES_H_

mlir/include/mlir/IR/AffineMap.h

public:
   /// when the caller knows it is safe to do so.
   unsigned getDimPosition(unsigned idx) const;
 
   /// Extracts the first result position where `input` dimension resides.
   /// Returns `llvm::None` if `input` is not a dimension expression or cannot be
   /// found in results.
   Optional<unsigned> getResultPosition(AffineExpr input) const;
 
+  /// Extracts the permuted position where the given input index resides.
+  /// Returns `llvm::None` if the input index is projected. Asserts on
+  /// non-projected permutation maps.
+  Optional<unsigned>
+  getPermutedPositionOfProjectedPermutation(unsigned input) const;
+
   /// Return true if any affine expression involves AffineDimExpr `position`.
   bool isFunctionOfDim(unsigned position) const {
     return llvm::any_of(getResults(), [&](AffineExpr e) {
       return e.isFunctionOfDim(position);
     });
   }
 
   /// Return true if any affine expression involves AffineSymbolExpr `position`.

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

 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformUtils.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 using namespace mlir::linalg;

 struct VectorizationPattern : public RewritePattern {
   explicit VectorizationPattern(MLIRContext *context)
       : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
     if (!linalgOp)
       return failure();
     return vectorize(rewriter, linalgOp);
   }
 };
 } // namespace
 
 DiagnosedSilenceableFailure
 transform::VectorizeOp::applyToOne(Operation *target,
                                    SmallVectorImpl<Operation *> &results,
                                    transform::TransformState &state) {
transform::VectorizeOp::applyToOne(Operation *target,
   if (failed(applyPatternsAndFoldGreedily(target, std::move(patterns))))
     return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
 
   results.push_back(target);
   return DiagnosedSilenceableFailure(success());
 }
 
 //===----------------------------------------------------------------------===//
+// MaskedVectorizeOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::MaskedVectorizeOp::apply(
+    mlir::transform::TransformResults &transformResults,
+    mlir::transform::TransformState &state) {
+  IRRewriter rewriter(getContext());
+  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
+  if (targets.empty())
+    return DiagnosedSilenceableFailure(success());
+
+  SmallVector<int64_t> vectorSizes;
+  for (OpFoldResult sz : getMixedVectorSizes()) {
+    if (sz.is<Attribute>()) {
+      auto attr = sz.get<Attribute>();
+      vectorSizes.push_back(attr.cast<IntegerAttr>().getInt());
+      continue;
+    }
+
+    ArrayRef<Operation *> szPayloads = state.getPayloadOps(sz.get<Value>());
+    if (szPayloads.size() != 1) {
+      auto diag = this->emitOpError(
+          "requires vector size handle that is mapped to 1 payload op");
+      diag.attachNote(sz.get<Value>().getLoc())
+          << "mapped to " << szPayloads.size() << " payload ops";
+      return DiagnosedSilenceableFailure::definiteFailure();
+    }
+
+    Operation *szPayloadOp = szPayloads[0];
+    if (szPayloadOp->getNumResults() != 1 ||
+        !szPayloadOp->getResult(0).getType().isIndex()) {
+      auto diag = this->emitOpError(
+          "requires vector size payload op with 1 index result");
+      diag.attachNote(szPayloadOp->getLoc()) << "vector size payload op";
+      return DiagnosedSilenceableFailure::definiteFailure();
+    }
+
+    IntegerAttr attr;
+    if (!matchPattern(szPayloadOp->getResult(0), m_Constant(&attr))) {
+      auto diag = this->emitOpError("requires constant vector size");
+      diag.attachNote(szPayloadOp->getLoc()) << "vector size payload op";
+      return DiagnosedSilenceableFailure::definiteFailure();
+    }
+
+    vectorSizes.push_back(attr.getInt());
+  }
+
+  // TODO: Check that the correct number of vectorSizes was provided.
+
+  for (Operation *target : targets) {
+    auto linalgOp = dyn_cast<LinalgOp>(target);
+    if (!linalgOp) {
+      Diagnostic diag(target->getLoc(), DiagnosticSeverity::Error);
+      diag << "cannot vectorize non-Linalg op";
+      return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+    }
+
+    if (failed(linalg::vectorize(rewriter, linalgOp, vectorSizes))) {
+      Diagnostic diag(target->getLoc(), DiagnosticSeverity::Error);
+      diag << "failed to vectorize op";
+      return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
+    }
+  }
+
+  return DiagnosedSilenceableFailure(success());
+}
+
+void transform::MaskedVectorizeOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  consumesHandle(getTarget(), effects);
+  onlyReadsHandle(getVectorSizes(), effects);
+}
+
+SmallVector<OpFoldResult> MaskedVectorizeOp::getMixedVectorSizes() {
+  OpBuilder b(getContext());
+  return getMixedValues(getStaticVectorSizes(), getVectorSizes(), b);
+}
+
+//===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Registers new ops and declares PDL as dependent dialect since the
 /// additional ops are using PDL types for operands and results.
 class LinalgTransformDialectExtension
     : public transform::TransformDialectExtension<

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

 //===- Vectorization.cpp - Implementation of linalg Vectorization ---------===//
 //
 // 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 the linalg dialect Vectorization transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/SliceAnalysis.h"
-#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
+#include "mlir/Dialect/Vector/Interfaces/MaskableOpInterface.h"
 #include "mlir/IR/AffineExpr.h"
-#include "mlir/IR/Matchers.h"
 #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;
if (res) {
       return WalkResult::interrupt();
     }
     res = op;
     return WalkResult::advance();
   });
   return res;
 }
 
+/// Contains the vectorization state and related methods used across the
+/// vectorization process of a given operation.
+struct VectorizationState {
+  VectorizationState(RewriterBase &rewriter) : rewriterGuard(rewriter) {}
+
+  /// Initializes the vectorization state, including the computation of the
+  /// canonical vector shape for vectorization.
+  LogicalResult initState(RewriterBase &rewriter, LinalgOp linalgOp,
+                          ArrayRef<int64_t> inputVectorSizes);
+
+  /// Returns the canonical vector shape used to vectorize the iteration space.
+  ArrayRef<int64_t> getCanonicalVecShape() const { return canonicalVecShape; }
+
+  /// Masks an operation with the canonical vector mask if the operation needs
+  /// masking. Returns the masked operation or the original operation if masking
+  /// is not needed. If provided, the canonical mask for this operation is
+  /// permuted using `maybeMaskingMap`.
+  Operation *maskOperation(RewriterBase &rewriter, Operation *opToMask,
+                           LinalgOp linalgOp,
+                           Optional<AffineMap> maybeMaskingMap = llvm::None);
+
+private:
+  /// Initializes the iteration space static sizes using the Linalg op
+  /// information. This may become more complicated in the future.
+  void initIterSpaceStaticSizes(LinalgOp linalgOp) {
+    iterSpaceStaticSizes.append(linalgOp.getStaticLoopRanges());
+  }
+
+  /// Generates 'tensor.dim' operations for all the dynamic dimensions of the
+  /// iteration space to be vectorized and store them in
+  /// `iterSpaceDynamicSizes`.
+  LogicalResult precomputeIterSpaceDynamicSizes(RewriterBase &rewriter,
+                                                LinalgOp linalgOp);
+
+  /// Create or retrieve an existing mask value to mask `opToMask` in the
+  /// canonical vector iteration space. If `maybeMaskingMap` the mask is
+  /// permuted using that permutation map. If a new mask is created, it will be
+  /// cached for future users.
+  Value getOrCreateMaskFor(RewriterBase &rewriter, Operation *opToMask,
+                           LinalgOp linalgOp,
+                           Optional<AffineMap> maybeMaskingMap);
+
+  // Holds the compile-time static sizes of the iteration space to vectorize.
+  // Dynamic dimensions are represented using ShapedType::kDynamicSize.
+  SmallVector<int64_t> iterSpaceStaticSizes;
+
+  /// Holds the runtime sizes of the iteration spaces to vectorize. Static
+  /// dimensions are represented with a empty value.
+  SmallVector<Value> iterSpaceDynamicSizes;
+
+  /// Holds the canonical vector shape used to vectorize the iteration space.
+  SmallVector<int64_t> canonicalVecShape;
+
+  /// Holds the active masks for permutations of the canonical vector iteration
+  /// space.
+  DenseMap<AffineMap, Value> activeMaskCache;
+
+  /// Global vectorization guard for the incoming rewriter. It's initialized
+  /// when the vectorization state is initialized.
+  OpBuilder::InsertionGuard rewriterGuard;
+};
+
+/// Generates 'tensor.dim' operations for all the dynamic dimensions of the
+/// iteration space to be vectorized and store them in
+/// `iterSpaceDynamicSizes`.
+LogicalResult
+VectorizationState::precomputeIterSpaceDynamicSizes(RewriterBase &rewriter,
+                                                    LinalgOp linalgOp) {
+  // TODO: Support 0-d vectors.
+  for (int vecDim = 0, end = canonicalVecShape.size(); vecDim < end; ++vecDim) {
+    if (!ShapedType::isDynamic(iterSpaceStaticSizes[vecDim])) {
+      // Add a empty value for static dimensions.
+      iterSpaceDynamicSizes.push_back(Value());
+      continue;
+    }
+
+    // Find an operand defined on this dimension of the iteration space to
+    // extract the runtime dimension size.
+    Value operand;
+    unsigned operandDimPos;
+    if (failed(linalgOp.mapIterationSpaceDimToOperandDim(vecDim, operand,
+                                                         operandDimPos)))
+      return failure();
+
+    Value dynamicDim = linalgOp.hasTensorSemantics()
+                           ? (Value)rewriter.create<tensor::DimOp>(
+                                 linalgOp.getLoc(), operand, operandDimPos)
+                           : (Value)rewriter.create<memref::DimOp>(
+                                 linalgOp.getLoc(), operand, operandDimPos);
+    iterSpaceDynamicSizes.push_back(dynamicDim);
+  }
+
+  return success();
+}
+
+/// Initializes the vectorization state, including the computation of the
+/// canonical vector shape for vectorization.
+// TODO: Move this to the constructor when we can remove the failure cases.
+LogicalResult
+VectorizationState::initState(RewriterBase &rewriter, LinalgOp linalgOp,
+                              ArrayRef<int64_t> inputVectorSizes) {
+  // Initialize the insertion point.
+  rewriter.setInsertionPoint(linalgOp);
+
+  if (!inputVectorSizes.empty()) {
+    // Get the canonical vector shape from the input vector sizes provided. This
+    // path should be taken to vectorize code with dynamic shapes and when using
+    // vector sizes greater than the iteration space sizes.
+    canonicalVecShape.append(inputVectorSizes.begin(), inputVectorSizes.end());
+  } else {
+    // Compute the canonical vector shape from the operation shape. If there are
+    // dynamic shapes, the operation won't be vectorized.
+    canonicalVecShape = linalgOp.getStaticLoopRanges();
+  }
+
+  LDBG("Canonical vector shape: ");
+  LLVM_DEBUG(llvm::interleaveComma(canonicalVecShape, llvm::dbgs()));
+  LLVM_DEBUG(llvm::dbgs() << "\n");
+
+  // Initialize iteration space static sizes.
+  initIterSpaceStaticSizes(linalgOp);
+
+  // Extract and register the runtime value of any potential dynamic shape
+  // needed to compute a mask during vectorization.
+  if (failed(precomputeIterSpaceDynamicSizes(rewriter, linalgOp)))
+    return failure();
+
+  if (ShapedType::isDynamicShape(canonicalVecShape))
+    return failure();
+  return success();
+}
+
+/// Create or retrieve an existing mask value to mask `opToMask` in the
+/// canonical vector iteration space. If `maybeMaskingMap` the mask is permuted
+/// using that permutation map. If a new mask is created, it will be cached for
+/// future users.
+Value VectorizationState::getOrCreateMaskFor(
+    RewriterBase &rewriter, Operation *opToMask, LinalgOp linalgOp,
+    Optional<AffineMap> maybeMaskingMap) {
+  // No mask is needed if the operation is not maskable.
+  auto maskableOp = dyn_cast<vector::MaskableOpInterface>(opToMask);
+  if (!maskableOp)
+    return Value();
+
+  assert(!maskableOp.isMasked() &&
+         "Masking an operation that is already masked");
+
+  // If no masking map was provided, use an identity map with the loop dims.
+  assert((!maybeMaskingMap || *maybeMaskingMap) &&
+         "Unexpected null mask permutation map");
+  AffineMap maskingMap =
+      maybeMaskingMap ? *maybeMaskingMap
+                      : AffineMap::getMultiDimIdentityMap(
+                            linalgOp.getNumLoops(), rewriter.getContext());
+
+  LDBG("Masking map: " << maskingMap << "\n");
+
+  // Return the active mask for the masking map of this operation if it was
+  // already created.
+  if (activeMaskCache.count(maskingMap)) {
+    Value mask = activeMaskCache[maskingMap];
+    LDBG("Reusing mask: " << mask << "\n");
+    return mask;
+  }
+
+  // Compute permuted projection of the iteration space to be masked and the
+  // corresponding mask shape. If the resulting iteration space dimensions are
+  // static and identical to the mask shape, masking is not needed for this
+  // operation.
+  // TODO: Improve this check. Only projected permutation indexing maps are
+  // supported.
+  SmallVector<int64_t> permutedStaticSizes =
+      applyPermutationMap(maskingMap, ArrayRef<int64_t>(iterSpaceStaticSizes));
+  SmallVector<int64_t> maskShape =
+      applyPermutationMap(maskingMap, ArrayRef<int64_t>(canonicalVecShape));
+  LDBG("Mask shape: ");
+  LLVM_DEBUG(llvm::interleaveComma(maskShape, llvm::dbgs()));
+  LLVM_DEBUG(llvm::dbgs() << "\n");
+
+  if (permutedStaticSizes == maskShape) {
+    LDBG("Masking is not needed for masking map: " << maskingMap << "\n");
+    activeMaskCache[maskingMap] = Value();
+    return Value();
+  }
+
+  // Compute the mask upper bound values by combining the permuted iteration
+  // space static sizes and the dynamic values.
+  SmallVector<Value> permutedDynamicSizes =
+      applyPermutationMap(maskingMap, ArrayRef<Value>(iterSpaceDynamicSizes));
+  SmallVector<Value> upperBounds;
+  for (auto [staticBound, dynBound] :
+       llvm::zip(permutedStaticSizes, permutedDynamicSizes))
+    upperBounds.push_back(ShapedType::isDynamic(staticBound)
+                              ? dynBound
+                              : rewriter.create<arith::ConstantIndexOp>(
+                                    linalgOp.getLoc(), staticBound));
+
+  assert(!maskShape.empty() && !upperBounds.empty() &&
+         "Masked 0-d vectors are not supported yet");
+
+  // Create the mask based on the value uppermension size values.
+  auto maskType = VectorType::get(maskShape, rewriter.getI1Type());
+  Value mask = rewriter.create<vector::CreateMaskOp>(linalgOp.getLoc(),
+                                                     maskType, upperBounds);
+  LDBG("Creating new mask: " << mask << "\n");
+  activeMaskCache[maskingMap] = mask;
+  return mask;
+}
+
+/// Masks an operation with the canonical vector mask if the operation needs
+/// masking. Returns the masked operation or the original operation if masking
+/// is not needed. If provided, the canonical mask for this operation is
+/// permuted using `maybeMaskingMap`.
+Operation *
+VectorizationState::maskOperation(RewriterBase &rewriter, Operation *opToMask,
+                                  LinalgOp linalgOp,
+                                  Optional<AffineMap> maybeMaskingMap) {
+  LDBG("Trying to mask: " << *opToMask << "\n");
+
+  // Create or retrieve mask for this operation.
+  Value mask =
+      getOrCreateMaskFor(rewriter, opToMask, linalgOp, maybeMaskingMap);
+
+  if (!mask) {
+    LDBG("No mask required\n");
+    return opToMask;
+  }
+
+  // Wrap the operation with a new `vector.mask` and update D-U chain.
+  assert(opToMask && "Expected a valid operation to mask");
+  auto opResults = opToMask->getResultTypes();
+  auto createRegionMask = [opToMask](OpBuilder &builder, Location loc) {
+    Block *insBlock = builder.getInsertionBlock();
+    // Create a block, put an op in that block. Look for a utility.
+    // Maybe in conversion pattern rewriter. Way to avoid splice.
+    // Set insertion point.
+    insBlock->getOperations().splice(
+        insBlock->begin(), opToMask->getBlock()->getOperations(), opToMask);
+    builder.create<vector::YieldOp>(loc, opToMask->getResults());
+  };
+  // TODO: Allow multiple results in vector.mask.
+  auto maskOp =
+      opResults.empty()
+          ? rewriter.create<vector::MaskOp>(opToMask->getLoc(), mask,
+                                            createRegionMask)
+          : rewriter.create<vector::MaskOp>(opToMask->getLoc(),
+                                            opToMask->getResultTypes().front(),
+                                            mask, createRegionMask);
+
+  Operation *maskOpTerminator = &maskOp.getMaskRegion().front().back();
+
+  for (auto [resIdx, resVal] : llvm::enumerate(opToMask->getResults()))
+    rewriter.replaceAllUsesExcept(resVal, maskOp.getResult(resIdx),
+                                  maskOpTerminator);
+
+  LDBG("Masked operation: " << *maskOp << "\n");
+  return maskOp;
+}
+
 /// Given an indexing `map` coming from a LinalgOp indexing, restricted to a
 /// projectedPermutation, compress the unused dimensions to serve as a
 /// permutation_map for a vector transfer operation.
 /// For example, given a linalg op such as:
 ///
 /// ```
 ///   %0 = linalg.generic {
 ///        indexing_maps = affine_map<(d0, d1, d2, d3, d4) -> (d4, d0, d2)>,

 }
 
 /// Build a vector.transfer_write of `value` into `outputOperand` at indices set
 /// to all `0`; where `outputOperand` is an output operand of the LinalgOp
 /// currently being vectorized. If `dest` has null rank, build an memref.store.
 /// Return the produced value or null if no value is produced.
 // Note: this is a true builder that notifies the OpBuilder listener.
 // TODO: Consider moving as a static helper on the ReduceOp.
-static Value buildVectorWrite(OpBuilder &b, Value value,
-                              OpOperand *outputOperand) {
-  Operation *write;
+static Value buildVectorWrite(RewriterBase &rewriter, Value value,
+                              OpOperand *outputOperand,
+                              VectorizationState &state) {
   Location loc = value.getLoc();
   auto linalgOp = cast<LinalgOp>(outputOperand->getOwner());
-  ArrayRef<int64_t> shape = linalgOp.getShape(outputOperand);
-  auto vectorType = VectorType::get(
-      shape, getElementTypeOrSelf(outputOperand->get().getType()));
+  AffineMap opOperandMap = linalgOp.getMatchingIndexingMap(outputOperand);
+  auto vectorType =
+      VectorType::get(opOperandMap.compose(state.getCanonicalVecShape()),
+                      getElementTypeOrSelf(outputOperand->get().getType()));
+
+  Operation *write;
   if (vectorType.getRank() > 0) {
-    // 0-d case is still special: do not invert the reindexing map.
-    AffineMap map =
-        reindexIndexingMap(linalgOp.getMatchingIndexingMap(outputOperand));
-    SmallVector<int64_t> transposeShape =
-        applyPermutationMap(inversePermutation(map), vectorType.getShape());
-    assert(!transposeShape.empty() && "unexpected empty transpose shape");
-    vectorType = VectorType::get(transposeShape, vectorType.getElementType());
+    AffineMap writeMap = reindexIndexingMap(opOperandMap);
     SmallVector<Value> indices(linalgOp.getRank(outputOperand),
-                               b.create<arith::ConstantIndexOp>(loc, 0));
-    value = broadcastIfNeeded(b, value, vectorType.getShape());
-    write = b.create<vector::TransferWriteOp>(
-        loc, value, outputOperand->get(), indices, map);
+                               rewriter.create<arith::ConstantIndexOp>(loc, 0));
+    value = broadcastIfNeeded(rewriter, value, vectorType.getShape());
+    write = rewriter.create<vector::TransferWriteOp>(
+        loc, value, outputOperand->get(), indices, writeMap);
   } else {
+    // 0-d case is still special: do not invert the reindexing writeMap.
     if (!value.getType().isa<VectorType>())
-      value = b.create<vector::BroadcastOp>(loc, vectorType, value);
+      value = rewriter.create<vector::BroadcastOp>(loc, vectorType, value);
     assert(value.getType() == vectorType && "incorrect type");
-    write = b.create<vector::TransferWriteOp>(
+    write = rewriter.create<vector::TransferWriteOp>(
         loc, value, outputOperand->get(), ValueRange{});
   }
-  LDBG("vectorized op: " << *write);
+
+  write = state.maskOperation(rewriter, write, linalgOp, opOperandMap);
+
+  // If masked, set in-bounds to true. Masking guarantees that the access will
+  // be in-bounds.
+  if (auto maskOp = dyn_cast<vector::MaskingOpInterface>(write)) {
+    auto maskedWriteOp = cast<vector::TransferWriteOp>(maskOp.getMaskableOp());
+    SmallVector<bool> inBounds(maskedWriteOp.getVectorType().getRank(), true);
+    maskedWriteOp.setInBoundsAttr(rewriter.getBoolArrayAttr(inBounds));
+  }
+
+  LDBG("vectorized op: " << *write << "\n");
   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.

 /// vector values are appended to `newResults`. Return
 /// VectorizationStatus::NoReplace to signal the vectorization algorithm that it
 /// should not try to map produced operations and instead return the results
 /// using the `newResults` vector making them available to the vectorization
 /// algorithm for RAUW. This function is meant to be used as a
 /// CustomVectorizationHook.
 static VectorizationResult
 vectorizeLinalgYield(RewriterBase &rewriter, Operation *op,
-                     const BlockAndValueMapping &bvm, LinalgOp linalgOp,
-                     SmallVectorImpl<Value> &newResults) {
+                     const BlockAndValueMapping &bvm, VectorizationState &state,
+                     LinalgOp linalgOp, SmallVectorImpl<Value> &newResults) {
   auto yieldOp = dyn_cast<linalg::YieldOp>(op);
   if (!yieldOp)
     return VectorizationResult{VectorizationStatus::Failure, nullptr};
-  for (const auto &outputs : llvm::enumerate(yieldOp.getValues())) {
+  for (const auto &output : llvm::enumerate(yieldOp.getValues())) {
     // TODO: Scan for an opportunity for reuse.
     // TODO: use a map.
-    Value vectorValue = bvm.lookup(outputs.value());
-    Value newResult = buildVectorWrite(
-        rewriter, vectorValue, linalgOp.getDpsInitOperand(outputs.index()));
+    Value vectorValue = bvm.lookup(output.value());
+    Value newResult =
+        buildVectorWrite(rewriter, vectorValue,
+                         linalgOp.getDpsInitOperand(output.index()), state);
     if (newResult)
       newResults.push_back(newResult);
   }
+
   return VectorizationResult{VectorizationStatus::NoReplace, nullptr};
 }
 
 /// Helper function to vectorize the index operations of a `linalgOp`. Return
 /// 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
 vectorizeLinalgIndex(RewriterBase &rewriter, Operation *op, LinalgOp linalgOp) {
   IndexOp indexOp = dyn_cast<linalg::IndexOp>(op);
   if (!indexOp)

 /// the `bvm` mapping. As a consequence, this function is meant to be called  on
 /// a topologically-sorted list of ops.
 /// This function does not update `bvm` but returns a VectorizationStatus that
 /// instructs the caller what `bvm` update needs to occur.
 static VectorizationResult
 vectorizeOneOp(RewriterBase &rewriter, LinalgOp linalgOp, Operation *op,
                const BlockAndValueMapping &bvm,
                ArrayRef<CustomVectorizationHook> customVectorizationHooks) {
-  LDBG("vectorize op " << *op);
+  LDBG("vectorize op " << *op << "\n");
 
   // 1. Try to apply any CustomVectorizationHook.
   if (!customVectorizationHooks.empty()) {
     for (auto &customFunc : customVectorizationHooks) {
       VectorizationResult result = customFunc(op, bvm);
       if (result.status == VectorizationStatus::Failure)
         continue;
       return result;

 /// iteration space. This eager broadcasting is introduced in the
 /// permutation_map of the vector.transfer_read operations. The eager
 /// broadcasting makes it trivial to detrmine where broadcast, transposes and
 /// reductions should occur, without any bookkeeping. The tradeoff is that, in
 /// the absence of good canonicalizations, the amount of work increases.
 /// This is not deemed a problem as we expect canonicalizations and foldings to
 /// aggressively clean up the useless work.
 static LogicalResult
-vectorizeAsLinalgGeneric(RewriterBase &rewriter, LinalgOp linalgOp,
+vectorizeAsLinalgGeneric(RewriterBase &rewriter, VectorizationState &state,
+                         LinalgOp linalgOp,
                          SmallVectorImpl<Value> &newResults) {
+  LDBG("Vectorizing operation as linalg generic\n");
   Block *block = linalgOp.getBlock();
 
   // 2. Values defined above the region can only be broadcast for now. Make them
   // map to themselves.
   BlockAndValueMapping bvm;
   SetVector<Value> valuesSet;
   mlir::getUsedValuesDefinedAbove(linalgOp->getRegion(0), valuesSet);
   bvm.map(valuesSet.getArrayRef(), valuesSet.getArrayRef());
 
   if (linalgOp.getNumDpsInits() == 0)
     return failure();
 
-  // TODO: the common vector shape is equal to the static loop sizes only when
-  // all indexing maps are projected permutations. For convs and stencils the
-  // logic will need to evolve.
-  SmallVector<int64_t> commonVectorShape = linalgOp.computeStaticLoopSizes();
-
   // 3. Turn all BBArgs into vector.transfer_read / load.
   Location loc = linalgOp.getLoc();
   Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
   for (OpOperand *opOperand : linalgOp.getOpOperandsMatchingBBargs()) {
     BlockArgument bbarg = linalgOp.getMatchingBlockArgument(opOperand);
     if (linalgOp.isScalar(opOperand)) {
       bvm.map(bbarg, opOperand->get());
       continue;
     }
-    VectorType readType;
-    AffineMap map;
-    // TODO: can we keep this simplification?
-    // if (linalgOp.getShape(&opOperand).empty()) {
-    //   readType = VectorType::get({}, bbarg.getType());
-    // } else {
-    if (opOperand->getOperandNumber() < linalgOp.getNumDpsInputs()) {
-      map = inverseAndBroadcastProjectedPermutation(
-          linalgOp.getMatchingIndexingMap(opOperand));
-      readType = VectorType::get(commonVectorShape,
-                                 getElementTypeOrSelf(opOperand->get()));
+
+    // 3.a. Convert the indexing map for this input/output to a transfer read
+    // permutation map and masking map.
+    AffineMap indexingMap = linalgOp.getMatchingIndexingMap(opOperand);
+
+    // Remove zeros from indexing map to use it as masking map.
+    SmallVector<int64_t> zeroPos;
+    auto results = indexingMap.getResults();
+    for (auto result : llvm::enumerate(results)) {
+      if (result.value().isa<AffineConstantExpr>()) {
+        zeroPos.push_back(result.index());
+      }
+    }
+    AffineMap maskingMap = indexingMap.dropResults(zeroPos);
+
+    AffineMap readMap;
+    SmallVector<int64_t> readVecShape;
+    if (linalgOp.isDpsInput(opOperand)) {
+      // 3.a.i. For input reads we use the canonical vector shape.
+      readMap = inverseAndBroadcastProjectedPermutation(indexingMap);
+      readVecShape = llvm::to_vector(state.getCanonicalVecShape());
     } else {
-      map = inversePermutation(
-          reindexIndexingMap(linalgOp.getMatchingIndexingMap(opOperand)));
-      readType = VectorType::get(map.compose(linalgOp.getShape(opOperand)),
-                                 getElementTypeOrSelf(opOperand->get()));
+      // 3.a.ii. For output reads (iteration-carried dependence, e.g.,
+      // reductions), the vector shape is computed by mapping the canonical
+      // vector shape to the output domain and back to the canonical domain.
+      readMap = inversePermutation(reindexIndexingMap(indexingMap));
+      readVecShape =
+          readMap.compose(indexingMap.compose(state.getCanonicalVecShape()));
     }
-    // }
 
-    auto shape = linalgOp.getShape(opOperand);
-    SmallVector<Value> indices(shape.size(), zero);
-    Value readValue = rewriter.create<vector::TransferReadOp>(
-        loc, readType, opOperand->get(), indices, map);
-    // Not all ops support 0-d vectors, extract the scalar for now.
+    auto readType =
+        VectorType::get(readVecShape, getElementTypeOrSelf(opOperand->get()));
+    SmallVector<Value> indices(linalgOp.getShape(opOperand).size(), zero);
+
+    Operation *read = rewriter.create<vector::TransferReadOp>(
+        loc, readType, opOperand->get(), indices, readMap);
+    read = state.maskOperation(rewriter, read, linalgOp, maskingMap);
+    Value readValue = read->getResult(0);
+
+    // 3.b. If masked, set in-bounds to true. Masking guarantees that the access
+    // will be in-bounds.
+    if (auto maskOp = dyn_cast<vector::MaskingOpInterface>(read)) {
+      SmallVector<bool> inBounds(readType.getRank(), true);
+      cast<vector::TransferReadOp>(maskOp.getMaskableOp())
+          .setInBoundsAttr(rewriter.getBoolArrayAttr(inBounds));
+    }
+
+    // 3.c. Not all ops support 0-d vectors, extract the scalar for now.
     // TODO: remove this.
     if (readValue.getType().cast<VectorType>().getRank() == 0)
       readValue = rewriter.create<vector::ExtractElementOp>(loc, readValue);
 
-    LDBG("new vectorized bbarg(" << bbarg.getArgNumber() << "): " << readValue);
+    LDBG("New vectorized bbarg(" << bbarg.getArgNumber() << "): " << readValue
+                                 << "\n");
     bvm.map(bbarg, readValue);
     bvm.map(opOperand->get(), readValue);
   }
 
   SmallVector<CustomVectorizationHook> hooks;
   // 4a. Register CustomVectorizationHook for yieldOp.
   CustomVectorizationHook vectorizeYield =
       [&](Operation *op,
           const BlockAndValueMapping &bvm) -> VectorizationResult {
-    return vectorizeLinalgYield(rewriter, op, bvm, linalgOp, newResults);
+    return vectorizeLinalgYield(rewriter, op, bvm, state, 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.
   CustomVectorizationHook vectorizeExtract =
       [&](Operation *op,
           const BlockAndValueMapping &bvm) -> VectorizationResult {
     return vectorizeTensorExtract(rewriter, op, linalgOp, bvm);
   };
   hooks.push_back(vectorizeExtract);
 
   // 5. Iteratively call `vectorizeOneOp` to each op in the slice.
   for (Operation &op : block->getOperations()) {
     VectorizationResult result =
         vectorizeOneOp(rewriter, linalgOp, &op, bvm, hooks);
     if (result.status == VectorizationStatus::Failure) {
-      LDBG("failed to vectorize: " << op);
+      LDBG("failed to vectorize: " << op << "\n");
       return failure();
     }
     if (result.status == VectorizationStatus::NewOp) {
-      LDBG("new vector op: " << *result.newOp;);
-      bvm.map(op.getResults(), result.newOp->getResults());
+      Operation *maybeMaskedOp =
+          state.maskOperation(rewriter, result.newOp, linalgOp);
+      LDBG("New vector op: " << *maybeMaskedOp << "\n");
+      bvm.map(op.getResults(), maybeMaskedOp->getResults());
     }
   }
 
   return success();
 }
 
 // TODO: probably need some extra checks for reduction followed by consumer
 // ops that may not commute (e.g. linear reduction + non-linear instructions).
 static LogicalResult reductionPreconditions(LinalgOp op) {
   if (llvm::none_of(op.getIteratorTypesArray(), isReductionIterator)) {
-    LDBG("reduction precondition failed: no reduction iterator");
+    LDBG("reduction precondition failed: no reduction iterator\n");
     return failure();
   }
   for (OpOperand *opOperand : op.getDpsInitOperands()) {
     AffineMap indexingMap = op.getMatchingIndexingMap(opOperand);
     if (indexingMap.isPermutation())
       continue;
 
     Operation *reduceOp = matchLinalgReduction(opOperand);
     if (!reduceOp || !getCombinerOpKind(reduceOp)) {
-      LDBG("reduction precondition failed: reduction detection failed");
+      LDBG("reduction precondition failed: reduction detection failed\n");
       return failure();
     }
   }
   return success();
 }
 
-static LogicalResult vectorizeStaticLinalgOpPrecondition(
-    linalg::LinalgOp op,
-    ArrayRef<CustomVectorizationPrecondition> customPreconditions) {
+static LogicalResult vectorizeDynamicLinalgOpPrecondition(linalg::LinalgOp op) {
+  // TODO: Masking only supports dynamic generic ops without reductions for now.
+  if (!isElementwise(op) &&
+      llvm::any_of(op.getIteratorTypesArray(), [](utils::IteratorType itType) {
+        return itType != utils::IteratorType::parallel;
+      }))
+    return failure();
+
+  // TODO: 0-d vectors are not supported yet.
+  if (llvm::any_of(op.getIndexingMapsArray(), [](AffineMap map) {
+        return map.isEmpty() || map.getResults().empty();
+      }))
+    return failure();
+
+  LDBG("Dynamically-shaped op meets vectorization pre-conditions\n");
+  return success();
+}
+
+LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(
+    LinalgOp linalgOp, ArrayRef<int64_t> inputVectorSizes) {
+  // Check API contract for input vector sizes.
+  if (!inputVectorSizes.empty()) {
+    assert(inputVectorSizes.size() == linalgOp.getNumLoops() &&
+           "Input vector sizes don't match the number of loops");
+    assert(!ShapedType::isDynamicShape(inputVectorSizes) &&
+           "Input vector sizes can't have dynamic dimensions");
+    assert(llvm::all_of(
+               llvm::zip(linalgOp.getStaticLoopRanges(), inputVectorSizes),
+               [](std::tuple<int64_t, int64_t> sizePair) {
+                 int64_t staticSize = std::get<0>(sizePair);
+                 int64_t inputSize = std::get<1>(sizePair);
+                 return ShapedType::isDynamic(staticSize) ||
+                        staticSize <= inputSize;
+               }) &&
+           "Input vector sizes must be smaller or equal than iteration space "
+           "static sizes");
+  }
+
+  // TODO: Masking is only supported for dynamic shapes so input vector sizes
+  // must be empty if the op is not dynamic.
+  if (!linalgOp.hasDynamicShape() && !inputVectorSizes.empty())
+    return failure();
+
+  if (linalgOp.hasDynamicShape() &&
+      failed(vectorizeDynamicLinalgOpPrecondition(linalgOp)))
+    return failure();
+
+  SmallVector<CustomVectorizationPrecondition> customPreconditions;
+
+  // Register CustomVectorizationPrecondition for extractOp.
+  customPreconditions.push_back(tensorExtractVectorizationPrecondition);
 
   // All types in the body should be a supported element type for VectorType.
-  for (Operation &innerOp : op->getRegion(0).front()) {
+  for (Operation &innerOp : linalgOp->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));
             })) {
       continue;
     }
     if (llvm::any_of(innerOp.getOperandTypes(), [](Type type) {
           return !VectorType::isValidElementType(type);
         })) {
       return failure();
     }
     if (llvm::any_of(innerOp.getResultTypes(), [](Type type) {
           return !VectorType::isValidElementType(type);
         })) {
       return failure();
     }
   }
-  if (isElementwise(op))
+  if (isElementwise(linalgOp))
     return success();
   // TODO: isaConvolutionOpInterface that can also infer from generic features.
   // But we will still need stride/dilation attributes that will be annoying to
   // reverse-engineer...
-  if (isa<ConvolutionOpInterface>(op.getOperation()))
+  if (isa<ConvolutionOpInterface>(linalgOp.getOperation()))
     return success();
   // TODO: the common vector shape is equal to the static loop sizes only when
   // all indexing maps are projected permutations. For convs and stencils the
   // logic will need to evolve.
-  if (!allIndexingsAreProjectedPermutation(op)) {
-    LDBG("precondition failed: not projected permutations");
+  if (!allIndexingsAreProjectedPermutation(linalgOp)) {
+    LDBG("precondition failed: not projected permutations\n");
     return failure();
   }
-  if (failed(reductionPreconditions(op))) {
-    LDBG("precondition failed: reduction preconditions");
+  if (failed(reductionPreconditions(linalgOp))) {
+    LDBG("precondition failed: reduction preconditions\n");
     return failure();
   }
   return success();
 }
 
-LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(LinalgOp linalgOp) {
-  // All types must be static shape to go to vector.
-  if (linalgOp.hasDynamicShape()) {
-    LDBG("precondition failed: dynamic shape");
+/// Emit a suitable vector form for a Linalg op. If provided, `inputVectorSizes`
+/// are used to vectorize this operation. `inputVectorSizes` must match the rank
+/// of the iteration space of the operation and the sizes must be smaller or
+/// equal than their counterpart interation space sizes, if static.
+/// `inputVectorShapes` also allows the vectorization of operations with dynamic
+/// shapes.
+LogicalResult mlir::linalg::vectorize(RewriterBase &rewriter, LinalgOp linalgOp,
+                                      ArrayRef<int64_t> inputVectorSizes) {
+  LDBG("Attempting to vectorize:\n" << linalgOp << "\n");
+  LDBG("Input vector sizes: ");
+  LLVM_DEBUG(llvm::interleaveComma(inputVectorSizes, llvm::dbgs()));
+  LLVM_DEBUG(llvm::dbgs() << "\n");
+
+  if (failed(vectorizeLinalgOpPrecondition(linalgOp, inputVectorSizes)))
+    return failure();
+
+  // Initialize vectorization state.
+  VectorizationState state(rewriter);
+  if (failed(state.initState(rewriter, linalgOp, inputVectorSizes))) {
+    LDBG("Vectorization state couldn't be initialized\n");
     return failure();
   }
 
-  SmallVector<CustomVectorizationPrecondition> customPreconditions;
-
-  // Register CustomVectorizationPrecondition for extractOp.
-  customPreconditions.push_back(tensorExtractVectorizationPrecondition);
-
-  return vectorizeStaticLinalgOpPrecondition(linalgOp, customPreconditions);
-}
-
-LogicalResult mlir::linalg::vectorize(RewriterBase &rewriter,
-                                      LinalgOp linalgOp) {
-  if (failed(vectorizeLinalgOpPrecondition(linalgOp)))
-    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)))
       return failure();
-    LDBG("Vectorize generic by broadcasting to a common shape: " << linalgOp);
-    if (failed(vectorizeAsLinalgGeneric(rewriter, linalgOp, results)))
+    LDBG("Vectorize generic by broadcasting to the canonical vector shape\n");
+    // TODO: 'vectorize' takes in a 'RewriterBase' which is up-casted to
+    // 'OpBuilder' when it is passed over to some methods like
+    // 'vectorizeAsLinalgGeneric'. This is highly problematic: if we erase an op
+    // within these methods, the actual rewriter won't be notified and we will
+    // end up with read-after-free issues!
+    if (failed(vectorizeAsLinalgGeneric(rewriter, state, linalgOp, results)))
       return failure();
   }
 
   if (!results.empty())
     rewriter.replaceOp(linalgOp, results);
   else
     rewriter.eraseOp(linalgOp);
 

 /// Check whether there is any interleaved use of any `values` between
 /// `firstOp` and `secondOp`. Conservatively return `true` if any op or value
 /// is in a different block.
 static bool mayExistInterleavedUses(Operation *firstOp, Operation *secondOp,
                                     ValueRange values) {
   if (firstOp->getBlock() != secondOp->getBlock() ||
       !firstOp->isBeforeInBlock(secondOp)) {
     LDBG("interleavedUses precondition failed, firstOp: "
-         << *firstOp << ", second op: " << *secondOp);
+         << *firstOp << ", second op: " << *secondOp << "\n");
     return true;
   }
   for (auto v : values) {
     for (auto &u : v.getUses()) {
       Operation *owner = u.getOwner();
       if (owner == firstOp || owner == secondOp)
         continue;
       // TODO: this is too conservative, use dominance info in the future.
       if (owner->getBlock() == firstOp->getBlock() &&
           (owner->isBeforeInBlock(firstOp) || secondOp->isBeforeInBlock(owner)))
         continue;
       LDBG(" found interleaved op " << *owner << ", firstOp: " << *firstOp
-                                    << ", second op: " << *secondOp);
+                                    << ", second op: " << *secondOp << "\n");
       return true;
     }
   }
   return false;
 }
 
 /// Return the unique subview use of `v` if it is indeed unique, null
 /// otherwise.

mlir/lib/Dialect/Vector/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRVectorDialect
   VectorOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Vector/IR
 
   DEPENDS
+  MLIRMaskableOpInterfaceIncGen
+  MLIRMaskingOpInterfaceIncGen
   MLIRVectorOpsIncGen
   MLIRVectorOpsEnumsIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces
   MLIRDestinationStyleOpInterface

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

void ReductionOp::print(OpAsmPrinter &p) {
   p << " ";
   getKindAttr().print(p);
   p << ", " << getVector();
   if (getAcc())
     p << ", " << getAcc();
   p << " : " << getVector().getType() << " into " << getDest().getType();
 }
 
+// MaskableOpInterface methods.
+
+/// Returns the mask type expected by this operation.
+Type ReductionOp::getExpectedMaskType() {
+  auto vecType = getVectorType();
+  return vecType.cloneWith(llvm::None,
+                           IntegerType::get(vecType.getContext(), /*width=*/1));
+}
+
 Value mlir::vector::getVectorReductionOp(arith::AtomicRMWKind op,
                                          OpBuilder &builder, Location loc,
                                          Value vector) {
   switch (op) {
   case arith::AtomicRMWKind::addf:
   case arith::AtomicRMWKind::addi:
     return builder.create<vector::ReductionOp>(vector.getLoc(),
                                                CombiningKind::ADD, vector);
if (paddingType != sourceElementType)
       return emitOpError(
           "requires formal padding and source of the same elemental type");
   }
 
   return verifyPermutationMap(permutationMap,
                               [&](Twine t) { return emitOpError(t); });
 }
 
+// MaskableOpInterface methods.
+
+/// Returns the mask type expected by this operation. Mostly used for
+/// verification purposes. It requires the operation to be vectorized."
+Type TransferReadOp::getExpectedMaskType() {
+  return inferTransferReadMaskType(getVectorType(), getPermutationMap());
+}
+
 template <typename TransferOp>
 static bool isInBounds(TransferOp op, int64_t resultIdx, int64_t indicesIdx) {
   // TODO: support more aggressive createOrFold on:
   // `op.indices()[indicesIdx] + vectorType < dim(op.source(), indicesIdx)`
   if (op.getShapedType().isDynamicDim(indicesIdx))
     return false;
   Value index = op.getIndices()[indicesIdx];
   auto cstOp = index.getDefiningOp<arith::ConstantIndexOp>();
if (failed(verifyTransferOp(cast<VectorTransferOpInterface>(getOperation()),
                               inferredMaskType, permutationMap,
                               getInBounds() ? *getInBounds() : ArrayAttr())))
     return failure();
 
   return verifyPermutationMap(permutationMap,
                               [&](Twine t) { return emitOpError(t); });
 }
 
+// MaskableOpInterface methods.
+
+/// Returns the mask type expected by this operation. Mostly used for
+/// verification purposes.
+Type TransferWriteOp::getExpectedMaskType() {
+  return inferTransferWriteMaskType(getVectorType(), getPermutationMap());
+}
+
 /// Fold:
 /// ```
 ///    %t1 = ...
 ///    %v = vector.transfer_read %t0[%c0...], {in_bounds = [true...]} :
 ///      tensor<static_sizesxf32>, vector<static_sizesxf32>
 ///    %t2 = vector.transfer_write %v, %t1[%c0...] {in_bounds = [true...]} :
 ///      vector<static_sizesxf32>, tensor<static_sizesxf32>
 /// ```
if (maskableOp->getNumResults() != getNumResults())
     return emitOpError("expects number of results to match maskable operation "
                        "number of results");
 
   if (!llvm::equal(maskableOp->getResultTypes(), getResultTypes()))
     return emitOpError(
         "expects result type to match maskable operation result type");
 
   // Mask checks.
-  if (getMask().getType() != maskableOp.getExpectedMaskType())
-    return emitOpError("expects a ") << maskableOp.getExpectedMaskType()
-                                     << " mask for the maskable operation";
+  Type expectedMaskType = maskableOp.getExpectedMaskType();
+  if (getMask().getType() != expectedMaskType)
+    return emitOpError("expects a ")
+           << expectedMaskType << " mask for the maskable operation";
 
   // Passthru checks.
   Value passthru = getPassthru();
   if (passthru) {
     if (!maskableOp.supportsPassthru())
       return emitOpError(
           "doesn't expect a passthru argument for this maskable operation");
 

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

matchAndRewriteMaskableOp(TransferWriteOp writeOp,
     rewriter.replaceOpWithNewOp<TransferWriteOp>(
         maskingOp.getOperation(), resultType, writeOp.getVector(),
         writeOp.getSource(), writeOp.getIndices(), writeOp.getPermutationMap(),
         maskingOp.getMask(), writeOp.getInBounds().value_or(ArrayAttr()));
     return success();
   }
 };
 
-/// Populates instances of `MaskOpRewritePattern` to lower masked operations
-/// with `vector.mask`. Patterns should rewrite the `vector.mask` operation and
-/// not its nested `MaskableOpInterface`.
-void populateVectorMaskLoweringPatternsForSideEffectingOps(
-    RewritePatternSet &patterns) {
-  patterns.add<MaskedTransferReadOpPattern, MaskedTransferWriteOpPattern>(
-      patterns.getContext());
-}
-
 struct LowerVectorMaskPass
     : public vector::impl::LowerVectorMaskPassBase<LowerVectorMaskPass> {
   using Base::Base;
 
   void runOnOperation() override {
     Operation *op = getOperation();
     MLIRContext *context = op->getContext();
 
     RewritePatternSet loweringPatterns(context);
     populateVectorMaskLoweringPatternsForSideEffectingOps(loweringPatterns);
 
     if (failed(applyPatternsAndFoldGreedily(op->getRegions(),
                                             std::move(loweringPatterns))))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<vector::VectorDialect>();
   }
 };
 
 } // namespace
 
+/// Populates instances of `MaskOpRewritePattern` to lower masked operations
+/// with `vector.mask`. Patterns should rewrite the `vector.mask` operation and
+/// not its nested `MaskableOpInterface`.
+void vector::populateVectorMaskLoweringPatternsForSideEffectingOps(
+    RewritePatternSet &patterns) {
+  patterns.add<MaskedTransferReadOpPattern, MaskedTransferWriteOpPattern>(
+      patterns.getContext());
+}
+
 std::unique_ptr<Pass> mlir::vector::createLowerVectorMaskPass() {
   return std::make_unique<LowerVectorMaskPass>();
 }

mlir/lib/IR/AffineMap.cpp

Optional<unsigned> AffineMap::getResultPosition(AffineExpr input) const {
   for (unsigned i = 0, numResults = getNumResults(); i < numResults; i++) {
     if (getResult(i) == input)
       return i;
   }
 
   return llvm::None;
 }
 
+/// Extracts the permuted position where the given input index resides.
+/// Returns `llvm::None` if the input index is projected. Asserts on
+/// non-projected permutation maps.
+Optional<unsigned>
+AffineMap::getPermutedPositionOfProjectedPermutation(unsigned input) const {
+  assert(isProjectedPermutation() && "invalid projected permutation request");
+  for (unsigned i = 0, numResults = getNumResults(); i < numResults; i++)
+    if (getDimPosition(i) == input)
+      return i;
+  return llvm::None;
+}
+
 /// Folds the results of the application of an affine map on the provided
 /// operands to a constant if possible. Returns false if the folding happens,
 /// true otherwise.
 LogicalResult
 AffineMap::constantFold(ArrayRef<Attribute> operandConstants,
                         SmallVectorImpl<Attribute> &results) const {
   // Attempt partial folding.
   SmallVector<int64_t, 2> integers;

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
   linalg.dot ins(%A, %B: memref<1584xf32>, memref<1584xf32>)
             outs(%C: memref<f32>)
   return
indexing_maps = [
     affine_map<(m, n, k) -> (k, n)>,
     affine_map<(m, n, k) -> (n, m)>
   ],
   iterator_types = ["parallel", "parallel", "reduction"]
 }
 
 // CHECK-LABEL: func @generic_output_transpose
 func.func @generic_output_transpose(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
                                     %C: memref<32x8xf32>) {
   //       CHECK: vector.transfer_read %{{.*}} : memref<8x16xf32>, vector<8x32x16xf32>
   //       CHECK: vector.transfer_read %{{.*}} : memref<16x32xf32>, vector<8x32x16xf32>
   //       CHECK: %[[ACC:.*]] = vector.transfer_read %{{.*}} : memref<32x8xf32>, vector<8x32xf32>
   //       CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32>
   //       CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[ACC]] [2] : vector<8x32x16xf32> to vector<8x32xf32>
   //       CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xf32>, memref<32x8xf32>
   linalg.generic #matmul_transpose_out_trait
     ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>)

 // CHECK-SAME:    : vector<16x32x64xf32> to vector<16x64xf32>
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.reduce"]} in %arg1
   %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
   %2 = transform.structured.vectorize %1
 }
+
+// -----
+
+func.func @vectorize_dynamic_identity(%arg0: tensor<?xf32>,
+                                      %arg1: tensor<?xf32>,
+                                      %arg2: tensor<?xf32>) -> tensor<?xf32> {
+  %0 = linalg.generic { indexing_maps = [affine_map<(d0) -> (d0)>,
+                                         affine_map<(d0) -> (d0)>,
+                                         affine_map<(d0) -> (d0)>],
+                   iterator_types = ["parallel"] }
+    ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
+    outs(%arg2 : tensor<?xf32>) {
+    ^bb(%in0: f32, %in1: f32, %out: f32) :
+      %0 = arith.addf %in0, %in1 : f32
+      linalg.yield %0 : f32
+    } -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL:   @vectorize_dynamic_identity
+// CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?xf32>
+// CHECK:           %[[VAL_7:.*]] = vector.create_mask %[[VAL_4]] : vector<4xi1>
+// CHECK:           %[[VAL_8:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+// CHECK:           %[[VAL_10:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+// CHECK:           %[[VAL_12:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+// CHECK:           %[[VAL_13:.*]] = arith.addf %[[VAL_8]], %[[VAL_10]] : vector<4xf32>
+// CHECK:           %[[VAL_14:.*]] = vector.mask %[[VAL_7]] { vector.transfer_write %{{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<?xf32> } : vector<4xi1> -> tensor<?xf32>
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+  transform.structured.masked_vectorize %0 vector_sizes [4]
+}
+
+// -----
+
+func.func @vectorize_dynamic_1d_broadcast(%arg0: tensor<?xf32>,
+                                          %arg1: tensor<?xf32>,
+                                          %arg2: tensor<?xf32>) -> tensor<?xf32> {
+  %0 = linalg.generic { indexing_maps = [affine_map<(d0) -> (0)>,
+                                         affine_map<(d0) -> (d0)>,
+                                         affine_map<(d0) -> (d0)>],
+                        iterator_types = ["parallel"] }
+    ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
+    outs(%arg2 : tensor<?xf32>) {
+    ^bb(%in0: f32, %in1: f32, %out: f32) :
+      %0 = arith.addf %in0, %in1 : f32
+      linalg.yield %0 : f32
+    } -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL:   @vectorize_dynamic_1d_broadcast
+// CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?xf32>
+// CHECK:           %[[VAL_7:.*]] = vector.transfer_read %{{.*}} {permutation_map = #{{.*}}} : tensor<?xf32>, vector<4xf32>
+// CHECK:           %[[VAL_9:.*]] = vector.create_mask %[[VAL_4]] : vector<4xi1>
+// CHECK:           %[[VAL_10:.*]] = vector.mask %[[VAL_9]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+// CHECK:           %[[VAL_12:.*]] = vector.mask %[[VAL_9]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+// CHECK:           %[[VAL_13:.*]] = arith.addf %[[VAL_7]], %[[VAL_10]] : vector<4xf32>
+// CHECK:           %[[VAL_14:.*]] = vector.mask %{{.*}} { vector.transfer_write %[[VAL_13]], {{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<?xf32> } : vector<4xi1> -> tensor<?xf32>
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+  transform.structured.masked_vectorize %0 vector_sizes [4]
+}
+
+// -----
+
+func.func @vectorize_dynamic_2d_transpose(%arg0: tensor<?x?xf32>,
+                                          %arg1: tensor<?x?xf32>,
+                                          %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>,
+                                         affine_map<(d0, d1) -> (d0, d1)>,
+                                         affine_map<(d0, d1) -> (d0, d1)>],
+                        iterator_types = ["parallel", "parallel"] }
+    ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+    outs(%arg2 : tensor<?x?xf32>) {
+    ^bb(%in0: f32, %in1: f32, %out: f32) :
+      %0 = arith.addf %in0, %in1 : f32
+      linalg.yield %0 : f32
+    } -> tensor<?x?xf32>
+    return %0 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL:   @vectorize_dynamic_2d_transpose
+// CHECK:           %[[VAL_3:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?x?xf32>
+// CHECK:           %[[VAL_5:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_6:.*]] = tensor.dim %{{.*}}, %[[VAL_5]] : tensor<?x?xf32>
+// CHECK:           %[[VAL_9:.*]] = vector.create_mask %[[VAL_6]], %[[VAL_4]] : vector<8x4xi1>
+// CHECK:           %[[VAL_10:.*]] = vector.mask %[[VAL_9]] { vector.transfer_read %{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : tensor<?x?xf32>, vector<4x8xf32> } : vector<8x4xi1> -> vector<4x8xf32>
+// CHECK:           %[[VAL_12:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_6]] : vector<4x8xi1>
+// CHECK:           %[[VAL_13:.*]] = vector.mask %[[VAL_12]] { vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>
+// CHECK:           %[[VAL_14:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_15:.*]] = vector.mask %[[VAL_12]] { vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>
+// CHECK:           %[[VAL_16:.*]] = arith.addf %[[VAL_10]], %[[VAL_13]] : vector<4x8xf32>
+// CHECK:           %[[VAL_17:.*]] = vector.mask %[[VAL_12]] { vector.transfer_write %[[VAL_16]], %{{.*}} {in_bounds = [true, true]} : vector<4x8xf32>, tensor<?x?xf32> } : vector<4x8xi1> -> tensor<?x?xf32>
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+  transform.structured.masked_vectorize %0 vector_sizes [4, 8]
+}
+
+// -----
+
+func.func @vectorize_dynamic_generic_2d_broadcast(%arg0: tensor<?x?xf32>,
+                                                  %arg1: tensor<?x?xf32>,
+                                                  %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (0, d1)>,
+                                         affine_map<(d0, d1) -> (d0, d1)>,
+                                         affine_map<(d0, d1) -> (d0, d1)>],
+                        iterator_types = ["parallel", "parallel"] }
+    ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+    outs(%arg2 : tensor<?x?xf32>) {
+    ^bb(%in0: f32, %in1: f32, %out: f32) :
+      %0 = arith.addf %in0, %in1 : f32
+      linalg.yield %0 : f32
+    } -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL:   @vectorize_dynamic_generic_2d_broadcast
+// CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?x?xf32>
+// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_6:.*]] = tensor.dim %{{.*}}, %[[VAL_5]] : tensor<?x?xf32>
+// CHECK:           %[[VAL_9:.*]] = vector.create_mask %[[VAL_6]] : vector<8xi1>
+// CHECK:           %[[VAL_10:.*]] = vector.mask %[[VAL_9]] { vector.transfer_read %{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : tensor<?x?xf32>, vector<4x8xf32> } : vector<8xi1> -> vector<4x8xf32>
+// CHECK:           %[[VAL_12:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_6]] : vector<4x8xi1>
+// CHECK:           %[[VAL_13:.*]] = vector.mask %[[VAL_12]] { vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>
+// CHECK:           %[[VAL_15:.*]] = vector.mask %[[VAL_12]] { vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>
+// CHECK:           %[[VAL_16:.*]] = arith.addf %[[VAL_10]], %[[VAL_13]] : vector<4x8xf32>
+// CHECK:           %[[VAL_18:.*]] = vector.mask %[[VAL_12]] { vector.transfer_write %{{.*}} {in_bounds = [true, true]} : vector<4x8xf32>, tensor<?x?xf32> } : vector<4x8xi1> -> tensor<?x?xf32>
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+  transform.structured.masked_vectorize %0 vector_sizes [4, 8]
+}
+

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":FuncDialect",
         ":FuncTransforms",
         ":IR",
         ":LinalgAnalysis",
         ":LinalgDialect",
         ":LinalgPassIncGen",
         ":LinalgStructuredOpsIncGen",
         ":LinalgUtils",
+        ":MaskableOpInterface",
         ":MathDialect",
         ":MemRefDialect",
         ":MemRefTransforms",
         ":Pass",
         ":SCFDialect",
         ":SCFTransforms",
         ":SCFUtils",
         ":SparseTensorDialect",