[Transform][Fusion] Introduce an iterative tiling and fusion pass in forward and backward fashion

[Yun-Fly]

Track issue: Support fine-grain fusion #54

• Matmul fusion with element-wise/reduce/broadcast ops.
• Pre-op and post-op fusion(a.k.a. producer and consumer fusion respectively).
• Multi-consumer and multi-producer support.
• Support multiple level tiling and fusion.
• Add flexible options to control the boundary of iterative fusion.
• Enable default tiling when no op is tiled before fusion.
• Cost-model to determine whether to fuse or not.

[ZhennanQin]

#122 is to update LLVM version. Please rebase after that PR is merged.

[dchigarev]

@Yun-Fly thanks for your PR!

I've been able to test your pass in the GPU pipeline (linalg->xegpu->gpu exe) by replacing tile-consumer-and-fuse-producers pass from tpp-mlir repo and everything works just fine (if change the default tiling value in the code to something greater than 16).

What prevents us from changing 1 in this line to a value from a parameter to the pass?

graph-compiler/lib/gc/Transforms/IterativeTilingAndFusion.cpp

Line 625 in 111d276

defaultTileSize[en] = rewriter.getIndexAttr(1);

[Yun-Fly]

@Yun-Fly thanks for your PR!

I've been able to test your pass in the GPU pipeline (linalg->xegpu->gpu exe) by replacing tile-consumer-and-fuse-producers pass from tpp-mlir repo and everything works just fine (if change the default tiling value in the code to something greater than 16).

Glad to hear this progress!

What prevents us from changing 1 in this line to a value from a parameter to the pass?

graph-compiler/lib/gc/Transforms/IterativeTilingAndFusion.cpp

Line 625 in 111d276

defaultTileSize[en] = rewriter.getIndexAttr(1);

Yes, its a quite great point! Firstly, I would like to share some context with you:

1. In CPU pipeline, this fusion pass is placed after another deep-tiled-matmul pass. In another word, we expects matmuls already tiled by best config(a.k.a tileSize).
2. As a fallback, the code line you pointed out is default tileSize for all kinds of linalg ops in case of no op is tiled before fusion as your said. However, in general, it is hard to designate specific tileSize to a certain op by a value from a parameter to the pass. E.g. if there exist two matmul, how do we pass different tileSize to them? In MLIR, although it maybe customized by transform.match in UT, it is not feasible in automatic pass.

I would still try to add a parameter to pass tileSize for certain kind of linalg op. i.e. -tile-size={"matmul", {32, 32}}. But in real case(like MLP), different shapes of matmuls may prefer different tileSize for better performance. And that is why we have a standalone pass to tile matmul before fusion in CPU pipeline. Would you like to share your plan about this in GPU pipeline?

[kurapov-peter]

@Yun-Fly, could you please add test cases that showcase the described scenarios? This would really help in understanding pre/post conditions.

[Yun-Fly]

I would still try to add a parameter to pass tileSize for certain kind of linalg op. i.e. -tile-size={"matmul", {32, 32}}. But in real case(like MLP), different shapes of matmuls may prefer different tileSize for better performance. And that is why we have a standalone pass to tile matmul before fusion in CPU pipeline. Would you like to share your plan about this in GPU pipeline?

Hi, @dchigarev . I have added new parameter called default-tile-size. You can specify tileSize by gc-opt -iterative-tiling-and-fusion="default-tile-size=matmul:{32,32}". This behavior will affect all matmuls in your input .mlir, which sounds similar to TPP. The difference is that we can also designate different tileSize for other kinds of ops except for contraction op, saying ="default-tile-size=matmul:{32,32},reduce:{16,16}". NOTE that, all of these default tileSize only make sense when the certain kind of op need self tiling rather than fused with any other already tiled op/group.

[Yun-Fly]

@Yun-Fly, could you please add test cases that showcase the described scenarios? This would really help in understanding pre/post conditions.

Sure, I will arrange them later. pre/post-op fusion is as known as producer/consumer fusion respectively.

[dchigarev]

@Yun-Fly

I have added new parameter called default-tile-size.

Thank you! I think this will be enough for our first gpu-pipeline prototype

[kurapov-peter]

The first pack of the comments/questions. Tried to go inside-out on the tiling using interface.

[kurapov-peter]

Second portion.

[kurapov-peter]

Suggested change

	FailureOr<int64_t> cstTileSizes =
	ValueBoundsConstraintSet::computeConstantBound(
	presburger::BoundType::UB, tileSizes[resultExpr.index()], nullptr,
	true);
	FailureOr<int64_t> cstTileSizes =
	ValueBoundsConstraintSet::computeConstantBound(
	presburger::BoundType::UB, tileSizes[resultExpr.index()], /*stopCondition=*/nullptr,
	/*closedUB=*/true);

Btw, are there cases when we know the tile sizes statically?

[Yun-Fly]

We assume all the tileSize is static rather than dymamic. BTW, do we need to support dynamic shape at this moment?

[kurapov-peter]

It's OK to start with the static case. We'll need the dynamic as well of course.

[kurapov-peter]

What does this check exactly?

[Yun-Fly]

As the name of this filter(noTilingOnReductionFilter) said, we need to ensure there exist no tiling on any reduction dimension. Otherwise, it will lead to wrong calculation result. I.e. tileSize should equal to IterationDomain. Certainly, if either of them is not dynamic, they are incomparable until RUNTIME , so just return failure for such case.

[kurapov-peter]

Is this to cover some weird uneven tiling/generic case? I mean, what's the reason a tail check is not sufficient?

[Yun-Fly]

As the name of this filter(exactTilingOnPackUnPackFilter) said, uneven tiling/generic case is not expected so far because it will involve more complex imperfect tiling case.

[kurapov-peter]

Wouldn't it be easier and safer to use the type ID for it?

[Yun-Fly]

Yeah, I agree. But the problem is that the default tileSize is passed by argument, like gc-opt -iterative-tiling-and-fusion="default-tile-size=matmul:{32,32}". By MLIR parser, they will be converted to std::string in fact.

[kurapov-peter]

Thanks, looks better now. It this still needs even more testing. All the minor things we can clean up later on.