[schedule][mlir] consumer fusion - #112
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
faa2da7 to
07f8b08
Compare
07f8b08 to
5f3b31b
Compare
| for i, loop_dim in enumerate(schedule.permutation[fuse_root]): | ||
| transform.AnnotateOp(op.new_loops[i], loop_dim) |
| consumer_handles = find_consumer_handles( | ||
| self._mlir_program.mlir_module, schedule.node_ident | ||
| ) | ||
| consumer_id = consumer_handles[0] |
There was a problem hiding this comment.
What if consumer_handles is empty? aka the current op is the last one of the graph
| assert self._named_sequence is not None | ||
|
|
||
| fuse_root = parent_name(schedule.fused_consumers[0]) | ||
| for fuse_axis in schedule.fused_consumers: |
There was a problem hiding this comment.
The loop over schedule.fused_consumers suggests multiple axes are supported, but it's only correct for one: fuse_root and the consumer handle are derived from [0] for every iteration, and sched_state.all_loops isn't refreshed after FuseConsumerOp, so a second iteration would reuse loop handles the first fusion may have invalidated (the fresh ones are in op.new_loops).
qaco
left a comment
qaco
left a comment
There was a problem hiding this comment.
Great work, thanks for doing this! See the requested changes above
3 participants
Motivation
Adds consumer fusion to the mlir backend.
Description
Implements fuse_consumer_at() for the mlir backend, the main use case being fusing relus that come after matmuls or conv2ds. It makes use of the new xtc-mlir python bindings, and a new xtc-mlir transform op FuseConsumerOp.
It only works when using the tensor dialect, the same as producer fusion.
Originally it took the name of the consumer as input, but I modified it to only fuse the first output to match how the tvm fuse_consumer_at is being implemented.