[mlir][Vector] Add a VectorUnrollInterface and expose UnrollVectorPattern.

[nicolasvasilache]

Summary

The UnrollVectorPattern is can be used in a programmable fashion by:

OwningRewritePatternList patterns;
    patterns.insert<UnrollVectorPattern<AddFOp>>(ArrayRef<int64_t>{2, 2}, ctx);
    patterns.insert<UnrollVectorPattern<vector::ContractionOp>>(
        ArrayRef<int64_t>{2, 2, 2}, ctx);
    ...
    applyPatternsAndFoldGreedily(getFunction(), patterns);

[ftynse]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

3185

Are there any ops derived from IntArithmeticOp that shouldn't have the interface methods declared? If not, let's add the trait in IntArithmeticOp instead. Same for FloatArithmeticOp and FloatUnaryOp.

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

109

Nit: remove the empty line

134

Given the "matching type" in the name, should there be a check for op types being same, or a static assert on some trait?

mlir/include/mlir/Interfaces/VectorUnrollInterface.h

2

Wrong header. Also, -*- C++ -*-

10

Wrong description

[nicolasvasilache]

Address review.

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

134

Refactored things a bit and renamed where appropriate.

[nicolasvasilache]

Cleanup spurious additions.

[nicolasvasilache]

Format

[nicolasvasilache]

CMake

[nicolasvasilache]

More CMake + liveness bugfix

[ftynse]

Please make sure clang-format is happy

[nicolasvasilache]

Format

[rriddle]

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

14

nit: Where is the Function include used?

mlir/include/mlir/Interfaces/VectorUnrollInterface.h

18

nit: Is this one necessary?

mlir/docs/OpDefinitions.md

InterfaceMethod<"/*insert doc here*/",
       "unsigned", "getNumInputsAndOutputs", (ins), [{
         return op.getNumInputs() + op.getNumOutputs();
     }]>,
 
     // Provide only a default definition of the method.
     // Note: `ConcreteOp` corresponds to the derived operation typename.
     InterfaceMethod<"/*insert doc here*/",
       "unsigned", "getNumWithDefault", (ins), /*methodBody=*/[{}], [{
-        ConcreteOp op = cast<ConcreteOp>(getOperation());
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         return op.getNumInputs() + op.getNumOutputs();
     }]>,
   ];
 }
 
 // Operation interfaces can optionally be wrapped inside
 // DeclareOpInterfaceMethods. This would result in autogenerating declarations
 // for members `foo`, `bar` and `fooStatic`. Methods with bodies are not

