[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]

Format.

[merge_guards_bot]

Unit tests: pass. 61930 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

[merge_guards_bot]

Unit tests: pass. 61930 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

[hanchung]

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

363

s/depth_multiplier/depthMultiplier

Same to others below.

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

51

Please fix it.

82

Please fix it

134–137

Do we also check if all types in resultTensorTypes are RankedTensorType?

144

Do we need to update maxPos with resultTensorTypes?

229

If this is a unary operation, why do we pass two input to the function?
I was expecting something like makeGenericLinalgOp(iterTypes, {I}, {}, resultTensorTypes, fun);
In addition, is there a test for this path?

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

910

Remove one "/" in the beginning?

959

nit: add comment for function argument, /*resultTensorTypes=*/{}

[ftynse]

I'm not sure I understand what happens to the Op operands that are supposed to be tensors in individual linalg_ constructors. The comment says outputs should not contain tensors, but in individual constructors there's no generic outputs, and some StructuredIndexed *must* be passed in.

Nit: typo in the commit description "identity identity"

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

186

There are no arguments called inputs and outputs here and below. I can suggest the formulation "StructuredIndexed for inputs may wrap a buffer or a tensor, those for outputs may wrap only buffers". Same below.

I'd drop the "TODO", it's the same as above.

195–197

It's not clear what happens to O, which I presume is the output, if you want tensor outputs.

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

235

Let's clean this semicolon :)

292

Please add a string with description to the assert

293

You can do something like

ArrayRef<StructuredIndexed> all(allIndexed)
bool isResultBuffer = C.getType().isa<MemRefType>();
auto inputs = all.drop_back(isResultBuffer);
auto outputs = all.take_back(isResultBuffer)

or even

ArrayRef<StructuredIndexed> all({A({m, k}), B({k, n}), C({m, n})});

340

Same as above

[nicolasvasilache]

Refocus on extending EDSC to tensors for pointwise ops.

[nicolasvasilache]

Thanks for your comments, I was a bit too ambitious with the generalization and missed some things.
I refocused on extending StructuredIndexed and limit this first diff to pointwise ops only.

