[mlir][Vector] WIP - Port vector-contract-to-outerproduct to use PDLL

[nicolasvasilache]

Summary

This is for gathering feedback atm, do not review as if it were intended to land.

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

LogicalResult matchAndRewrite(vector::ContractionOp op,
                                 PatternRewriter &rewriter) const override;
 
 private:
   /// Options to control the vector patterns.
   vector::VectorTransformsOptions vectorTransformOptions;
   FilterConstraintType filter;
 };
 
+FailureOr<Value> rewriteContractionToOuterProduct(PatternRewriter &rewriter,
+                                                  vector::ContractionOp op);
+
 /// Progressive lowering of a `vector.contract %a, %b, %c` with row-major matmul
 /// semantics to an output-size-unrolled sequence:
 /// ```
 ///    %out = arith.constant ... : vector<MxNxelt_type>
 ///    %bt = vector.transpose %b, [1, 0]
 ///    %aRow0 = vector.extract %a[0]
 ///    %btRow0 = vector.extract %bt[0]
 ///    %c00 = vector.reduce %atRow0, %bRow0

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

 
 private:
   vector::CombiningKind kind;
   Value lhs, rhs, res;
   VectorType lhsType;
 };
 } // namespace
 
+FailureOr<Value>
+mlir::vector::rewriteContractionToOuterProduct(PatternRewriter &rewriter,
+                                               vector::ContractionOp op) {
+  UnrolledOuterProductGenerator e(rewriter, op);
+  FailureOr<Value> matmatRes = e.matmat();
+  if (succeeded(matmatRes))
+    return matmatRes;
+  FailureOr<Value> matvecRes = e.matvec();
+  if (succeeded(matvecRes))
+    return matvecRes;
+  FailureOr<Value> tmatvecRes = e.tmatvec();
+  if (succeeded(tmatvecRes))
+    return tmatvecRes;
+  return failure();
+}
+
 /// Progressively lower a `vector.contract %a, %b, %c` with row-major matmul
 /// semantics to a reduction_size-unrolled sequence:
 /// ```
 ///    %at = vector.transpose %a, [1, 0]
 ///    %bRow0 = vector.extract %b[0]
 ///    %atRow0 = vector.extract %at[0]
 ///    %c0 = vector.outerproduct %atRow0, %bRow0, %c
 ///    ...
LogicalResult ContractionOpToOuterProductOpLowering::matchAndRewrite(
 
   if (vectorTransformOptions.vectorContractLowering !=
       vector::VectorContractLowering::OuterProduct)
     return failure();
 
   if (failed(filter(op)))
     return failure();
 
-  UnrolledOuterProductGenerator e(rewriter, op);
-  FailureOr<Value> matmatRes = e.matmat();
-  if (succeeded(matmatRes)) {
-    rewriter.replaceOp(op, *matmatRes);
-    return success();
-  }
-  FailureOr<Value> matvecRes = e.matvec();
-  if (succeeded(matvecRes)) {
-    rewriter.replaceOp(op, *matvecRes);
+  FailureOr<Value> res = rewriteContractionToOuterProduct(rewriter, op);
+  if (succeeded(res)) {
+    rewriter.replaceOp(op, *res);
     return success();
   }
-  FailureOr<Value> tmatvecRes = e.tmatvec();
-  if (succeeded(tmatvecRes)) {
-    rewriter.replaceOp(op, *tmatvecRes);
-    return success();
-  }
-
   return failure();
 }
 
 LogicalResult
 ContractionOpToDotLowering::matchAndRewrite(vector::ContractionOp op,
                                             PatternRewriter &rewriter) const {
   // TODO: implement masks
   if (llvm::size(op.masks()) != 0)

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

-// RUN: mlir-opt %s -test-vector-contraction-lowering | FileCheck %s
-// RUN: mlir-opt %s -test-vector-contraction-lowering=vector-lower-matrix-intrinsics=1 | FileCheck %s --check-prefix=MATRIX
-// RUN: mlir-opt %s -test-vector-contraction-lowering=vector-outerproduct=1 | FileCheck %s --check-prefix=OUTERPRODUCT
-// RUN: mlir-opt %s -test-vector-contraction-lowering=vector-filter-outerproduct=1 | FileCheck %s --check-prefix=FILTEROUTERPRODUCT
+// R-UN: mlir-opt %s -test-vector-contraction-lowering | FileCheck %s
+// R-UN: mlir-opt %s -test-vector-contraction-lowering=vector-lower-matrix-intrinsics=1 | FileCheck %s --check-prefix=MATRIX
+// R-UN: mlir-opt %s -test-vector-contraction-lowering=vector-outerproduct=1 | FileCheck %s --check-prefix=OUTERPRODUCT
+
+// RUN: mlir-pdll %p/vector-contract-transforms.pdll -x mlir -o %t | \
+// RUN: cat %s %t | mlir-opt -test-vector-contraction-to-outerproduct-pdll-lowering | FileCheck %s --check-prefix=OUTERPRODUCT
+
+// R-UN: mlir-opt %s -test-vector-contraction-lowering=vector-filter-outerproduct=1 | FileCheck %s --check-prefix=FILTEROUTERPRODUCT
 
 #dotp_accesses = [
   affine_map<(i) -> (i)>,
   affine_map<(i) -> (i)>,
   affine_map<(i) -> ()>
 ]
 #dotp_trait = {
   indexing_maps = #dotp_accesses,
   iterator_types = ["reduction"]
 }
 
