Page MenuHomePhabricator

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

Authored by hanchung on Wed, Nov 23, 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.Wed, Nov 23, 6:11 PM
hanchung requested review of this revision.Wed, Nov 23, 6:11 PM
Herald added a project: Restricted Project. · View Herald TranscriptWed, Nov 23, 6:11 PM
rengolin added inline comments.Thu, Nov 24, 6:18 AM
mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp
103

typo: dimensions

220

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.Fri, Nov 25, 3:41 PM

address comments

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

good point, done!

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

update undoPermutationToVector signature to follow applyPermutationToVector method

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

LGTM, thanks! @chelini any comments?

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

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