Merge branch 'main' into fix/cuda-graph-memfree-dep-tracking

Author: theonlychant

docs/user_guide/configuration.rst

@@ -159,7 +159,7 @@ Kernel-level settings can be passed as arguments to the :func:`@wp.kernel <warp.
 
 
     @wp.kernel(module_options={"fast_math": True}, module="unique")
-    def fast_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
+    def fast_kernel(a: wp.array[float], b: wp.array[float]):
         # fast_math is applied to this kernel's unique module
         tid = wp.tid()
         b[tid] = a[tid] + 1.0

docs/user_guide/tiles.rst

@@ -78,7 +78,7 @@ In Warp, tile objects are arrays of data where the tile elements may be scalars,
     TILE_THREADS = 64
 
     @wp.kernel
-    def compute(a: array2d(dtype=float)):
+    def compute(a: wp.array2d[float]):
         
         # obtain our 2d block index
         i, j = wp.tid()
@@ -483,7 +483,7 @@ The ``capacity`` must be a compile-time constant and ``dtype`` specifies the ele
     CAPACITY = wp.constant(256)
 
     @wp.kernel
-    def my_kernel(data: wp.array(dtype=float)):
+    def my_kernel(data: wp.array[float]):
         i, j = wp.tid()
         s = wp.tile_stack(capacity=CAPACITY, dtype=float)
         ...
@@ -554,7 +554,7 @@ output array:
     CAPACITY = wp.constant(BLOCK_DIM)  # at most one output per thread
 
     @wp.kernel
-    def compact_kernel(data: wp.array(dtype=float), out: wp.array(dtype=float)):
+    def compact_kernel(data: wp.array[float], out: wp.array[float]):
         _i, j = wp.tid()
 
         val = data[j]
@@ -982,8 +982,8 @@ then moves on. The next load cannot begin until the current store completes:
 
     @wp.kernel
     def sequential(
-        inp: wp.array2d(dtype=float),
-        out: wp.array2d(dtype=float),
+        inp: wp.array2d[float],
+        out: wp.array2d[float],
     ):
         for i in range(N_ROWS):
             a = wp.tile_load(inp, shape=(1, TILE_N), offset=(i, 0), storage="register")
@@ -995,8 +995,8 @@ then moves on. The next load cannot begin until the current store completes:
 
     @wp.kernel
     def pipelined(
-        inp: wp.array2d(dtype=float),
-        out: wp.array2d(dtype=float),
+        inp: wp.array2d[float],
+        out: wp.array2d[float],
     ):
         # Load first tile
         a = wp.tile_load(inp, shape=(1, TILE_N), offset=(0, 0), storage="register")

warp/__init__.pyi

@@ -3378,7 +3378,7 @@ def tile_scatter_add(a: Tile[Any, tuple[int, ...]], i: int32, value: Any, has_va
         .. code-block:: python
 
             @wp.kernel
-            def histogram(data: wp.array(dtype=float), out: wp.array(dtype=float)):
+            def histogram(data: wp.array[float], out: wp.array[float]):
 
                 bins = wp.tile_zeros(dtype=float, shape=4, storage="shared")
                 i = wp.tid()
@@ -4316,7 +4316,7 @@ def tile_stack(capacity: int32, dtype: Any) -> TileStack[Any, Any]:
             CAP = wp.constant(8)
 
             @wp.kernel
-            def compact_kernel(data: wp.array(dtype=int), out: wp.array(dtype=int), out_count: wp.array(dtype=int)):
+            def compact_kernel(data: wp.array[int], out: wp.array[int], out_count: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
 
@@ -4365,7 +4365,7 @@ def tile_stack_push(s: Any, value: Any, has_value: bool) -> int:
             CAP = wp.constant(8)
 
             @wp.kernel
-            def push_kernel(out_idx: wp.array(dtype=int)):
+            def push_kernel(out_idx: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
                 idx = wp.tile_stack_push(s, j * 10, j < 4)
@@ -4408,7 +4408,7 @@ def tile_stack_pop(s: Any) -> tuple[Any, int]:
             CAP = wp.constant(8)
 
             @wp.kernel
-            def pop_kernel(out: wp.array(dtype=int)):
+            def pop_kernel(out: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
                 wp.tile_stack_push(s, j * 10, j < 4)
@@ -4443,7 +4443,7 @@ def tile_stack_clear(s: Any) -> None:
             CAP = wp.constant(8)
 
             @wp.kernel
-            def clear_kernel(before: wp.array(dtype=int), after: wp.array(dtype=int)):
+            def clear_kernel(before: wp.array[int], after: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
                 wp.tile_stack_push(s, j, True)
@@ -4489,7 +4489,7 @@ def tile_stack_count(s: Any) -> int:
             CAP = wp.constant(8)
 
             @wp.kernel
-            def count_kernel(out_count: wp.array(dtype=int)):
+            def count_kernel(out_count: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
                 wp.tile_stack_push(s, j, j % 2 == 0)

warp/_src/builtins.py

@@ -4758,7 +4758,7 @@ def tile_scatter_add_dispatch_func(input_types, return_type, args):
         .. code-block:: python
 
             @wp.kernel
-            def histogram(data: wp.array(dtype=float), out: wp.array(dtype=float)):
+            def histogram(data: wp.array[float], out: wp.array[float]):
 
                 bins = wp.tile_zeros(dtype=float, shape=4, storage="shared")
                 i = wp.tid()
@@ -7018,7 +7018,7 @@ def tile_stack_dispatch_func(arg_types, return_type, arg_values):
             CAP = wp.constant(8)
 
             @wp.kernel
-            def compact_kernel(data: wp.array(dtype=int), out: wp.array(dtype=int), out_count: wp.array(dtype=int)):
+            def compact_kernel(data: wp.array[int], out: wp.array[int], out_count: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
 
@@ -7098,7 +7098,7 @@ def tile_stack_push_dispatch_func(arg_types, return_type, arg_values):
             CAP = wp.constant(8)
 
             @wp.kernel
-            def push_kernel(out_idx: wp.array(dtype=int)):
+            def push_kernel(out_idx: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
                 idx = wp.tile_stack_push(s, j * 10, j < 4)
@@ -7166,7 +7166,7 @@ def tile_stack_pop_dispatch_func(arg_types, return_type, arg_values):
             CAP = wp.constant(8)
 
             @wp.kernel
-            def pop_kernel(out: wp.array(dtype=int)):
+            def pop_kernel(out: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
                 wp.tile_stack_push(s, j * 10, j < 4)
@@ -7226,7 +7226,7 @@ def tile_stack_clear_dispatch_func(arg_types, return_type, arg_values):
             CAP = wp.constant(8)
 
             @wp.kernel
-            def clear_kernel(before: wp.array(dtype=int), after: wp.array(dtype=int)):
+            def clear_kernel(before: wp.array[int], after: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
                 wp.tile_stack_push(s, j, True)
@@ -7297,7 +7297,7 @@ def tile_stack_count_dispatch_func(arg_types, return_type, arg_values):
             CAP = wp.constant(8)
 
             @wp.kernel
-            def count_kernel(out_count: wp.array(dtype=int)):
+            def count_kernel(out_count: wp.array[int]):
                 _i, j = wp.tid()
                 s = wp.tile_stack(capacity=CAP, dtype=int)
                 wp.tile_stack_push(s, j, j % 2 == 0)