Skip to content

[WIP] Check that matmul epilogue contains only pointwise ops#2225

Closed
jacobhinkle wants to merge 4 commits intomainfrom
matmul_check_pointwise_epilogue
Closed

[WIP] Check that matmul epilogue contains only pointwise ops#2225
jacobhinkle wants to merge 4 commits intomainfrom
matmul_check_pointwise_epilogue

Conversation

@jacobhinkle
Copy link
Collaborator

This PR restricts the accepted matmul segments for the nvfuser matmul scheduler to only those containing pointwise epilogues. Additionally, it rules out cases for which we cannot yet reliably determine epilogue input vectorization due to transposes (TODO, see #2169).

Note that this check can be lifted when more epilogue cases are supported, e.g. #2213.

Fixes #2167.

This is stacked on #2175 and follow-up PR to that introducing LinearOp because currently segmentation fails for matmuls unless the complete fusion can be scheduled (see #1707). The MatmulOp and LinearOp IR nodes remove the need to inspect operand producer branches, so segmentation should work fine once that work is merged. This PR will be marked as draft until then.

This PR restricts the accepted matmul segments for the nvfuser matmul
scheduler to only those containing pointwise epilogues. Additionally, it
rules out cases for which we cannot yet reliably determine epilogue input
vectorization due to transposes (TODO, see #2169).

Note that this check can be lifted when more epilogue cases are
supported, e.g. #2213.

Fixes #2167.

This is stacked on #2175 and follow-up PR to that introducing LinearOp
because currently segmentation fails for matmuls unless the complete
fusion can be scheduled (see #1707). The MatmulOp and LinearOp IR nodes
remove the need to inspect operand producer branches, so segmentation
should work fine once that work is merged. This PR will be marked as
draft until then.
@jacobhinkle jacobhinkle added Matmuls Segmentation Issues related to nvFuser Segmentation Top-Down Matmul Dev labels May 9, 2024
@jacobhinkle
Copy link
Collaborator Author

This is outdated now

@jacobhinkle jacobhinkle closed this Oct 1, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Matmuls Segmentation Issues related to nvFuser Segmentation Top-Down Matmul Dev

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Restrict matmul scheduler to accepting only segments that are supported

1 participant