+#matvec_accesses = [
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (j)>,
+  affine_map<(i, j) -> (i)>
+]
+#matvec_trait = {
+  indexing_maps = #matvec_accesses,
+  iterator_types = ["parallel", "reduction"]
+}
+
+#vecmat_accesses = [
+  affine_map<(i, j) -> (j)>,
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (i)>
+]
+#vecmat_trait = {
+  indexing_maps = #vecmat_accesses,
+  iterator_types = ["parallel", "reduction"]
+}
+
+#matmat_accesses_0 = [
+  affine_map<(m, n, k) -> (m, k)>,
+  affine_map<(m, n, k) -> (k, n)>,
+  affine_map<(m, n, k) -> (m, n)>
+]
+#matmat_trait_0 = {
+  indexing_maps = #matmat_accesses_0,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+#matmat_accesses_1 = [
+  affine_map<(m, n, k) -> (m, k)>,
+  affine_map<(m, n, k) -> (n, k)>,
+  affine_map<(m, n, k) -> (m, n)>
+]
+#matmat_trait_1 = {
+  indexing_maps = #matmat_accesses_1,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+#matmat_accesses_2 = [
+  affine_map<(m, n, k) -> (k, m)>,
+  affine_map<(m, n, k) -> (k, n)>,
+  affine_map<(m, n, k) -> (m, n)>
+]
+#matmat_trait_2 = {
+  indexing_maps = #matmat_accesses_2,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+#matmat_accesses_3 = [
+  affine_map<(m, n, k) -> (k, m)>,
+  affine_map<(m, n, k) -> (n, k)>,
+  affine_map<(m, n, k) -> (m, n)>
+]
+#matmat_trait_3 = {
+  indexing_maps = #matmat_accesses_3,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+#matmat_accesses_4 = [
+  affine_map<(m, n, k) -> (m, k)>,
+  affine_map<(m, n, k) -> (k, n)>,
+  affine_map<(m, n, k) -> (n, m)>
+]
+#matmat_trait_4 = {
+  indexing_maps = #matmat_accesses_4,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+#matmat_accesses_5 = [
+  affine_map<(m, n, k) -> (m, k)>,
+  affine_map<(m, n, k) -> (k, n)>,
+  affine_map<(m, n, k) -> (n, m)>
+]
+#matmat_trait_5 = {
+  indexing_maps = #matmat_accesses_5,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+#matmat_accesses_6 = [
+  affine_map<(m, n, k) -> (m, k)>,
+  affine_map<(m, n, k) -> (k, n)>,
+  affine_map<(m, n, k) -> (n, m)>
+]
+#matmat_trait_6 = {
+  indexing_maps = #matmat_accesses_6,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+#matmat_accesses_7 = [
+  affine_map<(m, n, k) -> (m, k)>,
+  affine_map<(m, n, k) -> (k, n)>,
+  affine_map<(m, n, k) -> (n, m)>
+]
+#matmat_trait_7 = {
+  indexing_maps = #matmat_accesses_7,
+  iterator_types = ["parallel", "parallel", "reduction"]
+}
+
+#contraction2d_accesses = [
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> ()>
+]
+#contraction2d_trait = {
+  indexing_maps = #contraction2d_accesses,
+  iterator_types = ["reduction", "reduction"]
+}
+
+#contraction2d_trans_accesses = [
+  affine_map<(i, j) -> (i, j)>,
+  affine_map<(i, j) -> (j, i)>,
+  affine_map<(i, j) -> ()>
+]
+#contraction2d_trans_trait = {
+  indexing_maps = #contraction2d_trans_accesses,
+  iterator_types = ["reduction", "reduction"]
+}
+
+module @ir {
+
 // CHECK-LABEL: func @extract_contract1
 // CHECK-SAME: %[[A:.*0]]: vector<4xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<4xf32>,
 // CHECK-SAME: %[[C:.*2]]: f32
 // CHECK:      %[[F:.*]] = arith.mulf %[[A]], %[[B]] : vector<4xf32>
 // CHECK:      %[[R:.*]] = vector.reduction <add>, %[[F]] : vector<4xf32> into f32
 // CHECK:      %[[ACC:.*]] = arith.addf %[[R]], %[[C]] : f32
 // CHECK:      return %[[ACC]] : f32

 // CHECK:      return %[[ACC]] : i32
 
 func @extract_contract1_int(%arg0: vector<4xi32>, %arg1: vector<4xi32>, %arg2: i32) -> i32 {
   %0 = vector.contract #dotp_trait %arg0, %arg1, %arg2
     : vector<4xi32>, vector<4xi32> into i32
   return %0 : i32
 }
 
-#matvec_accesses = [
-  affine_map<(i, j) -> (i, j)>,
-  affine_map<(i, j) -> (j)>,
-  affine_map<(i, j) -> (i)>
-]
-#matvec_trait = {
-  indexing_maps = #matvec_accesses,
-  iterator_types = ["parallel", "reduction"]
-}
-
 // CHECK-LABEL: func @extract_contract2
 // CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<3xf32>,
 // CHECK-SAME: %[[C:.*2]]: vector<2xf32>
 // CHECK:      %[[R:.*]] = arith.constant dense<0.000000e+00> : vector<2xf32>
 // CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : vector<2x3xf32>
 // CHECK:      %[[T2:.*]] = arith.mulf %[[T0]], %[[B]] : vector<3xf32>
 // CHECK:      %[[T3:.*]] = vector.reduction <add>, %[[T2]] : vector<3xf32> into f32

 func @extract_contract2_int(%arg0: vector<2x3xi32>,
                         %arg1: vector<3xi32>,
       %arg2: vector<2xi32>) -> vector<2xi32> {
   %0 = vector.contract #matvec_trait %arg0, %arg1, %arg2
     : vector<2x3xi32>, vector<3xi32> into vector<2xi32>
   return %0 : vector<2xi32>
 }
 
