Skip to content

Commit

Permalink
[XLA:GPU][Emitters] Restrict the inliner.
Browse files Browse the repository at this point in the history 
At the moment we inline a callee that calls one of the functions called by the caller. This PR adjusts the logic to inline only such callees that call a subset of the caller's functions. That way we can be sure that after inlining the caller calls the same set of functions excluding the inlined one.

Background: jax-ml/jax#26162 contains an example of a MoF fusion that takes forever to compile.

The [indexing-based partitioner](44bc816) in combination with this change fix the issue.

PiperOrigin-RevId: 729079659
  • Loading branch information
@pifon2a @Google-ML-Automation
pifon2a authored and Google-ML-Automation committed Feb 21, 2025
1 parent a8e58ed commit 2cd0332
Show file tree
Hide file tree
Showing 2 changed files with 345 additions and 11 deletions.
Loading

0 comments on commit 2cd0332

Please sign in to comment.