warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/examples/interop/example_jax_callable.py

@@ -0,0 +1,146 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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
+#
+# http://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.
+
+###########################################################################
+# Example jax_callable()
+#
+# Examples of calling annotated Python functions from JAX.
+###########################################################################
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+import warp as wp
+from warp.jax_experimental import jax_callable
+
+
+@wp.kernel
+def scale_kernel(a: wp.array(dtype=float), s: float, output: wp.array(dtype=float)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+@wp.kernel
+def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtype=wp.vec2)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+# The Python function to call.
+# Note the argument annotations, just like Warp kernels.
+def scale_func(
+    # inputs
+    a: wp.array(dtype=float),
+    b: wp.array(dtype=wp.vec2),
+    s: float,
+    # outputs
+    c: wp.array(dtype=float),
+    d: wp.array(dtype=wp.vec2),
+):
+    wp.launch(scale_kernel, dim=a.shape, inputs=[a, s], outputs=[c])
+    wp.launch(scale_vec_kernel, dim=b.shape, inputs=[b, s], outputs=[d])
+
+
+@wp.kernel
+def accum_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
+    tid = wp.tid()
+    b[tid] += a[tid]
+
+
+def in_out_func(
+    a: wp.array(dtype=float),  # input only
+    b: wp.array(dtype=float),  # input and output
+    c: wp.array(dtype=float),  # output only
+):
+    wp.launch(scale_kernel, dim=a.size, inputs=[a, 2.0], outputs=[c])
+    wp.launch(accum_kernel, dim=a.size, inputs=[a, b])  # modifies `b`
+
+
+def example1():
+    jax_func = jax_callable(scale_func, num_outputs=2)
+
+    @jax.jit
+    def f():
+        # inputs
+        a = jnp.arange(10, dtype=jnp.float32)
+        b = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # wp.vec2
+        s = 2.0
+
+        # output shapes
+        output_dims = {"c": a.shape, "d": b.shape}
+
+        c, d = jax_func(a, b, s, output_dims=output_dims)
+
+        return c, d
+
+    r1, r2 = f()
+    print(r1)
+    print(r2)
+
+
+def example2():
+    jax_func = jax_callable(scale_func, num_outputs=2)
+
+    # NOTE: scalar arguments must be static compile-time constants
+    @partial(jax.jit, static_argnames=["s"])
+    def f(a, b, s):
+        # output shapes
+        output_dims = {"c": a.shape, "d": b.shape}
+
+        c, d = jax_func(a, b, s, output_dims=output_dims)
+
+        return c, d
+
+    # inputs
+    a = jnp.arange(10, dtype=jnp.float32)
+    b = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # wp.vec2
+    s = 3.0
+
+    r1, r2 = f(a, b, s)
+    print(r1)
+    print(r2)
+
+
+def example3():
+    # Using input-output arguments
+
+    jax_func = jax_callable(in_out_func, num_outputs=2, in_out_argnames=["b"])
+
+    f = jax.jit(jax_func)
+
+    a = jnp.ones(10, dtype=jnp.float32)
+    b = jnp.arange(10, dtype=jnp.float32)
+
+    b, c = f(a, b)
+    print(b)
+    print(c)
+
+
+def main():
+    wp.init()
+    wp.load_module(device=wp.get_device())
+
+    examples = [example1, example2, example3]
+
+    for example in examples:
+        print("\n===========================================================================")
+        print(f"{example.__name__}:")
+        example()
+
+
+if __name__ == "__main__":
+    main()

warp/examples/interop/example_jax_ffi_callback.py

