warp-lang 1.7.2rc1__py3-none-macosx_10_13_universal2.whl → 1.8.1__py3-none-macosx_10_13_universal2.whl

warp/jax_experimental/ffi.py

@@ -16,7 +16,8 @@
 import ctypes
 import threading
 import traceback
-from typing import Callable
+from enum import IntEnum
+from typing import Callable, Optional
 
 import jax
 
@@ -28,10 +29,17 @@ from warp.types import array_t, launch_bounds_t, strides_from_shape, type_to_war
 from .xla_ffi import *
 
 
+class GraphMode(IntEnum):
+    NONE = 0  # don't capture a graph
+    JAX = 1  # let JAX capture a graph
+    WARP = 2  # let Warp capture a graph
+
+
 class FfiArg:
-    def __init__(self, name, type):
+    def __init__(self, name, type, in_out=False):
         self.name = name
         self.type = type
+        self.in_out = in_out
         self.is_array = isinstance(type, wp.array)
 
         if self.is_array:
@@ -65,7 +73,7 @@ class FfiLaunchDesc:
 
 
 class FfiKernel:
-    def __init__(self, kernel, num_outputs, vmap_method, launch_dims, output_dims):
+    def __init__(self, kernel, num_outputs, vmap_method, launch_dims, output_dims, in_out_argnames):
         self.kernel = kernel
         self.name = generate_unique_name(kernel.func)
         self.num_outputs = num_outputs
@@ -76,17 +84,28 @@ class FfiKernel:
         self.launch_id = 0
         self.launch_descriptors = {}
 
+        in_out_argnames_list = in_out_argnames or []
+        in_out_argnames = set(in_out_argnames_list)
+        if len(in_out_argnames_list) != len(in_out_argnames):
+            raise AssertionError("in_out_argnames must not contain duplicate names")
+
         self.num_kernel_args = len(kernel.adj.args)
-        self.num_inputs = self.num_kernel_args - num_outputs
+        self.num_in_out = len(in_out_argnames)
+        self.num_inputs = self.num_kernel_args - num_outputs + self.num_in_out
         if self.num_outputs < 1:
             raise ValueError("At least one output is required")
         if self.num_outputs > self.num_kernel_args:
             raise ValueError("Number of outputs cannot be greater than the number of kernel arguments")
+        if self.num_outputs < self.num_in_out:
+            raise ValueError("Number of outputs cannot be smaller than the number of in_out_argnames")
 
         # process input args
         self.input_args = []
         for i in range(self.num_inputs):
-            arg = FfiArg(kernel.adj.args[i].label, kernel.adj.args[i].type)
+            arg_name = kernel.adj.args[i].label
+            arg = FfiArg(arg_name, kernel.adj.args[i].type, arg_name in in_out_argnames)
+            if arg_name in in_out_argnames:
+                in_out_argnames.remove(arg_name)
             if arg.is_array:
                 # keep track of the first input array argument
                 if self.first_array_arg is None:
@@ -96,11 +115,30 @@ class FfiKernel:
         # process output args
         self.output_args = []
         for i in range(self.num_inputs, self.num_kernel_args):
-            arg = FfiArg(kernel.adj.args[i].label, kernel.adj.args[i].type)
+            arg_name = kernel.adj.args[i].label
+            if arg_name in in_out_argnames:
+                raise AssertionError(
+                    f"Expected an output-only argument for argument {arg_name}."
+                    " in_out arguments should be placed before output-only arguments."
+                )
+            arg = FfiArg(arg_name, kernel.adj.args[i].type, False)
             if not arg.is_array:
                 raise TypeError("All output arguments must be arrays")
             self.output_args.append(arg)
 
+        if in_out_argnames:
+            raise ValueError(f"in_out_argnames: '{in_out_argnames}' did not match any function argument names.")
+
+        # Build input output aliases.
+        out_id = 0
+        input_output_aliases = {}
+        for in_id, arg in enumerate(self.input_args):
+            if not arg.in_out:
+                continue
+            input_output_aliases[in_id] = out_id
+            out_id += 1
+        self.input_output_aliases = input_output_aliases
+
         # register the callback
         FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
         self.callback_func = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame))