[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

[nicolasvasilache]

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

186

Done by virtue of refocusing the diff.

195–197

Done by virtue of refocusing the diff.

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

134–137

That is done by the verifier of linalg.generic.

144

Good point, I refocused the diff though so I don't need it anymore.

229

There was not and this was a mistake.
Now there is, thanks for surfacing!

292

obsolete, I'll revisit in a later diff, thanks!

293

obsolete, I'll revisit in a later diff, thanks!

340

obsolete, I'll revisit in a later diff, thanks!

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

910

obsolete, I'll revisit in a later diff, thanks!

959

obsolete, I'll revisit in a later diff, thanks!

[nicolasvasilache]

Reformat.

[merge_guards_bot]

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

clang-tidy: unknown.

clang-format: pass.

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

[asaadaldien]

LGTM!

[nicolasvasilache]

This was reverted due to an improper update.
See https://reviews.llvm.org/D73149 for the proper version.

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

 struct StructuredIndexed {
   StructuredIndexed(Value v) : value(v) {}
   StructuredIndexed operator()(ArrayRef<AffineExpr> indexings) {
     return StructuredIndexed(value, indexings);
   }
 
   operator Value() const /* implicit */ { return value; }
   ArrayRef<AffineExpr> getExprs() { return exprs; }
+  Type getType() { return value.getType(); }
 
 private:
   StructuredIndexed(Value v, ArrayRef<AffineExpr> indexings)
       : value(v), exprs(indexings.begin(), indexings.end()) {
-    assert(v.getType().isa<MemRefType>() && "MemRefType expected");
+    assert((v.getType().isa<MemRefType>() ||
+            v.getType().isa<RankedTensorType>()) &&
+           "MemRef or RankedTensor expected");
   }
   StructuredIndexed(ValueHandle v, ArrayRef<AffineExpr> indexings)
       : StructuredIndexed(v.getValue(), indexings) {}
 
   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.
+///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
+/// which case, the canonical identity indexing_map is assumed.
+//
+// TODO(ntv) In the future we may want to relax this identity assumption (e.g.
+// for automatic differentiation purposes). In that case we will want to make
+// StructuredIndexed work with ValueHandle to encode type or value.
 Operation *makeGenericLinalgOp(
     ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
-    ArrayRef<StructuredIndexed> outputs,
+    ArrayRef<StructuredIndexed> outputs, ArrayRef<Type> resultTensorTypes = {},
     function_ref<void(ArrayRef<BlockArgument>)> regionBuilder =
         defaultRegionBuilder,
     ArrayRef<Value> otherValues = {}, ArrayRef<Attribute> otherAttributes = {});
 
 namespace ops {
 using edsc::StructuredIndexed;
 using edsc::ValueHandle;
 using edsc::intrinsics::linalg_yield;

 /// build a perfect loop nest of parallel loops with arbitrarily complex
 /// innermost loop code and whatever (explicit) broadcast semantics.
 ///
 /// This can be used with both out-of-place and in-place semantics.
 /// The client is responsible for ensuring the region operations are compatible
 /// with in-place semantics and parallelism.
 
 /// Unary pointwise operation (with broadcast) entry point.
+///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
+/// which case, the canonical identity indexing_map is assumed.
+//
+// TODO(ntv) In the future we may want to relax this identity assumption (e.g.
+// for automatic differentiation purposes). In that case we will want to make
+// StructuredIndexed work with ValueHandle to encode type or value.
 using UnaryPointwiseOpBuilder = function_ref<Value(ValueHandle)>;
 Operation *linalg_pointwise(UnaryPointwiseOpBuilder unaryOp,
-                            StructuredIndexed I, StructuredIndexed O);
+                            StructuredIndexed I, StructuredIndexed O,
+                            ArrayRef<Type> resultTensorTypes = {});
 
 /// Build a linalg.pointwise with all `parallel` iterators and a region that
 /// computes `O = tanh(I)`. The client is responsible for specifying the proper
 /// indexings when creating the StructuredIndexed.
-Operation *linalg_pointwise_tanh(StructuredIndexed I, StructuredIndexed O);
+///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
+/// which case, the canonical identity indexing_map is assumed.
+//
+// TODO(ntv) In the future we may want to relax this identity assumption (e.g.
+// for automatic differentiation purposes). In that case we will want to make
+// StructuredIndexed work with ValueHandle to encode type or value.
+Operation *linalg_pointwise_tanh(StructuredIndexed I, StructuredIndexed O,
+                                 ArrayRef<Type> resultTensorTypes = {});
 
 /// Binary pointwise operation (with broadcast) entry point.
+///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
+/// which case, the canonical identity indexing_map is assumed.
+//
+// TODO(ntv) In the future we may want to relax this identity assumption (e.g.
+// for automatic differentiation purposes). In that case we will want to make
+// StructuredIndexed work with ValueHandle to encode type or value.
 using BinaryPointwiseOpBuilder = function_ref<Value(ValueHandle, ValueHandle)>;
 Operation *linalg_pointwise(BinaryPointwiseOpBuilder binaryOp,
                             StructuredIndexed I1, StructuredIndexed I2,
-                            StructuredIndexed O);
+                            StructuredIndexed O,
+                            ArrayRef<Type> resultTensorTypes = {});
 
 /// 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.
+///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
+/// which case, the canonical identity indexing_map is assumed.
+//
+// TODO(ntv) In the future we may want to relax this identity assumption (e.g.
+// for automatic differentiation purposes). In that case we will want to make
+// StructuredIndexed work with ValueHandle to encode type or value.
 Operation *linalg_pointwise_add(StructuredIndexed I1, StructuredIndexed I2,
-                                StructuredIndexed O);
+                                StructuredIndexed O,
+                                ArrayRef<Type> resultTensorTypes = {});
 
 /// Build a linalg.pointwise with all `parallel` iterators and a region that
 /// computes `O = max(I!, I2)`. The client is responsible for specifying the
 /// proper indexings when creating the StructuredIndexed.
