flydsl-tile-programming

name	flydsl-tile-programming
description	Guided step-by-step wizard for producing a new FlyDSL GPU kernel from a requirement: classify the kernel type, pick a skeleton, fill in compute, add control flow / sync / LDS, then test on GPU. Use when the user wants to WRITE a new kernel, port a Triton kernel to FlyDSL, or learn tile programming by following a procedure. For API/layout-algebra lookups, per-op reference tables, and troubleshooting, use the flydsl-kernel-authoring skill instead.
allowed-tools	Read Edit Bash Grep Glob Agent

FlyDSL Tile Programming

Guide users through writing GPU kernels using FlyDSL's tile programming model (CuTe-style layout algebra). This skill is a step-by-step wizard that takes a kernel requirement and produces a correct, tested FlyDSL kernel.

Trigger: User wants to write a new FlyDSL kernel, port a Triton kernel to FlyDSL, or learn tile programming patterns.

Prerequisites: FlyDSL installed (editable mode via pip install -e .). GPU access required for testing.

Scope (read this first): This skill is the procedure — follow the steps in order to produce a kernel. It is the companion to the flydsl-kernel-authoring skill, which is the reference (the full layout-algebra API surface, per-op tables, environment variables, and an exhaustive troubleshooting list). When you need to look something up rather than follow a step, go to flydsl-kernel-authoring. This wizard links there instead of duplicating those tables.

---

Step 1: Classify the Kernel Type

Ask the user what kind of kernel they need. Map to one of these patterns:

Pattern	Examples	Key Primitives
Elementwise	vecadd, scale, relu, abs	logical_divide + fx.copy
Reduction	sum, max, softmax, layernorm	make_buffer_tensor + warp shuffle + LDS
Tiled Copy	transpose, permute, gather	zipped_divide + TiledCopy
GEMM	matmul, batched gemm	TiledMma + TiledCopy + LDS
Fused	fused attention, GEMM+epilogue	Combine GEMM + elementwise

---

Step 2: Generate Kernel Skeleton

Based on the pattern, generate the appropriate skeleton. Every FlyDSL kernel has two parts:

import torch
import flydsl.compiler as flyc
import flydsl.expr as fx

@flyc.kernel
def my_kernel(A: fx.Tensor, B: fx.Tensor, ...):
    tid = fx.thread_idx.x
    bid = fx.block_idx.x
    # ... kernel body ...

@flyc.jit
def my_launch(A: fx.Tensor, B: fx.Tensor, ...,
              stream: fx.Stream = fx.Stream(None)):
    my_kernel(A, B, ...).launch(
        grid=(grid_x, grid_y, grid_z),
        block=(block_x, 1, 1),
        stream=stream
    )

Pattern A: Elementwise Kernel

The simplest pattern. Each thread processes VEC_WIDTH elements independently.

Data flow: Global -> Register -> Compute -> Register -> Global

import torch
import flydsl.compiler as flyc
import flydsl.expr as fx
from flydsl.expr.typing import Vector as Vec

BLOCK_DIM = 256
VEC_WIDTH = 4

@flyc.kernel
def elementwise_kernel(
    A: fx.Tensor,
    Out: fx.Tensor,
    BLOCK_DIM: fx.Constexpr[int],
    VEC_WIDTH: fx.Constexpr[int],
):
    bid = fx.block_idx.x
    tid = fx.thread_idx.x

    # === Step 1: Divide global tensor into block-sized tiles ===
    tile_size = BLOCK_DIM * VEC_WIDTH
    tA = fx.logical_divide(A, fx.make_layout(tile_size, 1))
    tOut = fx.logical_divide(Out, fx.make_layout(tile_size, 1))

    # === Step 2: Select this block's tile ===
    tA = fx.slice(tA, (None, bid))
    tOut = fx.slice(tOut, (None, bid))

    # === Step 3: Divide tile for per-thread vectorized access ===
    tA = fx.logical_divide(tA, fx.make_layout(VEC_WIDTH, 1))
    tOut = fx.logical_divide(tOut, fx.make_layout(VEC_WIDTH, 1))

    # === Step 4: Allocate register and set up copy atom ===
    copy_bits = VEC_WIDTH * 32
    copy_atom = fx.make_copy_atom(fx.UniversalCopy(copy_bits), fx.Float32)
    rA = fx.make_rmem_tensor(VEC_WIDTH, fx.Float32)
    rOut = fx.make_rmem_tensor(VEC_WIDTH, fx.Float32)

    # === Step 5: Load -> Compute -> Store ===
    fx.copy(copy_atom, fx.slice(tA, (None, tid)), rA)

    vA = Vec(fx.memref_load_vec(rA))
    # --- YOUR COMPUTE HERE ---
    vOut = vA * vA  # example: square
    # --- END COMPUTE ---
    fx.memref_store_vec(vOut, rOut)

    fx.copy(copy_atom, rOut, fx.slice(tOut, (None, tid)))

