[mlir][EDSC][Linalg] Compose linalg_matmul and vector.contract

[nicolasvasilache]

Summary

This revision allows model builder to create a linalg_matmul whose body
is a vector.contract. This shows the abstractions compose nicely.

[nicolasvasilache]

Format.

[aartbik]

mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h

136

typo: scalar

[aartbik]

trusting that you will fix both typos, good to land

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

218

expected 2 block arguments!

[nicolasvasilache]

ddress review comments.

mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h

 #ifndef MLIR_DIALECT_LINALG_EDSC_BUILDERS_H_
 #define MLIR_DIALECT_LINALG_EDSC_BUILDERS_H_
 
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 // TODO(ntv): Needed for SubViewOp::Range, clean this up.
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/EDSC/Builders.h"
+#include "mlir/EDSC/Intrinsics.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/Builders.h"
 
 namespace mlir {
 class AffineForOp;
 class BlockArgument;
 class SubViewOp;
 
 namespace loop {
 class ParallelOp;

 namespace ops {
 using edsc::StructuredIndexed;
 using edsc::ValueHandle;
 
 //===----------------------------------------------------------------------===//
 // EDSC builders for linalg generic operations.
 //===----------------------------------------------------------------------===//
 
-/// Build the body of a region to compute a multiply-accumulate, under the
-/// current ScopedContext, at the current insert point.
+/// Build the body of a region to compute a scalar multiply, under the current
+/// ScopedContext, at the current insert point.
+void mulRegionBuilder(ArrayRef<BlockArgument> args);
+
+/// Build the body of a region to compute a scalar multiply-accumulate, under
+/// the current ScopedContext, at the current insert point.
 void macRegionBuilder(ArrayRef<BlockArgument> args);
 
 /// TODO(ntv): In the future we should tie these implementations to something in
 /// Tablegen that generates the proper interfaces and the proper sugared named
 /// ops.
 
 /// Build a linalg.pointwise, under the current ScopedContext, at the current
 /// insert point, that computes:

 /// Build a linalg.pointwise with all `parallel` iterators and a region that
 /// computes `O = max(I1, I2)`. The client is responsible for specifying the
 /// proper indexings when creating the StructuredIndexed.
 Operation *linalg_pointwise_max(StructuredIndexed I1, StructuredIndexed I2,
                                 StructuredIndexed O);
 
 // TODO(ntv): Implement more useful pointwise operations on a per-need basis.
 
+using MatmulRegionBuilder = function_ref<void(ArrayRef<BlockArgument> args)>;
+
 /// Build a linalg.generic, under the current ScopedContext, at the current
 /// insert point, that computes:
 /// ```
 ///    (m, n, k) = (par, par, seq)
 ///    |
 ///    |  C(m, n) += A(m, k) * B(k, n)
 /// ```
-Operation *linalg_matmul(ValueHandle vA, ValueHandle vB, ValueHandle vC);
+Operation *linalg_matmul(ValueHandle vA, ValueHandle vB, ValueHandle vC,
+                         MatmulRegionBuilder regionBuilder = macRegionBuilder);
 
 /// Build a linalg.generic, under the current ScopedContext, at the current
 /// insert point, that computes:
 /// ```
 ///    (m, n, k) = (par, par, seq)
 ///    |
 ///    |  C(m, n) = sum_k(A(m, k) * B(k, n))
 /// ```
 /// and returns the tensor `C`.
-Operation *linalg_matmul(ValueHandle vA, ValueHandle vB, RankedTensorType tC);
+Operation *linalg_matmul(ValueHandle vA, ValueHandle vB, RankedTensorType tC,
+                         MatmulRegionBuilder regionBuilder = mulRegionBuilder);
 
 /// Build a linalg.generic, under the current ScopedContext, at the current
 /// insert point, that computes:
 /// ```
 ///    (m, n, k) = (par, par, seq)
 ///    |
 ///    |  D(m, n) = C(m, n) + sum_k(A(m, k) * B(k, n))
 /// ```
 /// and returns the tensor `D`.
 Operation *linalg_matmul(ValueHandle vA, ValueHandle vB, ValueHandle vC,
-                         RankedTensorType tD);
+                         RankedTensorType tD,
+                         MatmulRegionBuilder regionBuilder = macRegionBuilder);
 
-template <typename Container> Operation *linalg_matmul(Container values) {
+template <typename Container>
+Operation *linalg_matmul(Container values,
+                         MatmulRegionBuilder regionBuilder = macRegionBuilder) {
   assert(values.size() == 3 && "Expected exactly 3 values");
-  return linalg_matmul(values[0], values[1], values[2]);
+  return linalg_matmul(values[0], values[1], values[2], regionBuilder);
 }
 
 /// Build a linalg.generic, under the current ScopedContext, at the current
 /// insert point, that computes:
 /// ```
 ///    (batch, f, [h, w, ...], [kh, kw, ...], c) =
 ///    |  (par, par, [par, par, ...], [red, red, ...], red)
 ///    |

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

 
 static void mulRegionBuilder(ArrayRef<BlockArgument> args) {
   using edsc::op::operator*;
   assert(args.size() == 2 && "expected 2 block arguments");
   ValueHandle a(args[0]), b(args[1]);
   linalg_yield((a * b).getValue());
 }
 
+void mlir::edsc::ops::mulRegionBuilder(ArrayRef<BlockArgument> args) {
+  using edsc::op::operator+;
+  using edsc::op::operator*;
+  assert(args.size() == 2 && "expected 2 block arguments");
+  ValueHandle a(args[0]), b(args[1]);
+  linalg_yield((a * b).getValue());
+}
+
 void mlir::edsc::ops::macRegionBuilder(ArrayRef<BlockArgument> args) {
   using edsc::op::operator+;
   using edsc::op::operator*;
   assert(args.size() == 3 && "expected 3 block arguments");
   ValueHandle a(args[0]), b(args[1]), c(args[2]);
   linalg_yield((c + a * b).getValue());
 }
 
Operation *mlir::edsc::ops::linalg_pointwise_max(StructuredIndexed I1,
   BinaryPointwiseOpBuilder binOp([](ValueHandle a, ValueHandle b) -> Value {
     using edsc::op::operator>;
     return std_select(a > b, a, b).getValue();
   });
   return linalg_pointwise(binOp, I1, I2, O);
 }
 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
-                                          ValueHandle vC) {
+                                          ValueHandle vC,
+                                          MatmulRegionBuilder regionBuilder) {
   // clang-format off
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(vC);
   return makeGenericLinalgOp(
     {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
     {A({m, k}), B({k, n})},
     {C({m, n})},
-    macRegionBuilder);
+    regionBuilder);
   // clang-format on
 }
 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
-                                          RankedTensorType tC) {
+                                          RankedTensorType tC,
+                                          MatmulRegionBuilder regionBuilder) {
   // clang-format off
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(tC);
   return makeGenericLinalgOp(
     {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
     {A({m, k}), B({k, n})},
     {C({m, n})},
-    mulRegionBuilder);
+    regionBuilder);
   // clang-format on
 }
 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
-                                          ValueHandle vC, RankedTensorType tD) {
+                                          ValueHandle vC, RankedTensorType tD,
+                                          MatmulRegionBuilder regionBuilder) {
   // clang-format off
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(vC), D(tD);
   return makeGenericLinalgOp(
     {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
     {A({m, k}), B({k, n}), C({m, n})},
     {D({m, n})},
-    macRegionBuilder);
+    regionBuilder);
   // clang-format on
 }
 
 Operation *mlir::edsc::ops::linalg_conv_nhwc(ValueHandle vI, ValueHandle vW,
                                              ValueHandle vO,
                                              ArrayRef<int> strides,
                                              ArrayRef<int> dilations) {
   MLIRContext *ctx = ScopedContext::getContext();

mlir/test/EDSC/builder-api-test.cpp

TEST_FUNC(linalg_tensors_test) {
   linalg_pointwise_tanh(SA({i, j}), SC({i, j}));
   Value o1 = linalg_matmul(A, B, tensorType)->getResult(0);
   linalg_matmul(A, B, ValueHandle(o1), tensorType);
 
   f.print(llvm::outs());
   f.erase();
 }
 
-// CHECK-LABEL: func @vector_matmul_test(
-//  CHECK-SAME:   %[[A:.*]]: vector<4x16xf32>,
-//  CHECK-SAME:   %[[B:.*]]: vector<16x8xf32>,
-//  CHECK-SAME:   %[[C:.*]]: vector<4x8xf32>)
-//  CHECK:   vector.contract {{.*}}[affine_map<(d0, d1, d2) -> (d0, d2)>,
-//  CHECK-SAME:                     affine_map<(d0, d1, d2) -> (d2, d1)>,
-//  CHECK-SAME:                     affine_map<(d0, d1, d2) -> (d0, d1)>],
-//  CHECK-SAME:              {{.*}}["parallel", "parallel", "reduction"]
-//  CHECK-SAME: %[[A]], %[[B]], %[[C]]
-//  CHECK-SAME:   vector<4x16xf32>, vector<16x8xf32> into vector<4x8xf32>
-TEST_FUNC(vector_matmul_test) {
+// CHECK-LABEL: func @memref_vector_matmul_test(
+//  CHECK-SAME:   %[[A:.*]]: memref<?x?xvector<4x16xf32>>,
+//  CHECK-SAME:   %[[B:.*]]: memref<?x?xvector<16x8xf32>>,
+//  CHECK-SAME:   %[[C:.*]]: memref<?x?xvector<4x8xf32>>)
+//       CHECK:   linalg.generic {{.*}} %[[A]], %[[B]], %[[C]]
+//       CHECK:     vector.contract{{.*}}[affine_map<(d0, d1, d2) -> (d0,
+//  d2)>,
+//  CHECK-SAME:                       affine_map<(d0, d1, d2) -> (d2, d1)>,
+//  CHECK-SAME:                       affine_map<(d0, d1, d2) -> (d0, d1)>],
+//  CHECK-SAME:                {{.*}}["parallel", "parallel", "reduction"]
+//  CHECK-SAME:     vector<4x16xf32>, vector<16x8xf32> into vector<4x8xf32>
+//       CHECK:   memref<?x?xvector<4x16xf32>>, memref<?x?xvector<16x8xf32>>,
+//  CHECK-SAME:   memref<?x?xvector<4x8xf32>>
+TEST_FUNC(memref_vector_matmul_test) {
   using namespace edsc;
   using namespace edsc::ops;
 
   int64_t M = 4, N = 8, K = 16;
   auto f32Type = FloatType::getF32(&globalContext());
   auto mkVectorType = VectorType::get({M, K}, f32Type);
   auto knVectorType = VectorType::get({K, N}, f32Type);
   auto mnVectorType = VectorType::get({M, N}, f32Type);
-  auto f = makeFunction("vector_matmul_test", {},
-                        {mkVectorType, knVectorType, mnVectorType});
+  auto typeA =
+      MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
+                      mkVectorType, {}, 0);
+  auto typeB =
+      MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
+                      knVectorType, {}, 0);
+  auto typeC =
+      MemRefType::get({ShapedType::kDynamicSize, ShapedType::kDynamicSize},
+                      mnVectorType, {}, 0);
+  auto f = makeFunction("memref_vector_matmul_test", {}, {typeA, typeB, typeC});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   ValueHandle A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
-  vector_matmul(A, B, C);
+  auto contractionBuilder = [](ArrayRef<BlockArgument> args) {
+    assert(args.size() == 3 && "expected 3 block arguments");
+    (linalg_yield(vector_matmul(args[0], args[1], args[2])));
+  };
+  linalg_matmul(A, B, C, contractionBuilder);
+
   f.print(llvm::outs());
   f.erase();
 }
 
 int main() {
   RUN_TESTS();
   return 0;
 }