-#vecmat_accesses = [
-  affine_map<(i, j) -> (j)>,
-  affine_map<(i, j) -> (i, j)>,
-  affine_map<(i, j) -> (i)>
-]
-#vecmat_trait = {
-  indexing_maps = #vecmat_accesses,
-  iterator_types = ["parallel", "reduction"]
-}
-
 // CHECK-LABEL: func @extract_contract3
 // CHECK-SAME: %[[A:.*0]]: vector<3xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<2x3xf32>,
 // CHECK-SAME: %[[C:.*2]]: vector<2xf32>
 // CHECK:      %[[R:.*]] = arith.constant dense<0.000000e+00> : vector<2xf32>
 // CHECK:      %[[T0:.*]] = vector.extract %[[B]][0] : vector<2x3xf32>
 // CHECK:      %[[T2:.*]] = arith.mulf %[[T0]], %[[A]] : vector<3xf32>
 // CHECK:      %[[T3:.*]] = vector.reduction <add>, %[[T2]] : vector<3xf32> into f32
 // CHECK:      %[[T4:.*]] = vector.insert %[[T3]], %[[R]] [0] : f32 into vector<2xf32>
 // CHECK:      %[[T5:.*]] = vector.extract %[[B]][1] : vector<2x3xf32>
 // CHECK:      %[[T7:.*]] = arith.mulf %[[T5]], %[[A]] : vector<3xf32>
 // CHECK:      %[[T8:.*]] = vector.reduction <add>, %[[T7]] : vector<3xf32> into f32
 // CHECK:      %[[T9:.*]] = vector.insert %[[T8]], %[[T4]] [1] : f32 into vector<2xf32>
 // CHECK:      %[[T10:.*]] = arith.addf %[[T9]], %[[C]] : vector<2xf32>
 // CHECK:      return %[[T10]] : vector<2xf32>
 
 func @extract_contract3(%arg0: vector<3xf32>,
                         %arg1: vector<2x3xf32>,
                         %arg2: vector<2xf32>) -> vector<2xf32> {
   %0 = vector.contract #vecmat_trait %arg0, %arg1, %arg2
     : vector<3xf32>, vector<2x3xf32> into vector<2xf32>
   return %0 : vector<2xf32>
 }
 
-#matmat_accesses = [
-  affine_map<(i, j, k) -> (i, k)>,
-  affine_map<(i, j, k) -> (k, j)>,
-  affine_map<(i, j, k) -> (i, j)>
-]
-#matmat_trait = {
-  indexing_maps = #matmat_accesses,
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-
 // CHECK-LABEL: func @extract_contract4
 // CHECK-SAME: %[[A:.*0]]: vector<2x2xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<2x2xf32>,
 // CHECK-SAME: %[[C:.*2]]: vector<2x2xf32>
 // CHECK:    %[[R:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf32>
 // CHECK:    %[[Bt:.*]] = vector.transpose %arg1, [1, 0] : vector<2x2xf32> to vector<2x2xf32>
 // CHECK:    %[[T0:.*]] = vector.extract %[[A]][0] : vector<2x2xf32>
 // CHECK:    %[[T2:.*]] = vector.extract %[[Bt]][0] : vector<2x2xf32>

 // CHECK:    %[[T43:.*]] = vector.insert %[[T42]], %[[T34]] [1, 1] : f32 into vector<2x2xf32>
 //
 // CHECK:    %[[T52:.*]] = arith.addf %[[T43]], %[[C]] : vector<2x2xf32>
 // CHECK:    return %[[T52]] : vector<2x2xf32>
 
 func @extract_contract4(%arg0: vector<2x2xf32>,
                         %arg1: vector<2x2xf32>,
                         %arg2: vector<2x2xf32>) -> vector<2x2xf32> {
-  %0 = vector.contract #matmat_trait %arg0, %arg1, %arg2
+  %0 = vector.contract #matmat_trait_0 %arg0, %arg1, %arg2
     : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>
   return %0 : vector<2x2xf32>
 }
 
-#contraction2d_accesses = [
-  affine_map<(i, j) -> (i, j)>,
-  affine_map<(i, j) -> (i, j)>,
-  affine_map<(i, j) -> ()>
-]
-#contraction2d_trait = {
-  indexing_maps = #contraction2d_accesses,
-  iterator_types = ["reduction", "reduction"]
-}
-
 // CHECK-LABEL: func @full_contract1
 // CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<2x3xf32>,
 // CHECK-SAME: %[[C:.*2]]: f32
 // CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : vector<2x3xf32>
 // CHECK:      %[[T1:.*]] = vector.extract %[[B]][0] : vector<2x3xf32>
 // CHECK:      %[[T2:.*]] = arith.mulf %[[T0]], %[[T1]] : vector<3xf32>
 // CHECK:      %[[T3:.*]] = vector.reduction <add>, %[[T2]] : vector<3xf32> into f32
 // CHECK:      %[[T4:.*]] = arith.addf %[[T3]], %[[C]] : f32
 // CHECK:      %[[T5:.*]] = vector.extract %[[A]][1] : vector<2x3xf32>
 // CHECK:      %[[T6:.*]] = vector.extract %[[B]][1] : vector<2x3xf32>
 // CHECK:      %[[T7:.*]] = arith.mulf %[[T5]], %[[T6]] : vector<3xf32>
 // CHECK:      %[[T8:.*]] = vector.reduction <add>, %[[T7]] : vector<3xf32> into f32
 // CHECK:      %[[T9:.*]] = arith.addf %[[T8]], %[[T4]] : f32
 // CHECK:      return %[[T9]] : f32
 
 func @full_contract1(%arg0: vector<2x3xf32>,
                      %arg1: vector<2x3xf32>,
          %arg2: f32) -> f32 {
   %0 = vector.contract #contraction2d_trait %arg0, %arg1, %arg2
     : vector<2x3xf32>, vector<2x3xf32> into f32
   return %0 : f32
 }
 
