Node fusion via leaf nodes

torch.compile FX (Functional Transformations)
1

Hello everyone!
I’m perfoming node fusions and am thinking about using the following procedure:

  1. Take two subsequent nodes that perform an operation (say linear → relu or add → mul)
  2. Implement a fused module (LinearWithRelu with appropriate forward or addAndMul respectively)
  3. Set new fused module as leaf (or wrap function, in the case of addAndMul)
  4. Replace old nodes by fused ones.
  5. Recompile the GraphModule

Could this new IR have an impact (positive or negative) on the way the module is evaluated? Will the fused nodes be compiled and/or scheduled differently? Or this only affect the symbolic tracing?

2

Ok. Got no answer here but I already found one in the PyTorch Dev Forum, so I’m leaving it in case someones has a similar question:

The nuances of PyTorch Graph Capture

TL,DR
Could this new IR have an impact (positive or negative) on the way the module is evaluated? Yes, and that is the main purpose of changing the IR.
Will the fused nodes be compiled and/or scheduled differently? Or this only affect the symbolic tracing? Changing the graph will probably alter the scheduling of the node since “fused” nodes will be scheduled as a whole kernel, while the folding of such fusions will have separate schedules. Symbolic tracing is one procedure to capture the graph and thus changing the graph does not affect the symbolic tracing procedure per-se, but it does affect the result of the symbolic tracing.