warp-lang 1.8.0__py3-none-macosx_10_13_universal2.whl → 1.9.0__py3-none-macosx_10_13_universal2.whl

warp/examples/interop/example_jax_callable.py

@@ -42,7 +42,7 @@ def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtyp
 
 # The Python function to call.
 # Note the argument annotations, just like Warp kernels.
-def example_func(
+def scale_func(
     # inputs
     a: wp.array(dtype=float),
     b: wp.array(dtype=wp.vec2),
@@ -55,8 +55,23 @@ def example_func(
     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(example_func, num_outputs=2, vmap_method="broadcast_all")
+    jax_func = jax_callable(scale_func, num_outputs=2)
 
     @jax.jit
     def f():
@@ -78,7 +93,7 @@ def example1():
 
 
 def example2():
-    jax_func = jax_callable(example_func, num_outputs=2, vmap_method="broadcast_all")
+    jax_func = jax_callable(scale_func, num_outputs=2)
 
     # NOTE: scalar arguments must be static compile-time constants
     @partial(jax.jit, static_argnames=["s"])
@@ -100,11 +115,26 @@ def example2():
     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]
+    examples = [example1, example2, example3]
 
     for example in examples:
         print("\n===========================================================================")

warp/examples/interop/example_jax_kernel.py

@@ -72,6 +72,17 @@ def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtyp
     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)
@@ -189,11 +200,26 @@ def example7():
     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]
+    examples = [example1, example2, example3, example4, example5, example6, example7, example8]
 
     for example in examples:
         print("\n===========================================================================")

warp/fabric.py

@@ -211,7 +211,7 @@ class fabricarray(noncontiguous_array_base[T]):
                     allocator = self.device.get_allocator()
                     buckets_ptr = allocator.alloc(buckets_size)
                     cuda_stream = self.device.stream.cuda_stream