@@ -121,6 +159,9 @@ class FfiKernel:
         if vmap_method is None:
             vmap_method = self.vmap_method
 
+        # output types
+        out_types = []
+
         # process inputs
         static_inputs = {}
         for i in range(num_inputs):
@@ -150,6 +191,10 @@ class FfiKernel:
                 # stash the value to be retrieved by callback
                 static_inputs[input_arg.name] = input_arg.type(input_value)
 
+            # append in-out arg to output types
+            if input_arg.in_out:
+                out_types.append(get_jax_output_type(input_arg, input_value.shape))
+
         # launch dimensions
         if launch_dims is None:
             # use the shape of the first input array
@@ -162,8 +207,7 @@ class FfiKernel:
         else:
             launch_dims = tuple(launch_dims)
 
-        # output types
-        out_types = []
+        # output shapes
         if isinstance(output_dims, dict):
             # assume a dictionary of shapes keyed on argument name
             for output_arg in self.output_args:
@@ -185,6 +229,7 @@ class FfiKernel:
             self.name,
             out_types,
             vmap_method=vmap_method,
+            input_output_aliases=self.input_output_aliases,
         )
 
         # ensure the kernel module is loaded before the callback, otherwise graph capture may fail
@@ -238,9 +283,8 @@ class FfiKernel:
 
             arg_refs = []
 
-            # inputs
-            for i in range(num_inputs):
-                input_arg = self.input_args[i]
+            # input and in-out args
+            for i, input_arg in enumerate(self.input_args):
                 if input_arg.is_array:
                     buffer = inputs[i].contents
                     shape = buffer.dims[: input_arg.type.ndim]
@@ -255,10 +299,9 @@ class FfiKernel:
                     kernel_params[i + 1] = ctypes.addressof(arg)
                     arg_refs.append(arg)  # keep a reference
 
-            # outputs
-            for i in range(num_outputs):
-                output_arg = self.output_args[i]
-                buffer = outputs[i].contents
+            # pure output args (skip in-out FFI buffers)
+            for i, output_arg in enumerate(self.output_args):
+                buffer = outputs[i + self.num_in_out].contents
                 shape = buffer.dims[: output_arg.type.ndim]
                 strides = strides_from_shape(shape, output_arg.type.dtype)
                 arg = array_t(buffer.data, 0, output_arg.type.ndim, shape, strides)
@@ -295,30 +338,38 @@ class FfiKernel:
 class FfiCallDesc:
     def __init__(self, static_inputs):
         self.static_inputs = static_inputs
+        self.captures = {}
 
 
 class FfiCallable:
-    def __init__(self, func, num_outputs, graph_compatible, vmap_method, output_dims):
+    def __init__(self, func, num_outputs, graph_mode, vmap_method, output_dims, in_out_argnames):
         self.func = func
         self.name = generate_unique_name(func)
         self.num_outputs = num_outputs
         self.vmap_method = vmap_method
-        self.graph_compatible = graph_compatible
+        self.graph_mode = graph_mode
         self.output_dims = output_dims
         self.first_array_arg = None
-        self.has_static_args = False
         self.call_id = 0
         self.call_descriptors = {}
 
+        in_out_argnames_list = in_out_argnames or []
+        in_out_argnames = set(in_out_argnames_list)
+        if len(in_out_argnames_list) != len(in_out_argnames):
+            raise AssertionError("in_out_argnames must not contain duplicate names")
+
         # get arguments and annotations
         argspec = get_full_arg_spec(func)
 
         num_args = len(argspec.args)
-        self.num_inputs = num_args - num_outputs
+        self.num_in_out = len(in_out_argnames)
+        self.num_inputs = num_args - num_outputs + self.num_in_out
         if self.num_outputs < 1:
             raise ValueError("At least one output is required")
         if self.num_outputs > num_args:
             raise ValueError("Number of outputs cannot be greater than the number of kernel arguments")
+        if self.num_outputs < self.num_in_out:
+            raise ValueError("Number of outputs cannot be smaller than the number of in_out_argnames")
 
         if len(argspec.annotations) < num_args:
             raise RuntimeError(f"Incomplete argument annotations on function {self.name}")