+///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
+/// which case, the canonical identity indexing_map is assumed.
+//
+// TODO(ntv) In the future we may want to relax this identity assumption (e.g.
+// for automatic differentiation purposes). In that case we will want to make
+// StructuredIndexed work with ValueHandle to encode type or value.
 Operation *linalg_pointwise_max(StructuredIndexed I1, StructuredIndexed I2,
-                                StructuredIndexed O);
+                                StructuredIndexed O,
+                                ArrayRef<Type> resultTensorTypes = {});
 
 // 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:
 /// ```
 ///    (m, n, k) = (par, par, seq)
 ///    |
 ///    |  C(m, n) += A(m, k) * B(k, n)
 /// ```
-Operation *linalg_matmul(ValueHandle vA, ValueHandle vB, ValueHandle vC);
+///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
+/// which case, the canonical identity indexing_map is assumed.
+//
+// TODO(ntv) In the future we may want to relax this identity assumption (e.g.
+// for automatic differentiation purposes). In that case we will want to make
+// StructuredIndexed work with ValueHandle to encode type or value.
+Operation *linalg_matmul(ValueHandle vA, ValueHandle vB, ValueHandle vC,
+                         ArrayRef<Type> resultTensorTypes = {});
 
-template <typename Container> Operation *linalg_matmul(Container values) {
+template <typename Container>
+Operation *linalg_matmul(Container values,
+                         ArrayRef<Type> resultTensorTypes = {}) {
   assert(values.size() == 3 && "Expected exactly 3 values");
-  return linalg_matmul(values[0], values[1], values[2]);
+  assert(resultTensorTypes.size() <= 1 && "Expected at most 1 result tensor");
+  return linalg_matmul(values[0], values[1], values[2], resultTensorTypes);
 }
 
 /// 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)
 ///    |
 ///    | O(batch, [h, w, ...], f) +=
 ///    |   I(batch,
 ///    |     [
 ///    |       stride[0] * h + dilations[0] * kh,
 ///    |       stride[1] * w + dilations[1] * kw, ...
 ///          ],
 ///    |     c)
 ///    |   *
 ///    |   W([kh, kw, ...], c, f)
 /// ```
 /// If `dilations` or `strides` are left empty, the default value of `1` is used
 /// along each relevant dimension.
 ///
 /// For now `...` must be empty (i.e. only 2-D convolutions are supported).
 ///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
+/// which case, the canonical identity indexing_map is assumed.
+//
+// TODO(ntv) In the future we may want to relax this identity assumption (e.g.
+// for automatic differentiation purposes). In that case we will want to make
+// StructuredIndexed work with ValueHandle to encode type or value.
+//
 // TODO(ntv) Extend convolution rank with some template magic.
 Operation *linalg_conv_nhwc(ValueHandle vI, ValueHandle vW, ValueHandle vO,
+                            ArrayRef<Type> resultTensorTypes = {},
                             ArrayRef<int> strides = {},
                             ArrayRef<int> dilations = {});
 
 template <typename Container>
-Operation *linalg_conv_nhwc(Container values, ArrayRef<int> strides = {},
-                            ArrayRef<int> dilations = {}) {
+Operation *
+linalg_conv_nhwc(Container values, ArrayRef<Type> resultTensorTypes = {},
+                 ArrayRef<int> strides = {}, ArrayRef<int> dilations = {}) {
   assert(values.size() == 3 && "Expected exactly 3 values");
-  return linalg_conv_nhwc(values[0], values[1], values[2], strides, dilations);
+  assert(resultTensorTypes.size() <= 1 && "Expected at most 1 result tensor");
+  return linalg_conv_nhwc(values[0], values[1], values[2], resultTensorTypes,
+                          strides, dilations);
 }
 
 /// Build a linalg.generic, under the current ScopedContext, at the current
 /// insert point, that computes:
 /// ```
 ///    (batch, dm, c, [h, w, ...], [kh, kw, ...]) =
 ///    |  (par, par, par, [par, par, ...], [red, red, ...])
 ///    |
-///    | O(batch, [h, w, ...], c * depth_multiplier) +=
+///    | O(batch, [h, w, ...], c * depthMultiplier) +=
 ///    |   I(batch,
 ///    |     [
 ///    |       stride[0] * h + dilations[0] * kh,
 ///    |       stride[1] * w + dilations[1] * kw, ...
 ///          ],
 ///    |     c)
 ///    |   *
-///    |   W([kh, kw, ...], c, depth_multiplier)
+///    |   W([kh, kw, ...], c, depthMultiplier)
 /// ```
 /// If `dilations` or `strides` are left empty, the default value of `1` is used
 /// along each relevant dimension.
 ///
 /// For now `...` must be empty (i.e. only 2-D convolutions are supported).
 ///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
+/// which case, the canonical identity indexing_map is assumed.
+//
+// TODO(ntv) In the future we may want to relax this identity assumption (e.g.
+// for automatic differentiation purposes). In that case we will want to make
+// StructuredIndexed work with ValueHandle to encode type or value.
+//
 // TODO(ntv) Extend convolution rank with some template magic.
 Operation *linalg_dilated_conv_nhwc(ValueHandle vI, ValueHandle vW,
-                                    ValueHandle vO, int depth_multiplier = 1,
+                                    ValueHandle vO,
+                                    ArrayRef<Type> resultTensorTypes = {},
+                                    int depthMultiplier = 1,
                                     ArrayRef<int> strides = {},
                                     ArrayRef<int> dilations = {});
 
 template <typename Container>
-Operation *linalg_dilated_conv_nhwc(Container values, int depth_multiplier,
+Operation *linalg_dilated_conv_nhwc(Container values,
+                                    ArrayRef<Type> resultTensorTypes = {},
+                                    int depthMultiplier = 1,
                                     ArrayRef<int> strides = {},
                                     ArrayRef<int> dilations = {}) {
   assert(values.size() == 3 && "Expected exactly 3 values");
+  assert(resultTensorTypes.size() <= 1 && "Expected at most 1 result tensor");
   return linalg_dilated_conv_nhwc(values[0], values[1], values[2],
-                                  depth_multiplier, strides, dilations);
+                                  resultTensorTypes, depthMultiplier, strides,
+                                  dilations);
 }
 
 } // namespace ops
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_EDSC_BUILDERS_H_

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

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) {
   if (fun)
     fun();
   exit();
   return ValueHandle::null();
 }
 
 mlir::edsc::LoopNestRangeBuilder::LoopNestRangeBuilder(
     ArrayRef<ValueHandle *> ivs, ArrayRef<SubViewOp::Range> ranges) {
   loops.reserve(ranges.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) {
   if (fun)
     fun();
   for (auto &lit : reverse(loops)) {
     lit({});
   }
   return ValueHandle::null();
 }
 
 namespace mlir {
for (auto expr : sidx.getExprs()) {
           pos = std::max(pos, d.getPosition());
       });
     }
   }
 }
 
 Operation *mlir::edsc::makeGenericLinalgOp(
     ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
-    ArrayRef<StructuredIndexed> outputs,
+    ArrayRef<StructuredIndexed> outputBuffers, ArrayRef<Type> resultTensorTypes,
     function_ref<void(ArrayRef<BlockArgument>)> regionBuilder,
     ArrayRef<Value> otherValues, ArrayRef<Attribute> otherAttributes) {
+  assert(
+      llvm::all_of(llvm::make_range(outputBuffers.begin(), outputBuffers.end()),
+                   [](Value v) { return v.getType().isa<MemRefType>(); }) &&
+      "output operands must all be buffers.");
   auto &builder = edsc::ScopedContext::getBuilder();
   auto *ctx = builder.getContext();
   unsigned nInputs = inputs.size();
-  unsigned nOutputs = outputs.size();
+  unsigned nOutputs = outputBuffers.size() + resultTensorTypes.size();
   unsigned maxPos = 0;
   getMaxDimIndex(inputs, maxPos);
-  getMaxDimIndex(outputs, maxPos);
+  getMaxDimIndex(outputBuffers, 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)
+  for (auto out : outputBuffers)
     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());
