[PATCH] D153421: [mlir][Linalg] Implement the tiling interface for softmax

[Quentin Colombet via Phabricator via llvm-commits]

qcolombet marked 2 inline comments as done.
qcolombet added a subscriber: springerm.
qcolombet added inline comments.


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:55
+  if (!type.isDynamicDim(dim))
+    return builder.getIndexAttr(type.getDimSize(dim));
+
----------------
FYI, the "fix" I mentioned in my previous update is here. We use an attribute instead of a value here and that allows the fusing to work properly. I'll highlight what was broken without that.


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:65
+          }));
+}
+
----------------
@springerm Do you know if such utility function already exists somewhere?


================
Comment at: mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp:2239
+  Operation *tiledOp =
+      mlir::clone(builder, getOperation(), resultTypes, tiledOperands);
+
----------------
jpienaar wrote:
> s/mlir:://
I had to keep that otherwise clang grab the clone from `Operation::` which doesn't take any argument.


================
Comment at: mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp:531
+        << "\nSliceOp: " << sliceOpToTile.getOperation() << '\n';
+    return {};
+  }
----------------
This change is not technically required for this PR, but when I "screwed up" my tiling by returning a value instead of an attribute, this led to `slice` ops with dynamic dimensions and these were breaking `ExtractSliceOp::rankReduceIfNeeded`.

By changing this assert into a note, this makes the compiler more robust in case someone does the same mistake as I did. I.e., it will not fuse instead of crashing the compiler.


================
Comment at: mlir/test/Dialect/Linalg/tile-softmax.mlir:39
+
+transform.sequence failures(propagate) {
+  ^bb0(%arg1: !transform.any_op):
----------------
qcolombet wrote:
> rengolin wrote:
> > I'm assuming this also works with the tile-and-fuse pass. Could there be a test for that, too?
> Let me give a try!
@rengolin I added a test for tile and fuse and it exposed a weakness in the fuse algorithm, which I fixed here as well (it was asserting.)


Repository:
  rG LLVM Github Monorepo

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

https://reviews.llvm.org/D153421