@@ -331,17 +382,43 @@ class FfiCallable:
                 if arg_type is not None:
                     raise TypeError("Function must not return a value")
             else:
-                arg = FfiArg(arg_name, arg_type)
+                arg = FfiArg(arg_name, arg_type, arg_name in in_out_argnames)
+                if arg_name in in_out_argnames:
+                    in_out_argnames.remove(arg_name)
                 if arg.is_array:
                     if arg_idx < self.num_inputs and self.first_array_arg is None:
                         self.first_array_arg = arg_idx
-                else:
-                    self.has_static_args = True
                 self.args.append(arg)
+
+            if arg.in_out and arg_idx >= self.num_inputs:
+                raise AssertionError(
+                    f"Expected an output-only argument for argument {arg_name}."
+                    " in_out arguments should be placed before output-only arguments."
+                )
+
             arg_idx += 1
 
-        self.input_args = self.args[: self.num_inputs]
-        self.output_args = self.args[self.num_inputs :]
+        if in_out_argnames:
+            raise ValueError(f"in_out_argnames: '{in_out_argnames}' did not match any function argument names.")
+
+        self.input_args = self.args[: self.num_inputs]  # includes in-out args
+        self.output_args = self.args[self.num_inputs :]  # pure output args
+
+        # Buffer indices for array arguments in callback.
+        # In-out buffers are the same pointers in the XLA call frame,
+        # so we only include them for inputs and skip them for outputs.
+        self.array_input_indices = [i for i, arg in enumerate(self.input_args) if arg.is_array]
+        self.array_output_indices = list(range(self.num_in_out, self.num_outputs))
+
+        # Build input output aliases.
+        out_id = 0
+        input_output_aliases = {}
+        for in_id, arg in enumerate(self.input_args):
+            if not arg.in_out:
+                continue
+            input_output_aliases[in_id] = out_id
+            out_id += 1
+        self.input_output_aliases = input_output_aliases
 
         # register the callback
         FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
@@ -353,7 +430,9 @@ class FfiCallable:
     def __call__(self, *args, output_dims=None, vmap_method=None):
         num_inputs = len(args)
         if num_inputs != self.num_inputs:
-            raise ValueError(f"Expected {self.num_inputs} inputs, but got {num_inputs}")
+            input_names = ", ".join(arg.name for arg in self.input_args)
+            s = "" if self.num_inputs == 1 else "s"
+            raise ValueError(f"Expected {self.num_inputs} input{s} ({input_names}), but got {num_inputs}")
 
         # default argument fallback
         if vmap_method is None:
@@ -361,6 +440,9 @@ class FfiCallable:
         if output_dims is None:
             output_dims = self.output_dims
 
+        # output types
+        out_types = []
+
         # process inputs
         static_inputs = {}
         for i in range(num_inputs):
@@ -390,12 +472,11 @@ class FfiCallable:
                 # stash the value to be retrieved by callback
                 static_inputs[input_arg.name] = input_arg.type(input_value)
 
-        if output_dims is None and self.first_array_arg is not None:
-            # use the shape of the first input array
-            output_dims = get_warp_shape(self.input_args[self.first_array_arg], args[self.first_array_arg].shape)
+            # append in-out arg to output types
+            if input_arg.in_out:
+                out_types.append(get_jax_output_type(input_arg, input_value.shape))
 
-        # output types
-        out_types = []
+        # output shapes
         if isinstance(output_dims, dict):
             # assume a dictionary of shapes keyed on argument name
             for output_arg in self.output_args:
@@ -405,7 +486,9 @@ class FfiCallable:
                 out_types.append(get_jax_output_type(output_arg, dims))
         else:
             if output_dims is None:
-                raise ValueError("Unable to determine output dimensions")
+                if self.first_array_arg is None:
+                    raise ValueError("Unable to determine output dimensions")
+                output_dims = get_warp_shape(self.input_args[self.first_array_arg], args[self.first_array_arg].shape)
             elif isinstance(output_dims, int):
                 output_dims = (output_dims,)
             # assume same dimensions for all outputs
@@ -416,6 +499,7 @@ class FfiCallable:
             self.name,
             out_types,
             vmap_method=vmap_method,