-#contraction2d_trans_accesses = [
-  affine_map<(i, j) -> (i, j)>,
-  affine_map<(i, j) -> (j, i)>,
-  affine_map<(i, j) -> ()>
-]
-#contraction2d_trans_trait = {
-  indexing_maps = #contraction2d_trans_accesses,
-  iterator_types = ["reduction", "reduction"]
-}
-
 // CHECK-LABEL: func @full_contract2
 // CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<3x2xf32>,
 // CHECK-SAME: %[[C:.*2]]: f32
 // CHECK:      %[[Z:.*]] = arith.constant dense<0.000000e+00> : vector<3xf32>
 // CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : vector<2x3xf32>
 // CHECK:      %[[T1:.*]] = vector.extract %[[B]][0, 0] : vector<3x2xf32>
 // CHECK:      %[[T3:.*]] = vector.insert %[[T1]], %[[Z]] [0] : f32 into vector<3xf32>

 // REDUCE-NEXT: %[[ab12:.*]] = mul %[[a1]], %[[b02]] : vector<4xf32>
 // REDUCE-NEXT: %[[s12:.*]] = vector.reduction <add>, %[[ab12]] : vector<4xf32> into f32
 // REDUCE-NEXT: %[[r12:.*]] = vector.insert %[[s12]], %{{.*}} [1, 2] : f32 into vector<2x3xf32>
 //
 //      REDUCE: return %[[c3]] : vector<2x3xf32>
 func @matmul(%arg0: vector<2x4xf32>,
                           %arg1: vector<4x3xf32>,
                           %arg2: vector<2x3xf32>) -> vector<2x3xf32> {
-  %0 = vector.contract #matmat_trait %arg0, %arg1, %arg2
+  %0 = vector.contract #matmat_trait_0 %arg0, %arg1, %arg2
     : vector<2x4xf32>, vector<4x3xf32> into vector<2x3xf32>
   return %0 : vector<2x3xf32>
 }
 
 // CHECK-LABEL: func @broadcast_vec1d_from_scalar
 // CHECK-SAME: %[[A:.*0]]: f32
 // CHECK:      %[[T0:.*]] = vector.splat %[[A]] : vector<2xf32>
 // CHECK:      return %[[T0]] : vector<2xf32>

 // CHECK:      %[[T9:.*]] = vector.insert %[[T8]], %[[T6]] [1] : vector<1x7xi1> into vector<2x1x7xi1>
 // CHECK:      return %[[T9]] : vector<2x1x7xi1>
 
 func @genbool_var_3d(%arg0: index, %arg1: index, %arg2: index) -> vector<2x1x7xi1> {
   %0 = vector.create_mask %arg0, %arg1, %arg2 : vector<2x1x7xi1>
   return %0 : vector<2x1x7xi1>
 }
 
-#matmat_accesses_0 = [
-  affine_map<(m, n, k) -> (m, k)>,
-  affine_map<(m, n, k) -> (k, n)>,
-  affine_map<(m, n, k) -> (m, n)>
-]
-#matmat_trait_0 = {
-  indexing_maps = #matmat_accesses_0,
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-
 // OUTERPRODUCT-LABEL: func @matmul_0
 // OUTERPRODUCT-SAME: %[[A:[a-zA-Z0-9]*]]: vector<2x1xf32>,
 // OUTERPRODUCT-SAME: %[[B:[a-zA-Z0-9]*]]: vector<1x3xf32>,
 // OUTERPRODUCT-SAME: %[[C:[a-zA-Z0-9]*]]: vector<2x3xf32>
 //      OUTERPRODUCT: %[[At:.*]] = vector.transpose %[[A]], [1, 0]
 //      OUTERPRODUCT: %[[a0:.*]] = vector.extract %[[At]][0] : vector<1x2xf32>
 //      OUTERPRODUCT: %[[b0:.*]] = vector.extract %[[B]][0] : vector<1x3xf32>
 //      OUTERPRODUCT: %[[c0:.*]] = vector.outerproduct %[[a0]], %[[b0]], %[[C]]
 //      OUTERPRODUCT: return %[[c0]] : vector<2x3xf32>
 func @matmul_0(%arg0: vector<2x1xf32>, %arg1: vector<1x3xf32>, %arg2: vector<2x3xf32>)
 -> vector<2x3xf32>
 {
   %0 = vector.contract #matmat_trait_0 %arg0, %arg1, %arg2
     : vector<2x1xf32>, vector<1x3xf32> into vector<2x3xf32>
   return %0 : vector<2x3xf32>
 }
 
