Support arange with start,stop,step

Signed-off-by: Ziheng Deng <zihengd@nvidia.com>

Author: gunnersdeng

changelog.d/arange.md

@@ -0,0 +1,4 @@
+- Extend ``ct.arange()`` with optional ``start`` and ``step`` arguments:
+  ``ct.arange(size, start=0, step=1, dtype=...)``. ``size`` must be a constant
+  integer, while ``start`` and ``step`` may be dynamic numbers. For example, 
+  ``ct.arange(8, start=7, step=-1, dtype=ct.int32)`` creates a reversed range.

src/cuda/tile/_ir/ops.py

@@ -2891,12 +2891,17 @@ def arange(size: int, dtype: DType) -> Var:
 
 
 @impl(ct.arange)
-def arange_impl(size: Var, dtype: Var) -> Var:
+def arange_impl(size: Var, dtype: Var, start: Var, step: Var) -> Var:
     size_val = require_constant_int(size)
     dtype_val = require_dtype_spec(dtype)
     if not _is_power_of_2(size_val):
         raise TileTypeError(f"Result tile shape must be power of 2, got {size_val}")
-    return arange(size_val, dtype_val)
+    result = arange(size_val, dtype_val)
+    if not (step.is_constant() and step.get_constant() == 1):
+        result = binary_arithmetic_tensorlike("mul", result, astype(step, dtype_val))
+    if not (start.is_constant() and start.get_constant() == 0):
+        result = binary_arithmetic_tensorlike("add", result, astype(start, dtype_val))
+    return result
 
 
 @impl(ct.reshape)

src/cuda/tile/_stub.py

