Page MenuHomePhabricator

[mlir] Structured transforms: introduce op splitting
ClosedPublic

Authored by ftynse on Jul 4 2022, 9:40 AM.

Details

Summary

Introduce a new transformation on structured ops that splits the iteration
space into two parts along the specified dimension. The index at which the
splitting happens may be static or dynamic. This transformation can be seen as
a rudimentary form of index-set splitting that only supports the splitting
along hyperplanes parallel to the iteration space hyperplanes, and is therefore
decomposable into per-dimension application.

It is a key low-level transformation that enables independent scheduling for
different parts of the iteration space of the same op, which hasn't been
possible previously. It may be used to implement, e.g., multi-sized tiling. In
future, peeling can be implemented as a combination of split-off amount
computation and splitting.

The transformation is conceptually close to tiling in its separation of the
iteration and data spaces, but cannot be currently implemented on top of
TilingInterface as the latter does not properly support linalg.index
offsetting.

Note that the transformation intentionally bypasses folding of
tensor.extract_slice operations when creating them as this folding was found
to prevent repeated splitting of the same operation because due to internal
assumptions about extract/insert_slice combination in dialect utilities.

Diff Detail

Event Timeline

ftynse created this revision.Jul 4 2022, 9:40 AM
Herald added a project: Restricted Project. · View Herald TranscriptJul 4 2022, 9:40 AM
ftynse requested review of this revision.Jul 4 2022, 9:40 AM
Herald added a project: Restricted Project. · View Herald TranscriptJul 4 2022, 9:40 AM
nicolasvasilache accepted this revision.Jul 5 2022, 8:31 AM
nicolasvasilache added inline comments.
mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
160

Consider adding "two *complementary* parts".

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
202

seems like this would belong to the StructuredOpInterface

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
406

Seems we're past overdue of having applyToOne extended, also with the NoOp stuff.
I'll ping you offline to try and take a crack at it.

449

Add a // TODO: Use NoOp ?

454

ah ok, yes typo is -> in

456

Add a // TODO: Use NoOp ?

mlir/lib/Dialect/Linalg/Transforms/Split.cpp
21

Plz Reuse/adopt/evolve/move to StaticValueUtils.h

46

Some comment plz

58

You could also add a version of applyMapToValues that takes an ArrayRef<OpFoldResult> and folds the constants/compress to avoid this.

103

Seems like the code in l102-114 could be refactored in a properly named function + doc and reused?

141

I would take the tradeoff of bailing here with an error message and reusing TilingOpInterface if this is all it takes.
We could then extend when it becomes critical.

mlir/lib/Dialect/Linalg/Utils/Utils.cpp
926

no either

962

add it directly to the StructuredOpInterface ?

mlir/test/Dialect/Linalg/transform-op-split.mlir
350

error message is weird, is this a typo is -> in?

This revision is now accepted and ready to land.Jul 5 2022, 8:31 AM
ftynse updated this revision to Diff 442619.Jul 6 2022, 10:06 AM
ftynse marked 10 inline comments as done.

Review.

ftynse marked an inline comment as done.Jul 7 2022, 4:05 AM
ftynse added inline comments.
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
202

This is rather specific to some transformations, in particular because it takes operands distinct from the op operands. Putting that in the structured op interface sounds confusing to me. This rather belongs to TilingInterface when we get to using that.

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
406

OTOH, this only takes one loop and a couple of vectors per op. I am not entirely convinced the SFINAE and nested vector transposition complexity required to make applyToOne work in such cases is worth it.

449

Sounds irrelevant at this level, we are not in applyToOne, we don't want to return dummy ops here.

mlir/lib/Dialect/Linalg/Transforms/Split.cpp
21

This will make DialectUtils, and likely a bunch of transitive users, depend on the arithmetic dialect. Are you sure?

58

I added that, which expectedly turned out to be a non-negligible amount of code...

The issue is more fundamental with the functioning mode and the API of the folding subsystem. It materializes the constant result for each successfully folded operation and does not expose the "unmaterialized" OpFoldResult. The same is true for dimension computations that emit useless constants. This is reasonable for rewriting (and the main constant folder is a rewriting transform) because the folded op may have users that expect values and are not necessarily foldable themselves. This is less reasonable for the case here when we want to have folded-by-construction ops, and we can guarantee that there are no users (yet) so the materialization is premature.

141

I would rather not feature-creep this. The original patch was 800 lines, now I have at least 5 with at least 3k lines. We can port this separately and intentionally.

mlir/lib/Dialect/Linalg/Utils/Utils.cpp
962

I just carried over the comment that existed elsewhere in the code.

ftynse updated this revision to Diff 442857.Jul 7 2022, 4:11 AM
ftynse marked an inline comment as done.

More review.

This revision was landed with ongoing or failed builds.Jul 7 2022, 4:20 AM
This revision was automatically updated to reflect the committed changes.