[mlir][Linalg] Introduce a ContractionOpInterface

This revision takes advantage of recent extensions to vectorization to refactor contraction detection into a bona fide Linalg interface.
The mlit-linalg-ods-gen parser is extended to support adding such interfaces.
The detection that was originally enabling vectorization is refactored to serve as both a test on a generic LinalgOp as well as to verify ops that declare to conform to that interface.

This is plugged through Linalg transforms and strategies but it quickly becomes evident that the complexity and rigidity of the C++ class based templating does not pay for itself.
Therefore, this revision changes the API for vectorization patterns to get rid of templates as much as possible.
Variadic templates are relegated to the internals of LinalgTransformationFilter as much as possible and away from the user-facing APIs.

It is expected other patterns / transformations will follow the same path and drop as much C++ templating as possible from the class definition.

Differential revision: https://reviews.llvm.org/D95973

Author: nicolasvasilache

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

@@ -22,13 +22,21 @@
 
 namespace mlir {
 namespace linalg {
+class LinalgOp;
 
 /// Returns the values obtained by applying `map` to the list of values.
 SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
                                        AffineMap map, ValueRange values);
 
+/// Checks whether `linalgOp` conforms to ContractionOpInterface.
+// TODO: embed within `isa<ContractionOpInterface>` if possible / natural.
+bool isaContractionOpInterface(LinalgOp linalgOp);
+
 namespace detail {
 
+/// Verify that `op` conforms to ContractionOpInterface.
+LogicalResult verifyContractionInterface(Operation *op);
+
 /// Verify that `op` conforms to the invariants of StructuredOpInterface
 LogicalResult verifyStructuredOpInterface(Operation *op);
 

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

@@ -15,8 +15,28 @@
 
 include "mlir/IR/OpBase.td"
 
-// The linalg 'LinalgStructuredInterface' provides access to the 'LinalgOp'
-// interface.
+// The 'LinalgContractionOpInterface' provides access to the
+// 'ContractionOpInterface'.
+def LinalgContractionOpInterface : OpInterface<"ContractionOpInterface"> {
+  let description = [{
+   A Linalg contraction is defined in general terms:
+     1. Has 2 input and 1 output shapes.
+     2. Has at least one reduction dimension.
+     3. Has only projected permutation indexing maps.
+     4. its body computes `u5(u1(c) + u2(u3(a) * u4(b)))` on some field
+     (AddOpType, MulOpType), where u1, u2, u3, u4 and u5 represent scalar unary
+     operations that may change the type (e.g. for mixed-precision).
+   As a consequence, when vectorization of such an op occurs, the only special
+   behavior is that the (unique) MulOpType is vectorized into a
+   `vector.contract`. All other ops are handled in a generic fashion.
+   In the future, we may wish to allow more input arguments and elementwise and
+   constant operations that do not involve the reduction dimension(s).
+  }];
+  let cppNamespace = "::mlir::linalg";
+  let verify = [{ return detail::verifyContractionInterface($_op); }];
+}
+
+// The 'LinalgStructuredInterface' provides access to the 'LinalgOp' interface.
 def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
   let cppNamespace = "::mlir::linalg";
   let methods = [

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

@@ -1,36 +1,43 @@
-ods_def<MatmulOp>:
+ods_def<MatmulOp>
+implements_interface<LinalgContractionOpInterface> :
 def matmul(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
   C(m, n) = std_addf<k>(std_mulf(A(m, k), B(k, n)));
 }
 
-ods_def<MatmulColumnMajorOp>:
+ods_def<MatmulColumnMajorOp>
+implements_interface<LinalgContractionOpInterface> :
 def matmul_column_major(A: f32(K, M), B: f32(N, K)) -> (C: f32(N, M)) {
   C(n, m) = std_addf<k>(std_mulf(A(k, m), B(n, k)));
 }
 
-ods_def<MatmulI8I8I32Op>:
+ods_def<MatmulI8I8I32Op>
+implements_interface<LinalgContractionOpInterface> :
 def matmul_i8_i8_i32(A: i8(M, K), B: i8(K, N)) -> (C: i32(M, N)) {
   // TODO: ideally something closer to
   //   C(m, n) += cast<i32>(A(m, k)) * cast<i32>(B(k, n))
   C(m, n) = std_addi<k>(std_sexti32(std_muli(A(m, k), B(k, n))));
 }
 
-ods_def<MatvecOp>:
+ods_def<MatvecOp>
+implements_interface<LinalgContractionOpInterface> :
 def matvec(A: f32(M, N), y: f32(N)) -> (x: f32(M)) {
   x(m) = std_addf<n>(std_mulf(A(m, n), y(n)));
 }
 
-ods_def<VecmatOp>:
+ods_def<VecmatOp>
+implements_interface<LinalgContractionOpInterface> :
 def vecmat(y: f32(M), A: f32(M, N)) -> (x: f32(N)) {
   x(n) = std_addf<m>(std_mulf(y(m), A(m, n)));
 }
 
-ods_def<DotOp>:
+ods_def<DotOp>
+implements_interface<LinalgContractionOpInterface> :
 def dot(A: f32(M), B: f32(M)) -> (C: f32()) {
   C() = std_addf<m>(std_mulf(A(m), B(m)));
 }
 
-ods_def<BatchMatmulOp>:
+ods_def<BatchMatmulOp>
+implements_interface<LinalgContractionOpInterface> :
 def batch_matmul(A: f32(Batch, M, K), B: f32(Batch, K, N)) -> (C: f32(Batch, M, N)) {
   C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(b, k, n)));
 }

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

@@ -35,23 +35,33 @@ template <template <typename> class PatternType, typename ConcreteOpType,
           typename OptionsType,
           typename = std::enable_if_t<std::is_member_function_pointer<
               decltype(&ConcreteOpType::getOperationName)>::value>>
-void sfinae_enqueue(OwningRewritePatternList &patterList, OptionsType options,
+void sfinae_enqueue(OwningRewritePatternList &patternList, OptionsType options,
                     MLIRContext *context, StringRef opName,
                     linalg::LinalgTransformationFilter m) {
   assert(opName == ConcreteOpType::getOperationName() &&
          "explicit name must match ConcreteOpType::getOperationName");
-  patterList.insert<PatternType<ConcreteOpType>>(context, options, m);
+  patternList.insert<PatternType<ConcreteOpType>>(context, options, m);
 }
 
 /// SFINAE: Enqueue helper for OpType that do not have a `getOperationName`
 /// (e.g. LinalgOp, other interfaces, Operation*).
 template <template <typename> class PatternType, typename OpType,
           typename OptionsType>
-void sfinae_enqueue(OwningRewritePatternList &patterList, OptionsType options,
+void sfinae_enqueue(OwningRewritePatternList &patternList, OptionsType options,
                     MLIRContext *context, StringRef opName,
                     linalg::LinalgTransformationFilter m) {
   assert(!opName.empty() && "opName must not be empty");
-  patterList.insert<PatternType<OpType>>(opName, context, options, m);
+  patternList.insert<PatternType<OpType>>(opName, context, options, m);
+}
+
+template <typename PatternType, typename OpType, typename OptionsType>
+void enqueue(OwningRewritePatternList &patternList, OptionsType options,
+             MLIRContext *context, StringRef opName,
+             linalg::LinalgTransformationFilter m) {
+  if (!opName.empty())
+    patternList.insert<PatternType>(opName, context, options, m);
+  else
+    patternList.insert<PatternType>(m.addOpFilter<OpType>(), options);
 }
 
 /// Promotion transformation enqueues a particular stage-1 pattern for
@@ -112,13 +122,12 @@ struct Promote : public Transformation {
 /// Vectorization transformation enqueues a particular stage-1 pattern for
 /// `LinalgVectorizationPattern<LinalgOpType>` as well as copy to vector
 /// transfer rewrite forwarding patterns.
-template <typename LinalgOpType>
+template <typename LinalgOpType = LinalgOp>
 struct Vectorize : public Transformation {
   explicit Vectorize(
       linalg::LinalgVectorizationOptions options,
       linalg::LinalgTransformationFilter::FilterFunction f = nullptr)
-      : Transformation(f), opName(LinalgOpType::getOperationName()),
-        options(options) {}
+      : Transformation(f), opName(), options(options) {}
 
   Vectorize(StringRef name, linalg::LinalgVectorizationOptions options,
             linalg::LinalgTransformationFilter::FilterFunction f = nullptr)
@@ -128,7 +137,7 @@ struct Vectorize : public Transformation {
   buildRewritePatterns(MLIRContext *context,
                        linalg::LinalgTransformationFilter m) override {
     OwningRewritePatternList vectorizationPatterns;
-    sfinae_enqueue<linalg::LinalgVectorizationPattern, LinalgOpType>(
+    enqueue<linalg::LinalgVectorizationPattern, LinalgOpType>(
         vectorizationPatterns, options, context, opName, m);
     vectorizationPatterns.insert<linalg::LinalgCopyVTRForwardingPattern,
                                  linalg::LinalgCopyVTWForwardingPattern>(
@@ -235,16 +244,6 @@ struct CodegenStrategy {
             linalg::LinalgVectorizationOptions(), f));
     return *this;
   }
-  /// Append a pattern to rewrite `LinalgOpType` as a vector operation.
-  template <typename LinalgOpType>
-  CodegenStrategy &
-  vectorize(StringRef opName,
-            linalg::LinalgTransformationFilter::FilterFunction f = nullptr) {
-    transformationSequence.emplace_back(
-        std::make_unique<Vectorize<LinalgOpType>>(
-            opName, linalg::LinalgVectorizationOptions(), f));
-    return *this;
-  }
   /// Conditionally append a pattern to rewrite `LinalgOpType` as a vector
   /// operation.
   template <typename LinalgOpType>
@@ -254,13 +253,21 @@ struct CodegenStrategy {
     return b ? vectorize<LinalgOpType>(f) : *this;
     return *this;
   }
+  /// Append a pattern to rewrite `LinalgOpType` as a vector operation.
+  CodegenStrategy &
+  vectorize(StringRef opName,
+            linalg::LinalgTransformationFilter::FilterFunction f = nullptr) {
+    assert(!opName.empty() && "expected an op name");
+    transformationSequence.emplace_back(std::make_unique<Vectorize<LinalgOp>>(
+        opName, linalg::LinalgVectorizationOptions(), f));
+    return *this;
+  }
   /// Conditionally append a pattern to rewrite `LinalgOpType` as a vector
   /// operation.
-  template <typename LinalgOpType>
   CodegenStrategy &
   vectorizeIf(bool b, StringRef opName,
               linalg::LinalgTransformationFilter::FilterFunction f = nullptr) {
-    return b ? vectorize<LinalgOpType>(opName, f) : *this;
+    return b ? vectorize(opName, f) : *this;
     return *this;
   }
   /// Configure the post staged-patterns late vector transformations.