search

NVIDIA / Fuser

A Fusion Code Generator for NVIDIA GPUs (commonly known as "nvFuser")
Other
229 stars 42 forks source link

Fix MatmulOp IterDomain mapping #2246

Closed jacobhinkle closed 4 weeks ago

jacobhinkle commented 1 month ago

This PR does the following:

  1. Add MatmulOp to ir_utils::isTvOp so that its IterDomains will be automatically propagated by IdModel.
  2. Updates the tests to check that all non-Broadcast axes are properly mapped by IdModel through the MatmulOp.
  3. Changes the output of MatmulOp to have an IterType::Reduction axis in the last position of its root domain to represent the K dimension. This change was motivated by needing a way to have both operand K dimensions exact mapped together, as they would be if the op were translated to a mul+sum+cast.
  4. Updates the matmul op to translate trivial cases where K=1 to simple multiply+cast patterns.

Fixes #1707. In fact, that test was actually fixed by #2175 but the test validation was failing because isTvOp was not picking up the matmul as a reduction.

jacobhinkle commented 1 month ago

!build --diff

jacobhinkle commented 4 weeks ago

I agree about symbolic domains and I think we might need to handle this at concretization. Namely if K is a broadcast dim then we should translate to a mul op before segmentation. I'll work on an example for a case like that.