This is an archive of the discontinued LLVM Phabricator instance.

[mlir][TilingInterface] Modify `TilingInterface` methods to better return the state of the transformed IR.
ClosedPublic

Authored by mravishankar on Mar 1 2023, 5:10 PM.

Details

Summary

Currently the getTiledImplementation and generateResultTileValue
return just SmallVector<Operation *> and FailureOr<Value>.

  • For getTiledImplementation returning empty implies tiling wasnt done. There is also an implicit assumption that the tiled operation results correspond to the tiled values of the result of the original operation. This cannot handle cases where the tiled implementation might use multiple operations to compute the tiled value for the results of the untiled operation. Sometimes, the tiled operation might not directly give the tiled values, and might require casts, etc to get a replacement.
  • For generateResultTileValue, it is assumed that the op defining the returned Value is the operation that represents the tiled computation. Again presence of casts, etc violate this.

Instead make these methods return

struct TilingResult {
  SmallVector<Operation *> tiledOps;
  SmallVector<Value> tiledValues;
};

The tiledOps represent the operations generated that are relevant
for subsequent transformations. The tiledValues represent the tiled
values for the results of the original operation. This better
transmits the state of the transformed IR.

As a consequence the following methods also return FailureOr<TilingResult>

  • tensor::replaceExtractSliceWithTiledProducer
  • tensor::bubbleUpPadSlice

Diff Detail

Event Timeline

mravishankar created this revision.Mar 1 2023, 5:10 PM
Herald added a project: Restricted Project. · View Herald TranscriptMar 1 2023, 5:10 PM
mravishankar requested review of this revision.Mar 1 2023, 5:10 PM

Fix lit tests.

pifon2a accepted this revision.Mar 1 2023, 11:33 PM
pifon2a added inline comments.
mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
53–58

return tensor::bubbleUpPadSlice(b, cast<PadOp>(op), offsets, sizes);?

This revision is now accepted and ready to land.Mar 1 2023, 11:33 PM
pifon2a added inline comments.Mar 2 2023, 2:34 AM
mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
440–444

and here

Overall looks good to me, thanks!

I have some inline comments but carrying out the information makes the code more readable and robust.

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
531–532

assert(tileValues.size() == 1)?

596–599

Is this tested with more than one tiled Ops?

The return below just takes the first value?

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
53–58

Then why not just call bubblePadOp directly? That cast<PadOp> is looking like it could crash on incorrect wrapper usage anyway.

416–417

This is a common pattern, could you have a constructor for TilingResult(Operation*) and avoid the repetition?

440–444

Are these just wrappers on top of wrappers? Do we need them all?

667

And here... :) etc.

hanchung accepted this revision.Mar 6 2023, 10:04 AM
hanchung added inline comments.
mlir/include/mlir/Interfaces/TilingInterface.h
26–34

Can we add a document about the order in tiledOps? In SCFTilingResult, we have the doc about the order:

/// Transformation information returned after tiling.
struct SCFTilingResult {
  /// Tiled operations that are generated during tiling. The order does not
  /// matter except the last op. The replacements are expected to be the results
  /// of the last op.
  SmallVector<Operation *> tiledOps;
  /// The `scf.for` operations that iterate over the tiles.
  SmallVector<scf::ForOp> loops;
  /// Values to use as replacements for the untiled op. Is the same size as the
  /// number of results of the untiled op.
  SmallVector<Value> replacements;
};
mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
395–404

not related to this PR itself.. I've seen this code and feel that the assertion is not necessary. The casting to TilingInterface is not a requirement for the check. I can send out a fix for this part and the others if people think it's reasonable.

Rebase again.

ThomasRaoux added inline comments.
mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
479

This should be tileAndFuseResult->tiledValues[0] right?

It looks like we are missing a test where we fuse an op with multiple results used.
There is the following test but only the result 0 is used so this missed to catch this problems:
https://github.com/llvm/llvm-project/blob/ff11d6b6f6e27f5de389002b8f6102b6cf3ed474/mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir#L208