This is an archive of the discontinued LLVM Phabricator instance.

[mlir] fold more eagerly in structured op splitting
ClosedPublic

Authored by ftynse on Jul 8 2022, 9:13 AM.

Details

Summary

Existing implementation of structured op splitting creates several
affine.apply and affine.min operations in its subshape computation.
As these shapes are further used in data slice extraction, this may lead
to slice shapes being dynamic even when the original shapes and the
splitting point are static. This is particularly visible when splitting
is combined with further subsetting transformations such as tiling. Use
composition and folding more aggressively in splitting to avoid this.

In particular, introduce a createComposedAffineMin function that the
affine map used in "min" with the maps used by any affine.apply that
may be feeding the operands to the "min". This enables production of
more static shapes. Also introduce a createComposedFoldedAffineApply
function that combines the existing createComposedAffineApply with
in-place folding to propagate constants produced by zero-input affine
maps. Using these when splitting allows the subsequent canonicalizer
pass to recover static shapes for structured ops.

Diff Detail

Event Timeline

ftynse created this revision.Jul 8 2022, 9:13 AM
ftynse requested review of this revision.Jul 8 2022, 9:13 AM
Herald added a project: Restricted Project. · View Herald TranscriptJul 8 2022, 9:13 AM
bondhugula added inline comments.Jul 10 2022, 12:42 AM
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
738

This looks suspicious. Can you ever pass in a null map here? Even if you don't, if you pass in a map with > 1 results and which is already composed, the assertion message would incorrectly say "could not compose maps". It looks like this should be asserting in a different or perhaps earlier if at all.

819–826

Doc comment.

ftynse updated this revision to Diff 443601.Jul 11 2022, 4:39 AM
ftynse marked 2 inline comments as done.

Address review.

ftynse updated this revision to Diff 443624.Jul 11 2022, 6:37 AM

Fold even more eagerly and recover static shapes in splitting.

ftynse edited the summary of this revision. (Show Details)Jul 11 2022, 6:38 AM
nicolasvasilache requested changes to this revision.Jul 12 2022, 3:29 AM
nicolasvasilache added inline comments.
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
398

Is there no need for a folded version like in the above ?

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
739

Great, thanks, this has been missing since forever.
I wish we also renamed OpFoldResult to ValueOrAttr or something more descriptive but that was somehow contentious at the time ..

747

This function reinvents non trivial amounts of logic around normalization, dims/sym compression etc. and additionally gets clever with inplace logic.

How about just linearizing and shifting instead:

SmallVector<Value> dims, symbols;
SmallVector<AffineExpr> exprs;
for (unsigned i : llvm::seq<unsigned>(0, map.getNumResults())) {
  SmallVector<Value> allOperands(operands.begin(), operands.end());
  AffineMap map = map.getSubMap({i});
  fullyComposeAffineMapAndOperands(&map, &allOperands);
  unsigned numNewDims = map.getNumDims();
  map = map.shiftDims(dims.size()).shiftSymbols(symbols.size());
  llvm::append_range(dims, ArrayRef<Value>(allOperands).take_front(numNewDims));
  llvm::append_range(symbols, ArrayRef<Value>(allOperands).drop_front(numNewDims));
  exprs.push_back(map.getResult(0));
}

// At this point exprs, dims and symbols are aligned potentially with duplicates.
// Run one last normalization step to get rid of duplicates and then call createOrFold.

?

This revision now requires changes to proceed.Jul 12 2022, 3:29 AM
ftynse updated this revision to Diff 443918.Jul 12 2022, 5:15 AM
ftynse marked 3 inline comments as done.

Address review.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
739

I would generally like to have better support for folding IR on construction, but we need only a small subset of what would be required for a full-fledged AutoFoldIRBuilder.

747

Great idea, thanks! I wish I thought about that before spelling out all the logic here.

nicolasvasilache accepted this revision.Jul 12 2022, 5:34 AM

Nice!

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
739

+1

This revision is now accepted and ready to land.Jul 12 2022, 5:34 AM
This revision was automatically updated to reflect the committed changes.