[mlir][Linalg] Adding support for linalg_matmul with tensors.

[nicolasvasilache]

Summary

This revision provides 2 versions of matmul with tensors to account for the differences in buffer vs value semantics:

1. C(i, j) = sum_{r_k} A(i, r_k) * B(r_k, j)
2. D(i, j) = C(i, j) + sum_{r_k} A(i, r_k) * B(r_k, j)

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[antiagainst]

Neat!

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

1032

Do we want to use ShapedType::kDynamicSize for these magic numbers?

1042

I'm not sure about the conventions here but I'd typically prefer to have small and focused tests so it's easier to maintain and upate.

[nicolasvasilache]

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

1032

good point, I'll update everything in the test file in a followup CL.

1042

Right, so the tradeoff is bigger test vs boilerplate copy-pasta.
I'd prefer bigger test, up to some (undetermined) reasonable limit, at this time the test goes from 3 to 5 ops.
This falls into my reasonable category but I'll split if you have a strong preference for the copy-pasta.

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

 /// 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);
 
+/// 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);
+
+/// 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);
+
 template <typename Container> Operation *linalg_matmul(Container values) {
   assert(values.size() == 3 && "Expected exactly 3 values");
   return linalg_matmul(values[0], values[1], values[2]);
 }
 
 /// Build a linalg.generic, under the current ScopedContext, at the current
 /// insert point, that computes:
 /// ```

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

mlir::edsc::LoopRangeBuilder::LoopRangeBuilder(ValueHandle *iv,
                                                SubViewOp::Range range) {
   auto forOp =
       OperationHandle::createOp<ForOp>(range.offset, range.size, range.stride);
   *iv = ValueHandle(forOp.getInductionVar());
   auto *body = forOp.getBody();
   enter(body, /*prev=*/1);
 }
 