+            input_output_aliases=self.input_output_aliases,
             # has_side_effect=True,  # force this function to execute even if outputs aren't used
         )
 
@@ -425,22 +509,18 @@ class FfiCallable:
         module = wp.get_module(self.func.__module__)
         module.load(device)
 
-        if self.has_static_args:
-            # save call data to be retrieved by callback
-            call_id = self.call_id
-            self.call_descriptors[call_id] = FfiCallDesc(static_inputs)
-            self.call_id += 1
-            return call(*args, call_id=call_id)
-        else:
-            return call(*args)
+        # save call data to be retrieved by callback
+        call_id = self.call_id
+        self.call_descriptors[call_id] = FfiCallDesc(static_inputs)
+        self.call_id += 1
+        return call(*args, call_id=call_id)
 
     def ffi_callback(self, call_frame):
         try:
-            # TODO Try-catch around the body and return XLA_FFI_Error on error.
-            extension = call_frame.contents.extension_start
             # On the first call, XLA runtime will query the API version and traits
             # metadata using the |extension| field. Let us respond to that query
             # if the metadata extension is present.
+            extension = call_frame.contents.extension_start
             if extension:
                 # Try to set the version metadata.
                 if extension.contents.type == XLA_FFI_Extension_Type.Metadata:
@@ -448,17 +528,20 @@ class FfiCallable:
                     metadata_ext.contents.metadata.contents.api_version.major_version = 0
                     metadata_ext.contents.metadata.contents.api_version.minor_version = 1
                     # Turn on CUDA graphs for this handler.
-                    if self.graph_compatible:
+                    if self.graph_mode is GraphMode.JAX:
                         metadata_ext.contents.metadata.contents.traits = (
                             XLA_FFI_Handler_TraitsBits.COMMAND_BUFFER_COMPATIBLE
                         )
                     return None
 
-            if self.has_static_args:
-                # retrieve call info
-                attrs = decode_attrs(call_frame.contents.attrs)
-                call_id = int(attrs["call_id"])
-                call_desc = self.call_descriptors[call_id]
+            # retrieve call info
+            # NOTE: this assumes that there's only one attribute - call_id (int64).
+            # A more general but slower approach is this:
+            #   attrs = decode_attrs(call_frame.contents.attrs)
+            #   call_id = int(attrs["call_id"])
+            attr = ctypes.cast(call_frame.contents.attrs.attrs[0], ctypes.POINTER(XLA_FFI_Scalar)).contents
+            call_id = ctypes.cast(attr.value, ctypes.POINTER(ctypes.c_int64)).contents.value
+            call_desc = self.call_descriptors[call_id]
 
             num_inputs = call_frame.contents.args.size
             inputs = ctypes.cast(call_frame.contents.args.args, ctypes.POINTER(ctypes.POINTER(XLA_FFI_Buffer)))
@@ -469,16 +552,42 @@ class FfiCallable:
             assert num_inputs == self.num_inputs
             assert num_outputs == self.num_outputs
 
-            device = wp.device_from_jax(get_jax_device())
             cuda_stream = get_stream_from_callframe(call_frame.contents)
+
+            if self.graph_mode == GraphMode.WARP:
+                # check if we already captured an identical call
+                ip = [inputs[i].contents.data for i in self.array_input_indices]
+                op = [outputs[i].contents.data for i in self.array_output_indices]
+                buffer_hash = hash((*ip, *op))
+                capture = call_desc.captures.get(buffer_hash)
+
+                # launch existing graph
+                if capture is not None:
+                    # NOTE: We use the native graph API to avoid overhead with obtaining Stream and Device objects in Python.
+                    # This code should match wp.capture_launch().
+                    graph = capture.graph
+                    if graph.graph_exec is None:
+                        g = ctypes.c_void_p()
+                        if not wp.context.runtime.core.wp_cuda_graph_create_exec(
+                            graph.device.context, cuda_stream, graph.graph, ctypes.byref(g)
+                        ):
+                            raise RuntimeError(f"Graph creation error: {wp.context.runtime.get_error_string()}")
+                        graph.graph_exec = g
+
+                    if not wp.context.runtime.core.wp_cuda_graph_launch(graph.graph_exec, cuda_stream):
+                        raise RuntimeError(f"Graph launch error: {wp.context.runtime.get_error_string()}")
+
+                    # early out
+                    return
+
+            device = wp.device_from_jax(get_jax_device())
             stream = wp.Stream(device, cuda_stream=cuda_stream)
 
             # reconstruct the argument list
             arg_list = []
 
