Add cudnn_fusion decorator lowering computations to XLA cuDNN fusions.

jax/_src/cudnn/__init__.py

@@ -11,3 +11,5 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from .fusion import cudnn_fusion

jax/_src/cudnn/fusion.py

@@ -0,0 +1,91 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import functools
+import jax
+from jax import core as jax_core
+from jax.interpreters import mlir
+from jax.interpreters.mlir import hlo
+from jax.interpreters.mlir import ir
+
+
+
+def _cudnn_fusion_impl(*args, jaxpr, **unused_kwargs):
+  del unused_kwargs
+  return jax_core.jaxpr_as_fun(jaxpr)(*args)
+
+
+def _custom_abstract_eval(*args, jaxpr, **unused_kwargs):
+  del unused_kwargs
+  del args
+  return jaxpr.out_avals
+
+
+cudnn_fusion_p = jax_core.Primitive("cudnn_fusion")
+cudnn_fusion_p.multiple_results = True
+cudnn_fusion_p.def_abstract_eval(_custom_abstract_eval)
+cudnn_fusion_p.def_impl(_cudnn_fusion_impl)
+
+
+def call_cudnn_fusion(f, *args, **kwargs):
+  """Creates a new cudnn_fusion corresponding to calling
+  the given function f with args and kwargs."""
+  jaxpr, out_shapes = jax.make_jaxpr(
+    functools.partial(f, **kwargs), return_shape=True
+  )(*args)
+  flat_args = jax.tree.leaves(args)
+  out_tree = jax.tree.structure(out_shapes)
+  out_flat = cudnn_fusion_p.bind(*flat_args, name=f.__name__, jaxpr=jaxpr)
+  return jax.tree.unflatten(out_tree, out_flat)
+
+
+def _cudnn_fusion_stablehlo_lowering(
+  ctx,
+  *args,
+  name,
+  jaxpr,
+):
+  """Make cudnn_fusion which calls the implementation function.
+  Currently this leaks a CallOp since we're using the `core_call_lowering`
+  function, but this should get cleaned up by DCE easily.
+  """
+  impl = mlir.core_call_lowering(
+    ctx, *args, name=name + ".impl", call_jaxpr=jaxpr
+  )
+  call_op = impl[0].owner
+  called_fn = call_op.attributes["callee"]
+  cudnn_fusion = hlo.CustomCallOp(
+    [r.type for r in call_op.results],
+    call_op.operands,
+    call_target_name="__cudnn$fusion",
+    called_computations=ir.ArrayAttr.get([called_fn]),
+  )
+  return cudnn_fusion.results
+
+
+mlir.register_lowering(
+    cudnn_fusion_p, _cudnn_fusion_stablehlo_lowering, platform="cuda"
+  )
+
+
+def cudnn_fusion(f):
+  """Makes a function become a cuDNN kernel. Relies on XLA's handling of
+  custom fusions with __cudnn$fusion backend. Currently limited to GEMM
+  fusions. For example - batch matmul with mixed types and addition:
+
+  @cudnn_fusion
+  def fn(x, y, z):
+      return jnp.float32(jax.lax.batch_matmul(jnp.bfloat16(x), y)) + z
+  """
+  return functools.partial(call_cudnn_fusion, f)

tests/BUILD

@@ -1523,6 +1523,20 @@ py_test(
     ],
 )
 
+jax_test(
+    name = "cudnn_fusion_test",
+    srcs = ["cudnn_fusion_test.py"],
+    disable_backends = [
+        "cpu",
+        "tpu",
+    ],
+    enable_configs = [
+        "gpu_a100",
+        "gpu_h100",
+    ],
+    tags = ["multiaccelerator"],
+)
+
 exports_files(
     [
         "api_test.py",

tests/cudnn_fusion_test.py

@@ -0,0 +1,69 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from absl.testing import absltest, parameterized
+from unittest import SkipTest
+from jax._src import test_util as jtu
+import jax
+import jax.numpy as jnp
+from jax._src.cudnn import cudnn_fusion
+
+
+jax.config.parse_flags_with_absl()
+
+
+class CudnnFusionTest(jtu.JaxTestCase):
+  def setUp(self):
+    if (not jtu.test_device_matches(["cuda"]) or
+        not jtu.is_cuda_compute_capability_at_least("8.0")):
+      self.skipTest("Only works on >= sm80 GPUs")
+    super().setUp()
+
+  @parameterized.parameters(["", "pmap"])
+  @jtu.run_on_devices("cuda")
+  def test_cudnn_fusion(self, mode):
+    batch_size = 2
+    if mode == "pmap" and jax.device_count() < batch_size:
+        raise SkipTest("pmap test requires 2 GPUs")
+
+    @cudnn_fusion
+    def comp1(x, y, z):
+        return jnp.float32(jax.lax.batch_matmul(jnp.bfloat16(x), y)) + z
+
+    k = jax.random.key(0)
+    s = batch_size, 16, 16
+    x = jnp.int8(jax.random.normal(k, shape=s))
+    y = jnp.bfloat16(jax.random.normal(k, shape=s))
+    z = jnp.float32(jax.random.normal(k, shape=s))
+
+    fn = jax.pmap(comp1) if mode == "pmap" else comp1
+    jitted = jax.jit(comp1)
+    lowered = jitted.lower(x, y, z)
+    stablehlo = lowered.as_text("stablehlo")
+    self.assertIn("func.func private @comp1", stablehlo)
+    self.assertIn("__cudnn$fusion", stablehlo)
+
+    hlo = lowered.as_text("hlo")
+    self.assertIn('custom_call_target="__cudnn$fusion"', hlo)
+    self.assertIn("called_computations=", hlo)
+
+    hlo_after_opt = lowered.compile().as_text()
+    self.assertIn("kind=kCustom", hlo_after_opt)
+    self.assertIn("plan_id", hlo_after_opt)
+
+    self.assertAllClose(jitted(x, y, z), fn(x, y, z))
+
+
+if __name__ == '__main__':
+  absltest.main(testLoader=jtu.JaxTestLoader())