-#matmat_accesses_1 = [
-  affine_map<(m, n, k) -> (m, k)>,
-  affine_map<(m, n, k) -> (n, k)>,
-  affine_map<(m, n, k) -> (m, n)>
-]
-#matmat_trait_1 = {
-  indexing_maps = #matmat_accesses_1,
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-
 // OUTERPRODUCT-LABEL: func @matmul_1
 // OUTERPRODUCT-SAME: %[[A:[a-zA-Z0-9]*]]: vector<2x1xf32>,
 // OUTERPRODUCT-SAME: %[[B:[a-zA-Z0-9]*]]: vector<3x1xf32>,
 // OUTERPRODUCT-SAME: %[[C:[a-zA-Z0-9]*]]: vector<2x3xf32>
 //      OUTERPRODUCT: %[[At:.*]] = vector.transpose %[[A]], [1, 0]
 //      OUTERPRODUCT: %[[Bt:.*]] = vector.transpose %[[B]], [1, 0]
 //      OUTERPRODUCT: %[[a0:.*]] = vector.extract %[[At]][0] : vector<1x2xf32>
 //      OUTERPRODUCT: %[[b0:.*]] = vector.extract %[[Bt]][0] : vector<1x3xf32>
 //      OUTERPRODUCT: %[[c0:.*]] = vector.outerproduct %[[a0]], %[[b0]], %[[C]]
 //      OUTERPRODUCT: return %[[c0]] : vector<2x3xf32>
 func @matmul_1(%arg0: vector<2x1xf32>, %arg1: vector<3x1xf32>, %arg2: vector<2x3xf32>)
 -> vector<2x3xf32>
 {
   %0 = vector.contract #matmat_trait_1 %arg0, %arg1, %arg2
     : vector<2x1xf32>, vector<3x1xf32> into vector<2x3xf32>
   return %0 : vector<2x3xf32>
 }
 
-#matmat_accesses_2 = [
-  affine_map<(m, n, k) -> (k, m)>,
-  affine_map<(m, n, k) -> (k, n)>,
-  affine_map<(m, n, k) -> (m, n)>
-]
-#matmat_trait_2 = {
-  indexing_maps = #matmat_accesses_2,
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-
 // OUTERPRODUCT-LABEL: func @matmul_2
 // OUTERPRODUCT-SAME: %[[A:[a-zA-Z0-9]*]]: vector<1x2xf32>,
 // OUTERPRODUCT-SAME: %[[B:[a-zA-Z0-9]*]]: vector<1x3xf32>,
 // OUTERPRODUCT-SAME: %[[C:[a-zA-Z0-9]*]]: vector<2x3xf32>
 //      OUTERPRODUCT: %[[a0:.*]] = vector.extract %[[A]][0] : vector<1x2xf32>
 //      OUTERPRODUCT: %[[b0:.*]] = vector.extract %[[B]][0] : vector<1x3xf32>
 //      OUTERPRODUCT: %[[c0:.*]] = vector.outerproduct %[[a0]], %[[b0]], %[[C]]
 //      OUTERPRODUCT: return %[[c0]] : vector<2x3xf32>
 func @matmul_2(%arg0: vector<1x2xf32>, %arg1: vector<1x3xf32>, %arg2: vector<2x3xf32>)
 -> vector<2x3xf32>
 {
   %0 = vector.contract #matmat_trait_2 %arg0, %arg1, %arg2
     : vector<1x2xf32>, vector<1x3xf32> into vector<2x3xf32>
   return %0 : vector<2x3xf32>
 }
 
-#matmat_accesses_3 = [
-  affine_map<(m, n, k) -> (k, m)>,
-  affine_map<(m, n, k) -> (n, k)>,
-  affine_map<(m, n, k) -> (m, n)>
-]
-#matmat_trait_3 = {
-  indexing_maps = #matmat_accesses_3,
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-
 // OUTERPRODUCT-LABEL: func @matmul_3
 // OUTERPRODUCT-SAME: %[[A:[a-zA-Z0-9]*]]: vector<1x2xf32>,
 // OUTERPRODUCT-SAME: %[[B:[a-zA-Z0-9]*]]: vector<3x1xf32>,
 // OUTERPRODUCT-SAME: %[[C:[a-zA-Z0-9]*]]: vector<2x3xf32>
 //      OUTERPRODUCT: %[[Bt:.*]] = vector.transpose %[[B]], [1, 0]
 //      OUTERPRODUCT: %[[a0:.*]] = vector.extract %[[A]][0] : vector<1x2xf32>
 //      OUTERPRODUCT: %[[b0:.*]] = vector.extract %[[Bt]][0] : vector<1x3xf32>
 //      OUTERPRODUCT: %[[c0:.*]] = vector.outerproduct %[[a0]], %[[b0]], %[[C]]
 //      OUTERPRODUCT: return %[[c0]] : vector<2x3xf32>
 func @matmul_3(%arg0: vector<1x2xf32>, %arg1: vector<3x1xf32>, %arg2: vector<2x3xf32>)
 -> vector<2x3xf32>
 {
   %0 = vector.contract #matmat_trait_3 %arg0, %arg1, %arg2
     : vector<1x2xf32>, vector<3x1xf32> into vector<2x3xf32>
   return %0 : vector<2x3xf32>
 }
 
