cuTile.jl 1.0: Tile windows, atomics, and sparse views
Tim Besard
cuTile.jl has reached 1.0! The release adds tile windows via eachtile, masked and view-based atomics, an @atomic macro, sparse views, compiler remarks, and support for Tile IR 13.4, alongside a dedicated documentation site.
The package started as an experiment in expressing NVIDIA's tile-based programming model in Julia. Six months and three releases later, we're confident tagging an initial stable release. This comes with a dedicated documentation site.
Compared to v0.3, the 1.0 release adds the following features.
Tile windows
Previously, walking an array tile by tile meant passing the array, index, and shape to every ct.load and ct.store. The new ct.eachtile instead returns an indexable collection of fixed-shape windows. A blocked matrix multiplication shows the difference:
using CUDA, cuTile
import cuTile as ct
function matmul!(C, A, B)
a_tiles = ct.eachtile(A, (64, 32))
b_tiles = ct.eachtile(B, (32, 64))
c_tiles = ct.eachtile(C, (64, 64))
m, n = ct.bid(1), ct.bid(2)
acc = zeros(Float32, (64, 64))
for k in Int32(1):Int32(size(a_tiles, 2))
acc = muladd(a_tiles[m, k], b_tiles[k, n], acc)
end
c_tiles[m, n] = acc
return
end
A = CUDA.rand(Float16, 256, 128)
B = CUDA.rand(Float16, 128, 256)
C = CUDA.zeros(Float32, 256, 256)
@cuda backend=cuTile blocks=(4, 4) matmul!(C, A, B)
size(a_tiles, 2) returns the number of windows along that dimension, avoiding a separate trip-count calculation. step controls the distance between window origins: a smaller value produces overlap, while a larger one leaves gaps.
adjacent = ct.eachtile(a, (8, 8)) # step defaults to the shape
overlap = ct.eachtile(a, (8, 8); step=(4, 8)) # neighboring windows overlap
Partial edge windows are handled by the padding mode, as with a normal load. Unequal shape and step require Tile IR bytecode v13.3 or newer.
Atomics
cuTile now has three atomic-operation families, with different return values and ordering options.
The read-modify-write functions (ct.atomic_add and friends) return the old value and take a configurable memory order. In 1.0 they also accept a mask, useful for the tail block of a grid that does not divide the data evenly:
function histogram!(counts, data, n::Int32)
pid = ct.bid(1)
offs = (pid - Int32(1)) * Int32(128) .+ ct.arange(128)
vals = ct.load(data; index=pid, shape=(128,))
active = offs .<= n # mask off the tail
ct.atomic_add(counts, vals, Int32(1); mask=active)
return
end
The new ct.atomic_store_* family lowers to Tile IR's view-based atomic reductions. These reduce a tile into an array or an eachtile window and return nothing, using relaxed device-wide ordering:
function accumulate_tiles!(out, src)
tiles = ct.eachtile(out, (128,))
pid = ct.bid(1)
ct.atomic_store_add(tiles, 1, ct.load(src; index=pid, shape=(128,)))
return
end
ct.@atomic provides Base-style statement and value forms:
ct.@atomic counters[i] += update
ct.@atomic counters[i] = max(counters[i], value)
old_new = ct.@atomic counters[i] + value # returns old => new
Statement forms default to relaxed ordering, while value forms default to acquire-release. View-based reductions and ct.@atomic require Tile IR 13.3.
Sparse views
view and @view on a TileArray now accept positive step ranges. On arrays with two or more dimensions, one dimension may instead use a 1D integer tile, creating a sparse view that ct.load and ct.store lower to a Tile IR gather/scatter view:
function pick_rows!(dst, src)
rows = ct.arange(4; start=1, step=2) # rows 1, 3, 5, 7
selected = @view src[rows, 1:8]
tile = ct.load(selected, (4, 8))
ct.store(dst, (1, 1), tile)
return
end
The load shape is explicit and static, while the range starts may be runtime values. Sparse loads apply the requested padding and stores clip partially out-of-bounds elements; repeated indices are fine for loads, but conflicting stores are undefined. Step ranges and sparse views require Tile IR 13.3.
Compiler remarks
tileiras can report whether it selected tensor cores, vector loads, and other optimizations. code_tiled and @device_code_tiled now print those diagnostics with remarks=true.
Compile the matmul above with Float32 inputs:
A = CUDA.rand(Float32, 256, 128)
B = CUDA.rand(Float32, 128, 256)
C = CUDA.zeros(Float32, 256, 256)
ct.@device_code_tiled remarks=true @cuda backend=cuTile blocks=(4, 4) matmul!(C, A, B)
// tileiras optimization remarks
// Name: RemarkMemoryLoadInstructionSelected
// - RemarkId: 3
// - Remark: Load instruction selected
// Name: RemarkTensorCoreMMA
// - RemarkId: 1
// - Remark: MMA operation failed to optimize to use Tensor Cores, it is using FMA instructions instead
For this kernel, the Float32 multiply uses FMA instructions. With Float16 inputs and a Float32 accumulator, the compiler instead reports:
// Name: RemarkTensorCoreMMA
// - Remark: MMA operation successfully optimized to use Tensor Cores
Remarks require tileiras 13.4 or newer, which is still in early-access.
Programmatic dependent launch
Programmatic dependent launch can overlap the tail of a producer kernel with an independent preamble in the next kernel on the same stream. The producer signals when its dependents may start; the consumer is launched with dependent=true and waits before reading the producer's results:
function producer(a, producer_out)
ct.grid_dependency_control_launch_dependents()
tile = ct.load(a, 1, (32,)) # may overlap with the consumer
ct.store(producer_out, 1, tile)
return
end
function consumer(b, producer_out, out)
tile = ct.load(b, 1, (32,)) # independent preamble
ct.grid_dependency_control_wait()
ct.store(out, 1, tile + ct.load(producer_out, 1, (32,)))
return
end
stream = CUDA.stream()
@cuda backend=cuTile blocks=1 stream producer(a, producer_out)
@cuda backend=cuTile blocks=1 dependent=true stream consumer(b, producer_out, out)
The overlap is opportunistic, so correctness must never depend on the two kernels actually running concurrently. The feature requires Tile IR 13.4 and compute capability 9.0 or newer.
Other changes
Tile IR 13.4 is supported and emitted by default when
tileirasaccepts it. It bringsct.insert, the inverse ofct.extract, andcheck_bounds=falseonct.load/ct.store, an explicit promise that the whole tile is in bounds which drops the padding and selects Tile IR's unchecked encoding.Explicit rounding modes on float-to-float conversions:
Float32.(tile, RoundDown), withRoundNearest,RoundToZero,RoundDown,RoundUpandRoundNearestTiesAway. Supported modes and source/target pairs depend on the Tile IR version. Directed rounding generally requires 13.4; some conversions toFloat8_E8M0FNUare available in 13.3.64-bit indexing.
TileArraygained an index-type parameter, and arrays whose sizes or strides exceed the 32-bit range automatically useInt64. Smaller arrays continue to useInt32.ct.TileArray(a; index=Int64)selects wide indexing explicitly. Wide indexing requires Tile IR 13.3.Array construction syntax works in kernels:
[a, b, c], typed forms likeFloat32[a, b], bracket concatenation ([A; B],[a b; C],[A;;; B]) andcat(A, B...; dims).ct.catwas removed in favor of these.Multi-dimensional reductions.
dimsnow takes an integer, an iterable of integers, or:, for both tile-level and host-levelct.Tiledreductions.Memory ordering on plain loads and stores, not just atomics:
ct.loadandct.storeacceptmemory_orderandmemory_scope.Configuration moved to preferences. The
JULIA_CUTILE_CACHE_DIRandJULIA_CUTILE_CACHE_SIZEenvironment variables were replaced by thedisk_cache,cache_dir, andcache_size_bytespreferences. A newcompiler_timeout_secondspreference bounds eachtileirasinvocation.
cuTile.jl 1.0 requires CUDA.jl 6.3. See NEWS.md for the user-facing release history and the release notes for the merged pull requests. Please file an issue if you run into a problem.