-            # inputs
-            for i in range(num_inputs):
-                arg = self.input_args[i]
+            # input and in-out args
+            for i, arg in enumerate(self.input_args):
                 if arg.is_array:
                     buffer = inputs[i].contents
                     shape = buffer.dims[: buffer.rank - arg.dtype_ndim]
@@ -489,10 +598,9 @@ class FfiCallable:
                     value = call_desc.static_inputs[arg.name]
                     arg_list.append(value)
 
-            # outputs
-            for i in range(num_outputs):
-                arg = self.output_args[i]
-                buffer = outputs[i].contents
+            # pure output args (skip in-out FFI buffers)
+            for i, arg in enumerate(self.output_args):
+                buffer = outputs[i + self.num_in_out].contents
                 shape = buffer.dims[: buffer.rank - arg.dtype_ndim]
                 arr = wp.array(ptr=buffer.data, dtype=arg.type.dtype, shape=shape, device=device)
                 arg_list.append(arr)
@@ -500,9 +608,20 @@ class FfiCallable:
             # call the Python function with reconstructed arguments
             with wp.ScopedStream(stream, sync_enter=False):
                 if stream.is_capturing:
-                    with wp.ScopedCapture(stream=stream, external=True):
+                    # capturing with JAX
+                    with wp.ScopedCapture(external=True) as capture:
+                        self.func(*arg_list)
+                    # keep a reference to the capture object to prevent required modules getting unloaded
+                    call_desc.capture = capture
+                elif self.graph_mode == GraphMode.WARP:
+                    # capturing with WARP
+                    with wp.ScopedCapture() as capture:
                         self.func(*arg_list)
+                    wp.capture_launch(capture.graph)
+                    # keep a reference to the capture object and reuse it with same buffers
+                    call_desc.captures[buffer_hash] = capture
                 else:
+                    # not capturing
                     self.func(*arg_list)
 
         except Exception as e:
@@ -520,7 +639,9 @@ _FFI_KERNEL_REGISTRY: dict[str, FfiKernel] = {}
 _FFI_REGISTRY_LOCK = threading.Lock()
 
 