+  values.append(outputBuffers.begin(), outputBuffers.end());
 
   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
+              resultTensorTypes,
               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) {
   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(UnaryPointwiseOpBuilder unaryOp,
                                              StructuredIndexed I,
-                                             StructuredIndexed O) {
+                                             StructuredIndexed O,
+                                             ArrayRef<Type> resultTensorTypes) {
   SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
                                            edsc::IterType::Parallel);
   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);
+
+  // Distinguish between tensor and buffer semantics.
+  if (O.getType().isa<MemRefType>()) {
+    assert(resultTensorTypes.empty());
+    return makeGenericLinalgOp(iterTypes, {I}, {O}, {}, fun);
+  }
+  return makeGenericLinalgOp(iterTypes, {I, O}, {}, resultTensorTypes, fun);
 }
 
-Operation *mlir::edsc::ops::linalg_pointwise_tanh(StructuredIndexed I,
-                                                  StructuredIndexed O) {
+Operation *
+mlir::edsc::ops::linalg_pointwise_tanh(StructuredIndexed I, StructuredIndexed O,
+                                       ArrayRef<Type> resultTensorTypes) {
   ;
   using edsc::intrinsics::tanh;
   UnaryPointwiseOpBuilder unOp([](ValueHandle a) -> Value { return tanh(a); });
-  return linalg_pointwise(unOp, I, O);
+  return linalg_pointwise(unOp, I, O, resultTensorTypes);
 }
 
 /// Binary pointwise operation (with broadcast) entry point.
 Operation *mlir::edsc::ops::linalg_pointwise(BinaryPointwiseOpBuilder binaryOp,
                                              StructuredIndexed I1,
                                              StructuredIndexed I2,
-                                             StructuredIndexed O) {
+                                             StructuredIndexed O,
+                                             ArrayRef<Type> resultTensorTypes) {
   SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
                                            edsc::IterType::Parallel);
   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);
