Merge pull request #228 from JuliaGPU/tb/examples

Make examples slightly more idiomatic and aligned

Author: maleadt

README.md

@@ -98,17 +98,17 @@ Benchmarks comparing cuTile.jl against cuTile Python on an RTX 5080 (`tileiras`
 
 | Kernel | Size | Julia | Python | Status |
 |--------|------|-------|--------|--------|
-| Vector Addition | 2^27 f32 | 845 GB/s | 846 GB/s | OK (=) |
+| Vector Addition | 2^27 f32 | 844 GB/s | 845 GB/s | OK (=) |
 | Matrix Transpose | 8192² f32 | 812 GB/s | 814 GB/s | OK (=) |
-| Layer Norm fwd | 4096² f32 | 983 GB/s | 716 GB/s | +37% |
-| Layer Norm bwd | 4096² f32 | 248 GB/s | 251 GB/s | OK (-1%) |
-| Matrix Multiplication | 4096³ f32 | 47.5 TFLOPS | 43.5 TFLOPS | +9% |
-| Batch Matrix Multiply | 1024×512×2048 ×8 f32 | 34.0 TFLOPS | 30.8 TFLOPS | +10% |
-| FFT (3-stage Cooley-Tukey) | 512-pt ×64 c64 | 529 μs | 554 μs | +5% |
-| Mixture of Experts | 256tok 1024h 32e 2048i f16 | 27.0 TFLOPS | 20.1 TFLOPS | +34% |
-| Attention (FMHA) | 8×16×1024² ×64 f16 causal | 103.6 TFLOPS | 63.4 TFLOPS | +63% |
-| Softmax (TMA) | 4096² f32 | 849 GB/s | 857 GB/s | OK (-1%) |
-| Softmax (Chunked) | 4096² f32 | 1684 GB/s | 1640 GB/s | OK (+3%) |
+| Layer Norm fwd | 4096² f32 | 986 GB/s | 716 GB/s | +38% |
+| Layer Norm bwd | 4096² f32 | 246 GB/s | 251 GB/s | OK (-2%) |
+| Matrix Multiplication | 4096³ f32 | 47.4 TFLOPS | 43.5 TFLOPS | +9% |
+| Batch Matrix Multiply | 1024×512×2048 ×8 f32 | 34.2 TFLOPS | 30.9 TFLOPS | +11% |
+| FFT (3-stage Cooley-Tukey) | 512-pt ×64 c64 | 545 μs | 550 μs | OK (+1%) |
+| Mixture of Experts | 256tok 1024h 32e 2048i f16 | 27.7 TFLOPS | 20.3 TFLOPS | +36% |
+| Attention (FMHA) | 8×16×1024² ×64 f16 causal | 102.7 TFLOPS | 63.3 TFLOPS | +62% |
+| Softmax (TMA) | 4096² f32 | 838 GB/s | 843 GB/s | OK (-1%) |
+| Softmax (Chunked) | 4096² f32 | 1672 GB/s | 1636 GB/s | OK (+2%) |
 
 
 ## Supported Operations

examples/fmha.jl

@@ -28,13 +28,15 @@ function fmha_kernel(Q::ct.TileArray{T, 4}, K::ct.TileArray{T, 4},
                      CAUSAL::Bool, EVEN_K::Bool) where {T}
     ct.@compiler_options occupancy=2
 
-    # Map block IDs to batch and head indices
-    # Julia: bid(1) = x (seq tiles), bid(2) = y (batch * heads)
+    # Map block IDs to batch and head indices.
+    # Julia: bid(1) = x (seq tiles), bid(2) = y (batch * heads).
+    # `cld`/`mod1` give 1-indexed batch/head/kv-head straight from the
+    # 1-indexed bid_y, with no 0-indexed detour.
     bid_x = ct.bid(1)
-    bid_y = ct.bid(2) - Int32(1)  # 0-indexed for div/mod arithmetic
-    batch_idx = fld(bid_y, Int32(H)) + Int32(1)
-    head_idx = rem(bid_y, Int32(H)) + Int32(1)
-    off_kv_h = fld(head_idx - Int32(1), Int32(QUERY_GROUP_SIZE)) + Int32(1)
+    bid_y = ct.bid(2)
+    batch_idx = cld(bid_y, Int32(H))
+    head_idx = mod1(bid_y, Int32(H))
+    off_kv_h = cld(head_idx, Int32(QUERY_GROUP_SIZE))
 
     # Adjust qk_scale for exp2
     qk_scale = qk_scale * INV_LOG_2

examples/layernorm.jl

@@ -71,20 +71,14 @@ end
 #=============================================================================
  LayerNorm Backward Kernels
 
- Backward pass: computes gradients for LayerNorm.
- The full backward pass has two kernels:
- 1. layer_norm_bwd_dx - Computes dX (gradient with respect to input)
- 2. layer_norm_bwd_dwdb - Computes dW and dB (requires atomic accumulation)
-
- For now, we implement a simplified backward that just computes dX.
+ Backward pass: computes gradients for LayerNorm via two kernels:
+  1. layer_norm_bwd_dx_partial_dwdb - dX + per-group partial dW/dB (atomic).
+  2. layer_norm_bwd_dwdb            - final reduction of partial dW/dB.
 =============================================================================#
 
-"""
-Helper function for backward pass - loads data and computes common terms.
-This gets inlined by Julia's compiler.
-bid_m and j are 1-indexed (block ID and tile index).
-"""
-@inline function bwd_helper(X, W, DY, bid_m, j, mean, rstd, TILE_N, N)
+# Helper function for backward pass - loads data and computes common terms.
+# bid_m and j are 1-indexed (block ID and tile index).
+function bwd_helper(X, W, DY, bid_m, j, mean, rstd, TILE_N, N)
     tx = ct.load(X; index=(j, bid_m), shape=(TILE_N, 1), padding_mode=ct.PaddingMode.Zero)
     tw = reshape(ct.load(W; index=j, shape=(TILE_N,), padding_mode=ct.PaddingMode.Zero), (TILE_N, 1))
     tdy = ct.load(DY; index=(j, bid_m), shape=(TILE_N, 1), padding_mode=ct.PaddingMode.Zero)
@@ -103,53 +97,6 @@ bid_m and j are 1-indexed (block ID and tile index).
     return tdy, xhat_masked, wdy_masked
 end
 
-"""
-    layer_norm_bwd_dx(DX, DY, X, W, Mean, Rstd, TILE_N)
-
-Backward pass: computes gradient with respect to input X.
-
-Args:
-    DX: Output gradient with respect to X (N, M).
-    DY: Input gradient with respect to Y (N, M).
-    X: Input tensor (N, M).
-    W: Weight tensor (N,).
-    Mean: Mean tensor (M,).
-    Rstd: Reciprocal standard deviation tensor (M,).
-    TILE_N: Tile size along N dimension.
-"""
-function layer_norm_bwd_dx(DX::ct.TileArray{Float32, 2}, DY::ct.TileArray{Float32, 2},
-                           X::ct.TileArray{Float32, 2}, W::ct.TileArray{Float32, 1},
-                           Mean::ct.TileArray{Float32, 1}, Rstd::ct.TileArray{Float32, 1},
-                           TILE_N::Int)
-    bid_m = ct.bid(1)
-    num_tiles = ct.num_tiles(X, 1, (TILE_N, 1))
-    N = size(X, 1)
-
-    # Load mean and rstd for this row
-    mean = ct.load(Mean; index=bid_m, shape=(1,), padding_mode=ct.PaddingMode.Zero)
-    rstd = ct.load(Rstd; index=bid_m, shape=(1,), padding_mode=ct.PaddingMode.Zero)
-
-    # First pass: compute c1 and c2 reduction terms
-    c1 = zeros(Float32, (TILE_N, 1))
-    c2 = zeros(Float32, (TILE_N, 1))
-    for j in Int32(1):num_tiles
-        _, xhat, wdy = bwd_helper(X, W, DY, bid_m, j, mean, rstd, TILE_N, N)
-        c1 = c1 .+ (xhat .* wdy)
-        c2 = c2 .+ wdy
-    end
-    c1 = sum(c1; dims=1) / N
-    c2 = sum(c2; dims=1) / N
-
-    # Second pass: compute dX
-    for j in Int32(1):num_tiles
-        _, xhat, wdy = bwd_helper(X, W, DY, bid_m, j, mean, rstd, TILE_N, N)
-        tdx = (wdy .- (xhat .* c1 .+ c2)) .* rstd
-        ct.store(DX; index=(j, bid_m), tile=tdx)
-    end
-
-    return
-end
-
 """
     layer_norm_bwd_dx_partial_dwdb(DX, DY, DW, DB, X, W, Mean, Rstd, Locks, GROUP_SIZE_M, TILE_N)
 

examples/matmul.jl

@@ -8,17 +8,19 @@ using LinearAlgebra
 using cuTile: cuTile
 import cuTile as ct
 
-# 2D swizzle for better L2 cache locality
-# Groups blocks to access nearby memory regions together
-@inline function swizzle_2d(M, N, tm, tn, GROUP_SIZE_M, bid)
+# 2D swizzle for better L2 cache locality. Takes a 1-indexed bid and
+# returns 1-indexed (bid_m, bid_n). Modular arithmetic is done on the
+# 0-indexed bid internally; the conversion is contained in this helper.
+function swizzle_2d(M, N, tm, tn, GROUP_SIZE_M, bid)
     num_bid_m = cld(M, Int32(tm))
     num_bid_n = cld(N, Int32(tn))
     num_bid_in_group = Int32(GROUP_SIZE_M) * num_bid_n
-    group_id = fld(bid, num_bid_in_group)
+    bid0 = bid - Int32(1)
+    group_id = fld(bid0, num_bid_in_group)
     first_bid_m = group_id * Int32(GROUP_SIZE_M)
     group_size_m = min(num_bid_m - first_bid_m, Int32(GROUP_SIZE_M))
-    bid_m = first_bid_m + rem(bid, group_size_m)
-    bid_n = fld(rem(bid, num_bid_in_group), group_size_m)
+    bid_m = first_bid_m + rem(bid0, group_size_m) + Int32(1)
+    bid_n = fld(rem(bid0, num_bid_in_group), group_size_m) + Int32(1)
     return bid_m, bid_n
 end
 
@@ -31,11 +33,7 @@ function matmul_kernel(A::ct.TileArray{T,2}, B::ct.TileArray{T,2}, C::ct.TileArr
     bid = ct.bid(1)
     M = size(A, 1)
     N = size(B, 2)
-    # swizzle_2d expects 0-indexed bid, returns 0-indexed tile coords
-    bid_m_0, bid_n_0 = swizzle_2d(M, N, tm, tn, 8, bid - Int32(1))
-    # Convert to 1-indexed tile coordinates
-    bid_m = bid_m_0 + Int32(1)
-    bid_n = bid_n_0 + Int32(1)
+    bid_m, bid_n = swizzle_2d(M, N, tm, tn, 8, bid)
 
     # Number of K tiles to iterate over
     num_k = ct.num_tiles(A, 2, (tm, tk))

examples/moe.jl

@@ -18,15 +18,16 @@ import cuTile as ct
  Helper: 2D swizzle (same pattern as matmul.jl)
 =============================================================================#
 
-@inline function swizzle_2d(M, N, tm, tn, GROUP_SIZE_M, bid)
+function swizzle_2d(M, N, tm, tn, GROUP_SIZE_M, bid)
     num_bid_m = cld(M, Int32(tm))
     num_bid_n = cld(N, Int32(tn))
     num_bid_in_group = Int32(GROUP_SIZE_M) * num_bid_n
-    group_id = fld(bid, num_bid_in_group)
+    bid0 = bid - Int32(1)
+    group_id = fld(bid0, num_bid_in_group)
     first_bid_m = group_id * Int32(GROUP_SIZE_M)
     group_size_m = min(num_bid_m - first_bid_m, Int32(GROUP_SIZE_M))
-    bid_m = first_bid_m + rem(bid, group_size_m)
-    bid_n = fld(rem(bid, num_bid_in_group), group_size_m)
+    bid_m = first_bid_m + rem(bid0, group_size_m) + Int32(1)
+    bid_n = fld(rem(bid0, num_bid_in_group), group_size_m) + Int32(1)
     return bid_m, bid_n
 end
 
@@ -53,19 +54,18 @@ function fused_moe_kernel(A::ct.TileArray{T, 2}, B::ct.TileArray{T, 3},
     K = size(B, 1)
     N = size(B, 2)
 
-    bid = ct.bid(1) - Int32(1)  # 0-indexed for swizzle
+    bid = ct.bid(1)
     bid_m, bid_n = swizzle_2d(M, N, TILE_M, TILE_N, Int32(8), bid)
 
     # Gather 1-indexed token IDs for this block
-    token_id_indices = bid_m * Int32(TILE_M) .+ ct.arange(TILE_M)
+    token_id_indices = (bid_m - Int32(1)) * Int32(TILE_M) .+ ct.arange(TILE_M)
     token_ids = ct.gather(sorted_token_ids, token_id_indices)
 
-    # Map 1-indexed flat token_id to 1-indexed column in A
-    # token_id k → original token = (k-1) ÷ num_token_replicas + 1
-    a_tok_indices = (token_ids .- Int32(1)) .÷ Int32(num_token_replicas) .+ Int32(1)
+    # 1-indexed flat token_id → 1-indexed column in A. Each original token
+    # has `num_token_replicas` consecutive ids; ceil-divide recovers it.
+    a_tok_indices = cld.(token_ids, Int32(num_token_replicas))
 
-    # Expert for this block (scalar, 1-indexed tile index for load)
-    expert_id = sorted_expert_ids[bid_m + Int32(1)]
+    expert_id = sorted_expert_ids[bid_m]
 
     acc = zeros(Float32, TILE_N, TILE_M)
     num_k = cld(K, Int32(TILE_K))
@@ -81,7 +81,7 @@ function fused_moe_kernel(A::ct.TileArray{T, 2}, B::ct.TileArray{T, 3},
         # B is (K, N, num_experts): load (TILE_N, TILE_K) slice for this expert
         # order=(2,1,3) folds the transpose into the load via dim_map, matching
         # Python cuTile's order=(0,2,1) and avoiding an explicit permute in Tile IR.
-        b = ct.load(B; index=(bid_n + Int32(1), k, expert_id),
+        b = ct.load(B; index=(bid_n, k, expert_id),
                     shape=(TILE_N, TILE_K, 1), order=(2, 1, 3),
                     padding_mode=ct.PaddingMode.Zero)
         b = reshape(b, (TILE_N, TILE_K))
@@ -97,7 +97,7 @@ function fused_moe_kernel(A::ct.TileArray{T, 2}, B::ct.TileArray{T, 3},
 
     # Scatter result to C at token_id positions
     # C is (N, total_tokens): dim 1 = N, dim 2 = tokens
-    c_n_indices = bid_n * Int32(TILE_N) .+ ct.arange(TILE_N)  # 1-indexed
+    c_n_indices = (bid_n - Int32(1)) * Int32(TILE_N) .+ ct.arange(TILE_N)  # 1-indexed
     output = convert(ct.Tile{T}, acc)
     ct.scatter(C, (reshape(c_n_indices, (TILE_N, 1)),
                    reshape(token_ids, (1, TILE_M))), output)
@@ -204,24 +204,12 @@ function invoke_silu_and_mul_kernel(AB, C)
     inter = size(C, 1)  # C is (intermediate, total_tokens)
     total = size(AB, 2)
 
-    # Split AB(inter*2, total) into gate and up halves along dim 1.
-    #A_half = AB[1:inter, :]
-    #B_half = AB[inter+1:2*inter, :]
-    # FIXME: CUDA.jl's CuArray indexing (AB[1:inter, :]) uses a slow generic kernel.
-    # Use unsafe_copy2d! (cuMemcpy2D) for hardware-accelerated pitched 2D copy instead.
-    T = eltype(AB)
-    A_half = similar(AB, inter, total)
-    B_half = similar(AB, inter, total)
-    src_pitch = size(AB, 1) * sizeof(T)
-    dst_pitch = inter * sizeof(T)
-    CUDACore.unsafe_copy2d!(pointer(A_half), CUDACore.DeviceMemory,
-                        pointer(AB), CUDACore.DeviceMemory,
-                        inter, total; srcPitch=src_pitch, dstPitch=dst_pitch,
-                        async=true)
-    CUDACore.unsafe_copy2d!(pointer(B_half), CUDACore.DeviceMemory,
-                        pointer(AB) + inter * sizeof(T), CUDACore.DeviceMemory,
-                        inter, total; srcPitch=src_pitch, dstPitch=dst_pitch,
-                        async=true)
+    # Split AB(inter*2, total) into gate and up halves along dim 1 — mirrors
+    # cuTile Python's `AB.chunk(2, dim=-1)`. Views are non-contiguous along
+    # dim 2 but each block only loads a (TILE_N, 1) tile, so codegen is
+    # unaffected.
+    A_half = view(AB, 1:inter, :)
+    B_half = view(AB, (inter + 1):(2 * inter), :)
 
     tile_n = nextpow(2, inter)
     @cuda backend=cuTile blocks=total silu_and_mul_kernel(A_half, B_half, C, ct.Constant(tile_n))

src/launch.jl

@@ -27,12 +27,13 @@ struct KernelAdaptor end
 Adapt.adapt_storage(::KernelAdaptor, arr::AbstractArray) = TileArray(arr)
 Adapt.adapt_storage(::KernelAdaptor, t::Type) = Constant(t)
 
-# Adapt's default `adapt_structure(to, ::PermutedDimsArray)` recurses by
-# rebuilding `PermutedDimsArray(adapt(parent), perm)`. We can't follow that
+# Adapt's defaults for `PermutedDimsArray` and `SubArray` recurse by
+# rebuilding the wrapper around `adapt(parent)`. We can't follow that
 # pattern because `TileArray` isn't `<:AbstractArray` — strided-wrapper
 # state is absorbed into its `sizes`/`strides` fields directly. Short-circuit
 # the recursion so the whole wrapper becomes a single TileArray.
 Adapt.adapt_structure(::KernelAdaptor, arr::PermutedDimsArray) = TileArray(arr)
+Adapt.adapt_structure(::KernelAdaptor, arr::SubArray) = TileArray(arr)
 
 """
     cuTileconvert(x)