-#matmat_accesses_4 = [
-  affine_map<(m, n, k) -> (m, k)>,
-  affine_map<(m, n, k) -> (k, n)>,
-  affine_map<(m, n, k) -> (n, m)>
-]
-#matmat_trait_4 = {
-  indexing_maps = #matmat_accesses_4,
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-
 // OUTERPRODUCT-LABEL: func @matmul_4
 // OUTERPRODUCT-SAME: %[[A:[a-zA-Z0-9]*]]: vector<2x1xf32>,
 // OUTERPRODUCT-SAME: %[[B:[a-zA-Z0-9]*]]: vector<1x3xf32>,
 // OUTERPRODUCT-SAME: %[[C:[a-zA-Z0-9]*]]: vector<3x2xf32>
 //      OUTERPRODUCT: %[[At:.*]] = vector.transpose %[[A]], [1, 0]
 //      OUTERPRODUCT: %[[b0:.*]] = vector.extract %[[B]][0] : vector<1x3xf32>
 //      OUTERPRODUCT: %[[a0:.*]] = vector.extract %[[At]][0] : vector<1x2xf32>
 //      OUTERPRODUCT: %[[c0:.*]] = vector.outerproduct %[[b0]], %[[a0]], %[[C]]
 //      OUTERPRODUCT: return %[[c0]] : vector<3x2xf32>
 func @matmul_4(%arg0: vector<2x1xf32>, %arg1: vector<1x3xf32>, %arg2: vector<3x2xf32>)
 -> vector<3x2xf32>
 {
   %0 = vector.contract #matmat_trait_4 %arg0, %arg1, %arg2
     : vector<2x1xf32>, vector<1x3xf32> into vector<3x2xf32>
   return %0 : vector<3x2xf32>
 }
 
-#matmat_accesses_5 = [
-  affine_map<(m, n, k) -> (m, k)>,
-  affine_map<(m, n, k) -> (k, n)>,
-  affine_map<(m, n, k) -> (n, m)>
-]
-#matmat_trait_5 = {
-  indexing_maps = #matmat_accesses_5,
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-
 // OUTERPRODUCT-LABEL: func @matmul_5
 // OUTERPRODUCT-SAME: %[[A:[a-zA-Z0-9]*]]: vector<2x1xf32>,
 // OUTERPRODUCT-SAME: %[[B:[a-zA-Z0-9]*]]: vector<1x3xf32>,
 // OUTERPRODUCT-SAME: %[[C:[a-zA-Z0-9]*]]: vector<3x2xf32>
 //      OUTERPRODUCT: %[[At:.*]] = vector.transpose %[[A]], [1, 0]
 //      OUTERPRODUCT-DAG: %[[a0:.*]] = vector.extract %[[At]][0] : vector<1x2xf32>
 //      OUTERPRODUCT-DAG: %[[b0:.*]] = vector.extract %[[B]][0] : vector<1x3xf32>
 //      OUTERPRODUCT: %[[c0:.*]] = vector.outerproduct %[[b0]], %[[a0]], %[[C]]
 //      OUTERPRODUCT: return %[[c0]] : vector<3x2xf32>
 func @matmul_5(%arg0: vector<2x1xf32>, %arg1: vector<1x3xf32>, %arg2: vector<3x2xf32>)
 -> vector<3x2xf32>
 {
   %0 = vector.contract #matmat_trait_5 %arg0, %arg1, %arg2
     : vector<2x1xf32>, vector<1x3xf32> into vector<3x2xf32>
   return %0 : vector<3x2xf32>
 }
 
-#matmat_accesses_6 = [
-  affine_map<(m, n, k) -> (m, k)>,
-  affine_map<(m, n, k) -> (k, n)>,
-  affine_map<(m, n, k) -> (n, m)>
-]
-#matmat_trait_6 = {
-  indexing_maps = #matmat_accesses_6,
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-
 // OUTERPRODUCT-LABEL: func @matmul_6
 // OUTERPRODUCT-SAME: %[[A:[a-zA-Z0-9]*]]: vector<2x1xf32>,
 // OUTERPRODUCT-SAME: %[[B:[a-zA-Z0-9]*]]: vector<1x3xf32>,
 // OUTERPRODUCT-SAME: %[[C:[a-zA-Z0-9]*]]: vector<3x2xf32>
 //      OUTERPRODUCT: %[[At:.*]] = vector.transpose %[[A]], [1, 0]
 //      OUTERPRODUCT-DAG: %[[a0:.*]] = vector.extract %[[At]][0] : vector<1x2xf32>
 //      OUTERPRODUCT-DAG: %[[b0:.*]] = vector.extract %[[B]][0] : vector<1x3xf32>
 //      OUTERPRODUCT: %[[c0:.*]] = vector.outerproduct %[[b0]], %[[a0]], %[[C]]
 //      OUTERPRODUCT: return %[[c0]] : vector<3x2xf32>
 func @matmul_6(%arg0: vector<2x1xf32>, %arg1: vector<1x3xf32>, %arg2: vector<3x2xf32>)
 -> vector<3x2xf32>
 {
   %0 = vector.contract #matmat_trait_6 %arg0, %arg1, %arg2
     : vector<2x1xf32>, vector<1x3xf32> into vector<3x2xf32>
   return %0 : vector<3x2xf32>
 }
 