@flyc.jit
def elementwise_launch(
    A: fx.Tensor, Out: fx.Tensor, N: fx.Int32,
    stream: fx.Stream = fx.Stream(None),
):
    tile_size = BLOCK_DIM * VEC_WIDTH
    grid_x = (N + tile_size - 1) // tile_size
    elementwise_kernel(A, Out, BLOCK_DIM, VEC_WIDTH).launch(
        grid=(grid_x, 1, 1), block=(BLOCK_DIM, 1, 1), stream=stream
    )

# === Test ===
N = 1024
A = torch.randn(N, dtype=torch.float32, device="cuda")
Out = torch.empty(N, dtype=torch.float32, device="cuda")
elementwise_launch(A, Out, N, stream=torch.cuda.Stream())
torch.cuda.synchronize()
assert torch.allclose(Out, A * A, atol=1e-5)

Pattern B: Tiled 2D Copy (Transpose, Gather)

Uses zipped_divide + TiledCopy for 2D data movement with explicit thread-value mapping.

Data flow: Global[M,N] -> Fragment -> Global[M,N] (with layout change)

@flyc.kernel
def tiled_copy_kernel(A: fx.Tensor, B: fx.Tensor):
    tid = fx.thread_idx.x
    bid = fx.block_idx.x

    block_m, block_n = 8, 24
    # make_tile is varargs -- passing a list raises ValueError
    tile = fx.make_tile(
        fx.make_layout(block_m, 1),
        fx.make_layout(block_n, 1),
    )

    # Wrap as buffer tensors (AMD buffer descriptors)
    A = fx.rocdl.make_buffer_tensor(A)
    B = fx.rocdl.make_buffer_tensor(B)

    # Divide into tiles, select block's tile
    bA = fx.zipped_divide(A, tile)
    bB = fx.zipped_divide(B, tile)
    bA = fx.slice(bA, (None, bid))
    bB = fx.slice(bB, (None, bid))

    # Thread-value layout: how threads cooperate on the tile
    thr_layout = fx.make_layout((4, 1), (1, 1))   # 4 threads along M
    val_layout = fx.make_layout((1, 8), (1, 1))    # each loads 8 along N
    copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32)
    # Build the TV layout with make_layout_tv -- do NOT hand-build it from
    # raked_product alone: make_layout_tv also derives the TV mapping via
    # composition(right_inverse(layout_mn), ...). See examples/02-tiledCopy.py.
    tile_mn, tv_layout = fx.make_layout_tv(thr_layout, val_layout)

    # Build tiled copy and get thread's partition
    tiled_copy = fx.make_tiled_copy(copy_atom, tv_layout, tile_mn)
    thr_copy = tiled_copy.get_slice(tid)
    src = thr_copy.partition_S(bA)
    dst = thr_copy.partition_D(bB)
    frag = fx.make_fragment_like(src)

    # Copy: global A -> frag -> global B
    fx.copy(copy_atom, src, frag)
    fx.copy(copy_atom, frag, dst)

Pattern C: Tiled MMA (GEMM)

Uses TiledMma + TiledCopy for matrix multiply with AMD MFMA instructions.

Data flow: Global -> (TiledCopy) -> Fragment A,B -> (MFMA) -> Fragment C -> Global

block_m, block_n, block_k = 64, 64, 8