@@ -1893,12 +1893,15 @@ def atomic_xor(array, indices, update, /, *,
 
 
 @stub
-def arange(size, /, *, dtype) -> Tile:
-    """Creates a tile with value starting from 0 to `size - 1`.
+def arange(size, /, *, dtype, start=0, step=1) -> Tile:
+    """Creates a 1-D tile of length ``size`` with values
+    ``start, start + step, ..., start + (size - 1) * step``.
 
     Args:
-        size (const int): Size of the tile.
+        size (const int): Size of the tile. Must be a constant integer that is a power of two.
         dtype (DType): Datatype of the tile.
+        start: Value of the first element. Defaults to ``0``.
+        step: The gap between adjacent values. Defaults to ``1``.
 
     Returns:
         Tile:
@@ -1908,12 +1911,21 @@ def arange(size, /, *, dtype) -> Tile:
         .. testcode::
             :template: kernel_wrapper.py
 
-            tile = ct.arange(4, dtype=ct.int32)
-            print(tile)
+            tile_0 = ct.arange(4, dtype=ct.int32)
+            print(tile_0)
+            tile_1 = ct.arange(8, start=2, dtype=ct.int32)
+            print(tile_1)
+            tile_2 = ct.arange(4, start=2, step=2, dtype=ct.int32)
+            print(tile_2)
+            tile_3 = ct.arange(8, start=7, step=-1, dtype=ct.int32)
+            print(tile_3)
 
         .. testoutput::
 
             [0, 1, 2, 3]
+            [2, 3, 4, 5, 6, 7, 8, 9]
+            [2, 4, 6, 8]
+            [7, 6, 5, 4, 3, 2, 1, 0]
     """
 
 

test/test_generate.py

@@ -7,24 +7,85 @@
 
 from math import ceil
 import cuda.tile as ct
-from util import assert_equal
+from cuda.tile import TileTypeError
+from util import assert_equal, assert_close
 from conftest import int_dtypes, float_dtypes, dtype_id
 
 
 @ct.kernel
-def arange(x, TILE: ct.Constant[int]):
+def arange_dynamic_start_step(x, step, TILE: ct.Constant[int]):
     bid = ct.bid(0)
-    start = ct.astype(bid * TILE, x.dtype)
-    tx = start + ct.arange(TILE, dtype=x.dtype)
+    tx = ct.arange(TILE, start=bid * TILE, step=step, dtype=x.dtype)
     ct.store(x, index=(bid,), tile=tx)
 
 
 @pytest.mark.parametrize("shape", [(128,)])
 @pytest.mark.parametrize("tile", [64])
 @pytest.mark.parametrize("dtype", int_dtypes + float_dtypes, ids=dtype_id)
-def test_arange(shape, dtype, tile):
+def test_arange_dynamic_start_step(shape, dtype, tile):
     x = torch.zeros(shape, dtype=dtype, device='cuda')
     grid = (ceil(shape[0] / tile), 1, 1)
-    ct.launch(torch.cuda.current_stream(), grid, arange, (x, tile))
+    ct.launch(torch.cuda.current_stream(), grid, arange_dynamic_start_step, (x, 1, tile))
     ref = torch.arange(len(x), dtype=dtype, device=x.device)
     assert_equal(x, ref)
+
+
+@pytest.mark.parametrize("size,start,step", [
+    (128, None, None),  # arange(size)
+    (64, 8, None), (64, -16, None),  # arange(size, start)
+    (64, 0, 2), (64, 64, -1), (16, -8, -3), (4, 8.5, -1.1),  # arange(size, start, step)
+    (8, 10, 0)  # step=0
+])
+@pytest.mark.parametrize("dtype", int_dtypes + float_dtypes, ids=dtype_id)
+def test_arange(size, start, step, dtype):
+    @ct.kernel
+    def arange_kernel(x):
+        if start is None:
+            tx = ct.arange(size, dtype=x.dtype)
+        elif step is None:
+            tx = ct.arange(size, start=start, dtype=x.dtype)
+        else:
+            tx = ct.arange(size, start=start, step=step, dtype=x.dtype)
+        ct.store(x, index=(0,), tile=tx)
+
+    x = torch.zeros(size, dtype=dtype, device='cuda')
+    ct.launch(torch.cuda.current_stream(), (1, 1, 1), arange_kernel, (x,))
+    if step == 0:
+        ref = torch.full((size,), start, dtype=dtype, device=x.device)
+    else:
+        start = 0 if start is None else start
+        step = 1 if step is None else step
+        ref = torch.arange(start, start + size * step, step, dtype=dtype, device=x.device)
+    assert_close(x, ref)
+
+
+@pytest.mark.parametrize("size,start,step,error_message", [
+    (3, None, None, "Result tile shape must be power of 2"),
+    (5, 0, 2, "Result tile shape must be power of 2"),
+    (0.1, None, None, 'Expected an integer constant')
+])
+def test_arange_invalid_size(size, start, step, error_message):
+    @ct.kernel
+    def arange_kernel(x):
+        if start is None:
+            tx = ct.arange(size, dtype=x.dtype)
+        elif step is None:
+            tx = ct.arange(size, start=start, dtype=x.dtype)
+        else:
+            tx = ct.arange(size, start=start, step=step, dtype=x.dtype)
+        ct.store(x, index=(0,), tile=tx)
+
+    with pytest.raises(TileTypeError, match=error_message):
+        x = torch.zeros(1, dtype=torch.int32, device='cuda')
+        ct.launch(torch.cuda.current_stream(), (1, 1, 1), arange_kernel, (x,))
+
+
+def test_arange_reject_dynamic_size():
+    @ct.kernel
+    def arange_dynamic_size(x):
+        tx = ct.arange(ct.bid(0), dtype=x.dtype)
+        ct.store(x, index=(0,), tile=tx)
+
+    with pytest.raises(TileTypeError, match="Expected an integer constant"):
+        x = torch.zeros(1, dtype=torch.int32, device='cuda')
+        ct.launch(torch.cuda.current_stream(), (1, 1, 1), arange_dynamic_size, (x,))

test/test_token_order.py

@@ -22,12 +22,12 @@ def store_buffer(X, TILE: ct.Constant[int]):
         ct.store(X, index=(0,), tile=tx0)
         # reverse so that each SIMT thread is less likely to be assigned the same address
         X_ALIAS = X
-        reverse_offset = TILE - 1 - ct.arange(TILE, dtype=np.int32)
+        reverse_offset = ct.arange(TILE, start=TILE - 1, step=-1, dtype=np.int32)
         tx1 = ct.gather(X_ALIAS, reverse_offset)
         ct.store(X_ALIAS, index=(0,), tile=tx1)
 
     def store_buffer_alternative(X, TILE: ct.Constant[int]):
-        reverse_offset = TILE - 1 - ct.arange(TILE, dtype=np.int32)
+        reverse_offset = ct.arange(TILE, start=TILE - 1, step=-1, dtype=np.int32)
         tx2 = ct.arange(TILE, dtype=X.dtype)
         ct.scatter(X, reverse_offset, tx2)
         tx3 = ct.load(X, index=(0,), shape=(TILE,))
@@ -37,7 +37,7 @@ def serialized_for_loop(X, TILE: ct.Constant[int]):
         ct.store(X, index=(0,), tile=ct.arange(TILE, dtype=X.dtype))
         # flip the buffer 3 times
         for i in range(3):
-            reverse_offset = TILE - 1 - ct.arange(TILE, dtype=np.int32)
+            reverse_offset = ct.arange(TILE, start=TILE - 1, step=-1, dtype=np.int32)
             tx = ct.gather(X, reverse_offset)
             ct.store(X, index=(0,), tile=tx)
 
@@ -46,7 +46,7 @@ def serialized_while_loop(X, TILE: ct.Constant[int]):
         # flip the buffer 3 times
         i = 0
         while i < 3:
-            reverse_offset = TILE - 1 - ct.arange(TILE, dtype=np.int32)
+            reverse_offset = ct.arange(TILE, start=TILE - 1, step=-1, dtype=np.int32)
             tx = ct.gather(X, reverse_offset)
             ct.store(X, index=(0,), tile=tx)
             i += 1