/tilegym-converting-cutile-to-julia

cuTile Python ↔ cuTile.jl (Julia) API Mapping

Import & Setup

Kernel Definition

Kernel Launch

Grid & Block IDs

Memory Operations

Padding Modes

Tile Construction

Type Conversion

Type Names

Arithmetic & Element-wise

Comparisons & Logic

Math Functions

Reductions

Scans (Prefix Operations)

Shape Operations

Matrix Operations

Conditional / Selection

Atomic Operations

Control Flow

Host Harness

Memory Layout Considerations

Python uses row-major (C-order), Julia uses column-major (Fortran-order).

For 2D arrays, this is largely transparent since cuTile handles it via strides.

For batched operations (3D+), consider reordering dimensions:

• Python: (Batch, M, K) — batch is first (outermost in row-major)
• Julia: (M, K, Batch) — batch is last (outermost in column-major)

This gives optimal memory access patterns in each language.

When converting, either:

1. Transpose the layout (recommended for performance): change array shapes and adjust kernel indexing
2. Keep the layout and accept potentially suboptimal memory access patterns

Bitwise Operations

1D Element-wise Pattern (TMA load/store)

For simple 1D element-wise ops (dropout, activations), use the TMA ct.load/ct.store pattern. No to_col_major/from_col_major needed — 1D arrays have no row/col distinction.

# 1D kernel using TMA block indexing
function my_1d_kernel(x::ct.TileArray{T,1}, output::ct.TileArray{T,1},
                      BLOCK_SIZE::Int) where {T}
    bid = ct.bid(1)
    x_tile = ct.load(x; index=bid, shape=(BLOCK_SIZE,))
    # ... process ...
    ct.store(output; index=bid, tile=result_tile)
    return nothing
end

Host harness: flatten input, pad to BLOCK_SIZE multiple, launch kernel, trim output.

2D Persistent Scheduling Pattern (RoPE, etc.)

For row-per-block kernels with many rows, use persistent scheduling with ct.Constant for tile sizes:

function my_kernel(data::ct.TileArray{T,2}, TILE_HD::Int) where {T}
    ct.@compiler_options occupancy=2
    bid = ct.bid(1)
    num_programs = ct.num_blocks(1)
    n_rows = size(data, 2)
    row_idx = bid
    while row_idx <= n_rows
        tile = ct.load(data; index=(Int32(1), row_idx), shape=(TILE_HD, 1))
        # ... process row ...
        ct.store(data; index=(Int32(1), row_idx), tile=result)
        row_idx += num_programs
    end
    return
end

# Launch with ct.Constant for tile size
ct.launch(my_kernel, num_blocks, data_cu, ct.Constant(tile_hd))

Kernel Patterns for Large Tensors

When a single ct.load of the entire data exceeds hardware limits, use one of these patterns:

Pattern A: Column-loop with ct.load/ct.store (Online Algorithm)

Best for TMA-based kernels where each chunk is a contiguous tile along columns. Uses for loops for column iteration and ct.num_tiles for tile count.

function online_kernel(output::ct.TileArray{T, 2}, input::ct.TileArray{T, 2},
                       TILE_SIZE::Int) where {T}
    row_idx = ct.bid(1)
    num_col_tiles = ct.num_tiles(input, 2, (1, TILE_SIZE))

    m_prev = fill(-Inf32, (1, 1))
    l_prev = zeros(Float32, 1, 1)

    for col_idx in Int32(1):num_col_tiles
        tile = ct.load(input; index=(row_idx, col_idx), shape=(1, TILE_SIZE),
                      padding_mode=ct.PaddingMode.NegInf)
        tile = convert(ct.Tile{Float32}, tile)
        tile_max = maximum(tile; dims=2)
        m_curr = max.(tile_max, m_prev)
        l_prev = l_prev .* exp.(m_prev .- m_curr)
        l_prev = sum(exp.(tile .- m_curr); dims=2) .+ l_prev
        m_prev = m_curr
    end

    for col_idx in Int32(1):num_col_tiles
        tile = ct.load(input; index=(row_idx, col_idx), shape=(1, TILE_SIZE),
                      padding_mode=ct.PaddingMode.NegInf)
        tile = convert(ct.Tile{Float32}, tile)
        result = exp.(tile .- m_prev) ./ l_prev
        ct.store(output; index=(row_idx, col_idx), tile=convert(ct.Tile{T}, result))
    end
    return
end

Pattern B: Chunked with ct.gather/ct.scatter and ct.Constant (Preferred)

Use when you need multiple passes over column chunks. Pass tile sizes as ct.Constant at launch — no @eval needed.

function chunked_kernel(output::ct.TileArray{T, 2}, input::ct.TileArray{T, 2},
                        n_cols::Int, TILE_SIZE::Int) where {T}
    ct.@compiler_options occupancy=4
    row_idx = ct.bid(1)
    num_chunks = (n_cols + TILE_SIZE - Int32(1)) ÷ Int32(TILE_SIZE)
    col_offsets_base = ct.arange(TILE_SIZE)
    row_tile = ct.Tile(row_idx)

    row_max = fill(-Inf32, (1,))
    denominator = zeros(Float32, TILE_SIZE)

    for chunk_idx in Int32(0):num_chunks - Int32(1)
        col_indices = ct.broadcast_to(ct.Tile(chunk_idx * Int32(TILE_SIZE)), (TILE_SIZE,)) .+ col_offsets_base
        chunk = ct.gather(input, (row_tile, col_indices); check_bounds=true, padding_value=T(-Inf))
        chunk = convert(ct.Tile{Float32}, chunk)
        row_max = max.(row_max, ct.Tile(maximum(chunk)))
    end
    # ... pass 2 and 3 similarly ...
    return
end

function julia_chunked_softmax(output::CuMatrix{T}, input::CuMatrix{T};
                               tile_size::Int=1024) where {T}
    M, N = size(input)
    ct.launch(chunked_kernel, M, output, input, ct.Constant(N), ct.Constant(tile_size))
    CUDA.synchronize()
    return
end

Quick Conversion Reference