-def jax_kernel(kernel, num_outputs=1, vmap_method="broadcast_all", launch_dims=None, output_dims=None):
+def jax_kernel(
+    kernel, num_outputs=1, vmap_method="broadcast_all", launch_dims=None, output_dims=None, in_out_argnames=None
+):
     """Create a JAX callback from a Warp kernel.
 
     NOTE: This is an experimental feature under development.
@@ -528,6 +649,7 @@ def jax_kernel(kernel, num_outputs=1, vmap_method="broadcast_all", launch_dims=N
     Args:
         kernel: The Warp kernel to launch.
         num_outputs: Optional. Specify the number of output arguments if greater than 1.
+                     This must include the number of ``in_out_arguments``.
         vmap_method: Optional. String specifying how the callback transforms under ``vmap()``.
                      This argument can also be specified for individual calls.
         launch_dims: Optional. Specify the default kernel launch dimensions. If None, launch
@@ -536,12 +658,13 @@ def jax_kernel(kernel, num_outputs=1, vmap_method="broadcast_all", launch_dims=N
         output_dims: Optional. Specify the default dimensions of output arrays.  If None, output
                      dimensions are inferred from the launch dimensions.
                      This argument can also be specified for individual calls.
+        in_out_argnames: Optional. Names of input-output arguments.
 
     Limitations:
         - All kernel arguments must be contiguous arrays or scalars.
         - Scalars must be static arguments in JAX.
-        - Input arguments are followed by output arguments in the Warp kernel definition.
-        - There must be at least one output argument.
+        - Input and input-output arguments must precede the output arguments in the ``kernel`` definition.
+        - There must be at least one output or input-output argument.
         - Only the CUDA backend is supported.
     """
     key = (
@@ -554,7 +677,7 @@ def jax_kernel(kernel, num_outputs=1, vmap_method="broadcast_all", launch_dims=N
 
     with _FFI_REGISTRY_LOCK:
         if key not in _FFI_KERNEL_REGISTRY:
-            new_kernel = FfiKernel(kernel, num_outputs, vmap_method, launch_dims, output_dims)
+            new_kernel = FfiKernel(kernel, num_outputs, vmap_method, launch_dims, output_dims, in_out_argnames)
             _FFI_KERNEL_REGISTRY[key] = new_kernel
 
     return _FFI_KERNEL_REGISTRY[key]
@@ -563,9 +686,11 @@ def jax_kernel(kernel, num_outputs=1, vmap_method="broadcast_all", launch_dims=N
 def jax_callable(
     func: Callable,
     num_outputs: int = 1,
-    graph_compatible: bool = True,
-    vmap_method: str = "broadcast_all",
+    graph_compatible: Optional[bool] = None,  # deprecated
+    graph_mode: GraphMode = GraphMode.JAX,
+    vmap_method: Optional[str] = "broadcast_all",
     output_dims=None,
+    in_out_argnames=None,
 ):
     """Create a JAX callback from an annotated Python function.
 
@@ -576,31 +701,50 @@ def jax_callable(
     Args:
         func: The Python function to call.
         num_outputs: Optional. Specify the number of output arguments if greater than 1.
+            This must include the number of ``in_out_arguments``.
         graph_compatible: Optional. Whether the function can be called during CUDA graph capture.
+            This argument is deprecated, use ``graph_mode`` instead.
+        graph_mode: Optional. CUDA graph capture mode.
+            ``GraphMode.JAX`` (default): Let JAX capture the graph, which may be used as a subgraph in an enclosing capture.
+            ``GraphMode.WARP``: Let Warp capture the graph. Use this mode when the callable cannot be used as a subraph,
+            such as when the callable uses conditional graph nodes.
+            ``GraphMode.NONE``: Disable graph capture. Use when the callable performs operations that are not legal in a graph,
+            such as host synchronization.
         vmap_method: Optional. String specifying how the callback transforms under ``vmap()``.
             This argument can also be specified for individual calls.
         output_dims: Optional. Specify the default dimensions of output arrays.
             If ``None``, output dimensions are inferred from the launch dimensions.
             This argument can also be specified for individual calls.
+        in_out_argnames: Optional. Names of input-output arguments.
 
     Limitations:
         - All kernel arguments must be contiguous arrays or scalars.
         - Scalars must be static arguments in JAX.
-        - Input arguments are followed by output arguments in the Warp kernel definition.
-        - There must be at least one output argument.
+        - Input and input-output arguments must precede the output arguments in the ``func`` definition.
+        - There must be at least one output or input-output argument.
         - Only the CUDA backend is supported.
     """
+
+    if graph_compatible is not None:
+        wp.utils.warn(
+            "The `graph_compatible` argument is deprecated, use `graph_mode` instead.",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+        if graph_compatible is False:
+            graph_mode = GraphMode.NONE
+
     key = (
         func,
         num_outputs,
-        graph_compatible,
+        graph_mode,
         vmap_method,
         tuple(sorted(output_dims.items())) if output_dims else output_dims,
     )
 
     with _FFI_REGISTRY_LOCK:
         if key not in _FFI_CALLABLE_REGISTRY:
-            new_callable = FfiCallable(func, num_outputs, graph_compatible, vmap_method, output_dims)
+            new_callable = FfiCallable(func, num_outputs, graph_mode, vmap_method, output_dims, in_out_argnames)
             _FFI_CALLABLE_REGISTRY[key] = new_callable
 
     return _FFI_CALLABLE_REGISTRY[key]
@@ -631,7 +775,6 @@ def register_ffi_callback(name: str, func: Callable, graph_compatible: bool = Tr
 
     def ffi_callback(call_frame):
         try:
-            # TODO Try-catch around the body and return XLA_FFI_Error on error.
             extension = call_frame.contents.extension_start
             # On the first call, XLA runtime will query the API version and traits
             # metadata using the |extension| field. Let us respond to that query