Merge pull request jax-ml#25041 from dfm:ffi-example-refactor

PiperOrigin-RevId: 700093685

Author: Google-ML-Automation

examples/ffi/CMakeLists.txt

@@ -15,8 +15,7 @@ find_package(nanobind CONFIG REQUIRED)
 set(
   JAX_FFI_EXAMPLE_PROJECTS
   "rms_norm"
-  "attrs"
-  "counter"
+  "cpu_examples"
 )
 
 foreach(PROJECT ${JAX_FFI_EXAMPLE_PROJECTS})
@@ -27,9 +26,9 @@ endforeach()
 
 if(JAX_FFI_EXAMPLE_ENABLE_CUDA)
   enable_language(CUDA)
-  add_library(_cuda_e2e SHARED "src/jax_ffi_example/cuda_e2e.cu")
-  set_target_properties(_cuda_e2e PROPERTIES POSITION_INDEPENDENT_CODE ON
-                                             CUDA_STANDARD 17)
-  target_include_directories(_cuda_e2e PUBLIC ${XLA_DIR})
-  install(TARGETS _cuda_e2e LIBRARY DESTINATION ${SKBUILD_PROJECT_NAME})
+  add_library(_cuda_examples SHARED "src/jax_ffi_example/cuda_examples.cu")
+  set_target_properties(_cuda_examples PROPERTIES POSITION_INDEPENDENT_CODE ON
+                                                  CUDA_STANDARD 17)
+  target_include_directories(_cuda_examples PUBLIC ${XLA_DIR})
+  install(TARGETS _cuda_examples LIBRARY DESTINATION ${SKBUILD_PROJECT_NAME})
 endif()

examples/ffi/README.md

@@ -11,18 +11,19 @@ Within the example project, there are several example calls:
    demonstrates the most basic use of the FFI. It also includes customization of
    behavior under automatic differentiation using `jax.custom_vjp`.
 
-2. `counter`: This example demonstrates a common pattern for how an FFI call can
-   use global cache to maintain state between calls. This pattern is useful when
-   an FFI call requires an expensive initialization step which shouldn't be
-   run on every execution, or if there is other shared state that could be
-   reused between calls. In this simple example we just count the number of
-   times the call was executed.
+2. `cpu_examples`: This submodule includes several smaller examples:
 
-3. `attrs`: An example demonstrating the different ways that attributes can be
-   passed to the FFI. For example, we can pass arrays, variadic attributes, and
-   user-defined types. Full support of user-defined types isn't yet supported
-   by XLA, so that example will be added in the future.
+   * `counter`: This example demonstrates a common pattern for how an FFI call
+     can use global cache to maintain state between calls. This pattern is
+     useful when an FFI call requires an expensive initialization step which
+     shouldn't be run on every execution, or if there is other shared state
+     that could be reused between calls. In this simple example we just count
+     the number of times the call was executed.
+   * `attrs`: An example demonstrating the different ways that attributes can be
+     passed to the FFI. For example, we can pass arrays, variadic attributes,
+     and user-defined types. Full support of user-defined types isn't yet
+     supported by XLA, so that example will be added in the future.
 
-4. `cuda_e2e`: An end-to-end example demonstrating the use of the JAX FFI with
-   CUDA. The specifics of the kernels are not very important, but the general
-   structure, and packaging of the extension are useful for testing.
+3. `cuda_examples`: An end-to-end example demonstrating the use of the JAX FFI
+   with CUDA. The specifics of the kernels are not very important, but the
+   general structure, and packaging of the extension are useful for testing.

examples/ffi/src/jax_ffi_example/counter.cc

@@ -14,13 +14,27 @@ limitations under the License.
 ==============================================================================*/
 
 #include <cstdint>
+#include <mutex>
+#include <string_view>
+#include <unordered_map>
 
 #include "nanobind/nanobind.h"
 #include "xla/ffi/api/ffi.h"
 
 namespace nb = nanobind;
 namespace ffi = xla::ffi;
 
