Repository: jax-ml/jax
Pull request #29473 (2025-06-14 02:21:34)

Add custom-call ops to roofline.
Add custom-call ops to roofline.

Recursively calculates the roofline result for the primitives from the custom function.



Patch:
diff --git a/jax/experimental/roofline/roofline.py b/jax/experimental/roofline/roofline.py
index fcfe3ff4b..dbfe3c983 100644
--- a/jax/experimental/roofline/roofline.py
+++ b/jax/experimental/roofline/roofline.py
@@ -188,6 +188,16 @@ def _roofline_interpreter(
           pin_lhs_in_vmem=pin_lhs_in_vmem,
           pin_rhs_in_vmem=pin_rhs_in_vmem,
         )
+      elif "call_jaxpr" in eqn.params:
+        # Used for custom_jvp_call_p. Recursively calculates roofline result for
+        # all primitives in the custom function.
+        result += _roofline_interpreter(
+          util.wrap_name(f_name, eqn.primitive.name),
+          eqn.params['call_jaxpr'],
+          mesh,
+          pin_lhs_in_vmem=pin_lhs_in_vmem,
+          pin_rhs_in_vmem=pin_rhs_in_vmem,
+        )
       else:
         if eqn.primitive not in _rooflines:
           msg = f"No roofline rule for {eqn.primitive}."


Test patch:
diff --git a/tests/roofline_test.py b/tests/roofline_test.py
index 9ed85c506..a234c0fe4 100644
--- a/tests/roofline_test.py
+++ b/tests/roofline_test.py
@@ -46,6 +46,31 @@ def create_inputs(
   return mesh, tuple(arrays)
 
 
+def example_function(x):
+  return jnp.sin(x) + x**2
+
+
+@jax.custom_jvp
+def example_custom_function(x):
+  """Example custom function.
+
+  Small wrapper around `example_function`. We define `example_custom_function`
+  separately since we add the `@jax.custom_jvp` decorator and want to compare
+  its behavior to `example_function`'s in tests.
+  """
+  return example_function(x)
+
+
+@example_custom_function.defjvp
+def example_custom_function_jvp(primals, tangents):
+  """Example custom function jvp.
+
+  Normally this function would define a mathematically correct JVP, but its
+  definition has 0 effect on the roofline result, so we keep it very simple.
+  """
+  return example_custom_function(primals), tangents
+
+
 class RooflineTest(jtu.JaxTestCase):
 
   def setUp(self):
@@ -802,6 +827,33 @@ class RooflineTest(jtu.JaxTestCase):
           result.unfused_hbm_bytes, self._bytes_per_word * expected_memory
       )
 
+  def test_custom_jvp_call_p_roofline(self):
+    dummy_input = jnp.ones((3, 8))
+
+    _, base_result = roofline.roofline(example_function)(dummy_input)
+    _, custom_result = roofline.roofline(example_custom_function)(dummy_input)
+
+    self.assertEqual(custom_result.unfused_flops, base_result.unfused_flops)
+    self.assertEqual(
+        custom_result.unfused_hbm_bytes, base_result.unfused_hbm_bytes
+    )
+
+  def test_custom_jvp_call_p_roofline_with_neg(self):
+    dummy_input = jnp.ones((3, 8))
+
+    def with_neg(f):
+      return lambda x: jax.lax.neg(f(x))
+
+    _, base_result = roofline.roofline(with_neg(example_function))(dummy_input)
+    _, custom_result = roofline.roofline(with_neg(example_custom_function))(
+        dummy_input
+    )
+
+    self.assertEqual(custom_result.unfused_flops, base_result.unfused_flops)
+    self.assertEqual(
+        custom_result.unfused_hbm_bytes, base_result.unfused_hbm_bytes
+    )
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())