cuda.tile.arange#

cuda.tile.arange(size, /, *, dtype, start=0, step=1)#

Creates a 1-D tile of length size with values start, start + step, ..., start + (size - 1) * step.

Parameters:
  • 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.

Return type:

Tile

Examples

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)
import cuda.tile as ct
import torch

@ct.kernel
def kernel():
    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)


torch.cuda.init()
ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
torch.cuda.synchronize()

Output

[0, 1, 2, 3]
[2, 3, 4, 5, 6, 7, 8, 9]
[2, 4, 6, 8]
[7, 6, 5, 4, 3, 2, 1, 0]