+// ----------
+// Attributes
+// ----------
+//
+// An example demonstrating the different ways that attributes can be passed to
+// the FFI.
+//
+// For example, we can pass arrays, variadic attributes, and user-defined types.
+// Full support of user-defined types isn't yet supported by XLA, so that
+// example will be added in the future.
+
 ffi::Error ArrayAttrImpl(ffi::Span<const int32_t> array,
                          ffi::ResultBufferR0<ffi::S32> res) {
   int64_t total = 0;
@@ -54,13 +68,52 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(DictionaryAttr, DictionaryAttrImpl,
                                   .Ret<ffi::BufferR0<ffi::S32>>()
                                   .Ret<ffi::BufferR0<ffi::S32>>());
 
-NB_MODULE(_attrs, m) {
+// -------
+// Counter
+// -------
+//
+// An example demonstrating how an FFI call can maintain "state" between calls
+//
+// In this case, the ``Counter`` call simply accumulates the number of times it
+// was executed, but this pattern can also be used for more advanced use cases.
+// For example, this pattern is used in jaxlib for:
+//
+// 1. The GPU solver linear algebra kernels which require an expensive "handler"
+//    initialization, and
+// 2. The ``triton_call`` function which caches the compiled triton modules
+//    after their first use.
+
+ffi::Error CounterImpl(int64_t index, ffi::ResultBufferR0<ffi::S32> out) {
+  static std::mutex mutex;
+  static auto &cache = *new std::unordered_map<int64_t, int32_t>();
+  {
+    const std::lock_guard<std::mutex> lock(mutex);
+    auto it = cache.find(index);
+    if (it != cache.end()) {
+      out->typed_data()[0] = ++it->second;
+    } else {
+      cache.insert({index, 0});
+      out->typed_data()[0] = 0;
+    }
+  }
+  return ffi::Error::Success();
+}
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    Counter, CounterImpl,
+    ffi::Ffi::Bind().Attr<int64_t>("index").Ret<ffi::BufferR0<ffi::S32>>());
+
+// Boilerplate for exposing handlers to Python
+NB_MODULE(_cpu_examples, m) {
   m.def("registrations", []() {
     nb::dict registrations;
     registrations["array_attr"] =
         nb::capsule(reinterpret_cast<void *>(ArrayAttr));
     registrations["dictionary_attr"] =
         nb::capsule(reinterpret_cast<void *>(DictionaryAttr));
+
+    registrations["counter"] = nb::capsule(reinterpret_cast<void *>(Counter));
+
     return registrations;
   });
 }

examples/ffi/src/jax_ffi_example/counter.py

@@ -12,22 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""An example demonstrating the different ways that attributes can be passed to
-the FFI.
-
-For example, we can pass arrays, variadic attributes, and user-defined types.
-Full support of user-defined types isn't yet supported by XLA, so that example
-will be added in the future.
-"""
-
 import numpy as np
 
 import jax
 import jax.extend as jex
 
-from jax_ffi_example import _attrs
+from jax_ffi_example import _cpu_examples
 
-for name, target in _attrs.registrations().items():
+for name, target in _cpu_examples.registrations().items():
   jex.ffi.register_ffi_target(name, target)
 
 
@@ -43,3 +35,8 @@ def dictionary_attr(**kwargs):
       "dictionary_attr",
       (jax.ShapeDtypeStruct((), np.int32), jax.ShapeDtypeStruct((), np.int32)),
   )(**kwargs)
+
+
+def counter(index):
+  return jex.ffi.ffi_call(
+    "counter", jax.ShapeDtypeStruct((), jax.numpy.int32))(index=int(index))

examples/ffi/src/jax_ffi_example/attrs.cc renamed to examples/ffi/src/jax_ffi_example/cpu_examples.cc

@@ -27,7 +27,7 @@
 import jax.extend as jex
 
 # Load the shared library with the FFI target definitions
-SHARED_LIBRARY = os.path.join(os.path.dirname(__file__), "lib_cuda_e2e.so")
+SHARED_LIBRARY = os.path.join(os.path.dirname(__file__), "lib_cuda_examples.so")
 library = ctypes.cdll.LoadLibrary(SHARED_LIBRARY)
 
 jex.ffi.register_ffi_target("foo-fwd", jex.ffi.pycapsule(library.FooFwd),