Merge branch 'tile_rand' into 'main'

Support generating tiles of random numbers (GH-1010)

See merge request omniverse/warp!1818

Author: daedalus5

CHANGELOG.md

@@ -31,6 +31,7 @@
   compilation time. Lower levels may improve compile times but reduce run-time performance. It can be configured
   globally, or per-module using `wp.set_module_options({"optimization_level": #})`. Currently it only has an effect
   on GPU modules and only when using CUDA Toolkit 12.9 or newer ([GH-1084](https://github.com/NVIDIA/warp/issues/1084)).
+- Add `wp.tile_randi()` and `wp.tile_randf()` to support generating tiles of random numbers ([GH-1010](https://github.com/NVIDIA/warp/issues/1010))
 - Add `wp.mesh_query_ray_anyhit` for ray any-hit queries ([GH-1097](https://github.com/NVIDIA/warp/issues/1097)).
 - Add support for group-aware construction and queries for `wp.Mesh` ([GH-1097](https://github.com/NVIDIA/warp/issues/1097)):
   - New constructor argument: `groups` for per-face group IDs.

docs/modules/functions.rst

@@ -1917,6 +1917,187 @@ Tile Primitives
     :returns: A tile filled with the specified value
 
 
+.. py:function:: tile_randi(shape: tuple[int, ...], rng: uint32, storage: str) -> Tile[Any,tuple[int, ...]]
+
+    .. hlist::
+       :columns: 8
+
+       * Kernel
+
+    Generate a tile of random integers.
+
+    :param shape: Shape of the output tile
+    :param rng: Random number generator state, typically from :func:`warp.rand_init`
+    :param storage: The storage location for the tile: ``"register"`` for registers
+      (default) or ``"shared"`` for shared memory.
+    :returns: A tile of random integers with the specified shape
+
+    Example:
+
+    .. testcode::
+
+        TILE_M, TILE_N = 2, 2
+        M, N = 2, 2
+        seed = 42
+
+        @wp.kernel
+        def rand_kernel(seed: int, x: wp.array2d(dtype=int)):
+            i, j = wp.tid()
+            rng = wp.rand_init(seed, i * TILE_M + j)
+            t = wp.tile_randi(shape=(TILE_M, TILE_N), rng=rng)
+            wp.tile_store(x, t, offset=(i * TILE_M, j * TILE_N))
+
+        x = wp.zeros(shape=(M * TILE_M, N * TILE_N), dtype=int)
+        wp.launch_tiled(rand_kernel, dim=[M, N], inputs=[seed, x], block_dim=32)
+        print(x.numpy())
+
+    .. testoutput::
+
+        [[  798497746  1803297529  -955788638    17806966]
+         [ 1788185933  1320194893  2073257406 -2009156320]
+         [ -257534450 -1138585923  1145322783  -321794125]
+         [-2096177388 -1835610841  1159339128  -652221052]]
+    
+
+
+.. py:function:: tile_randi(shape: tuple[int, ...], rng: uint32, min: int32, max: int32, storage: str) -> Tile[Any,tuple[int, ...]]
+    :noindex:
+    :nocontentsentry:
+
+    .. hlist::
+       :columns: 8
+
+       * Kernel
+
+    Generate a tile of random integers within a specified range.
+
+    :param shape: Shape of the output tile
+    :param rng: Random number generator state, typically from :func:`warp.rand_init`
+    :param min: Minimum value (inclusive) for random integers
+    :param max: Maximum value (exclusive) for random integers
+    :param storage: The storage location for the tile: ``"register"`` for registers
+      (default) or ``"shared"`` for shared memory.
+    :returns: A tile of random integers in the range [min, max) with the specified shape
+
+    Example:
+
+    .. testcode::
+
+        TILE_M, TILE_N = 2, 2
+        M, N = 2, 2
+        seed = 42
+
+        @wp.kernel
+        def rand_range_kernel(seed: int, x: wp.array2d(dtype=int)):
+            i, j = wp.tid()
+            rng = wp.rand_init(seed, i * TILE_M + j)
+            t = wp.tile_randi(shape=(TILE_M, TILE_N), rng=rng, min=-5, max=5)
+            wp.tile_store(x, t, offset=(i * TILE_M, j * TILE_N))
+
+        x = wp.zeros(shape=(M * TILE_M, N * TILE_N), dtype=int)
+        wp.launch_tiled(rand_range_kernel, dim=[M, N], inputs=[seed, x], block_dim=32)
+        print(x.numpy())
+
+    .. testoutput::
+
+        [[ 1  4  3  1]
+         [-2 -2  1  1]
+         [ 1 -2 -2 -4]
+         [ 3  0  3 -1]]
+    
+
+
+.. py:function:: tile_randf(shape: tuple[int, ...], rng: uint32, storage: str) -> Tile[Any,tuple[int, ...]]
+
+    .. hlist::
+       :columns: 8
+
+       * Kernel
+
+    Generate a tile of random floats.
+
+    :param shape: Shape of the output tile
+    :param rng: Random number generator state, typically from :func:`warp.rand_init`
+    :param storage: The storage location for the tile: ``"register"`` for registers
+      (default) or ``"shared"`` for shared memory.
+    :returns: A tile of random floats in the range [0, 1) with the specified shape
+
+    Example:
+
+    .. testcode::
+
+        TILE_M, TILE_N = 2, 2
+        M, N = 2, 2
+        seed = 42
+
+        @wp.kernel
+        def rand_kernel(seed: int, x: wp.array2d(dtype=float)):
+            i, j = wp.tid()
+            rng = wp.rand_init(seed, i * TILE_M + j)
+            t = wp.tile_randf(shape=(TILE_M, TILE_N), rng=rng)
+            wp.tile_store(x, t, offset=(i * TILE_M, j * TILE_N))
+
+        x = wp.zeros(shape=(M * TILE_M, N * TILE_N), dtype=float)
+        wp.launch_tiled(rand_kernel, dim=[M, N], inputs=[seed, x], block_dim=32)
+        print(x.numpy())
+
+    .. testoutput::
+
+        [[0.1859147  0.41986287 0.7774631  0.00414598]
+         [0.41634446 0.3073818  0.4827178  0.53220683]
+         [0.9400381  0.73490226 0.26666623 0.9250764 ]
+         [0.51194566 0.57261354 0.26992965 0.8481429 ]]
+    
+
+
+.. py:function:: tile_randf(shape: tuple[int, ...], rng: uint32, min: float32, max: float32, storage: str) -> Tile[Any,tuple[int, ...]]
+    :noindex:
+    :nocontentsentry:
+
+    .. hlist::
+       :columns: 8
+
+       * Kernel
+
+    Generate a tile of random floats within a specified range.
+
+    :param shape: Shape of the output tile
+    :param rng: Random number generator state, typically from :func:`warp.rand_init`
+    :param min: Minimum value (inclusive) for random floats
+    :param max: Maximum value (exclusive) for random floats
+    :param storage: The storage location for the tile: ``"register"`` for registers
+      (default) or ``"shared"`` for shared memory.
+    :returns: A tile of random floats in the range [min, max) with the specified shape
+
+    Example:
+
+    .. testcode::
+
+        TILE_M, TILE_N = 2, 2
+        M, N = 2, 2
+        seed = 42
+
+        @wp.kernel
+        def rand_range_kernel(seed: int, x: wp.array2d(dtype=float)):
+            i, j = wp.tid()
+            rng = wp.rand_init(seed, i * TILE_M + j)
+            t = wp.tile_randf(shape=(TILE_M, TILE_N), rng=rng, min=-5.0, max=5.0)
+            wp.tile_store(x, t, offset=(i * TILE_M, j * TILE_N))
+
+        x = wp.zeros(shape=(M * TILE_M, N * TILE_N), dtype=float)
+        wp.launch_tiled(rand_range_kernel, dim=[M, N], inputs=[seed, x], block_dim=32)
+        print(x.numpy())
+
+    .. testoutput::
+
+        [[-3.140853   -0.80137134  2.7746308  -4.95854   ]
+         [-0.83655536 -1.9261819  -0.17282188  0.32206833]
+         [ 4.400381    2.3490226  -2.3333378   4.2507644 ]
+         [ 0.11945665  0.7261354  -2.3007035   3.481429  ]]
+
+    
+
+
 .. py:function:: tile_arange(*args: Scalar, dtype: Scalar, storage: str) -> Tile[Scalar,tuple[int]]
 
     .. hlist::

docs/modules/tiles.rst

@@ -204,6 +204,8 @@ Construction
 * :func:`tile_zeros`
 * :func:`tile_ones`
 * :func:`tile_full`
+* :func:`tile_randi`
+* :func:`tile_randf`
 * :func:`tile_arange`
 * :func:`tile`
 * :func:`untile`

warp/__init__.pyi

@@ -2624,6 +2624,169 @@ def tile_full(shape: tuple[int, ...], value: Any, dtype: Any, storage: str) -> T
     """
     ...
 
+@over
+def tile_randi(shape: tuple[int, ...], rng: uint32, storage: str) -> Tile[Any, tuple[int, ...]]:
+    """Generate a tile of random integers.
+
+    :param shape: Shape of the output tile
+    :param rng: Random number generator state, typically from :func:`warp.rand_init`
+    :param storage: The storage location for the tile: ``"register"`` for registers
+      (default) or ``"shared"`` for shared memory.
+    :returns: A tile of random integers with the specified shape
+
+    Example:
+
+    .. testcode::
+
+        TILE_M, TILE_N = 2, 2
+        M, N = 2, 2
+        seed = 42
+
+        @wp.kernel
+        def rand_kernel(seed: int, x: wp.array2d(dtype=int)):
+            i, j = wp.tid()
+            rng = wp.rand_init(seed, i * TILE_M + j)
+            t = wp.tile_randi(shape=(TILE_M, TILE_N), rng=rng)
+            wp.tile_store(x, t, offset=(i * TILE_M, j * TILE_N))
+
+        x = wp.zeros(shape=(M * TILE_M, N * TILE_N), dtype=int)
+        wp.launch_tiled(rand_kernel, dim=[M, N], inputs=[seed, x], block_dim=32)
+        print(x.numpy())
+
+    .. testoutput::
+
+        [[  798497746  1803297529  -955788638    17806966]
+         [ 1788185933  1320194893  2073257406 -2009156320]
+         [ -257534450 -1138585923  1145322783  -321794125]
+         [-2096177388 -1835610841  1159339128  -652221052]]
+
+    """
+    ...
+
+@over
+def tile_randi(shape: tuple[int, ...], rng: uint32, min: int32, max: int32, storage: str) -> Tile[Any, tuple[int, ...]]:
+    """Generate a tile of random integers within a specified range.
+
+    :param shape: Shape of the output tile
+    :param rng: Random number generator state, typically from :func:`warp.rand_init`
+    :param min: Minimum value (inclusive) for random integers
+    :param max: Maximum value (exclusive) for random integers
+    :param storage: The storage location for the tile: ``"register"`` for registers
+      (default) or ``"shared"`` for shared memory.
+    :returns: A tile of random integers in the range [min, max) with the specified shape
+
+    Example:
+
+    .. testcode::
+
+        TILE_M, TILE_N = 2, 2
+        M, N = 2, 2
+        seed = 42
+
+        @wp.kernel
+        def rand_range_kernel(seed: int, x: wp.array2d(dtype=int)):
+            i, j = wp.tid()
+            rng = wp.rand_init(seed, i * TILE_M + j)
+            t = wp.tile_randi(shape=(TILE_M, TILE_N), rng=rng, min=-5, max=5)
+            wp.tile_store(x, t, offset=(i * TILE_M, j * TILE_N))
+
+        x = wp.zeros(shape=(M * TILE_M, N * TILE_N), dtype=int)
+        wp.launch_tiled(rand_range_kernel, dim=[M, N], inputs=[seed, x], block_dim=32)
+        print(x.numpy())
+
+    .. testoutput::
+
+        [[ 1  4  3  1]
+         [-2 -2  1  1]
+         [ 1 -2 -2 -4]
+         [ 3  0  3 -1]]
+
+    """
+    ...
+
+@over
+def tile_randf(shape: tuple[int, ...], rng: uint32, storage: str) -> Tile[Any, tuple[int, ...]]:
+    """Generate a tile of random floats.
+
+    :param shape: Shape of the output tile
+    :param rng: Random number generator state, typically from :func:`warp.rand_init`
+    :param storage: The storage location for the tile: ``"register"`` for registers
+      (default) or ``"shared"`` for shared memory.
+    :returns: A tile of random floats in the range [0, 1) with the specified shape
+
+    Example:
+
+    .. testcode::
+
+        TILE_M, TILE_N = 2, 2
+        M, N = 2, 2
+        seed = 42
+
+        @wp.kernel
+        def rand_kernel(seed: int, x: wp.array2d(dtype=float)):
+            i, j = wp.tid()
+            rng = wp.rand_init(seed, i * TILE_M + j)
+            t = wp.tile_randf(shape=(TILE_M, TILE_N), rng=rng)
+            wp.tile_store(x, t, offset=(i * TILE_M, j * TILE_N))
+
+        x = wp.zeros(shape=(M * TILE_M, N * TILE_N), dtype=float)
+        wp.launch_tiled(rand_kernel, dim=[M, N], inputs=[seed, x], block_dim=32)
+        print(x.numpy())
+
+    .. testoutput::
+
+        [[0.1859147  0.41986287 0.7774631  0.00414598]
+         [0.41634446 0.3073818  0.4827178  0.53220683]
+         [0.9400381  0.73490226 0.26666623 0.9250764 ]
+         [0.51194566 0.57261354 0.26992965 0.8481429 ]]
+
+    """
+    ...
+
+@over
+def tile_randf(
+    shape: tuple[int, ...], rng: uint32, min: float32, max: float32, storage: str
+) -> Tile[Any, tuple[int, ...]]:
+    """Generate a tile of random floats within a specified range.
+
+    :param shape: Shape of the output tile
+    :param rng: Random number generator state, typically from :func:`warp.rand_init`
+    :param min: Minimum value (inclusive) for random floats
+    :param max: Maximum value (exclusive) for random floats
+    :param storage: The storage location for the tile: ``"register"`` for registers
+      (default) or ``"shared"`` for shared memory.
+    :returns: A tile of random floats in the range [min, max) with the specified shape
+
+    Example:
+
+    .. testcode::
+
+        TILE_M, TILE_N = 2, 2
+        M, N = 2, 2
+        seed = 42
+
+        @wp.kernel
+        def rand_range_kernel(seed: int, x: wp.array2d(dtype=float)):
+            i, j = wp.tid()
+            rng = wp.rand_init(seed, i * TILE_M + j)
+            t = wp.tile_randf(shape=(TILE_M, TILE_N), rng=rng, min=-5.0, max=5.0)
+            wp.tile_store(x, t, offset=(i * TILE_M, j * TILE_N))
+
+        x = wp.zeros(shape=(M * TILE_M, N * TILE_N), dtype=float)
+        wp.launch_tiled(rand_range_kernel, dim=[M, N], inputs=[seed, x], block_dim=32)
+        print(x.numpy())
+
+    .. testoutput::
+
+        [[-3.140853   -0.80137134  2.7746308  -4.95854   ]
+         [-0.83655536 -1.9261819  -0.17282188  0.32206833]
+         [ 4.400381    2.3490226  -2.3333378   4.2507644 ]
+         [ 0.11945665  0.7261354  -2.3007035   3.481429  ]]
+
+
+    """
+    ...
+
 @over
 def tile_arange(*args: Scalar, dtype: Scalar, storage: str) -> Tile[Scalar, tuple[int]]:
     """Generate a tile of linearly spaced elements.