+  // Distinguish between tensor and buffer semantics.
+  if (O.getType().isa<MemRefType>()) {
+    assert(resultTensorTypes.empty());
+    return makeGenericLinalgOp(iterTypes, {I1, I2}, {O}, {}, fun);
+  }
+  return makeGenericLinalgOp(iterTypes, {I1, I2, O}, {}, resultTensorTypes,
+                             fun);
 }
 
-Operation *mlir::edsc::ops::linalg_pointwise_add(StructuredIndexed I1,
-                                                 StructuredIndexed I2,
-                                                 StructuredIndexed O) {
+Operation *
+mlir::edsc::ops::linalg_pointwise_add(StructuredIndexed I1,
+                                      StructuredIndexed I2, StructuredIndexed O,
+                                      ArrayRef<Type> resultTensorTypes) {
   using edsc::op::operator+;
   BinaryPointwiseOpBuilder binOp(
       [](ValueHandle a, ValueHandle b) -> Value { return a + b; });
-  return linalg_pointwise(binOp, I1, I2, O);
+  return linalg_pointwise(binOp, I1, I2, O, resultTensorTypes);
 }
 
-Operation *mlir::edsc::ops::linalg_pointwise_max(StructuredIndexed I1,
-                                                 StructuredIndexed I2,
-                                                 StructuredIndexed O) {
+Operation *
+mlir::edsc::ops::linalg_pointwise_max(StructuredIndexed I1,
+                                      StructuredIndexed I2, StructuredIndexed O,
+                                      ArrayRef<Type> resultTensorTypes) {
   BinaryPointwiseOpBuilder binOp([](ValueHandle a, ValueHandle b) -> Value {
     using edsc::intrinsics::select;
     using edsc::op::operator>;
     return select(a > b, a, b).getValue();
   });
-  return linalg_pointwise(binOp, I1, I2, O);
+  return linalg_pointwise(binOp, I1, I2, O, resultTensorTypes);
 }
 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
-                                          ValueHandle vC) {
-  // clang-format off
+                                          ValueHandle vC,
+                                          ArrayRef<Type> resultTensorTypes) {
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(vC);
+
+  assert(!C.getType().isa<MemRefType>() || resultTensorTypes.empty());
+  StructuredIndexed allIndexed[3]{A({m, k}), B({k, n}), C({m, n})};
+  ArrayRef<StructuredIndexed> inputs =
+      (C.getType().isa<MemRefType>())
+          ? ArrayRef<StructuredIndexed>{allIndexed, allIndexed + 2}
+          : ArrayRef<StructuredIndexed>{allIndexed, allIndexed + 3};
+  ArrayRef<StructuredIndexed> outputs =
+      (C.getType().isa<MemRefType>())
+          ? ArrayRef<StructuredIndexed>{allIndexed + 2, allIndexed + 3}
+          : ArrayRef<StructuredIndexed>{};
   return makeGenericLinalgOp(
-    {IterType::Parallel, IterType::Parallel, IterType::Reduction},
-    {A({m, k}), B({k, n})},
-    {C({m, n})},
-    macRegionBuilder);
-  // clang-format on
+      {IterType::Parallel, IterType::Parallel, IterType::Reduction}, inputs,
+      outputs, resultTensorTypes, macRegionBuilder);
 }
 
 Operation *mlir::edsc::ops::linalg_conv_nhwc(ValueHandle vI, ValueHandle vW,
                                              ValueHandle vO,
+                                             ArrayRef<Type> resultTensorTypes,
                                              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");
 
   // Some short names.
   auto par = IterType::Parallel;
   auto red = IterType::Reduction;
   auto s = strides;
   auto d = dilations;
 
   AffineExpr b, f, h, w, kh, kw, c;
   bindDims(ctx, b, f, h, w, kh, kw, c);
   unsigned numDims = c.cast<AffineDimExpr>().getPosition() + 1;
   StructuredIndexed I(vI), W(vW), O(vO);
-  // clang-format off
-  return makeGenericLinalgOp(
-    {par, par, par, par, red, red, red}, {
-      I({b,
-         // Roundtrip to flattened form to serve as canonicalization and ensure
-         // consistent ordering of subexpressions.
+
+  assert(!O.getType().isa<MemRefType>() || resultTensorTypes.empty());
+  // clang-format off
+  StructuredIndexed allIndexed[3] = {
+      I({b,
          simplifyAffineExpr(s[0] * h + d[0] * kh, numDims, 0),
          simplifyAffineExpr(s[1] * w + d[1] * kw, numDims, 0),
          c}),
-      W({kh, kw, c, f})}, {
-      O({b, h, w, f})},
-    macRegionBuilder);
+      W({kh, kw, c, f}),
+      O({b, h, w, f})};
   // clang-format on
+  auto inputs = (O.getType().isa<MemRefType>())
+                    ? ArrayRef<StructuredIndexed>{allIndexed, allIndexed + 2}
+                    : ArrayRef<StructuredIndexed>{allIndexed, allIndexed + 3};
+  ArrayRef<StructuredIndexed> outputs =
+      (O.getType().isa<MemRefType>())
+          ? ArrayRef<StructuredIndexed>{allIndexed + 2, allIndexed + 3}
+          : ArrayRef<StructuredIndexed>{};
+  return makeGenericLinalgOp({par, par, par, par, red, red, red}, inputs,
+                             outputs, resultTensorTypes, macRegionBuilder);
 }
 
 Operation *mlir::edsc::ops::linalg_dilated_conv_nhwc(
-    ValueHandle vI, ValueHandle vW, ValueHandle vO, int depth_multiplier,
+    ValueHandle vI, ValueHandle vW, ValueHandle vO,
+    ArrayRef<Type> resultTensorTypes, int depthMultiplier,
     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");
 
   // Some short names.
   auto par = IterType::Parallel;
   auto red = IterType::Reduction;
   auto s = strides;
   auto d = dilations;
 
   // clang-format off
   AffineExpr b, dm, c, h, w, kh, kw;
   bindDims(ctx, b, dm, c, h, w, kh, kw);
   unsigned numDims = kw.cast<AffineDimExpr>().getPosition() + 1;
   StructuredIndexed I(vI), W(vW), O(vO);
-  return makeGenericLinalgOp(
-    {par, par, par, par, par, red, red}, {
+  // Roundtrip to flattened form to serve as canonicalization and ensure
+  // consistent ordering of subexpressions.
+  // clang-format off
+  StructuredIndexed allIndexed[3] = {
       I({b,
          // Roundtrip to flattened form to serve as canonicalization and ensure
          // consistent ordering of subexpressions.
          simplifyAffineExpr(s[0] * h + d[0] * kh, numDims, 0),
          simplifyAffineExpr(s[1] * w + d[1] * kw, numDims, 0),
          c}),
-      W({kh, kw, c, dm})}, {
-      O({b, h, w, simplifyAffineExpr(c * depth_multiplier + dm, numDims, 0)})},
-    macRegionBuilder);
+      W({kh, kw, c, dm}),
+      O({b, h, w, simplifyAffineExpr(c * depthMultiplier + dm, numDims, 0)})};
   // clang-format on
+  auto inputs = (O.getType().isa<MemRefType>())
+                    ? ArrayRef<StructuredIndexed>{allIndexed, allIndexed + 2}
+                    : ArrayRef<StructuredIndexed>{allIndexed, allIndexed + 3};
+  ArrayRef<StructuredIndexed> outputs =
+      (O.getType().isa<MemRefType>())
+          ? ArrayRef<StructuredIndexed>{allIndexed + 2, allIndexed + 3}
+          : ArrayRef<StructuredIndexed>{};
+  return makeGenericLinalgOp({par, par, par, par, par, red, red}, inputs,
+                             outputs, resultTensorTypes, macRegionBuilder);
 }

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

 
 TEST_FUNC(zero_and_sign_extendi_op_i1_to_i8) {
   using namespace edsc;
   using namespace edsc::intrinsics;
   using namespace edsc::op;
   auto i1Type = IntegerType::get(1, &globalContext());
   auto i8Type = IntegerType::get(8, &globalContext());
   auto memrefType = MemRefType::get({}, i1Type, {}, 0);
-  auto f = makeFunction("zero_and_sign_extendi_op", {}, {memrefType, memrefType});
+  auto f =
+      makeFunction("zero_and_sign_extendi_op", {}, {memrefType, memrefType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   IndexedValue A(f.getArgument(0));
   IndexedValue B(f.getArgument(1));
   // clang-format off
   edsc::intrinsics::zero_extendi(*A, i8Type);
   edsc::intrinsics::sign_extendi(*B, i8Type);
TEST_FUNC(empty_map_load_store) {
   // CHECK:  [[A:%.*]] = affine.load %{{.*}}[]
   // CHECK:  affine.store [[A]], %{{.*}}[]
   // clang-format on
   f.print(llvm::outs());
   f.erase();
 }
 
 // CHECK-LABEL: func @affine_if_op
-// CHECK:       affine.if affine_set<([[d0:.*]], [[d1:.*]]){{\[}}[[s0:.*]], [[s1:.*]]{{\]}}
+// CHECK:       affine.if affine_set<([[d0:.*]], [[d1:.*]]){{\[}}
+// CHECK-SAME:    [[s0:.*]], [[s1:.*]]{{\]}}
 // CHECK-NOT:   else
-// CHECK:       affine.if affine_set<([[d0:.*]], [[d1:.*]]){{\[}}[[s0:.*]], [[s1:.*]]{{\]}}
-// CHECK-NEXT: } else {
+// CHECK:       affine.if affine_set<([[d0:.*]], [[d1:.*]]){{\[}}
+// CHECK-SAME:    [[s0:.*]], [[s1:.*]]{{\]}}
+// CHECK-NEXT:  } else {
 TEST_FUNC(affine_if_op) {
   using namespace edsc;
   using namespace edsc::intrinsics;
   using namespace edsc::op;
   auto f32Type = FloatType::getF32(&globalContext());
   auto memrefType = MemRefType::get({-1, -1}, f32Type, {}, 0);
   auto f = makeFunction("affine_if_op", {}, {memrefType});
 
TEST_FUNC(linalg_matmul_test) {
   ScopedContext scope(builder, f.getLoc());
   linalg_matmul(makeValueHandles(llvm::to_vector<3>(f.getArguments())));
 
   f.print(llvm::outs());
   f.erase();
 }
 
 // clang-format off
+// CHECK-LABEL: func @linalg_matmul_mixed_tensors
+//       CHECK:   linalg.generic {args_in = 3 : 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
+//       CHECK:   }: tensor<?x?xf32>, memref<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
+// clang-format on
+TEST_FUNC(linalg_matmul_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_matmul_mixed_tensors", {},
+                        {tensorType, memrefType, tensorType});
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+  linalg_matmul(makeValueHandles(llvm::to_vector<3>(f.getArguments())),
+                tensorType);
+
+  f.print(llvm::outs());
+  f.erase();
+}
+
+// clang-format off
 // CHECK-LABEL: func @linalg_conv_nhwc
 //       CHECK:   linalg.generic {args_in = 2 : i64, args_out = 1 : i64,
 // CHECK-SAME: indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2 * 3 + d4 * 5, d3 * 4 + d5 * 6, d6)>,
 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d1)>,
 // CHECK-SAME: affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2, d3, d1)>],
 // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "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
 //       CHECK:   }: memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
 // clang-format on
 TEST_FUNC(linalg_conv_nhwc) {
   using namespace edsc;
   using namespace edsc::ops;
 
   auto f32Type = FloatType::getF32(&globalContext());
   auto memrefType = MemRefType::get({-1, -1, -1, -1}, f32Type, {}, 0);
   auto f = makeFunction("linalg_conv_nhwc", {},
                         {memrefType, memrefType, memrefType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
-  linalg_conv_nhwc(makeValueHandles(llvm::to_vector<3>(f.getArguments())),
+  linalg_conv_nhwc(makeValueHandles(llvm::to_vector<3>(f.getArguments())), {},
                    /*strides=*/{3, 4}, /*dilations=*/{5, 6});
 
   f.print(llvm::outs());
   f.erase();
 }
 
 // clang-format off
 // CHECK-LABEL: func @linalg_dilated_conv_nhwc
TEST_FUNC(linalg_dilated_conv_nhwc) {
   auto memrefType = MemRefType::get({-1, -1, -1, -1}, f32Type, {}, 0);
   auto f = makeFunction("linalg_dilated_conv_nhwc", {},
                         {memrefType, memrefType, memrefType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   linalg_dilated_conv_nhwc(
       makeValueHandles(llvm::to_vector<3>(f.getArguments())),
+      /*outputTensorTypes=*/{},
       /*depth_multiplier=*/7,
       /*strides=*/{3, 4}, /*dilations=*/{5, 6});
 
   f.print(llvm::outs());
   f.erase();
 }
 
 // clang-format off