@flyc.kernel
def gemm_kernel(A: fx.Tensor, B: fx.Tensor, C: fx.Tensor):
    tid = fx.thread_idx.x
    bid = fx.block_idx.x

    # Define tiles
    tileA = fx.make_tile(block_m, block_k)
    tileB = fx.make_tile(block_n, block_k)
    tileC = fx.make_tile(block_m, block_n)

    # Wrap as buffer tensors
    A = fx.rocdl.make_buffer_tensor(A)
    B = fx.rocdl.make_buffer_tensor(B)
    C = fx.rocdl.make_buffer_tensor(C)

    # Divide and select block's tile
    bA = fx.slice(fx.zipped_divide(A, tileA), (None, bid))
    bB = fx.slice(fx.zipped_divide(B, tileB), (None, bid))
    bC = fx.slice(fx.zipped_divide(C, tileC), (None, bid))

    # === MMA setup ===
    # MFMA(M, N, K, ABType[, AccType]) -- 4th arg is the A/B operand type;
    # the accumulator defaults to f32. Here operand and acc coincide.
    mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(16, 16, 4, fx.Float32))

    # Tile the MMA atom across threads: 2x2 = 4 MMA atoms per warp
    tiled_mma = fx.make_tiled_mma(
        mma_atom,
        fx.make_layout((2, 2, 1), (1, 2, 0))  # (M_rep, N_rep, K_rep)
    )
    thr_mma = tiled_mma.thr_slice(tid)

    # === Copy setup (matched to MMA layout) ===
    copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy32b(), fx.Float32)
    tiled_copy_A = fx.make_tiled_copy_A(copy_atom, tiled_mma)
    tiled_copy_B = fx.make_tiled_copy_B(copy_atom, tiled_mma)
    tiled_copy_C = fx.make_tiled_copy_C(copy_atom, tiled_mma)

    thr_copy_A = tiled_copy_A.get_slice(tid)
    thr_copy_B = tiled_copy_B.get_slice(tid)
    thr_copy_C = tiled_copy_C.get_slice(tid)

    # === Partition data ===
    # Copy partitions (for data movement)
    copy_src_A = thr_copy_A.partition_S(bA)
    copy_src_B = thr_copy_B.partition_S(bB)
    copy_dst_C = thr_copy_C.partition_S(bC)

    # MMA partitions (for compute)
    part_A = thr_mma.partition_A(bA)
    part_B = thr_mma.partition_B(bB)
    part_C = thr_mma.partition_C(bC)

    # === Allocate fragments (registers) ===
    frag_A = thr_mma.make_fragment_A(part_A)
    frag_B = thr_mma.make_fragment_B(part_B)
    frag_C = thr_mma.make_fragment_C(part_C)

    # Retile fragments for copy compatibility
    copy_frag_A = thr_copy_A.retile(frag_A)
    copy_frag_B = thr_copy_B.retile(frag_B)
    copy_frag_C = thr_copy_C.retile(frag_C)

    # === Execute: Load A,B -> GEMM -> Store C ===
    fx.copy(copy_atom, copy_src_A, copy_frag_A, pred=None)
    fx.copy(copy_atom, copy_src_B, copy_frag_B, pred=None)
    fx.gemm(mma_atom, frag_C, frag_A, frag_B, frag_C)
    fx.copy(copy_atom, copy_frag_C, copy_dst_C, pred=None)

Pattern D: Raw Buffer Load/Store (Low-level Escape Hatch)

Direct AMD buffer intrinsics, bypassing the layout algebra. Reach for this only when the access has no layout form — typically a scalar base with a per-thread element offset. For anything with a tile structure, use Pattern A/B/C, which go through make_buffer_tensor + copy atoms and get an OOB-checked V# descriptor built for you.

from kernels.common import buffer_ops   # moved out of flydsl.expr in #880

@flyc.kernel
def buffer_kernel(A: fx.Tensor, B: fx.Tensor, N: fx.Constexpr[int]):
    tid = fx.thread_idx.x
    bid = fx.block_idx.x
    gid = bid * 256 + tid

    rsrc_a = buffer_ops.create_buffer_resource(A)
    rsrc_b = buffer_ops.create_buffer_resource(B)

    # offset is in ELEMENTS (not bytes!) -- buffer_load converts internally
    # fx.T.* are properties on Types -- reference them, do not call them
    data = buffer_ops.buffer_load(rsrc_a, gid * 4, vec_width=4, dtype=fx.T.f32)
    # ... compute on data ...
    buffer_ops.buffer_store(data, rsrc_b, gid * 4)

The element-vs-byte offset is a classic source of bugs; see the kernel-code-cleanup skill to migrate an existing kernel onto the layout API.

---

Step 3: Fill in the Compute Logic

Common compute recipes (all work on vectors):

from flydsl.expr.typing import Vector as Vec

# Scale: C = A * scalar
scale = Vec.filled(VEC_WIDTH, 2.0, fx.Float32)
vC = Vec(vA) * scale

# Add: C = A + B
vC = Vec(vA) + Vec(vB)

# FMA: D = A * B + C
vC = Vec(vA) * Vec(vB) + Vec(vC)

# ReLU: C = max(A, 0)
zero = Vec.filled(VEC_WIDTH, 0.0, fx.Float32)
vC = Vec(vA).maximumf(zero)

# Abs: C = |A|
v = Vec(vA)
neg = -v
is_neg = v < zero