[mlir][Linalg] Add tensor support to Linalg EDSC Builders

[nicolasvasilache]

Summary

This diff extends the Linalg EDSC builders so we can easily create mixed
tensor/buffer linalg.generic ops. This is expected to be useful for
HLO -> Linalg lowering.

The StructuredIndexed struct is made to derive from ValueHandle and can
now capture a type + indexing expressions. This is used to represent return
tensors.

Pointwise unary and binary builders are extended to allow both output buffers
and return tensors. This has implications on the number of region arguments.

[nicolasvasilache]

https://reviews.llvm.org/D72863 landed in an improper state and was reverted, I must have mixed something when updating the diff and later patching again before pushing...
This is the proper update.

[merge_guards_bot]

Unit tests: pass. 62017 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[asaadaldien]

LGTM! clang-format + few nits...

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

161–165

How about directly inserting into values ? e.g

std::copy_if(s.begin(), s.end(), std::back_inserter(values), [](StructuredIndexed s) {return s.hasValue();});

165

same here

[nicolasvasilache]

Address review comment and clang-format.

[nicolasvasilache]

Typo.

[merge_guards_bot]

Unit tests: pass. 62017 tests passed, 0 failed and 783 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 62114 tests passed, 0 failed and 808 were skipped.

clang-tidy: fail. clang-tidy found 0 errors and 6 warnings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

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

inline StringRef toString(IterType t) {
   case IterType::Parallel:
     return getParallelIteratorTypeName();
   case IterType::Reduction:
     return getReductionIteratorTypeName();
   }
   llvm_unreachable("Unsupported IterType");
 }
 
-/// A StructuredIndexed represents a captured value that can be indexed and
-/// passed to the `makeGenericLinalgOp`. It allows writing intuitive index
-/// expressions such as:
+/// A StructuredIndexed represents an indexable quantity that is either:
+/// 1. a captured value, which is suitable for buffer and tensor operands, or;
+/// 2. a captured type, which is suitable for tensor return values.
+///
+/// A StructuredIndexed itself is indexed and passed to `makeGenericLinalgOp`.
+/// It enable an idiomatic syntax for index expressions such as:
 ///
 /// ```
-///      StructuredIndexed A(vA), B(vB), C(vC);
+///      StructuredIndexed A(buffer_or_tensor_value), B(buffer_or_tensor_value),
+///        C(buffer_value_or_tensor_type);
 ///      makeGenericLinalgOp({A({m, n}), B({k, n})}, {C({m, n})}, ... );
 /// ```