-#matmat_accesses_7 = [
-  affine_map<(m, n, k) -> (m, k)>,
-  affine_map<(m, n, k) -> (k, n)>,
-  affine_map<(m, n, k) -> (n, m)>
-]
-#matmat_trait_7 = {
-  indexing_maps = #matmat_accesses_7,
-  iterator_types = ["parallel", "parallel", "reduction"]
-}
-
 // OUTERPRODUCT-LABEL: func @matmul_7
 // OUTERPRODUCT-SAME: %[[A:[a-zA-Z0-9]*]]: vector<2x1xf32>,
 // OUTERPRODUCT-SAME: %[[B:[a-zA-Z0-9]*]]: vector<1x3xf32>,
 // OUTERPRODUCT-SAME: %[[C:[a-zA-Z0-9]*]]: vector<3x2xf32>
 //      OUTERPRODUCT: %[[At:.*]] = vector.transpose %[[A]], [1, 0]
 //      OUTERPRODUCT-DAG: %[[a0:.*]] = vector.extract %[[At]][0] : vector<1x2xf32>
 //      OUTERPRODUCT-DAG: %[[b0:.*]] = vector.extract %[[B]][0] : vector<1x3xf32>
 //      OUTERPRODUCT: %[[c0:.*]] = vector.outerproduct %[[b0]], %[[a0]], %[[C]]

 //      FILTEROUTERPRODUCT: %[[c0:.*]] = vector.contract {{{.*}}} %[[A]], %[[B]], %[[C]]
 func @matmul_4_not_filtered(%arg0: vector<3x4xf32>, %arg1: vector<4x4xf32>, %arg2: vector<3x4xf32>)
 -> vector<3x4xf32>
 {
   %0 = vector.contract #matmat_trait_0 %arg0, %arg1, %arg2
     : vector<3x4xf32>, vector<4x4xf32> into vector<3x4xf32>
   return %0 : vector<3x4xf32>
 }
+
+} // end module @ir

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

+// RUN: mlir-pdll %p/vector-contract-transforms.pdll -x mlir | FileCheck %s
+
+// TODO: Once https://reviews.llvm.org/D119900 lands, get I64 etc by just:
+// include "mlir/IR/OpBase.td" 
+
+//---------------------------------------------------------------------------//
+// Drop once https://reviews.llvm.org/D119900 lands.
+//---------------------------------------------------------------------------//
+
+// TODO: all this below is disabled because we apparently need C++ plumbing to 
+// register the constraints functions in the pdl module.
+// Atm, trying to use native inline constraints crashes with:
+// ```
+//   ByteCode.cpp:760: void {anonymous}::Generator::generate(
+//     mlir::pdl_interp::ApplyConstraintOp, {anonymous}::ByteCodeWriter&): 
+//   Assertion `constraintToMemIndex.count(op.name()) && 
+//     "expected index for constraint function"' failed.
+// ```
+// 
+//Constraint I64(attr: Attr) [{
+//  return success(attr.isSignlessInteger(64));
+//}];
+//
+//Constraint ShapedType(type: Type) [{
+//  return success(type.isa<ShapedType>());
+//}];
+//
+//Constraint ShapedValue(value: Value) [{
+//  return success(value.getType().isa<ShapedType>());
+//}];
+//
+//Constraint hasRank(type: ShapedType, target_rank: I64) [{
+//  return success(type.cast<ShapedType>().getRank() == 
+//                 target_rank.cast<IntegerAttr>().getInt());
+//}];
+
+//---------------------------------------------------------------------------//
+// END: Drop once https://reviews.llvm.org/D119900 lands.
+//---------------------------------------------------------------------------//
+
+// Import native constraints.
+Constraint IntegerAttr(attr: Attr);
+Constraint ShapedType(type: Type);
+
+// This fails with assertion: 
+// `values.size() == 1 && "expected values to have a single entity"' failed.
+// Constraint hasRank(type: ShapedType, target_rank: IntegerAttr);
+
+Constraint hasRank1(type: ShapedType);
+Constraint hasRank2(type: ShapedType);
+
+// Import native rewrite.
+Rewrite contractionToOuterProduct(op: Op) -> Value;
+
+Pattern VectorContractToOuterProduct222 {
+  // There is no `lhs.type` or `type(lhs)`; is it because we bind explicitly?
+  // Therefore we need to explicitly create a new type.
+  // Something like let `lhs: ShapedValue` + `type(lhs)` would be nice too.
+
+  // Explicitly filter only vector.contract with types of rank all 2.
+  // This is because the `contractionToOuterProduct` is not allowed to fail but
+  // it only supports a `matmul` and `matvec` flavors of contractions atm.
+  let lhsType : hasRank2;
+  let lhs: Value<lhsType>;
+  // hasRank2(lhsType);
+
+  let rhsType : ShapedType;
+  let rhs: Value<rhsType>;
+  hasRank2(rhsType);
+
+  let resType : ShapedType;
+  hasRank2(resType);
+
+  // This fails with assertion: 
+  // `values.size() == 1 && "expected values to have a single entity"' failed.
+  // hasRank(lhsType, attr<"2 : i64">);
+
+  let contract = op<vector.contract>(lhs, rhs, acc: Value) -> (resType);
+  replace contract with contractionToOuterProduct(contract);
+}
+
+// CHECK: pdl.pattern @VectorContractToOuterProduct222 : benefit(0)
 No newline at end of file

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

 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::vector;
