[PATCH] D72863: [mlir][Linalg] Add tensor support to Linalg EDSC Builders

[Nicolas Vasilache via Phabricator via llvm-commits]

nicolasvasilache marked 21 inline comments as done.
nicolasvasilache added inline comments.


================
Comment at: mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h:186
+///
+/// This accepts both buffers and tensors as `inputs` but only buffers as
+/// `outputs`. Output tensors can be specified with `resultTensorTypes`, in
----------------
ftynse wrote:
> 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.
Done by virtue of refocusing the diff.


================
Comment at: mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h:195-196
 Operation *linalg_pointwise(UnaryPointwiseOpBuilder unaryOp,
-                            StructuredIndexed I, StructuredIndexed O);
+                            StructuredIndexed I, StructuredIndexed O,
+                            ArrayRef<Type> resultTensorTypes = {});
 
----------------
ftynse wrote:
> It's not clear what happens to `O`, which I presume is the output, if you want tensor outputs.
Done by virtue of refocusing the diff.


================
Comment at: mlir/lib/Dialect/Linalg/EDSC/Builders.cpp:134-137
+  assert(
+      llvm::all_of(llvm::make_range(outputBuffers.begin(), outputBuffers.end()),
+                   [](Value v) { return v.getType().isa<MemRefType>(); }) &&
+      "output operands must all be buffers.");
----------------
hanchung wrote:
> Do we also check if all types in resultTensorTypes are RankedTensorType?
That is done by the verifier of linalg.generic.


================
Comment at: mlir/lib/Dialect/Linalg/EDSC/Builders.cpp:144
   getMaxDimIndex(inputs, maxPos);
-  getMaxDimIndex(outputs, maxPos);
+  getMaxDimIndex(outputBuffers, maxPos);
   // maxPos is 0 indexed, need to turn this into a count (i.e. +1)
----------------
hanchung wrote:
> Do we need to update maxPos with resultTensorTypes?
Good point, I refocused the diff though so I don't need it anymore.


================
Comment at: mlir/lib/Dialect/Linalg/EDSC/Builders.cpp:229
+  }
+  return makeGenericLinalgOp(iterTypes, {I, O}, {}, resultTensorTypes, fun);
 }
----------------
hanchung wrote:
> 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?
There was not and this was a mistake.
Now there is, thanks for surfacing!


================
Comment at: mlir/lib/Dialect/Linalg/EDSC/Builders.cpp:292
+
+  assert(!C.getType().isa<MemRefType>() || resultTensorTypes.empty());
+  StructuredIndexed allIndexed[3]{A({m, k}), B({k, n}), C({m, n})};
----------------
ftynse wrote:
> Please add a string with description to the assert
obsolete, I'll revisit in a later diff, thanks!


================
Comment at: mlir/lib/Dialect/Linalg/EDSC/Builders.cpp:293
+  assert(!C.getType().isa<MemRefType>() || resultTensorTypes.empty());
+  StructuredIndexed allIndexed[3]{A({m, k}), B({k, n}), C({m, n})};
+  ArrayRef<StructuredIndexed> inputs =
----------------
ftynse wrote:
> 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})});
> ```
> ```
obsolete, I'll revisit in a later diff, thanks!


================
Comment at: mlir/lib/Dialect/Linalg/EDSC/Builders.cpp:340
   // clang-format on
+  auto inputs = (O.getType().isa<MemRefType>())
+                    ? ArrayRef<StructuredIndexed>{allIndexed, allIndexed + 2}
----------------
ftynse wrote:
> Same as above
obsolete, I'll revisit in a later diff, thanks!


================
Comment at: mlir/test/EDSC/builder-api-test.cpp:910
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}
+///      CHECK:   ^bb0(%[[a0:.*]]: f32, %[[a1:.*]]: f32, %[[a2:.*]]: f32):
+//       CHECK:     %[[a3:.*]] = mulf %[[a0]], %[[a1]] : f32
----------------
hanchung wrote:
> Remove one "/" in the beginning?
obsolete, I'll revisit in a later diff, thanks!


================
Comment at: mlir/test/EDSC/builder-api-test.cpp:959
   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});
----------------
hanchung wrote:
> nit: add comment for function argument, /*resultTensorTypes=*/{}
obsolete, I'll revisit in a later diff, thanks!


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D72863/new/

https://reviews.llvm.org/D72863