-struct StructuredIndexed {
-  StructuredIndexed(Value v) : value(v) {}
+struct StructuredIndexed : public ValueHandle {
+  StructuredIndexed(Type type) : ValueHandle(type) {}
+  StructuredIndexed(Value value) : ValueHandle(value) {}
+  StructuredIndexed(ValueHandle valueHandle) : ValueHandle(valueHandle) {}
   StructuredIndexed operator()(ArrayRef<AffineExpr> indexings) {
-    return StructuredIndexed(value, indexings);
+    return StructuredIndexed(*this, indexings);
   }
 
-  operator Value() const /* implicit */ { return value; }
   ArrayRef<AffineExpr> getExprs() { return exprs; }
 
 private:
+  StructuredIndexed(Type t, ArrayRef<AffineExpr> indexings)
+      : ValueHandle(t), exprs(indexings.begin(), indexings.end()) {
+    assert(t.isa<RankedTensorType>() && "RankedTensor expected");
+  }
   StructuredIndexed(Value v, ArrayRef<AffineExpr> indexings)
-      : value(v), exprs(indexings.begin(), indexings.end()) {
-    assert(v.getType().isa<MemRefType>() && "MemRefType expected");
+      : ValueHandle(v), exprs(indexings.begin(), indexings.end()) {
+    assert((v.getType().isa<MemRefType>() ||
+            v.getType().isa<RankedTensorType>()) &&
+           "MemRef or RankedTensor expected");
   }
-  StructuredIndexed(ValueHandle v, ArrayRef<AffineExpr> indexings)
-      : StructuredIndexed(v.getValue(), indexings) {}
+  StructuredIndexed(ValueHandle vh, ArrayRef<AffineExpr> indexings)
+      : ValueHandle(vh), exprs(indexings.begin(), indexings.end()) {}
 
-  Value value;
   SmallVector<AffineExpr, 4> exprs;
 };
 
 inline void defaultRegionBuilder(ArrayRef<BlockArgument> args) {}
 
+/// Build a `linalg.generic` op with the specified `inputs`, `outputs` and
+/// `region`.
+///
+/// `otherValues` and `otherAttributes` may be passed and will be appended as
+/// operands and attributes respectively.
+///
+/// Prerequisites:
+/// =============
+///
+/// 1. `inputs` may contain StructuredIndexed that capture either buffer or
+/// tensor values.
+/// 2. `outputs` may contain StructuredIndexed that capture either buffer values
+/// or tensor types. If both buffer values and tensor types are present, then
+/// all buffer values must appear before any tensor type. Without this
+/// restriction output tensor results would need to be reordered, which would
+/// result in surprising behavior when combined with region definition.
 Operation *makeGenericLinalgOp(
     ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,
     function_ref<void(ArrayRef<BlockArgument>)> regionBuilder =
         defaultRegionBuilder,
     ArrayRef<Value> otherValues = {}, ArrayRef<Attribute> otherAttributes = {});
 
 namespace ops {

 
 /// Build a linalg.pointwise with all `parallel` iterators and a region that
 /// computes `O = I1 + I2`. The client is responsible for specifying the proper
 /// indexings when creating the StructuredIndexed.
 Operation *linalg_pointwise_add(StructuredIndexed I1, StructuredIndexed I2,
                                 StructuredIndexed O);
 
 /// Build a linalg.pointwise with all `parallel` iterators and a region that
-/// computes `O = max(I!, I2)`. The client is responsible for specifying the
+/// 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.
 
 /// Build a linalg.generic, under the current ScopedContext, at the current
 /// insert point, that computes:

mlir/include/mlir/EDSC/Builders.h

void swap(ValueHandle &other) {
     if (this == &other)
       return;
     std::swap(t, other.t);
     std::swap(v, other.v);
   }
 
   /// Implicit conversion useful for automatic conversion to Container<Value>.
   operator Value() const { return getValue(); }
+  operator Type() const { return getType(); }
   operator bool() const { return hasValue(); }
 
   /// Generic mlir::Op create. This is the key to being extensible to the whole
   /// of MLIR without duplicating the type system or the op definitions.
   template <typename Op, typename... Args>
   static ValueHandle create(Args... args);
 
   /// Generic mlir::Op create. This is the key to being extensible to the whole

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

static void getMaxDimIndex(ArrayRef<StructuredIndexed> structuredIndices,
   }
 }
 
 Operation *mlir::edsc::makeGenericLinalgOp(
     ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,
     function_ref<void(ArrayRef<BlockArgument>)> regionBuilder,
     ArrayRef<Value> otherValues, ArrayRef<Attribute> otherAttributes) {
+  for (unsigned i = 0, e = outputs.size(); i + 1 < e; ++i)
+    assert(!(outputs[i].getType().isa<RankedTensorType>() &&
+             outputs[i + 1].getType().isa<MemRefType>()) &&
+           "output tensors must be passed after output buffers");
   auto &builder = edsc::ScopedContext::getBuilder();
   auto *ctx = builder.getContext();
   unsigned nInputs = inputs.size();
   unsigned nOutputs = outputs.size();
   unsigned maxPos = 0;
   getMaxDimIndex(inputs, maxPos);
   getMaxDimIndex(outputs, maxPos);
   // maxPos is 0 indexed, need to turn this into a count (i.e. +1)
   unsigned nDims = maxPos + 1;
 
   SmallVector<AffineMap, 4> maps;
   maps.reserve(nInputs + nOutputs);
   for (auto in : inputs)
     maps.push_back(
         AffineMap::get(/*dimCount=*/nDims, /*symbolCount=*/0, in.getExprs()));
   for (auto out : outputs)
     maps.push_back(
         AffineMap::get(/*dimCount=*/nDims, /*symbolCount=*/0, out.getExprs()));
 
   unsigned nViews = nInputs + nOutputs;
   SmallVector<Value, 4> values;
   values.reserve(nViews);
   values.append(inputs.begin(), inputs.end());
-  values.append(outputs.begin(), outputs.end());
+  std::copy_if(outputs.begin(), outputs.end(), std::back_inserter(values),
+               [](StructuredIndexed s) { return s.hasValue(); });
+  SmallVector<Type, 4> types;
+  std::copy_if(outputs.begin(), outputs.end(), std::back_inserter(types),
+               [](StructuredIndexed s) { return !s.hasValue(); });
 
   auto iteratorStrTypes = functional::map(toString, iteratorTypes);
   // clang-format off
   auto *op =
       edsc::ScopedContext::getBuilder()
           .create<linalg::GenericOp>(
               edsc::ScopedContext::getLocation(),
-              ArrayRef<Type>{}, // TODO(ntv): support tensors
+              types,
               values,
               IntegerAttr::get(IntegerType::get(64, ctx), nInputs),
               IntegerAttr::get(IntegerType::get(64, ctx), nOutputs),
               builder.getAffineMapArrayAttr(maps),
               builder.getStrArrayAttr(iteratorStrTypes),
               StringAttr() /*doc*/,
               FlatSymbolRefAttr() /*fun*/,
               StringAttr() /*library_call*/
void mlir::edsc::ops::macRegionBuilder(ArrayRef<BlockArgument> args) {
   linalg_yield((c + a * b).getValue());
 }
 
 Operation *mlir::edsc::ops::linalg_pointwise(UnaryPointwiseOpBuilder unaryOp,
                                              StructuredIndexed I,
                                              StructuredIndexed O) {
   SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
                                            edsc::IterType::Parallel);
+  if (O.getType().isa<RankedTensorType>()) {
+    auto fun = [&unaryOp](ArrayRef<BlockArgument> args) {
+      assert(args.size() == 1 && "expected 1 block arguments");
+      ValueHandle a(args[0]);
+      linalg_yield(unaryOp(a));
+    };
+    return makeGenericLinalgOp(iterTypes, {I}, {O}, fun);
+  }
   auto fun = [&unaryOp](ArrayRef<BlockArgument> args) {
     assert(args.size() == 2 && "expected 2 block arguments");
     ValueHandle a(args[0]);
     linalg_yield(unaryOp(a));
   };
   return makeGenericLinalgOp(iterTypes, {I}, {O}, fun);
 }
 
 Operation *mlir::edsc::ops::linalg_pointwise_tanh(StructuredIndexed I,
                                                   StructuredIndexed O) {
-  ;
   using edsc::intrinsics::tanh;
   UnaryPointwiseOpBuilder unOp([](ValueHandle a) -> Value { return tanh(a); });
   return linalg_pointwise(unOp, I, O);
 }
 
 /// Binary pointwise operation (with broadcast) entry point.
 Operation *mlir::edsc::ops::linalg_pointwise(BinaryPointwiseOpBuilder binaryOp,
                                              StructuredIndexed I1,
                                              StructuredIndexed I2,
                                              StructuredIndexed O) {
   SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
                                            edsc::IterType::Parallel);
+  if (O.getType().isa<RankedTensorType>()) {
+    auto fun = [&binaryOp](ArrayRef<BlockArgument> args) {
+      assert(args.size() == 2 && "expected 2 block arguments");
+      ValueHandle a(args[0]), b(args[1]);
+      linalg_yield(binaryOp(a, b));
+    };
+    return makeGenericLinalgOp(iterTypes, {I1, I2}, {O}, fun);
+  }
   auto fun = [&binaryOp](ArrayRef<BlockArgument> args) {
     assert(args.size() == 3 && "expected 3 block arguments");
     ValueHandle a(args[0]), b(args[1]);
     linalg_yield(binaryOp(a, b));
   };
   return makeGenericLinalgOp(iterTypes, {I1, I2}, {O}, fun);
 }
 

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