Repository: google/jax
Pull request #21484 (2024-05-29 13:57:55)

Add a helper function for building custom call lowering rules
This function provides sensible defaults for custom call lowering rules with the goal of reducing the amount of boilerplate required for implementing custom calls. While all of the behavior can be overridden, this will be particularly useful when (a) the inputs and outputs are in row-major order because then the layouts and types can be evaluated from the context avals, and/or (b) the custom call supports `api_version=4` attributes provided via a `backend_config` dictionary.

There are a number of lowering rules in the core library that could be rewritten using this helper, but I think it probably wouldn't make sense to port those as part of this PR.

For now the tests only check that lowering succeeds for parameters/attributes of different types, but it might be useful to check that the output is appropriate. I haven't figured out a reasonable way to do that yet - thoughts would be appreciated!

@superbobry — I'd love to hear your feedback. Since this is closely based on code you wrote previously, I added you as a co-author on the commit, but I'm happy to remove that if you'd prefer.

dfm: Thanks for the reviews, @superbobry! I think this is getting close to ready. Here are the main updates:

1. After conversations with @superbobry and @froystig, I agreed that it makes sense for this to live somewhere under `jax.extend`. In particular, I think it actually makes the most sense for this to be in `jax.extend.ffi`, and then we can focus specifically on supporting the new XLA FFI interface, rather than general custom calls. With this in mind, I've renamed the function to `jax.extend.ffi.ffi_lowering`.
2. I updated the docstring to omit any discussion of primitives (or binding) for now. As we settle on a public API for primitives, this could be updated.

There's one inline question about supporting dynamic shapes out of the box, but I'm not very familiar with that use case. Otherwise, I think this is good to go!

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d47f59253..9a9beb412 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,10 +8,7 @@ Remember to align the itemized text with the first line of an item within a list
 
 ## jax 0.4.29
 
-* Changes
-  * We anticipate that this will be the last release of JAX and jaxlib
-    supporting a monolithic CUDA jaxlib. Future releases will use the CUDA
-    plugin jaxlib (e.g. `pip install jax[cuda12]`).
+* Breaking changes
   * JAX now requires ml_dtypes version 0.4.0 or newer.
 
 * Deprecations
@@ -37,11 +34,8 @@ Remember to align the itemized text with the first line of an item within a list
 ## jaxlib 0.4.29
 
 * Bug fixes