@@ -0,0 +1,132 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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
+#
+# http://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.
+
+###########################################################################
+# Example register_ffi_callback()
+#
+# Examples of calling Python functions from JAX.
+# Target functions must have the form func(inputs, outputs, attrs, ctx).
+###########################################################################
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import warp as wp
+from warp.jax import get_jax_device
+from warp.jax_experimental import register_ffi_callback
+
+
+@wp.kernel
+def scale_kernel(a: wp.array(dtype=float), s: float, output: wp.array(dtype=float)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+@wp.kernel
+def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtype=wp.vec2)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+def example1():
+    # the Python function to call
+    def print_args(inputs, outputs, attrs, ctx):
+        def buffer_to_string(b):
+            return f"{b.dtype}{list(b.shape)} @{b.data:x}"
+
+        print("Inputs:     ", ", ".join([buffer_to_string(b) for b in inputs]))
+        print("Outputs:    ", ", ".join([buffer_to_string(b) for b in outputs]))
+        print("Attributes: ", "".join([f"\n  {k}: {str(v)}" for k, v in attrs.items()]))  # noqa: RUF010
+
+    # register callback
+    register_ffi_callback("print_args", print_args)
+
+    # set up call
+    call = jax.ffi.ffi_call("print_args", jax.ShapeDtypeStruct((1, 2, 3), jnp.int8))
+
+    # call it
+    call(
+        jnp.arange(16),
+        jnp.arange(32.0).reshape((4, 8)),
+        str_attr="hi",
+        f32_attr=np.float32(4.2),
+        dict_attr={"a": 1, "b": 6.4},
+    )
+
+
+def example2():
+    # the Python function to call
+    def warp_func(inputs, outputs, attrs, ctx):
+        # input arrays
+        a = inputs[0]
+        b = inputs[1]
+
+        # scalar attributes
+        s = attrs["scale"]
+
+        # output arrays
+        c = outputs[0]
+        d = outputs[1]
+
+        device = wp.device_from_jax(get_jax_device())
+        stream = wp.Stream(device, cuda_stream=ctx.stream)
+
+        with wp.ScopedStream(stream):
+            # launch with arrays of scalars
+            wp.launch(scale_kernel, dim=a.shape, inputs=[a, s], outputs=[c])
+
+            # launch with arrays of vec2
+            # NOTE: the input shapes are from JAX arrays, we need to strip the inner dimension for vec2 arrays
+            wp.launch(scale_vec_kernel, dim=b.shape[0], inputs=[b, s], outputs=[d])
+
+    # register callback
+    register_ffi_callback("warp_func", warp_func)
+
+    n = 10
+
+    # inputs
+    a = jnp.arange(n, dtype=jnp.float32)
+    b = jnp.arange(n, dtype=jnp.float32).reshape((n // 2, 2))  # array of wp.vec2
+    s = 2.0
+
+    # set up call
+    out_types = [
+        jax.ShapeDtypeStruct(a.shape, jnp.float32),
+        jax.ShapeDtypeStruct(b.shape, jnp.float32),  # array of wp.vec2
+    ]
+    call = jax.ffi.ffi_call("warp_func", out_types)
+
+    # call it
+    c, d = call(a, b, scale=s)
+
+    print(c)
+    print(d)
+
+
+def main():
+    wp.init()
+    wp.load_module(device=wp.get_device())
+
+    examples = [example1, example2]
+
+    for example in examples:
+        print("\n===========================================================================")
+        print(f"{example.__name__}:")
+        example()
+
+
+if __name__ == "__main__":
+    main()

warp/examples/interop/example_jax_kernel.py

@@ -0,0 +1,232 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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
+#
+# http://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.
+
+###########################################################################
+# Example jax_kernel()
+#
+# Examples of calling a Warp kernel from JAX.
+###########################################################################
+
+import math
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+import warp as wp
+from warp.jax_experimental import jax_kernel
+
+
+@wp.kernel
+def add_kernel(a: wp.array(dtype=int), b: wp.array(dtype=int), output: wp.array(dtype=int)):
+    tid = wp.tid()
+    output[tid] = a[tid] + b[tid]
+
+
+@wp.kernel
+def sincos_kernel(angle: wp.array(dtype=float), sin_out: wp.array(dtype=float), cos_out: wp.array(dtype=float)):
+    tid = wp.tid()
+    sin_out[tid] = wp.sin(angle[tid])
+    cos_out[tid] = wp.cos(angle[tid])
+
+
+@wp.kernel
+def diagonal_kernel(output: wp.array(dtype=wp.mat33)):
+    tid = wp.tid()
+    d = float(tid + 1)
+    output[tid] = wp.mat33(d, 0.0, 0.0, 0.0, d * 2.0, 0.0, 0.0, 0.0, d * 3.0)
+
+
+@wp.kernel
+def matmul_kernel(
+    a: wp.array2d(dtype=float),  # NxK
+    b: wp.array2d(dtype=float),  # KxM
+    c: wp.array2d(dtype=float),  # NxM
+):
+    # launch dims should be (N, M)
+    i, j = wp.tid()
+    N = a.shape[0]
+    K = a.shape[1]
+    M = b.shape[1]
+    if i < N and j < M:
+        s = wp.float32(0)
+        for k in range(K):
+            s += a[i, k] * b[k, j]
+        c[i, j] = s
+
+
+@wp.kernel
+def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtype=wp.vec2)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+@wp.kernel
+def in_out_kernel(
+    a: wp.array(dtype=float),  # input only
+    b: wp.array(dtype=float),  # input and output
+    c: wp.array(dtype=float),  # output only
+):
+    tid = wp.tid()
+    b[tid] += a[tid]
+    c[tid] = 2.0 * a[tid]
+
+
+def example1():
+    # two inputs and one output
+    jax_add = jax_kernel(add_kernel)
+
+    @jax.jit
+    def f():
+        n = 10
+        a = jnp.arange(n, dtype=jnp.int32)
+        b = jnp.ones(n, dtype=jnp.int32)
+        return jax_add(a, b)
+
+    print(f())
+
+
+def example2():
+    # one input and two outputs
+    jax_sincos = jax_kernel(sincos_kernel, num_outputs=2)
+
+    @jax.jit
+    def f():
+        n = 32
+        a = jnp.linspace(0, 2 * math.pi, n)
+        return jax_sincos(a)
+
+    s, c = f()
+    print(s)
+    print()
+    print(c)
+
+
+def example3():
+    # multiply vectors by scalar
+    jax_scale_vec = jax_kernel(scale_vec_kernel)
+
+    @jax.jit
+    def f():
+        a = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # array of vec2
+        s = 2.0
+        return jax_scale_vec(a, s)
+
+    b = f()
+    print(b)
+
+
+def example4():
+    # multiply vectors by scalar (static arg)
+    jax_scale_vec = jax_kernel(scale_vec_kernel)
+
+    # NOTE: scalar arguments must be static compile-time constants
+    @partial(jax.jit, static_argnames=["s"])
+    def f(a, s):
+        return jax_scale_vec(a, s)
+
+    a = jnp.arange(10, dtype=jnp.float32).reshape((5, 2))  # array of vec2
+    s = 3.0
+
+    b = f(a, s)
+    print(b)
+
+
+def example5():
+    N, M, K = 3, 4, 2
+
+    # specify default launch dims
+    jax_matmul = jax_kernel(matmul_kernel, launch_dims=(N, M))
+
+    @jax.jit
+    def f():
+        a = jnp.full((N, K), 2, dtype=jnp.float32)
+        b = jnp.full((K, M), 3, dtype=jnp.float32)
+
+        # use default launch dims
+        return jax_matmul(a, b)
+
+    print(f())
+
+
+def example6():
+    # don't specify default launch dims
+    jax_matmul = jax_kernel(matmul_kernel)
+
+    @jax.jit
+    def f():
+        N1, M1, K1 = 3, 4, 2
+        a1 = jnp.full((N1, K1), 2, dtype=jnp.float32)
+        b1 = jnp.full((K1, M1), 3, dtype=jnp.float32)
+
+        # use custom launch dims
+        result1 = jax_matmul(a1, b1, launch_dims=(N1, M1))
+
+        N2, M2, K2 = 4, 3, 2
+        a2 = jnp.full((N2, K2), 2, dtype=jnp.float32)
+        b2 = jnp.full((K2, M2), 3, dtype=jnp.float32)
+
+        # use custom launch dims
+        result2 = jax_matmul(a2, b2, launch_dims=(N2, M2))
+
+        return result1, result2
+
+    r1, r2 = f()
+    print(r1)
+    print()
+    print(r2)
+
+
+def example7():
+    # no inputs and one output
+    jax_diagonal = jax_kernel(diagonal_kernel)
+
+    @jax.jit
+    def f():
+        # launch dimensions determine output size
+        return jax_diagonal(launch_dims=4)
+
+    print(f())
+
+
+def example8():
+    # Using input-output arguments
+
+    jax_func = jax_kernel(in_out_kernel, num_outputs=2, in_out_argnames=["b"])
+
+    f = jax.jit(jax_func)
+
+    a = jnp.ones(10, dtype=jnp.float32)
+    b = jnp.arange(10, dtype=jnp.float32)
+
+    b, c = f(a, b)
+    print(b)
+    print(c)
+
+
+def main():
+    wp.init()
+    wp.load_module(device=wp.get_device())
+
+    examples = [example1, example2, example3, example4, example5, example6, example7, example8]
+
+    for example in examples:
+        print("\n===========================================================================")
+        print(f"{example.__name__}:")
+        example()
+
+
+if __name__ == "__main__":
+    main()