[FEA]: Support Indexing & Slicing in Array

[ZhangZhiPku]

Is this a new feature, an improvement, or a change to existing functionality?

New Feature

How would you describe the priority of this feature request?

High

Please provide a clear description of problem this feature solves

I am writing a GEMM program with cutile. Before using cutile, I was already very used to structuring CUDA programs in three levels: device, block, and warp. However, when implementing GEMM with cutile, I ran into an issue I hadn’t anticipated.

In a standard GEMM (D = A * B + C), we launch the kernel with a 2D grid. Each block is responsible for computing a [tile_m, tile_n] tile of the output matrix. Therefore, the block needs to read an [M, K] tile from matrix A and an [N, K] tile from matrix B (assuming B is column-major).

I normally structure my program like this:

1. device_gemm: responsible for splitting the work at the block level, partitioning A and B into tiles that each block will handle.
2. block_gemm: responsible for the actual GEMM computation.

But when I tried to write device_gemm, I encountered a problem:

@ct.kernel
def device_gemm(
    x: ct.Array, y: ct.Array, out: ct.Array,
    tm: ct.Constant, tn: ct.Constant, tk: ct.Constant
):
    block_idx_x, block_idx_y = ct.bid(0), ct.bid(1)

    # can not create local_x here
    # cuda.tile._exception.TileTypeError: Non-constant slices are not supported
    local_x = x[block_idx_x * tm: block_idx_x * tm + tm, :]
    local_y = y[block_idx_y * tn: block_idx_y * tn + tn, :]

    acc = tile_gemm(local_x, local_y, tm, tn, tk)
    ct.store(out, (block_idx_x, block_idx_y), acc)

With the current cutile design, I’m unable to “relocate” an array. In fact, what I need is to obtain a new array via indexing and slicing, which is just a view of the original array.

While I could pass the entire array directly into tile_gemm and then rely on ct.load inside it to perform all data accesses at once, believe me, that’s not a good approach. Creating new arrays (as views) through indexing and slicing makes the program much more readable.

Feature Description

Support indexing & slicing in ct.Array

@ct.kernel
def kernel(
    x: ct.Array, tile_size: ct.Constant
):
    block_idx = ct.bid(0)

    # just like this:
    local_x: ct.ArrayView = x[block_idx * tile_size: block_idx * tile_size + tile_size]

Describe your ideal solution

Introduce a view mechanism to ct.Array, so that we can relocate ct.Array without change anything else in cutile.

Describe any alternatives you have considered

No response

Additional context

No response

Contributing Guidelines

• I agree to follow cuTile Python's contributing guidelines
• I have searched the open feature requests and have found no duplicates for this feature request

[haijieg]

Creating view on Array is useful but I'm not sure I understand why you need it for GEMM. In a standard GEMM program, i.e. here, a single array with ct.load is sufficient. Why do you want sub array view first?

[ZhangZhiPku]

Creating view on Array is useful but I'm not sure I understand why you need it for GEMM. In a standard GEMM program, i.e. here, a single array with ct.load is sufficient. Why do you want sub array view first?

Splitting an array into a series of smaller arrays makes the program more elegant. Let’s take the following code snippet as an example:

@ct.function
def tile_gemm(
    x: ct.Array, y: ct.Array, 
    tm: ct.Constant, tn: ct.Constant, tk: ct.Constant
) -> ct.Tile:
    k_iter = x.shape[-1] // tk
    acc = ct.zeros((tm, tn), dtype=ct.float32)

    for k in range(k_iter):
        tile_x = ct.load(x, (0, k), (tm, tk))
        tile_y = ct.load(y, (0, k), (tn, tk))
        # mma is not possible with transposed x and no-transposed y here!
        ct.mma(tile_x, tile_y, acc)

    return acc

@ct.kernel
def device_gemm(
    x: ct.Array, y: ct.Array, out: ct.Array,
    tm: ct.Constant, tn: ct.Constant, tk: ct.Constant
):
    block_idx_x, block_idx_y = ct.bid(0), ct.bid(1)

    # can not create local_x here
    # cuda.tile._exception.TileTypeError: Non-constant slices are not supported
    local_x = x[block_idx_x * tm: block_idx_x * tm + tm, :]
    local_y = y[block_idx_y * tn: block_idx_y * tn + tn, :]

    acc = tile_gemm(local_x, local_y, tm, tn, tk)
    ct.store(out, (block_idx_x, block_idx_y), acc)

The tile&block‑based programming model in CUDA is excellent.

However, once we decompose the work down to blocks, each block should have its own “working space” — it should only see the part of the data that belongs to it. This greatly improves readability and maintainability of the program. We shouldn’t pass the full x and y matrices into tile_gemm; instead, we should pass only the sub‑matrices that the current block is responsible for.

In the example GEMM code, the device and block execution levels are not clearly separated. For complex programs, this hurts maintainability and readability. Structuring the code in two levels (device‑level and block‑level) would solve these problems.

[haijieg]

I understand what are you trying to do. I might look slightly cleaner, but it will add unnecessary index calculation to create those sub array views.

[chfwinter]

@ZhangZhiPku, instead of passing arrays, you might consider passing tiles in tile_gemm. I noticed that ct.extract can split the input tiles to create smaller ones. Would that work for your situation?

[ZhangZhiPku]

@ZhangZhiPku, instead of passing arrays, you might consider passing tiles in tile_gemm. I noticed that ct.extract can split the input tiles to create smaller ones. Would that work for your situation?

It probably won’t solve the problem. Cutile requires the shape of a tile being statically determined by ct.Constant. Since only static shapes can enable optimizations such as TMA, loop unrolling, vectorized memory access, and vectorized type casting.
However, the parameter x passed to tile_gemm is not feasible: the tile shape there would be x: [M: int, tk: ct.Constant], where M is a runtime variable.

[Da1sypetals]

I believe these operations are not in the range of cuTile's abstraction level. Forcing them in can destroy DX, which is pretty much the only reason why cuTile exist.