struct TestVectorScanLowering
   }
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     populateVectorScanLoweringPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 
+// The contract between PDLL and C++ is implicit here, this does not seem
+// viable to import C++ native rewrites because we would need a lot of
+// manual error messages to make this usable without surprising crashes.
+void contractionToOuterProduct(llvm::ArrayRef<mlir::PDLValue> inputs,
+                               mlir::ArrayAttr, mlir::PatternRewriter &rewriter,
+                               mlir::PDLResultList &results) {
+  ContractionOpToOuterProductOpLowering lowering(
+      vector::VectorTransformsOptions(), rewriter.getContext());
+  // For now, this operation is not faillible so we must assert.
+  assert(inputs.size() == 1 && "inputs precondition failed");
+  Operation *op = inputs.front().dyn_cast<Operation *>();
+  assert(op && "op precondition failed");
+  op->dump();
+  auto res = vector::rewriteContractionToOuterProduct(
+      rewriter, dyn_cast<vector::ContractionOp>(op));
+  assert(succeeded(res) && "contractionToOuterProduct must not fail to apply");
+  results.push_back(*res);
+}
+
+struct TestVectorContractionToOuterProductWithPDLL
+    : public PassWrapper<TestVectorContractionToOuterProductWithPDLL,
+                         OperationPass<ModuleOp>> {
+  StringRef getArgument() const final {
+    return "test-vector-contraction-to-outerproduct-pdll-lowering";
+  }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<pdl_interp::PDLInterpDialect>();
+  }
+
+  TestVectorContractionToOuterProductWithPDLL() = default;
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+
+    ModuleOp irModule = getOperation().lookupSymbol<ModuleOp>(
+        StringAttr::get(&getContext(), "ir"));
+    ModuleOp patternModule;
+    getOperation().walk([&](Operation *op) {
+      if (op->getDialect()->getNamespace() == "pdl") {
+        patternModule = op->getParentOfType<ModuleOp>();
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    });
+    if (!irModule || !patternModule)
+      return;
+
+    RewritePatternSet patternList(&getContext());
+    PDLPatternModule pdlPatternModule(patternModule);
+    // The contract between PDLL and C++ is implicit here, this does not seem
+    // viable to import C++ native constraints because we would need a lot of
+    // manual error messages to make this usable without surprising crashes.
+    pdlPatternModule.registerConstraintFunction(
+        "IntegerAttr",
+        [](const mlir::PDLValue &ins, mlir::ArrayAttr &,
+           mlir::PatternRewriter &) -> LogicalResult {
+          llvm::errs() << "IntegerAttr\n";
+          return success(ins.cast<Attribute>().isa<IntegerAttr>());
+        });
+    // The contract between PDLL and C++ is implicit here, this does not seem
+    // viable to import C++ native constraints because we would need a lot of
+    // manual error messages to make this usable without surprising crashes.
+    pdlPatternModule.registerConstraintFunction(
+        "ShapedType",
+        [](const mlir::PDLValue &in, mlir::ArrayAttr &attrs,
+           mlir::PatternRewriter &) -> LogicalResult {
+          // TODO: atm both in and attrs can be NULL when this is called so we
+          // need to guard explicitly agains null in.
+          return success(in && in.cast<Type>().isa<ShapedType>());
+        });
+    // The contract between PDLL and C++ is implicit here, this does not seem
+    // viable to import C++ native constraints because we would need a lot of
+    // manual error messages to make this usable without surprising crashes.
+    // This fails with assertion:
+    // `values.size() == 1 && "expected values to have a single entity"' failed.
+    // pdlPatternModule.registerConstraintFunction(
+    //     "hasRank",
+    //     [](const mlir::PDLValue &in, mlir::ArrayAttr &attrs,
+    //        mlir::PatternRewriter &) -> LogicalResult {
+    //       return success(in.cast<Type>().cast<ShapedType>().getRank() ==
+    //                      attrs[0].cast<IntegerAttr>().getInt());
+    //     });
+    // The contract between PDLL and C++ is implicit here, this does not seem
+    // viable to import C++ native constraints because we would need a lot of
+    // manual error messages to make this usable without surprising crashes.
+    pdlPatternModule.registerConstraintFunction(
+        "hasRank",
+        [](const mlir::PDLValue &in, mlir::ArrayAttr &attrs,
+           mlir::PatternRewriter &) -> LogicalResult {
+          return success(in.cast<Type>().cast<ShapedType>().getRank() == 1);
+        });
+    // The contract between PDLL and C++ is implicit here, this does not seem
+    // viable to import C++ native constraints because we would need a lot of
+    // manual error messages to make this usable without surprising crashes.
+    pdlPatternModule.registerConstraintFunction(
+        "hasRank2",
+        [](const mlir::PDLValue &in, mlir::ArrayAttr &attrs,
+           mlir::PatternRewriter &) -> LogicalResult {
+          return success(in.cast<Type>().cast<ShapedType>().getRank() == 2);
+        });
+    pdlPatternModule.registerRewriteFunction("contractionToOuterProduct",
+                                             contractionToOuterProduct);
+    patternList.add(std::move(pdlPatternModule));
+
+    // Invoke the pattern driver with the provided patterns.
+    (void)applyPatternsAndFoldGreedily(irModule.getBodyRegion(),
+                                       std::move(patternList));
+  }
+};
+
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestVectorLowerings() {
   PassRegistration<TestVectorToVectorLowering>();
 
   PassRegistration<TestVectorContractionLowering>();
void registerTestVectorLowerings() {
 
   PassRegistration<TestVectorReduceToContractPatternsPatterns>();
 
   PassRegistration<TestVectorTransferDropUnitDimsPatterns>();
 
   PassRegistration<TestFlattenVectorTransferPatterns>();
 
   PassRegistration<TestVectorScanLowering>();
+
+  PassRegistration<TestVectorContractionToOuterProductWithPDLL>();
 }
 } // namespace test
 } // namespace mlir