mlir/include/mlir/Dialect/StandardOps/IR/Ops.h

 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/VectorUnrollInterface.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 // Pull in all enum type definitions and utility function declarations.
 #include "mlir/Dialect/StandardOps/IR/OpsEnums.h.inc"
 
 namespace mlir {
 class AffineMap;
 class Builder;

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 
 #ifndef STANDARD_OPS
 #define STANDARD_OPS
 
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/VectorUnrollInterface.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 
 def StandardOps_Dialect : Dialect {
   let name = "std";
   let cppNamespace = "";
   let hasConstantMaterializer = 1;
 }
 
class UnaryOpSameOperandAndResultType<string mnemonic,
                                       list<OpTrait> traits = []> :
     UnaryOp<mnemonic, !listconcat(traits, [SameOperandsAndResultType])> {
   let parser = [{
     return impl::parseOneResultSameOperandTypeOp(parser, result);
   }];
 }
 
 class FloatUnaryOp<string mnemonic, list<OpTrait> traits = []> :
-    UnaryOpSameOperandAndResultType<mnemonic, traits>,
+    UnaryOpSameOperandAndResultType<mnemonic,
+      !listconcat(traits,
+                  [DeclareOpInterfaceMethods<VectorUnrollOpInterface>])>,
     Arguments<(ins FloatLike:$operand)>;
 
 // Base class for standard arithmetic operations.  Requires operands and
 // results to be of the same type, but does not constrain them to specific
 // types.  Individual classes will have `lhs` and `rhs` accessor to operands.
 class ArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
     Op<StandardOps_Dialect, mnemonic,
        !listconcat(traits, [NoSideEffect, SameOperandsAndResultType])> {

 // tensors thereof.  This operation takes two operands and returns one result,
 // each of these is required to be of the same type.  This type may be an
 // integer scalar type, a vector whose element type is an integer type, or an
 // integer tensor.  The custom assembly form of the operation is as follows
 //
 //     <op>i %0, %1 : i32
 //
 class IntArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
-    ArithmeticOp<mnemonic, traits>,
+    ArithmeticOp<mnemonic,
+      !listconcat(traits,
+                  [DeclareOpInterfaceMethods<VectorUnrollOpInterface>])>,
     Arguments<(ins SignlessIntegerLike:$lhs, SignlessIntegerLike:$rhs)>;
 
 // Base class for standard arithmetic binary operations on floats, vectors and
 // tensors thereof.  This operation has two operands and returns one result,
 // each of these is required to be of the same type.  This type may be a
 // floating point scalar type, a vector whose element type is a floating point
 // type, or a floating point tensor.  The custom assembly form of the operation
 // is as follows
 //
 //     <op>f %0, %1 : f32
 //
 class FloatArithmeticOp<string mnemonic, list<OpTrait> traits = []> :
-    ArithmeticOp<mnemonic, traits>,
+    ArithmeticOp<mnemonic,
+      !listconcat(traits,
+                  [DeclareOpInterfaceMethods<VectorUnrollOpInterface>])>,
     Arguments<(ins FloatLike:$lhs, FloatLike:$rhs)>;
 
 // Base class for standard arithmetic operations on complex numbers with a
 // floating-point element type.
 // These operations take two operands and return one result, all of which must
 // be complex numbers of the same type.
 // The assembly format is as follows
 //
def ViewOp : Std_Op<"view", [
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // XOrOp
 //===----------------------------------------------------------------------===//
 
 def XOrOp : IntArithmeticOp<"xor", [Commutative]> {
   let summary = "integer binary xor";
   let description = [{
     The `xor` operation takes two operands and returns one result, each of these
     is required to be the same type. This type may be an integer scalar type, a
     vector whose element type is integer, or a tensor of integers. It has no
     standard attributes.
 
     Example:
 

mlir/include/mlir/Dialect/Vector/VectorOps.h

-//===- VectorOps.h - MLIR Super Vectorizer Operations -----------*- C++ -*-===//
+//===- VectorOps.h - MLIR Vector Dialect Operations -------------*- C++ -*-===//
 //
 // 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 defines the Vector dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_VECTOR_VECTOROPS_H
 #define MLIR_DIALECT_VECTOR_VECTOROPS_H
 
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/VectorUnrollInterface.h"
 
 namespace mlir {
 class MLIRContext;
 class OwningRewritePatternList;
 namespace vector {
 
 /// Collect a set of vector-to-vector canonicalization patterns.
 void populateVectorToVectorCanonicalizationPatterns(

mlir/include/mlir/Dialect/Vector/VectorOps.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef VECTOR_OPS
 #define VECTOR_OPS
 
 include "mlir/Dialect/Affine/IR/AffineOpsBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/VectorUnrollInterface.td"
 
 def Vector_Dialect : Dialect {
   let name = "vector";
   let cppNamespace = "vector";
   let hasConstantMaterializer = 1;
 }
 
 // Base class for Vector dialect ops.
 class Vector_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<Vector_Dialect, mnemonic, traits> {
   // For every vector op, there needs to be a:
   //   * void print(OpAsmPrinter &p, ${C++ class of Op} op)
   //   * LogicalResult verify(${C++ class of Op} op)
   //   * ParseResult parse${C++ class of Op}(OpAsmParser &parser,
   //                                         OperationState &result)
   // functions.
   let printer = [{ return ::print(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
 // TODO(andydavis, ntv) Add an attribute to specify a different algebra
 // with operators other than the current set: {*, +}.
 def Vector_ContractionOp :
-  Vector_Op<"contract", [NoSideEffect,
+  Vector_Op<"contract", [
+      NoSideEffect,
       PredOpTrait<"lhs and rhs have same element type", TCopVTEtIsSameAs<0, 1>>,
       PredOpTrait<"third operand acc and result have same element type",
-                 TCresVTEtIsSameAsOpBase<0, 2>>]>,
+                  TCresVTEtIsSameAsOpBase<0, 2>>,
+      DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
+    ]>,
     Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyType:$acc,
                Variadic<VectorOf<[I1]>>:$masks,
                AffineMapArrayAttr:$indexing_maps, ArrayAttr:$iterator_types)>,
     Results<(outs AnyType)> {
   let summary = "vector contraction operation";
   let description = [{
     Computes the sum of products of vector elements along contracting
     dimension pairs from 2 vectors of rank M and N respectively, adds this
code extraTransferDeclaration = [{
     }
     VectorType getVectorType() {
       return vector().getType().cast<VectorType>();
     }
   }];
 }
 
 def Vector_TransferReadOp :
-  Vector_Op<"transfer_read">,
+  Vector_Op<"transfer_read", [
+      DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
+    ]>,
     Arguments<(ins AnyMemRef:$memref, Variadic<Index>:$indices,
                AffineMapAttr:$permutation_map, AnyType:$padding,
                OptionalAttr<BoolArrayAttr>:$masked)>,
     Results<(outs AnyVector:$vector)> {
 
   let summary = "Reads a supervector from memory into an SSA vector value.";
 
   let description = [{
static AffineMap getTransferMinorIdentityMap(
         return impl::getTransferMinorIdentityMap(memRefType, vectorType);
       }
   }];
 
   let hasFolder = 1;
 }
 
 def Vector_TransferWriteOp :
-  Vector_Op<"transfer_write">,
+  Vector_Op<"transfer_write", [
+      DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
+    ]>,
     Arguments<(ins AnyVector:$vector, AnyMemRef:$memref,
                Variadic<Index>:$indices,
                AffineMapAttr:$permutation_map,
                OptionalAttr<BoolArrayAttr>:$masked)> {
 
   let summary = "The vector.transfer_write op writes a supervector to memory.";
 
   let description = [{

mlir/include/mlir/Dialect/Vector/VectorTransformPatterns.td

 
 include "mlir/IR/OpBase.td"
 
 class HasShape<list<int> shape> :
   CPred<"$0.getType().cast<ShapedType>().hasStaticShape({" #
     StrJoinInt<shape>.result # "})">;
 
 class UnrollVectorOp<list<int> factors> : NativeCodeCall<
-  "unrollSingleResultOpMatchingType($_builder, $0.getDefiningOp(), " #
+  "unrollSingleResultVectorOp($_builder, $0.getDefiningOp(), " #
     "{" # StrJoinInt<factors>.result # "})">;
 
 #endif // VECTOR_TRANSFORM_PATTERNS

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

 //===- VectorTransforms.h - Vector transformations as patterns --*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_VECTOR_VECTORTRANSFORMS_H_
 #define DIALECT_VECTOR_VECTORTRANSFORMS_H_
 
 #include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/Dialect/Vector/VectorUtils.h"
+#include "mlir/IR/Function.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 class MLIRContext;
 class OwningRewritePatternList;
 
 /// Collect a set of patterns to convert from the Vector dialect to itself.
 /// Should be merged with populateVectorToSCFLoweringPattern.
 void populateVectorToVectorConversionPatterns(
     MLIRContext *context, OwningRewritePatternList &patterns,
     ArrayRef<int64_t> coarseVectorShape = {},
     ArrayRef<int64_t> fineVectorShape = {});
 
 namespace vector {
 
-// Entry point for unrolling declarative pattern rewrites.
-// `op` is unrolled to the `targetShape` as follows, for each of its operands:
-//   1. the unrolled type `unrolledVectorType` and number of unrolled instances
-//   `numUnrolledInstances` are computed from the `targetShape`. For now it is
-//   assumed the unrolling factors divide the vector sizes.
-//   2. a fakeFork cast op is inserted that takes the operand and returns
-//   `numUnrolledInstances` results of type `unrolledVectorType`.
-//   3. the original op is cloned `numUnrolledInstances` times, once for each
-//   result of the fakeFork cast op.
-//   4. a fakeJoin cast op takes all these results and merges them into a single
-//   aggregate vector result whose size matches the original non-unrolled op
-//   operand types.
-//
-// Example:
-//
-//    opA(operand0, operand1)  // numUnrolledInstances = 3
-//
-//            operand0                   operand1
-//               |                          |
-//             fork                       fork
-//        <----------gather all fork ops --------->
-//              /|\                        /|\
-//          f00 f01 f02                f10 f11 f12
-//        <---------- clone op 3 times --------->
-//          opA0(f00, f10), opA1(f01, f11), opA2(f02, f12)
-//                 \            |            /
-//      <-------------------- join ------------------------->
-//
-// Other local patterns then kick in iteratively (including DCE) and compose
-// until all the fakeFork and fakeJoin ops are removed.
-//
-// This will be extended in the future to support more advanced use cases than
-// simple pointwise ops.
-SmallVector<Value, 1>
-unrollSingleResultOpMatchingType(OpBuilder &builder, Operation *op,
+/// Entry point for unrolling declarative pattern rewrites.
+/// `op` is unrolled to the `targetShape` as follows, for each of its operands:
+///   1. the unrolled type `unrolledVectorType` and number of unrolled instances
+///   `numUnrolledInstances` are computed from the `targetShape`. For now it is
+///   assumed the unrolling factors divide the vector sizes.
+///   2. a fakeFork cast op is inserted that takes the operand and returns
+///   `numUnrolledInstances` results of type `unrolledVectorType`.
+///   3. the original op is cloned `numUnrolledInstances` times, once for each
+///   result of the fakeFork cast op.
+///   4. a fakeJoin cast op takes all these results and merges them into a
+///   single aggregate vector result whose size matches the original
+///   non-unrolled op operand types.
+///
+/// Example:
+///
+///    opA(operand0, operand1)  // numUnrolledInstances = 3
+///
+///            operand0                   operand1
+///               |                          |
+///             fork                       fork
+///        <----------gather all fork ops --------->
+///              /|\                        /|\
+///          f00 f01 f02                f10 f11 f12
+///        <---------- clone op 3 times --------->
+///          opA0(f00, f10), opA1(f01, f11), opA2(f02, f12)
+///                 \            |            /
+///      <-------------------- join ------------------------->
+///
+/// Other local patterns then kick in iteratively (including DCE) and compose
+/// until all the fakeFork and fakeJoin ops are removed.
+///
+/// This will be extended in the future to support more advanced use cases than
+/// simple pointwise ops.
+SmallVector<Value, 1> unrollSingleResultVectorOp(OpBuilder &builder,
+                                                 Operation *op,
                                                  ArrayRef<int64_t> targetShape);
 
+/// Pattern to apply `unrollSingleResultVectorOp` to a `targetShape`
+/// declaratively.
+template <typename OpTy>
+struct UnrollVectorPattern : public OpRewritePattern<OpTy> {
+  using FilterConstraintType = std::function<LogicalResult(OpTy op)>;
+  UnrollVectorPattern(
+      ArrayRef<int64_t> targetShape, MLIRContext *context,
+      FilterConstraintType constraint = [](OpTy op) { return success(); })
+      : OpRewritePattern<OpTy>(context),
+        targetShape(targetShape.begin(), targetShape.end()),
+        filter(constraint) {}
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    if (failed(filter(op)))
+      return failure();
+    auto unrollableVectorOp =
+        dyn_cast<VectorUnrollOpInterface>(op.getOperation());
+    if (!unrollableVectorOp)
+      return failure();
+    auto maybeUnrollShape = unrollableVectorOp.getShapeForUnroll();
+    if (!maybeUnrollShape)
+      return failure();
+    auto maybeShapeRatio = shapeRatio(*maybeUnrollShape, targetShape);
+    if (!maybeShapeRatio ||
+        llvm::all_of(*maybeShapeRatio, [](int64_t v) { return v == 1; }))
+      return failure();
+    if (op.getOperation()->getNumResults() != 1)
+      return failure();
+    auto resultVector = unrollSingleResultVectorOp(rewriter, op, targetShape);
+    if (resultVector.size() != 1)
+      return failure();
+    rewriter.replaceOp(op, resultVector.front());
+    return success();
+  }
+
+private:
+  SmallVector<int64_t, 4> targetShape;
+  FilterConstraintType filter;
+};
+
 } // namespace vector
 
 //===----------------------------------------------------------------------===//
 // Finer-grained patterns exposed for more control over individual lowerings.
 //===----------------------------------------------------------------------===//
 
 /// Progressive lowering of a `vector.contract %a, %b, %c` with row-major matmul
 /// semantics to:
 /// ```
 ///    %flattened_a = vector.shape_cast %a
 ///    %flattened_b = vector.shape_cast %b
 ///    %flattened_d = vector.matmul %flattened_a, %flattened_b
 ///    %d = vector.shape_cast %%flattened_d
 ///    %e = add %c, %d
 /// ```
 /// `vector.matmul` later lowers to `llvm.matrix.multiply`.
 //
 /// This only kicks in when VectorTransformsOptions is set to OuterProduct and
 /// the vector.contract op is a row-major matrix multiply.
 class ContractionOpToMatmulOpLowering
     : public OpRewritePattern<vector::ContractionOp> {
 public:
   using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
 
   ContractionOpToMatmulOpLowering(
       vector::VectorTransformsOptions vectorTransformsOptions,
       MLIRContext *context)
       : OpRewritePattern<vector::ContractionOp>(context),
         vectorTransformsOptions(vectorTransformsOptions) {}
 
   LogicalResult match(vector::ContractionOp op) const override;
   void rewrite(vector::ContractionOp op,
                PatternRewriter &rewriter) const override;
 
 private:
   /// Options to control the vector patterns.

mlir/include/mlir/Interfaces/CMakeLists.txt

 add_mlir_interface(CallInterfaces)
 add_mlir_interface(ControlFlowInterfaces)
 add_mlir_interface(CopyOpInterface)
 add_mlir_interface(DerivedAttributeOpInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)
 add_mlir_interface(SideEffectInterfaces)
+add_mlir_interface(VectorUnrollInterface)
 add_mlir_interface(ViewLikeInterface)
 

mlir/include/mlir/Interfaces/VectorUnrollInterface.h

+//===- VectorUnrollInterface.h - Vector unrolling interface ---------------===//
+//
+// 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 operation interface for vector ops that can be
+// unrolled.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_VECTORUNROLLINTERFACE_H
+#define MLIR_INTERFACES_VECTORUNROLLINTERFACE_H
+
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/StandardTypes.h"
+
+namespace mlir {
+
+#include "mlir/Interfaces/VectorUnrollInterface.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_INTERFACES_VECTORUNROLLINTERFACE_H

mlir/include/mlir/Interfaces/VectorUnrollInterface.td

+//===- VectorUnrollInterface.td - VectorUnroll interface ---*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface for operations on vectors that can be unrolled.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_VECTORUNROLLINTERFACE
+#define MLIR_INTERFACES_VECTORUNROLLINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def VectorUnrollOpInterface : OpInterface<"VectorUnrollOpInterface"> {
+  let description = [{
+    Encodes properties of an operation on vectors that can be unrolled.
+  }];
+
+  let methods = [
+    InterfaceMethod<[{
+        Returns the shape ratio of unrolling to the target vector shape
+        `targetShape`. Returns `None` if the op cannot be unrolled to the target
+        vector shape.
+      }],
+      "Optional<SmallVector<int64_t, 4>>",
+      "getShapeForUnroll",
+      (ins),
+      /*methodBody=*/[{}],
+      [{
+        auto vt = this->getOperation()->getResult(0).getType().
+          template dyn_cast<VectorType>();
+        if (!vt)
+          return None;
+        SmallVector<int64_t, 4> res(vt.getShape().begin(), vt.getShape().end());
+        return res;
+      }]
+    >,
+  ];
+}
+
+#endif // MLIR_INTERFACES_VECTORUNROLLINTERFACE

mlir/lib/Dialect/StandardOps/CMakeLists.txt

add_mlir_dialect_library(MLIRStandardOps
   MLIRStandardOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIREDSC
   MLIRIR
   MLIRSideEffectInterfaces
+  MLIRVectorUnrollInterface
   MLIRViewLikeInterface
   )
 
 add_subdirectory(Transforms)

mlir/lib/Dialect/Vector/CMakeLists.txt

add_mlir_dialect_library(MLIRVector
   LINK_LIBS PUBLIC
   MLIREDSC
   MLIRIR
   MLIRStandardOps
   MLIRAffineOps
   MLIRSCF
   MLIRLoopAnalysis
   MLIRSideEffectInterfaces
+  MLIRVectorUnrollInterface
   )

mlir/lib/Dialect/Vector/VectorOps.cpp

SmallVector<AffineMap, 4> ContractionOp::getIndexingMaps() {
   SmallVector<AffineMap, 4> res;
   auto mapAttrs = indexing_maps().getValue();
   res.reserve(mapAttrs.size());
   for (auto mapAttr : mapAttrs)
     res.push_back(mapAttr.cast<AffineMapAttr>().getValue());
   return res;
 }
 
+Optional<SmallVector<int64_t, 4>> ContractionOp::getShapeForUnroll() {
+  SmallVector<int64_t, 4> shape;
+  getIterationBounds(shape);
+  return shape;
+}
+
 //===----------------------------------------------------------------------===//
 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(vector::ExtractElementOp op) {
   VectorType vectorType = op.getVectorType();
   if (vectorType.getRank() != 1)
     return op.emitOpError("expected 1-D vector");

 
 OpFoldResult TransferReadOp::fold(ArrayRef<Attribute>) {
   /// transfer_read(memrefcast) -> transfer_read
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return OpFoldResult();
 }
 
+Optional<SmallVector<int64_t, 4>> TransferReadOp::getShapeForUnroll() {
+  auto s = getVectorType().getShape();
+  return SmallVector<int64_t, 4>{s.begin(), s.end()};
+}
+
 //===----------------------------------------------------------------------===//
 // TransferWriteOp
 //===----------------------------------------------------------------------===//
 
 /// Builder that sets permutation map to 'getMinorIdentityMap'.
 void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
                             Value vector, Value memref, ValueRange indices,
                             ArrayRef<bool> maybeMasked) {
return verifyPermutationMap(permutationMap,
                               [&op](Twine t) { return op.emitOpError(t); });
 }
 
 LogicalResult TransferWriteOp::fold(ArrayRef<Attribute>,
                                     SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCast(*this);
 }
 
+Optional<SmallVector<int64_t, 4>> TransferWriteOp::getShapeForUnroll() {
+  auto s = getVectorType().getShape();
+  return SmallVector<int64_t, 4>{s.begin(), s.end()};
+}
+
 //===----------------------------------------------------------------------===//
 // ShapeCastOp
 //===----------------------------------------------------------------------===//
 
 /// Returns true if each element of 'a' is equal to the product of a contiguous
 /// sequence of the elements of 'b'. Returns false otherwise.
 static bool isValidShapeCast(ArrayRef<int64_t> a, ArrayRef<int64_t> b) {
   unsigned rankA = a.size();

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 #include "mlir/IR/Function.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Interfaces/VectorUnrollInterface.h"
 
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "vector-to-vector"
 
 using namespace mlir;

 // *) Add pattern which matches SourceOp -> StridedSlice -> UserOp which checks
 //    if there are duplicate identical StridedSlice ops from SourceOp, and
 //    rewrites itself to use the first duplicate. This transformation should
 //    cause users of identifical StridedSlice ops to reuse the same StridedSlice
 //    operation, and leave the duplicate StridedSlice ops with no users
 //    (removable with DCE).
 
 // TODO(andydavis) Generalize this to support structured ops beyond
-// vector ContractionOp, and merge it with 'unrollSingleResultOpMatchingType'
+// vector ContractionOp, and merge it with 'unrollSingleResultVectorOp'
 static Value unrollSingleResultStructuredOp(Operation *op,
                                             ArrayRef<int64_t> iterationBounds,
                                             std::vector<VectorState> &vectors,
                                             unsigned resultIndex,
                                             ArrayRef<int64_t> targetShape,
                                             OpBuilder &builder) {
   auto shapedType = op->getResult(0).getType().dyn_cast_or_null<ShapedType>();
   if (!shapedType || !shapedType.hasStaticShape())
static Value unrollSingleResultStructuredOp(Operation *op,
   Value insertSlicesOp = builder.create<vector::InsertSlicesOp>(
       op->getLoc(), resultVectorType, tupleOp, builder.getI64ArrayAttr(sizes),
       builder.getI64ArrayAttr(strides));
   return insertSlicesOp;
 }
 
 static void getVectorContractionOpUnrollState(
     vector::ContractionOp contractionOp, ArrayRef<int64_t> targetShape,
-    SmallVectorImpl<int64_t> &iterationBounds,
     std::vector<VectorState> &vectors, unsigned &resultIndex) {
-  // Get contraction op iteration bounds.
-  contractionOp.getIterationBounds(iterationBounds);
-  assert(iterationBounds.size() == targetShape.size());
   // Get map from iteration space index to lhs/rhs/result shape index.
   std::vector<DenseMap<int64_t, int64_t>> iterationIndexMapList;
   contractionOp.getIterationIndexMap(iterationIndexMapList);
   unsigned numIterators = iterationIndexMapList.size();
   vectors.resize(numIterators);
   unsigned accOperandIndex = vector::ContractionOp::getAccOperandIndex();
   for (unsigned i = 0; i < numIterators; ++i) {
     vectors[i].type = contractionOp.getOperand(i).getType().cast<VectorType>();
     vectors[i].indexMap = iterationIndexMapList[i];
     vectors[i].operandIndex = i;
     vectors[i].isAcc = i == accOperandIndex ? true : false;
   }
 
   if (llvm::size(contractionOp.masks()) == 2) {
     // Add vectors for lhs/rhs vector mask arguments. Masks have the
     // same vector shape lhs/rhs args, so copy their index maps.
     vectors.push_back({contractionOp.getLHSVectorMaskType(),
                        vectors[0].indexMap, accOperandIndex + 1, false});
     vectors.push_back({contractionOp.getRHSVectorMaskType(),
                        vectors[1].indexMap, accOperandIndex + 2, false});
   }
-  // Unroll 'op' 'iterationBounds' to 'targetShape'.
   // TODO(andydavis) Use linalg style 'args_in'/'args_out' to partition
   // 'vectors' instead of 'resultIndex'.
   resultIndex = accOperandIndex;
 }
 
-static void
-getVectorElementwiseOpUnrollState(Operation *op, ArrayRef<int64_t> targetShape,
-                                  SmallVectorImpl<int64_t> &iterationBounds,
+static void getVectorElementwiseOpUnrollState(Operation *op,
+                                              ArrayRef<int64_t> targetShape,
                                               std::vector<VectorState> &vectors,
                                               unsigned &resultIndex) {
   // Verify that operation and operands all have the same vector shape.
   auto resultType = op->getResult(0).getType().dyn_cast_or_null<VectorType>();
   assert(resultType && "Expected op with vector result type");
   auto resultShape = resultType.getShape();
   // Verify that all operands have the same vector type as result.
   assert(llvm::all_of(op->getOperandTypes(),
                       [=](Type type) { return type == resultType; }));
-  // Populate 'iterationBounds' with 'resultShape' for elementwise operations.
-  iterationBounds.assign(resultShape.begin(), resultShape.end());
 
   // Create trivial elementwise identity index map based on 'resultShape'.
   DenseMap<int64_t, int64_t> indexMap;
   indexMap.reserve(resultShape.size());
   for (unsigned i = 0; i < resultShape.size(); ++i)
     indexMap[i] = i;
 
   // Create VectorState each operand and single result.
   unsigned numVectors = op->getNumOperands() + op->getNumResults();
   vectors.resize(numVectors);
   for (unsigned i = 0; i < op->getNumOperands(); ++i)
     vectors[i] = {resultType, indexMap, i, false};
   vectors[numVectors - 1] = {resultType, indexMap, -1, false};
   resultIndex = numVectors - 1;
 }
 
 // Entry point for unrolling declarative pattern rewrites.
-SmallVector<Value, 1> mlir::vector::unrollSingleResultOpMatchingType(
-    OpBuilder &builder, Operation *op, ArrayRef<int64_t> targetShape) {
+SmallVector<Value, 1>
+mlir::vector::unrollSingleResultVectorOp(OpBuilder &builder, Operation *op,
+                                         ArrayRef<int64_t> targetShape) {
   assert(op->getNumResults() == 1 && "Expected single result operation");
 
   // Populate 'iterationBounds', 'vectors' and 'resultIndex' to unroll 'op'.
   SmallVector<int64_t, 6> iterationBounds;
+  auto unrollableVectorOp = cast<VectorUnrollOpInterface>(op);
+  auto maybeUnrollShape = unrollableVectorOp.getShapeForUnroll();
+  assert(maybeUnrollShape && "Trying to unroll an incorrect vector op");
+
   std::vector<VectorState> vectors;
   unsigned resultIndex;
 
   if (auto contractionOp = dyn_cast<vector::ContractionOp>(op)) {
     // Populate state for vector ContractionOp.
-    getVectorContractionOpUnrollState(contractionOp, targetShape,
-                                      iterationBounds, vectors, resultIndex);
+    getVectorContractionOpUnrollState(contractionOp, targetShape, vectors,
+                                      resultIndex);
   } else {
     // Populate state for vector elementwise op.
-    getVectorElementwiseOpUnrollState(op, targetShape, iterationBounds, vectors,
-                                      resultIndex);
+    getVectorElementwiseOpUnrollState(op, targetShape, vectors, resultIndex);
   }
 
   // Unroll 'op' with 'iterationBounds' to 'targetShape'.
   return SmallVector<Value, 1>{unrollSingleResultStructuredOp(
-      op, iterationBounds, vectors, resultIndex, targetShape, builder)};
+      op, *maybeUnrollShape, vectors, resultIndex, targetShape, builder)};
 }
 
 /// Generates slices of 'vectorType' according to 'sizes' and 'strides, and
 /// calls 'fn' with linear index and indices for each slice.
 static void generateTransferOpSlices(
     Type memrefElementType, VectorType vectorType, TupleType tupleType,
     ArrayRef<int64_t> sizes, ArrayRef<int64_t> strides, ArrayRef<Value> indices,
     OpBuilder &builder, function_ref<void(unsigned, ArrayRef<Value>)> fn) {

mlir/lib/Interfaces/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   CallInterfaces.cpp
   ControlFlowInterfaces.cpp
   CopyOpInterface.cpp
   DerivedAttributeOpInterface.cpp
   InferTypeOpInterface.cpp
   LoopLikeInterface.cpp
   SideEffectInterfaces.cpp
+  VectorUnrollInterface.cpp
   ViewLikeInterface.cpp
   )
 
 function(add_mlir_interface_library name)
   add_mlir_library(MLIR${name}
     ${name}.cpp
 
     ADDITIONAL_HEADER_DIRS

 
 add_mlir_interface_library(CallInterfaces)
 add_mlir_interface_library(ControlFlowInterfaces)
 add_mlir_interface_library(CopyOpInterface)
 add_mlir_interface_library(DerivedAttributeOpInterface)
 add_mlir_interface_library(InferTypeOpInterface)
 add_mlir_interface_library(LoopLikeInterface)
 add_mlir_interface_library(SideEffectInterfaces)
+add_mlir_interface_library(VectorUnrollInterface)
 add_mlir_interface_library(ViewLikeInterface)
 

mlir/lib/Interfaces/VectorUnrollInterface.cpp

+//===- VectorUnrollInterface.cpp - Unrollable vector operations -*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Interfaces/VectorUnrollInterface.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// VectorUnroll Interfaces
+//===----------------------------------------------------------------------===//
+
+/// Include the definitions of the VectorUntoll interfaces.
+#include "mlir/Interfaces/VectorUnrollInterface.cpp.inc"

mlir/test/Dialect/Vector/vector-transforms.mlir

 // RUN: mlir-opt %s -test-vector-to-vector-conversion | FileCheck %s
+// RUN: mlir-opt %s -test-vector-unrolling-patterns | FileCheck %s
 
 // CHECK-DAG: #[[MAP0:map[0-9]+]] = affine_map<(d0, d1) -> (d0, d1)>
 // CHECK-DAG: #[[MAP1:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
 
 // CHECK-LABEL: func @add4x2
 //      CHECK: %[[ES1:.*]] = vector.extract_slices %{{.*}}, [2, 2], [1, 1] : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
 // CHECK-NEXT: %[[ES2:.*]] = vector.extract_slices %{{.*}}, [2, 2], [1, 1] : vector<4x2xf32> into tuple<vector<2x2xf32>, vector<2x2xf32>>
 // CHECK-NEXT: %[[TG1:.*]] = vector.tuple_get %[[ES1]], 0 : tuple<vector<2x2xf32>, vector<2x2xf32>>

mlir/test/lib/Transforms/TestVectorTransforms.cpp

void runOnFunction() override {
     if (lowerToFlatTranspose)
       transposeLowering = VectorTransposeLowering::Flat;
     VectorTransformsOptions options{contractLowering, transposeLowering};
     populateVectorContractLoweringPatterns(patterns, &getContext(), options);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };
 
+struct TestVectorUnrollingPatterns
+    : public PassWrapper<TestVectorUnrollingPatterns, FunctionPass> {
+  void runOnFunction() override {
+    MLIRContext *ctx = &getContext();
+    OwningRewritePatternList patterns;
+    patterns.insert<UnrollVectorPattern<AddFOp>>(ArrayRef<int64_t>{2, 2}, ctx);
+    patterns.insert<UnrollVectorPattern<vector::ContractionOp>>(
+        ArrayRef<int64_t>{2, 2, 2}, ctx);
+    populateVectorToVectorCanonicalizationPatterns(patterns, ctx);
+    populateVectorToVectorTransformationPatterns(patterns, ctx);
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
+  }
+};
+
 } // end anonymous namespace
 
 namespace mlir {
 void registerTestVectorConversions() {
   PassRegistration<TestVectorToVectorConversion> vectorToVectorPass(
       "test-vector-to-vector-conversion",
       "Test conversion patterns between ops in the vector dialect");
 
   PassRegistration<TestVectorSlicesConversion> slicesPass(
       "test-vector-slices-conversion",
       "Test conversion patterns that lower slices ops in the vector dialect");
 
   PassRegistration<TestVectorContractionConversion> contractionPass(
       "test-vector-contraction-conversion",
       "Test conversion patterns that lower contract ops in the vector dialect");
+
+  PassRegistration<TestVectorUnrollingPatterns> contractionUnrollingPass(
+      "test-vector-unrolling-patterns",
+      "Test conversion patterns to unroll contract ops in the vector dialect");
 }
 } // namespace mlir