This is an archive of the discontinued LLVM Phabricator instance.

[mlir][tensor] Implement TilingInterface for tensor.pack op.
ClosedPublic

Authored by hanchung on Nov 23 2022, 6:11 PM.

Details

Summary

We can compute the offsets and sizes for the slice of input because the
iteration domain is defined over outer loops. If the dimension is tiled,
the i-th index is the product of offset_i and inner_tile_i.

Different from tiling a pad op, we do not have to deal with reading zero
data from input. Because the tiling sizes are indicated to packed outer
dimensions. We will read either the entire tile or partial tile for each
packed tile. The scf.if and tensor.generate ops are not needed in this
context.

Co-authored-by: Lorenzo Chelini <l.chelini@icloud.com>

Diff Detail

Event Timeline

hanchung created this revision.Nov 23 2022, 6:11 PM
hanchung requested review of this revision.Nov 23 2022, 6:11 PM
Herald added a project: Restricted Project. · View Herald TranscriptNov 23 2022, 6:11 PM
rengolin added inline comments.Nov 24 2022, 6:18 AM
mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
104

typo: dimensions

221

This seems to be created from line 209:

tiledOperands.push_back(b.create<ExtractSliceOp>(
    loc, packOp.getDest(), outputOffsets, outputSizes, strides));

Then some push_backs, then this constant indexing. It's confusing.

If it's always from the same op, then it would be much clearer if this was just:

  auto extractOp = b.create<ExtractSliceOp>(...);
  tiledOperands.push_back(extractOp);

  ...

Operation *tiledPackOp =
        b.create<PackOp>(loc, {extractOp.getType()}, tiledOperands, op->getAttrs());
hanchung updated this revision to Diff 478032.Nov 25 2022, 3:41 PM

address comments

hanchung marked 2 inline comments as done.Nov 25 2022, 3:41 PM
hanchung added inline comments.
mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
221

good point, done!

hanchung updated this revision to Diff 478034.Nov 25 2022, 5:58 PM
hanchung marked an inline comment as done.

update undoPermutationToVector signature to follow applyPermutationToVector method

rengolin accepted this revision.Nov 26 2022, 5:09 AM

LGTM, thanks! @chelini any comments?

This revision is now accepted and ready to land.Nov 26 2022, 5:09 AM

It's thanksgiving on US side, I'll wait for couple weekdays, and see if there are comments from others.

Build with BUILD_SHARED_LIBS=ON fails with a linking error: undefined reference to `mlir::tensor::createDimValues(mlir::OpBuilder&, mlir::Location, mlir::Value)'

It's thanksgiving on US side, I'll wait for couple weekdays, and see if there are comments from others.

yes, please.

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
75

Can we please rename this method to something like: computeInversePermutationForDimsPos? And please let's also update the comment.

131

(Optional) I find this method a bit too long. Maybe we should bring the lambdas out?

154

We can use an ArrayRef here.

205

single stmt body, drop braces?

231

I did not quite get this comment. Specifically this sentence: In this context, the outer dimensions of result tile position is the same.

242

I would just (void) here, like we did in line 119. reifyResultShapes does not fail.

244

Can you please add a comment here? It is not clear to me why we have these checks.

hanchung updated this revision to Diff 479045.Nov 30 2022, 12:13 PM
hanchung marked 4 inline comments as done.

address comments

hanchung updated this revision to Diff 479046.Nov 30 2022, 12:16 PM

fix builds

hanchung updated this revision to Diff 479060.Nov 30 2022, 1:22 PM

fix build + rebase

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
103

can you go grab AffineBuilder from iree-dialects/Dialect/LinalgExt/Transforms/Utils.h and put it in a reusable place?
This has been written one too many times.

mlir/include/mlir/Dialect/Utils/IndexingUtils.h
80 ↗(On Diff #479060)

please add an inversePermutation instead and just

applyPermutationToVector(vec, inversePermutation(permutation))
mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
79

I have strong doubts that this is needed .. can we compose existing stuff instead rather than create new one-off functions?
I already spotted this undoPermutationToVector and I suspect this is another instance.

Please increase the doc significantly because I cannot tell what this wants to do.

hanchung marked 3 inline comments as done.Dec 1 2022, 12:07 PM
hanchung added inline comments.
mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
79

it's not needed after I revisited the semantics of outerDimsPerm.

103

moved to Affine/Utils

hanchung updated this revision to Diff 479380.Dec 1 2022, 12:07 PM
hanchung marked 2 inline comments as done.

address comments

hanchung updated this revision to Diff 479407.Dec 1 2022, 1:25 PM

improve readability

hanchung updated this revision to Diff 479477.Dec 1 2022, 4:44 PM

rename method and variable

hanchung updated this revision to Diff 479481.Dec 1 2022, 5:13 PM

test stack

hanchung updated this revision to Diff 479483.Dec 1 2022, 5:17 PM

delete unrelated assert

hanchung updated this revision to Diff 479700.Dec 2 2022, 11:44 AM

fix bazel build

mravishankar accepted this revision.Dec 2 2022, 5:04 PM

I have reviewed this code previously in IREE. Looks the same to me. So fine by me

hanchung updated this revision to Diff 480193.Dec 5 2022, 12:03 PM

clang-format

This revision was landed with ongoing or failed builds.Dec 5 2022, 2:00 PM
This revision was automatically updated to reflect the committed changes.