[PATCH] D74172: [mlir][Linalg] Implement fusion of linalg.generic operation on tensors.

[Nicolas Vasilache via Phabricator via llvm-commits]

nicolasvasilache accepted this revision.
nicolasvasilache added a comment.
This revision is now accepted and ready to land.

This looks really great, thanks a lot for pushing this Mahesh!
I think it is an elegant solution that shows how proper usage of the StructuredOps abstraction avoid creating problems that should never have existed in the first place (e.g. carrying broadcast and all the phase ordering issues related to buffer allocation, fusion and erasing unnecessary buffers).

Approving conditioned on addressing all comments.



================
Comment at: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp:337
+  if (!producerOp || !consumerOp || producerOp.getNumOutputs() != 1 ||
+      producerOp.getResult(0) != consumerOp.getOperand(consumerIdx)) {
+    return false;
----------------
trivial braces here and below plz


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp:346
+  // must match the number of loops of the producer.
+  AffineMap consumerIndexMap = consumerOp.getIndexingMap(consumerIdx);
+  if (consumerIndexMap.getNumResults() != producerOp.getNumLoops()) {
----------------
re. AffineMap, note that the current "inversion" procedure is best effort atm and will fail in many cases.
You want to check that the inverse computation succeeeds and return false if not.
Otherwise I suspect you'll crash later.

Example that I think should fail today and may crash your transformation down the line: `(i, j) -> (i, i+j)`


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp:390
+  AffineMap consumerIndexMap = consumerOp.getIndexingMap(consumerIdx);
+  AffineMap invProducerResultIndexMap =
+      inversePermutation(producerOp.getOutputIndexingMap(0));
----------------
This can fail in a bunch of even simple cases.
Plz see my comment above re. checking it in the precondition.


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp:411
+                               consumerOp.indexing_maps().end());
+  fusedIndexingMapAttrs.erase(
+      std::next(fusedIndexingMapAttrs.begin(), consumerIdx));
----------------
Haven't checked deeply that you don't have an off by 1 here somewhere, please be extra careful but you know better than me :)


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp:521
+struct FuseGenericTensorOps : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
----------------
Thanks for writing this as a pattern that will compose with other patterns! 


================
Comment at: mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp:525
+                                     PatternRewriter &rewriter) const override {
+    if (!op.hasTensorSemantics()) {
+      return matchFailure();
----------------
trivial braces here and everywhere please


================
Comment at: mlir/test/Dialect/Linalg/fusion-tensor.mlir:91
+
+// CHECK-LABEL: @add_broadcast_mul_fusion
+func @add_broadcast_mul_fusion(%arg0: tensor<?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
----------------
Very cool, thanks Mahesh!


Repository:
  rG LLVM Github Monorepo

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

https://reviews.llvm.org/D74172