-  * Fixed a bug where XLA sharded some concatenation operations incorrectly,
+  * Fixes a bug where XLA sharded some concatenation operations incorrectly,
     which manifested as an incorrect output for cumulative reductions (#21403).
-  * Fixed a bug where XLA:CPU miscompiled certain matmul fusions
-    (https://github.com/openxla/xla/pull/13301).
-  * Fixes a compiler crash on GPU (https://github.com/google/jax/issues/21396).
 
 ## jax 0.4.28 (May 9, 2024)
 
diff --git a/jax/_src/extend/ffi.py b/jax/_src/extend/ffi.py
index c73e86058..479046b84 100644
--- a/jax/_src/extend/ffi.py
+++ b/jax/_src/extend/ffi.py
@@ -16,8 +16,16 @@ from __future__ import annotations
 
 import os
 import ctypes
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any
 
+import numpy as np
+
+from jax._src import dtypes
+from jax._src.interpreters import mlir
 from jax._src.lib import jaxlib
+from jax._src.lib.mlir import ir
+from jax._src.typing import DimSize
 
 
 def pycapsule(funcptr):
@@ -59,3 +67,77 @@ def include_dir() -> str:
   """Get the path to the directory containing header files bundled with jaxlib"""
   jaxlib_dir = os.path.dirname(os.path.abspath(jaxlib.__file__))
   return os.path.join(jaxlib_dir, "include")
+
+
+def ffi_lowering(
+    call_target_name: str,
+    *,
+    operand_layouts: Sequence[Sequence[DimSize]] | None = None,
+    result_layouts: Sequence[Sequence[DimSize]] | None = None,
+    backend_config: Mapping[str, ir.Attribute] | None = None,
+    **lowering_args: Any
+) -> mlir.LoweringRule:
+  """Build a lowering rule for an foreign function interface (FFI) target.
+
+  By default, this lowering rule can use the input and output abstract values to
+  compute the input and output types and shapes for the custom call, assuming
+  row-major layouts.
+
+  If keyword arguments are passed to the lowering rule, these are treated as
+  attributes, and added to `backend_config`.
+
+  Args:
+    call_target_name: The name of the custom call target.
+    operand_layouts: A sequence of layouts (dimension orders) for each operand.
+      By default, the operands are assumed to be row-major.
+    result_layouts: A sequence of layouts (dimension orders) for each result.
+      By default, the results are assumed to be row-major.
+    backend_config: Configuration data for the custom call. Any keyword
+      arguments passed to the lowering rule will added to this dictionary.
+    lowering_args: Any other arguments to :func:`mlir.custom_call` will also be
+      passed through if provided as extra arguments to this function.
+  """
+
+  def _lowering(
+    ctx: mlir.LoweringRuleContext, *operands: ir.Value, **params: Any
+  ) -> Sequence[ir.Value | Sequence[ir.Value]]:
+    kwargs = dict(lowering_args)
+    kwargs.setdefault("api_version", 4)
+    kwargs["backend_config"] = dict(
+      backend_config or {}, **{k: _ir_attribute(v) for k, v in params.items()})
+    if "result_types" not in kwargs:
+      kwargs["result_types"] = [mlir.aval_to_ir_type(aval) for aval in ctx.avals_out]
+    if operand_layouts is None:
+      kwargs["operand_layouts"] = _default_layouts(aval.shape for aval in ctx.avals_in)  # pytype: disable=attribute-error
+    if result_layouts is None:
+      kwargs["result_layouts"] = _default_layouts(aval.shape for aval in ctx.avals_out)
+
+    return mlir.custom_call(call_target_name, operands=operands, **kwargs).results  # type: ignore
+
+  return _lowering
+
+
+def _default_layouts(shapes: Iterable[Sequence[DimSize]]) -> list[list[DimSize]]:
+  return [list(reversed(range(len(shape)))) for shape in shapes]
+
+
+def _ir_attribute(obj: Any) -> ir.Attribute:
+  # TODO(dfm): Similar functions exist in Pallas and Mosaic GPU. Perhaps these
+  # could be consolidated into mlir or similar.
+  if isinstance(obj, str):
+    return ir.StringAttr.get(obj)
+  elif isinstance(obj, bool):
+    return ir.BoolAttr.get(obj)
+  elif isinstance(obj, int):
+    return mlir.i64_attr(obj)
+  elif isinstance(obj, float):
+    return ir.FloatAttr.get_f64(obj)
+  elif hasattr(obj, "dtype"):
+    if not (dtypes.is_python_scalar(obj) or np.isscalar(obj)):
+      raise TypeError("Only scalar attributes are supported")
+    mlir_type = mlir.dtype_to_ir_type(obj.dtype)
+    if isinstance(mlir_type, ir.IntegerType):
+      return ir.IntegerAttr.get(mlir_type, obj)
+    elif isinstance(mlir_type, ir.FloatType):
+      return ir.FloatAttr.get(mlir_type, obj)
+  raise TypeError(f"Unsupported attribute type: {type(obj)}")
diff --git a/jax/extend/ffi.py b/jax/extend/ffi.py
index 12db44f09..565b37cbb 100644
--- a/jax/extend/ffi.py
+++ b/jax/extend/ffi.py
@@ -16,6 +16,7 @@
 # See PEP 484 & https://github.com/google/jax/issues/7570
 
 from jax._src.extend.ffi import (
-  include_dir as include_dir,
-  pycapsule as pycapsule,
+    ffi_lowering as ffi_lowering,
+    include_dir as include_dir,
+    pycapsule as pycapsule,
 )


Test patch:
diff --git a/tests/extend_test.py b/tests/extend_test.py
index 86133c02b..8aee24224 100644
--- a/tests/extend_test.py
+++ b/tests/extend_test.py
@@ -14,7 +14,9 @@
 
 import os
 import unittest
-from absl.testing import absltest
+
+import numpy as np
+from absl.testing import absltest, parameterized
 
 import jax
 import jax.extend as jex
@@ -22,12 +24,15 @@ import jax.numpy as jnp
 
 from jax._src import abstract_arrays
 from jax._src import api
+from jax._src import core
 from jax._src import linear_util
 from jax._src import prng
 from jax._src import test_util as jtu
+from jax._src.interpreters import mlir
+from jax._src.lib.mlir import ir
+from jax._src.extend import ffi
 from jax._src.lib import xla_extension_version
 
-
 jax.config.parse_flags_with_absl()
 
 
@@ -94,6 +99,26 @@ class FfiTest(jtu.JaxTestCase):
     for header in ["c_api.h", "api.h", "ffi.h"]:
       self.assertTrue(os.path.exists(os.path.join(base_dir, header)))
 
+  @parameterized.parameters(
+      [True, int(1), float(5.0),
+       np.int32(-5), np.float32(0.5)])
+  def testIrAttribute(sel, value):
+    with mlir.make_ir_context(), ir.Location.unknown():
+      const = mlir.ir_constant(value)
+      attr = ffi._ir_attribute(value)
+      assert const.type.element_type == attr.type
+
+  @parameterized.parameters([True, 1, 5.0, "param", np.float32(0.5)])
+  def testParams(self, param):
+    prim = core.Primitive("test_ffi")
+    prim.def_abstract_eval(lambda *args, **kwargs: args[0])
+    mlir.register_lowering(prim, jex.ffi.ffi_lowering("test_ffi"))
+
+    # TODO(dfm): Currently testing that lowering works with different types of
+    # parameters, but we should probably actually check the emitted HLO.
+    func = jax.jit(lambda *args: prim.bind(*args, param=param))
+    func.lower(jnp.linspace(0, 5, 10))
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())