-                    runtime.core.memcpy_h2d(
+                    runtime.core.wp_memcpy_h2d(
                         self.device.context, buckets_ptr, ctypes.addressof(buckets), buckets_size, cuda_stream
                     )
                     self.deleter = allocator.deleter

warp/fem/cache.py

@@ -34,7 +34,7 @@ _key_re = re.compile("[^0-9a-zA-Z_]+")
 
 def _make_key(obj, suffix: str, options: Optional[Dict[str, Any]] = None):
     # human-readable part
-    key = _key_re.sub("", f"{obj.__name__}_{suffix}")
+    suffix = str(suffix)
 
     sorted_opts = sorted(options.items()) if options is not None else ()
     opts_str = "".join(
@@ -49,7 +49,7 @@ def _make_key(obj, suffix: str, options: Optional[Dict[str, Any]] = None):
     uid = hashlib.blake2b(bytes(opts_str, encoding="utf-8"), digest_size=4).hexdigest()
 
     # avoid long keys, issues on win
-    key = f"{key[:64]}_{uid}"
+    key = f"{obj.__name__}_{suffix[:32]}_{uid}"
 
     return key
 
@@ -62,7 +62,10 @@ def _arg_type_name(arg_type):
     return wp.types.get_type_code(wp.types.type_to_warp(arg_type))
 
 
-def _make_cache_key(func, key, argspec=None):
+def _make_cache_key(func, key, argspec=None, allow_overloads: bool = True):
+    if not allow_overloads:
+        return key
+
     if argspec is None:
         annotations = get_annotations(func)
     else:
@@ -80,6 +83,7 @@ def _register_function(
 ):
     # wp.Function will override existing func for a given key...
     # manually add back our overloads
+    key = _key_re.sub("", key)
     existing = module.functions.get(key)
     new_fn = wp.Function(
         func=func,
@@ -95,9 +99,9 @@ def _register_function(
     return module.functions[key]
 
 
-def get_func(func, suffix: str, code_transformers=None):
+def get_func(func, suffix: str, code_transformers=None, allow_overloads=False):
     key = _make_key(func, suffix)
-    cache_key = _make_cache_key(func, key)
+    cache_key = _make_cache_key(func, key, allow_overloads=allow_overloads)
 
     if cache_key not in _func_cache:
         module = wp.get_module(func.__module__)
@@ -111,9 +115,9 @@ def get_func(func, suffix: str, code_transformers=None):
     return _func_cache[cache_key]
 
 
-def dynamic_func(suffix: str, code_transformers=None):
+def dynamic_func(suffix: str, code_transformers=None, allow_overloads=False):
     def wrap_func(func: Callable):
-        return get_func(func, suffix=suffix, code_transformers=code_transformers)
+        return get_func(func, suffix=suffix, code_transformers=code_transformers, allow_overloads=allow_overloads)
 
     return wrap_func
 
@@ -122,46 +126,49 @@ def get_kernel(
     func,
     suffix: str,
     kernel_options: Optional[Dict[str, Any]] = None,
+    allow_overloads=False,
 ):
     if kernel_options is None:
         kernel_options = {}
 
     key = _make_key(func, suffix, kernel_options)
-    cache_key = _make_cache_key(func, key)
+    cache_key = _make_cache_key(func, key, allow_overloads=allow_overloads)
 
     if cache_key not in _kernel_cache:
-        module_name = f"{func.__module__}.dyn.{key}"
+        kernel_key = _key_re.sub("", key)
+        module_name = f"{func.__module__}.dyn.{kernel_key}"
         module = wp.get_module(module_name)
         module.options = dict(wp.get_module(func.__module__).options)
         module.options.update(kernel_options)
-        _kernel_cache[cache_key] = wp.Kernel(func=func, key=key, module=module, options=kernel_options)
+        _kernel_cache[cache_key] = wp.Kernel(func=func, key=kernel_key, module=module, options=kernel_options)
     return _kernel_cache[cache_key]
 
 
-def dynamic_kernel(suffix: str, kernel_options: Optional[Dict[str, Any]] = None):
+def dynamic_kernel(suffix: str, kernel_options: Optional[Dict[str, Any]] = None, allow_overloads=False):
     if kernel_options is None:
         kernel_options = {}
 
     def wrap_kernel(func: Callable):
-        return get_kernel(func, suffix=suffix, kernel_options=kernel_options)
+        return get_kernel(func, suffix=suffix, kernel_options=kernel_options, allow_overloads=allow_overloads)
 
     return wrap_kernel
 
 
 def get_struct(struct: type, suffix: str):
     key = _make_key(struct, suffix)
-    # used in codegen
-    struct.__qualname__ = key
+    cache_key = key
 
-    if key not in _struct_cache:
+    if cache_key not in _struct_cache:
+        # used in codegen
+        struct.__qualname__ = _key_re.sub("", key)
         module = wp.get_module(struct.__module__)
-        _struct_cache[key] = wp.codegen.Struct(
-            key=key,
+        _struct_cache[cache_key] = wp.codegen.Struct(
+            key=struct.__qualname__,
             cls=struct,
             module=module,
         )
 
-    return _struct_cache[key]
+    return _struct_cache[cache_key]
 
 
 def dynamic_struct(suffix: str):
@@ -293,12 +300,13 @@ def get_integrand_kernel(
         options.update(kernel_options)
 
     kernel_key = _make_key(integrand.func, suffix, options=options)
-    cache_key = _make_cache_key(integrand, kernel_key, integrand.argspec)
+    cache_key = _make_cache_key(integrand, kernel_key, integrand.argspec, allow_overloads=True)
 
     if cache_key not in _kernel_cache:
         if kernel_fn is None:
             return None
 
+        kernel_key = _key_re.sub("", kernel_key)
         module = wp.get_module(f"{integrand.module.name}.{kernel_key}")
         module.options = options
         _kernel_cache[cache_key] = wp.Kernel(

warp/fem/domain.py

@@ -237,11 +237,11 @@ class Cells(GeometryDomain):
         filter_target = True
         pos_type = cache.cached_vec_type(self.geometry.dimension, dtype=float)
 
-        @cache.dynamic_func(suffix=self.name)
+        @cache.dynamic_func(suffix=self.name, allow_overloads=True)
         def cell_partition_lookup(args: self.DomainArg, pos: pos_type, max_dist: float):
             return filtered_cell_lookup(args.geo, pos, max_dist, args.index, filter_target)
 
-        @cache.dynamic_func(suffix=self.name)
+        @cache.dynamic_func(suffix=self.name, allow_overloads=True)
         def cell_partition_lookup(args: self.DomainArg, pos: pos_type):
             max_dist = 0.0
             return filtered_cell_lookup(args.geo, pos, max_dist, args.index, filter_target)

warp/fem/field/nodal_field.py

@@ -232,7 +232,7 @@ class NodalFieldBase(DiscreteField):
             @cache.dynamic_func(suffix=self.name)
             def eval_grad_outer_world_space(args: self.ElementEvalArg, s: Sample):
                 grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
-                return eval_grad_outer_ref_space(args, s, grad_transform)
+                return eval_grad_outer(args, s, grad_transform)
 
             return eval_grad_outer_world_space
         else:
@@ -240,7 +240,7 @@ class NodalFieldBase(DiscreteField):
             @cache.dynamic_func(suffix=self.name)
             def eval_grad_outer_ref_space(args: self.ElementEvalArg, s: Sample):
                 grad_transform = 1.0
-                return eval_grad_outer_ref_space(args, s, grad_transform)
+                return eval_grad_outer(args, s, grad_transform)
 
             return eval_grad_outer_ref_space