-ValueHandle mlir::edsc::LoopRangeBuilder::
-operator()(std::function<void(void)> fun) {
+ValueHandle
+mlir::edsc::LoopRangeBuilder::operator()(std::function<void(void)> fun) {
   if (fun)
     fun();
   exit();
   return ValueHandle::null();
 }
 
 mlir::edsc::LoopNestRangeBuilder::LoopNestRangeBuilder(
     ArrayRef<ValueHandle *> ivs, ArrayRef<SubViewOp::Range> ranges) {
mlir::edsc::LoopNestRangeBuilder::LoopNestRangeBuilder(
   assert(loops.size() == ivs.size() && "Mismatch loops vs ivs size");
 }
 
 mlir::edsc::LoopNestRangeBuilder::LoopNestRangeBuilder(
     ArrayRef<ValueHandle *> ivs, ArrayRef<Value> ranges)
     : LoopNestRangeBuilder(
           ivs, SmallVector<ValueHandle, 4>(ranges.begin(), ranges.end())) {}
 
-ValueHandle LoopNestRangeBuilder::LoopNestRangeBuilder::
-operator()(std::function<void(void)> fun) {
+ValueHandle LoopNestRangeBuilder::LoopNestRangeBuilder::operator()(
+    std::function<void(void)> fun) {
   if (fun)
     fun();
   for (auto &lit : reverse(loops)) {
     lit({});
   }
   return ValueHandle::null();
 }
 
Operation *mlir::edsc::makeGenericLinalgOp(
   auto handles = makeValueHandles(blockTypes);
   BlockBuilder(&b, op->getRegion(0),
                makeHandlePointers(MutableArrayRef<ValueHandle>(handles)))(
       [&] { regionBuilder(b.getBlock()->getArguments()); });
   assert(op->getRegion(0).getBlocks().size() == 1);
   return op;
 }
 
+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::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_matmul(ValueHandle vA, ValueHandle vB,
   return makeGenericLinalgOp(
     {IterType::Parallel, IterType::Parallel, IterType::Reduction},
     {A({m, k}), B({k, n})},
     {C({m, n})},
     macRegionBuilder);
   // clang-format on
 }
 
+Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
+                                          RankedTensorType tC) {
+  // clang-format off
+  AffineExpr m, n, k;
+  bindDims(ScopedContext::getContext(), m, n, k);
+  StructuredIndexed A(vA), B(vB), C(tC);
+  return makeGenericLinalgOp(
+    {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+    {A({m, k}), B({k, n})},
+    {C({m, n})},
+    mulRegionBuilder);
+  // clang-format on
+}
+
+Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
+                                          ValueHandle vC, RankedTensorType tD) {
+  // clang-format off
+  AffineExpr m, n, k;
+  bindDims(ScopedContext::getContext(), m, n, k);
+  StructuredIndexed A(vA), B(vB), C(vC), D(tD);
+  return makeGenericLinalgOp(
+    {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+    {A({m, k}), B({k, n}), C({m, n})},
+    {D({m, n})},
+    macRegionBuilder);
+  // 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();
   // TODO(ntv) some template magic to make everything rank-polymorphic.
   assert((dilations.empty() || dilations.size() == 2) && "only 2-D conv atm");
   assert((strides.empty() || strides.size() == 2) && "only 2-D conv atm");

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

TEST_FUNC(linalg_pointwise_test) {
   linalg_pointwise_max(SA({i, j}), SB({i, j}), SC({i, j}));
   linalg_pointwise_tanh(SA({i, j}), SC({i, j}));
 
   f.print(llvm::outs());
   f.erase();
 }
 
 // clang-format off
-// CHECK-LABEL: func @linalg_pointwise_mixed_tensors
-//       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
-// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
-// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
-//       CHECK:       addf
-//       CHECK:     }: tensor<?x?xf32>, memref<?x?xf32> -> tensor<?x?xf32>
-//       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
-// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
-// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
-//       CHECK:       cmpf "ogt"
-//       CHECK:       select
-//       CHECK:   }: tensor<?x?xf32>, memref<?x?xf32> -> tensor<?x?xf32>
-//       CHECK:   linalg.generic {args_in = 1 : i64, args_out = 1 : i64,
-// CHECK-SAME:      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
-// CHECK-SAME:      iterator_types = ["parallel", "parallel"]}
-//       CHECK:     tanh
-//       CHECK:   }: tensor<?x?xf32> -> tensor<?x?xf32>
-// clang-format on
-TEST_FUNC(linalg_pointwise_mixed_tensors_test) {
-  using namespace edsc;
-  using namespace edsc::ops;
-
-  auto f32Type = FloatType::getF32(&globalContext());
-  auto memrefType = MemRefType::get({-1, -1}, f32Type, {}, 0);
-  auto tensorType = RankedTensorType::get({-1, -1}, f32Type);
-  auto f = makeFunction("linalg_pointwise_mixed_tensors", {},
-                        {tensorType, memrefType});
-
-  OpBuilder builder(f.getBody());
-  ScopedContext scope(builder, f.getLoc());
-  ValueHandle A(f.getArgument(0)), B(f.getArgument(1));
-  AffineExpr i, j;
-  bindDims(&globalContext(), i, j);
-  StructuredIndexed SA(A), SB(B), SC(tensorType);
-  linalg_pointwise_add(SA({i, j}), SB({i, j}), SC({i, j}));
-  linalg_pointwise_max(SA({i, j}), SB({i, j}), SC({i, j}));
-  linalg_pointwise_tanh(SA({i, j}), SC({i, j}));
-
-  f.print(llvm::outs());
-  f.erase();
-}
-
-// clang-format off
 // CHECK-LABEL: func @linalg_matmul
 //       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
 // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}
 ///      CHECK:   ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
 //       CHECK:     %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
 //       CHECK:     %[[a4:.*]] = addf %[[a2]], %[[a3]] : f32
 //       CHECK:     linalg.yield %[[a4]] : f32
TEST_FUNC(linalg_metadata_ops) {
   auto reshaped = linalg_reshape(v, ArrayRef<ArrayRef<AffineExpr>>{{i, j}, k});
   linalg_reshape(memrefType, reshaped,
                  ArrayRef<ArrayRef<AffineExpr>>{{i, j}, k});
 
   f.print(llvm::outs());
   f.erase();
 }
 
+// clang-format off
+// CHECK-LABEL: func @linalg_tensors
+//       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
+// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
+//       CHECK:       addf
+//       CHECK:     }: tensor<?x?xf32>, memref<?x?xf32> -> tensor<?x?xf32>
+//       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
+// CHECK-SAME: indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+// CHECK-SAME: iterator_types = ["parallel", "parallel"]}
+//       CHECK:       cmpf "ogt"
+//       CHECK:       select
+//       CHECK:   }: tensor<?x?xf32>, memref<?x?xf32> -> tensor<?x?xf32>
+//       CHECK:   linalg.generic {args_in = 1 : i64, args_out = 1 : i64,
+// CHECK-SAME:      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+// CHECK-SAME:      iterator_types = ["parallel", "parallel"]}
+//       CHECK:     tanh
+//       CHECK:   }: tensor<?x?xf32> -> tensor<?x?xf32>
+//       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
+//  CHECK-SAME:     indexing_maps = [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:     iterator_types = ["parallel", "parallel", "reduction"]}
+//       CHECK:     mulf
+//       CHECK:   }: tensor<?x?xf32>, memref<?x?xf32> -> tensor<?x?xf32>
+//       CHECK:   linalg.generic {args_in = 3 : i64, args_out = 1 : i64,
+//  CHECK-SAME:     indexing_maps = [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:                      affine_map<(d0, d1, d2) -> (d0, d1)>],
+//  CHECK-SAME:     iterator_types = ["parallel", "parallel", "reduction"]
+//       CHECK:     mulf
+//       CHECK:     addf
+//       CHECK:   }: tensor<?x?xf32>, memref<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
+// clang-format on
+TEST_FUNC(linalg_tensors_test) {
+  using namespace edsc;
+  using namespace edsc::ops;
+
+  auto f32Type = FloatType::getF32(&globalContext());
+  auto memrefType = MemRefType::get({-1, -1}, f32Type, {}, 0);
+  auto tensorType = RankedTensorType::get({-1, -1}, f32Type);
+  auto f = makeFunction("linalg_tensors", {}, {tensorType, memrefType});
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+  ValueHandle A(f.getArgument(0)), B(f.getArgument(1));
+  AffineExpr i, j;
+  bindDims(&globalContext(), i, j);
+  StructuredIndexed SA(A), SB(B), SC(tensorType);
+  linalg_pointwise_add(SA({i, j}), SB({i, j}), SC({i, j}));
+  linalg_pointwise_max(SA({i, j}), SB({i, j}), SC({i, j}));
+  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();
+}
+
 int main() {
   RUN_TESTS();
   return 0;
 }