Import NVIDIA/warp 1.13.0 from GitHub.

PiperOrigin-RevId: 923878216 Change-Id: If9df5e0672e956eb62874fb6f4f74ba3a02977b3
2026-05-30 04:47:29 -07:00
parent 3a6e0ee9b3
commit b9c7a4b81a
5 changed files with 29 additions and 26 deletions
@@ -179,8 +179,8 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
    wp_field.shape = wp_field.shape[1:]
  # warp scalars
  wp_dtype = type(wp_field)
-  if wp_dtype in wp.types.warp_type_to_np_dtype:
-    return wp.types.warp_type_to_np_dtype[wp_dtype](wp_field)
+  if wp_dtype in wp._src.types.warp_type_to_np_dtype:
+    return wp.dtype_to_numpy(wp_dtype)(wp_field)

  # warp arrays
  if isinstance(wp_field, wp.array):
@@ -131,13 +131,13 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):

    # Launch the kernel.
    wp._src.context.runtime.core.wp_cuda_launch_kernel(
-        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream
+        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream, None
    )


 def _create_jax_warp_primitive():
    import jax  # noqa: PLC0415
-    from jax._src.interpreters import batching  # noqa: PLC0415
+    from jax.interpreters import batching  # noqa: PLC0415
    from jax.interpreters import mlir  # noqa: PLC0415
    from jax.interpreters.mlir import ir  # noqa: PLC0415
    from jaxlib.hlo_helpers import custom_call  # noqa: PLC0415
@@ -170,9 +170,9 @@ def _create_jax_warp_primitive():
        # Figure out the number of outputs.
        wp_kernel = _registered_kernels[params["kernel"]]
        output_count = len(wp_kernel.adj.args) - len(args)
-        shape, dim = next((a.shape, d) for a, d in zip(args, dims) if d is not None)
+        shape, dim = next((a.shape, d) for a, d in zip(args, dims, strict=True) if d is not None)
        size = shape[dim]
-        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims)]
+        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims, strict=True)]
        # Create the batched primitive.
        return _jax_warp_p.bind(*args, **params), [dims[0]] * output_count

@@ -211,6 +211,7 @@ def _create_jax_warp_primitive():
    def base_type_to_jax_ir(warp_dtype):
        warp_to_jax_dict = {
            wp.float16: ir.F16Type.get(),
+            wp.bfloat16: ir.BF16Type.get(),
            wp.float32: ir.F32Type.get(),
            wp.float64: ir.F64Type.get(),
            wp.int8: ir.IntegerType.get_signless(8),
@@ -232,6 +233,7 @@ def _create_jax_warp_primitive():
    def base_type_is_compatible(warp_type, jax_ir_type):
        jax_ir_to_warp = {
            "f16": wp.float16,
+            "bf16": wp.bfloat16,
            "f32": wp.float32,
            "f64": wp.float64,
            "i8": wp.int8,
@@ -324,7 +326,7 @@ def _create_jax_warp_primitive():
        # Figure out the types and shapes of the input arrays.
        arg_strings = []
        operand_layouts = []
-        for actual, warg in zip(args, wp_kernel.adj.args):
+        for actual, warg in zip(args, wp_kernel.adj.args, strict=False):
            wtype = warg.type
            rtt = ir.RankedTensorType(actual.type)

@@ -8,8 +8,8 @@ import ctypes
 import inspect
 import threading
 import traceback
+from collections.abc import Callable
 from enum import IntEnum
-from typing import Callable

 import jax

@@ -25,6 +25,7 @@ from warp._src.types import (
    type_size_in_bytes,
    type_to_warp,
 )
+from warp._src.utils import warn

 from .xla_ffi import *

@@ -249,19 +250,18 @@ class FfiKernel:
                # check dtype
                if input_value.dtype != input_arg.jax_scalar_type:
                    raise TypeError(
-                            f"Invalid data type for array argument '{input_arg.name}',"
-                            f" expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
+                        f"Invalid data type for array argument '{input_arg.name}', expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
                    )
                # check ndim
                if input_value.ndim != input_arg.jax_ndim:
                    raise TypeError(
-                            f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
+                        f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
                    )
                # check inner dims
                for d in range(input_arg.dtype_ndim):
                    if input_value.shape[input_arg.type.ndim + d] != input_arg.dtype_shape[d]:
                        raise TypeError(
-                                f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
+                            f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
                        )
            else:
                # make sure scalar is not a traced variable, should be static
@@ -455,6 +455,7 @@ class FfiKernel:
                        hooks.forward_smem_bytes,
                        kernel_params,
                        stream,
+                        None,  # apic_info
                    )
                else:
                    wp._src.context.runtime.core.wp_cpu_launch_kernel(
@@ -467,9 +468,7 @@ class FfiKernel:
        except Exception as e:
            print(traceback.format_exc())
            return create_ffi_error(
-                call_frame.contents.api,
-                XLA_FFI_Error_Code.UNKNOWN,
-                f"FFI callback error: {type(e).__name__}: {e}",
+                call_frame.contents.api, XLA_FFI_Error_Code.UNKNOWN, f"FFI callback error: {type(e).__name__}: {e}"
            )


@@ -617,7 +616,9 @@ class FfiCallable:
        self.input_output_aliases = input_output_aliases

        # register the callback
-        FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
+        FFI_CCALLFUNC = ctypes.CFUNCTYPE(
+                ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame)
+        )

        self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
        ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
@@ -1071,7 +1072,7 @@ class FfiCallable:
                        call_desc.capture = capture

                    else:
-                        # not capturing or on CPU
+                        # not capturing
                        self.func(*arg_list)

        except Exception as e:
@@ -1359,7 +1360,7 @@ def jax_kernel(
                try:
                    gi.zero_()
                except Exception as e:
-                    wp.utils.warn(f"Failed to zero gradient array: {e}", stacklevel=2)
+                    warn(f"Failed to zero gradient array: {e}", stacklevel=2)
                    raise e

        # NOTE: We cannot use a passed launch_dims here, the backward rule doesn't receive it (and it could be wrong under pmap/vmap).
@@ -1445,7 +1446,7 @@ def jax_kernel(
        non_static_inputs, output_vals_tuple = residuals

        input_vals = list(non_static_inputs)
-        for i, v in zip(static_args, nondiff_vals):
+        for i, v in zip(static_args, nondiff_vals, strict=True):
            input_vals.insert(i, v)

        # Normalize grad outputs and handle nested containers (e.g., single tuple for multi-output)
@@ -1461,7 +1462,7 @@ def jax_kernel(

        out_dims_map = {}
        param_ann = {p.name: p.annotation for p in parameters[:num_inputs]}
-        for name, val in zip(differentiable_input_names, non_static_inputs):
+        for name, val in zip(differentiable_input_names, non_static_inputs, strict=True):
            ann = param_ann.get(name)
            if ann is None:
                continue