This guide covers writing GPU kernels with FlyDSL: the @flyc.jit and @flyc.kernel decorators, expression API, launch configuration, shared memory, and synchronization. It documents the @flyc.kernel/@flyc.jit API from flydsl.compiler and flydsl.expr (python/flydsl/).
| Concept | API | Description |
|---|---|---|
| JIT host func | @flyc.jit |
Emit host-side launcher with JIT compilation |
| GPU kernel | @flyc.kernel |
Define GPU kernel function |
| Launch | kernel(...).launch(grid=, block=) |
Configure and emit GPU launch |
| Thread ID | fx.gpu.thread_idx.x |
Get thread index in workgroup |
| Block ID | fx.gpu.block_idx.x |
Get block/workgroup index |
| Block dim | fx.gpu.block_dim.x |
Get block dimension size |
| Compile-time | fx.Constexpr[int] |
Compile-time constant parameter |
| Tensor arg | fx.Tensor |
GPU tensor argument (via DLPack) |
| Stream arg | fx.Stream |
CUDA/HIP stream argument |
| Barrier | fx.gpu.barrier() |
Workgroup synchronization |
| Constants | fx.Int32 / fx.Int64 / fx.Float32 |
Create typed DSL constants (fx.Int64 for index/offset values; fx.Index is deprecated) |
| Range loop | range_constexpr(n) |
Compile-time unrolled loop |
| Buffer load | buffer_ops.buffer_load(rsrc, off) |
AMD buffer load intrinsic |
| Scalar buffer load | fx.rocdl.s_buffer_load(rsrc, off, count=1) |
Wave-uniform AMD buffer load; off is in i32 elements and the result is fx.Int32 (or a four-word fx.Vector with count=4) |
import flydsl.compiler as flyc
import flydsl.expr as fx
from flydsl.expr import gpu
@flyc.kernel
def vec_add_kernel(
A: fx.Tensor,
B: fx.Tensor,
C: fx.Tensor,
N: fx.Constexpr[int],
):
tid = gpu.thread_idx.x
bid = gpu.block_idx.x
idx = bid * 256 + tid
# ... kernel body using fx.*, ArithValue, Vector, and buffer ops ...
@flyc.jit
def vec_add(
A: fx.Tensor,
B: fx.Tensor,
C: fx.Tensor,
N: fx.Constexpr[int],
stream: fx.Stream = fx.Stream(None),
):
vec_add_kernel(A, B, C, N).launch(
grid=(N // 256,),
block=(256,),
stream=stream,
)
# Usage:
import torch
A = torch.randn(1024, device="cuda", dtype=torch.float32)
B = torch.randn(1024, device="cuda", dtype=torch.float32)
C = torch.empty(1024, device="cuda", dtype=torch.float32)
vec_add(A, B, C, 1024)@flyc.kernelwraps the function as aKernelFunction@flyc.jitwraps the function as aJitFunction- On first call,
JitFunction.__call__triggers:- AST rewriting (Python loops/ifs → MLIR scf ops)
- MLIR module creation with
gpu.container_module - Tracing the jit function body to generate MLIR ops
- Calling
vec_add_kernel(...)captures a pending kernel call .launch()emits the specializedgpu.funcandgpu.launch_funcMlirCompiler.compile()runs the full pass pipelineCompiledArtifactstores the compiled module and lazily owns the resulting ExecutionEngine
- Subsequent calls with the same type signature use the cached binary
Maps a PyTorch tensor to an MLIR memref descriptor via DLPack:
@flyc.kernel
def my_kernel(input: fx.Tensor, output: fx.Tensor):
# input and output are Tensor wrappers around ir.Value (memref)
...At the host boundary, torch.Tensor is converted through the registered
PyTorch/DLPack argument adapter.
Compile-time constant. The value is embedded directly in the generated IR:
@flyc.kernel
def my_kernel(data: fx.Tensor, N: fx.Constexpr[int], dtype: fx.Constexpr[str]):
for i in range_constexpr(N // 64): # unrolled at compile time
...Different Constexpr values produce different compiled kernels (separate cache entries).
Runtime integer parameter (passed as i32):
@flyc.jit
def launch(data: fx.Tensor, size: fx.Int32, stream: fx.Stream = fx.Stream(None)):
...Python int values are converted to Int32 via the JitArgumentRegistry.
CUDA/HIP stream for asynchronous kernel launch:
@flyc.jit
def launch(data: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
my_kernel(data).launch(grid=(1,), block=(256,), stream=stream)
# Launch on specific stream:
stream = torch.cuda.Stream()
launch(data, stream=fx.Stream(stream))Register a raw Python type with a JitArgument adapter and the DslType that
will appear inside the traced function:
from flydsl.compiler import JitArgumentRegistry
@JitArgumentRegistry.register(MyCustomType, dsl_type=MyDslType)
class MyCustomAdaptor:
def __init__(self, value: MyCustomType):
self.value = value
def __get_ir_types__(self):
return [...] # MLIR types for this argument
def __cache_signature__(self):
return (...) # every property that can change generated code
def __c_abi_spec__(self):
return [...] # ordered (ctypes storage type, fill(argument, storage)) slotsThe adapter's C-ABI slots may outnumber its MLIR types (for example, a dynamic memref has data and layout slots).
from flydsl.expr import gpu
# Thread index within workgroup (returns Int32)
tid_x = gpu.thread_idx.x
tid_y = gpu.thread_idx.y
tid_z = gpu.thread_idx.z
# Block (workgroup) index within grid
bid_x = gpu.block_idx.x
bid_y = gpu.block_idx.y
# Block dimensions
bdim_x = gpu.block_dim.x
# Grid dimensions
gdim_x = gpu.grid_dim.x
# Low-level (returns raw ir.Value)
raw_tid = gpu.thread_id("x")
raw_bid = gpu.block_id("x")import flydsl.expr as fx
# Constants (prefer DSL numeric types)
c42 = fx.Int64(42) # 64-bit integer constant (prefer over the deprecated fx.Index)
c3_14 = fx.Float32(3.14) # f32 constant
mask = fx.Int32(0xFF) # i32 constant
# Arithmetic (operator overloading via ArithValue / Numeric)
result = a + b
result = a * 2
result = a // 4
result = a % 16
largest = fx.max(a, b)
smallest = fx.min(a, b)
tiles = fx.ceildiv(count, tile_size)
# Cast (prefer DSL numeric constructors)
i64_val = fx.Int64(int_val) # cast to 64-bit integer (fx.Index is deprecated)
i32_val = fx.Int32(i64_val) # cast to i32
# Select
result = cond.select(true_val, false_val) # when cond is an ArithValue
# Bitwise
result = a & b
result = a ^ b
result = a << 4Use direct arith.*FOp(..., fastmath=...) only where explicit fastmath flags are performance-critical.
Use fx.max / fx.min for type-dispatched extrema and fx.ceildiv for
overflow-safe integer ceil division. fx.maxnumf / fx.minnumf intentionally
retain non-NaN-wins semantics, and fx.ceil_div remains the layout/int-tuple API.
from flydsl.expr.typing import Vector as Vec
# Build vector from elements
vec = Vec.from_elements([a, b, c, d], fx.Float32)
# Vector store to memref
vec.store(memref, [idx])
# Extract, bitcast, and convert
elem = vec[idx]
as_i32 = vec.bitcast(fx.Int32)
as_bf16 = vec.to(fx.BFloat16)from flydsl.expr import rocdl
# Buffer tensor — wraps a Tensor with AMD buffer resource descriptor
A_buf = rocdl.make_buffer_tensor(A)
# RDNA bounded buffer — supplying the byte extent enables hardware OOB checks
A_buf = rocdl.make_buffer_tensor(
A,
num_records_bytes=problem_size * element_bytes,
)
# MFMA MMA atom constructor (CDNA3/CDNA4) — returns MmaAtomCDNA3_MFMAType
atom_type = rocdl.MFMA(m=16, n=16, k=32, elem_ty_ab=fx.Float8E4M3FNUZ)
# Buffer copy atom types
copy_op = rocdl.BufferCopy128b() # 128-bit buffer copy
copy_op = rocdl.BufferCopy64b() # 64-bit buffer copy
copy_op = rocdl.BufferCopy32b() # 32-bit buffer copyOn RDNA, supplying num_records_bytes selects the raw-buffer descriptor mode
where out-of-range reads return zero and writes are suppressed. The byte count
may be dynamic. Omitting it keeps the default unchecked descriptor; passing
max_size=False derives the count from the tensor layout and enables checking.
See {ref}gfx1250 WMMA & TDM atoms <wmma-tdm-atoms> below for the
gfx1250 WMMA (incl. MX-scaled) MMA atoms and the TDM async copy atom.
Signature: (result_type, [a, b, c, cbsz, abid, blgp]) — trailing ints default to 0.
result = rocdl.mfma_f32_16x16x16f16(result_type, [a, b, acc])
result = rocdl.mfma_f32_16x16x32_fp8_fp8(result_type, [a, b, acc])
result = rocdl.mfma_i32_16x16x32_i8(result_type, [a, b, acc])
result = rocdl.mfma_f32_16x16x16bf16_1k(result_type, [a, b, acc]) # BF16 1K variant
# GFX950 scaled MFMA (MXFP4/FP6/FP8)
result = rocdl.mfma_scale_f32_16x16x128_f8f6f4(
result_type, [a, b, acc, cbsz, blgp, opselA, scaleA, opselB, scaleB]
)Control instruction scheduling for performance tuning:
rocdl.sched_mfma(cnt) # wait for cnt MFMA instructions to complete
rocdl.sched_vmem(cnt) # wait for cnt VMEM reads to complete
rocdl.sched_dsrd(cnt) # wait for cnt DS (LDS) reads to complete
rocdl.sched_dswr(cnt) # wait for cnt DS (LDS) writes to completeSingle-instruction hardware math (guaranteed 1 VALU cycle, lower precision than math.*):
# Base-2 exponential (v_exp_f32)
result = rocdl.exp2(T.f32, x)
# Reciprocal (v_rcp_f32)
result = rocdl.rcp(T.f32, x)Generated upstream ROCDL builders and repository-only compatibility helpers are
not part of the stable expression API and are intentionally omitted here. New
kernels should prefer rocdl.make_buffer_tensor() plus fx.copy() or typed
Tensor/Pointer access.
(wmma-tdm-atoms)=
gfx1250 kernels run wave32 and use WMMA (not MFMA) for matrix math and the
TDM (Tensor Data Mover) engine for async whole-tile Global↔LDS copies. All of the
factories below live in flydsl.expr.rocdl (from flydsl.expr import rocdl).
WMMA MMA atom — rocdl.WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs)
is arch-dispatched (gfx11 v16 ABI; gfx12 / gfx1250 v8 ABI). On gfx1250 it builds
MmaOpGFX1250_WMMAType (M=N=16). Supported K / dtypes:
| dtype (A,B → Acc) | K | notes |
|---|---|---|
| f32 → f32 | 4 | |
| f16/bf16 → f32 or same | 32 | |
| fp8/bf8 (E4M3FN / E5M2, any mix) → f32 or f16 | 64, 128 | native OCP fp8 |
| i8 → i32 | 64 | sign_a / sign_b / clamp kwargs |
| i4 → i32 | 32 | sign_a / sign_b / clamp kwargs |
mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 32, fx.Float16)) # f16 → f32
mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 128, fx.Float8E4M3FN)) # fp8 → f32
# signed int4 with accumulator clamp:
mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 32, T.i4, T.i32, sign_a=True, sign_b=True, clamp=True))On RDNA4 (gfx1200 / gfx1201) the same factory builds MmaOpGFX120X_WMMAType
instead. RDNA4 shares the v8 register ABI, while its supported floating-point
forms use 16x16x16: f16/bf16 and every fp8(E4M3FN)/bf8(E5M2) A/B combination
accumulating to f32. Pass a different B type with the keyword-only elem_ty_b;
it defaults to the A type. The 16x16x32 BF16 and fp8 K=64/128 forms above are
gfx1250-only and are rejected by the atom's verifier. See
kernels/gemm/rdna_f16_gemm.py for a full pipelined f16 example and
kernels/gemm/rdna_fp8_preshuffle_gemm.py for an RDNA4 FP8 implementation.
mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 16, fx.BFloat16, fx.Float32)) # RDNA4
mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 16, fx.Float8E4M3FN, fx.Float32)) # RDNA4 FP8
mma = fx.make_mma_atom(rocdl.WMMA(16, 16, 16, fx.Float8E4M3FN, fx.Float32,
elem_ty_b=fx.Float8E5M2)) # RDNA4 FP8 x BF8MX-scaled WMMA — rocdl.WMMAScale(m, n, k, elem_ty_a, elem_ty_b=None, elem_ty_acc=None, *, opsel_a=0, opsel_b=0, mod_c=0, reuse_a=False, reuse_b=False, block_size=32) builds the E8M0 block-scaled WMMA (V_WMMA_SCALE /
V_WMMA_SCALE16) for the unified f8/f6/f4 operand format. Shapes: 16x16x128
(f8/f6/f4) and 32x16x128 (fp4-only). Per-operand E8M0 scales are atom state
(scale_a / scale_b); block_size 32 → i32 scale state, 16 → i64.
mma = fx.make_mma_atom(rocdl.WMMAScale(16, 16, 128, fx.Float8E4M3FN))
mma = fx.atom_set_value(mma, "scale_a", fx.Int32(scale_a)) # E8M0 scales
mma = fx.atom_set_value(mma, "scale_b", fx.Int32(scale_b))
fx.gemm(mma, frag_C, frag_A, frag_B, frag_C)For per-tile scales, pass each operand as a list or tuple containing its data and scale fragments:
fx.gemm(tiled_mma, frag_D, [frag_A, scales_A], [frag_B, scales_B], frag_C)If A has shape (V, M, K), its scale fragment has shape (1, M, K); B uses
(1, N, K). Rank-2 operands omit the K mode. Scales use i32 for CDNA4 MFMA
and block-32 WMMA, or i64 for block-16 WMMA. Use zero strides in tile modes
to broadcast a scale. The compiler slices all tensors in each operand group
together and fully expands the static M/N/K dimensions. All operands in each
group must be Tensors.
The first tensor is always the primary operand. Further tensors are defined by the MMA atom; Single
Tensor operands and scalar atom-state keywords remain supported for both bare atoms and tiled MMA.
tiled_mma.set_value("scale_a", scale) returns a new tiled MMA with the same layout and
permutation. Explicit scale tensors override the corresponding atom-state fields; omitted scale
tensors use those fields. The current scaled atoms accept data and an optional scale; they do not
consume sparsity metadata.
Tiled copies support the same state update API: tiled_copy.set_value("soffset", offset)
returns a new TiledCopy with its original tile and thread/value layout. Passing it to
fx.copy uses the updated atom state. Both tiled types also accept a dictionary of fields.
Use atom_callback to select a different atom for each static tile, for example
to choose a byte from packed CDNA4 scale words explicitly:
fx.gemm(
tiled_mma, frag_D, [frag_A, scales_A], [frag_B, scales_B], frag_C,
atom_callback=lambda atom, mnk: fx.make_mma_atom(
fx.rocdl.cdna4.MFMA_Scale(
16, 16, 128, fx.Float8E4M3FN,
opsel_a=mnk[0] % 2, opsel_b=mnk[1] % 4,
)
),
)When a callback is present, the callback receives the original atom with its runtime state and a
tuple of integer tile indices (m, n, k), and returns the atom for that tile. Rank-2 operands use
k=0; rank-1 calls use (0, 0, 0).
TDM tiled copy atom (preferred) — build the atom over the whole global
tensor with rocdl.cdna5.make_tiled_tdm_atom; it returns the atom plus a
coordinate tensor. fx.zipped_divide picks this block's tile and
rocdl.cdna5.tdm_partition cuts the LDS tile and the coordinate tile into one
warp's share. Padding is row slack in the LDS box layout, out-of-bounds comes from
the tensor's own extent plus boundary_check, and the K-loop advances by indexing
the coordinate rest with an i32 tile number:
lds = fx.SharedAllocator().allocate(fx.Array[fx.Float16, TM * TK]).peek()
box = fx.make_layout((TM, TK), (TK, 1)) # row slack here IS the padding
smem = fx.make_view(lds.ptr, fx.make_layout(((TM, TK), 1), ((TK, 1), TM * TK)))
gA = fx.make_view(fx.get_iter(A), fx.make_layout((M, K), (K, 1))) # FULL tensor
atom, coord = rocdl.cdna5.make_tiled_tdm_atom(rocdl.TensorLoad(), gA, box, (TM, TK))
blk = fx.zipped_divide(coord, (TM, TK))[None, (fx.block_idx.x, None)]
tAs, tAg = rocdl.cdna5.tdm_partition(atom, 0, fx.make_layout(1, 1), smem, blk)
fx.copy(atom, tAg[None, kt], tAs[None, 0]) # kt: i32 tile index, NOT a byte offset
rocdl.s_wait_tensorcnt(0)For a store pass rocdl.TensorStore(), flip the fx.copy operands, and keep the
LDS box packed — a store cannot de-pad LDS.
Older TDM atom — rocdl.make_tdm_atom still exists and still lowers, but it
is not what new code should reach for. It takes the extents, strides and padding
as explicit arguments and bumps the K loop with a byte imm_offset; the tiled
atom reads all of that off the tensors. The in-tree gfx1250 kernels are on the
tiled atom, and the only copies left on the old one are those it cannot express:
a non-contiguous innermost dim, and gather/scatter
(make_tensor_gather_descriptor / tensor_load_gather). See
docs/api/dsl.rst for its signature if you have to read one of those.
from flydsl.expr import gpu
# Barrier (workgroup synchronization)
gpu.barrier()
# Shared memory address space attribute
addrspace = gpu.smem_space()
addrspace_int = gpu.smem_space(int=True)The ASTRewriter transforms Python for loops:
@flyc.kernel
def my_kernel(data: fx.Tensor, N: fx.Constexpr[int]):
# Compile-time unrolled loop
for i in range_constexpr(N):
# This loop is fully unrolled in the generated IR
...
# Runtime loop (lowered by the AST rewriter)
for i in range(runtime_value):
...Mark a value as compile-time constant:
from flydsl.expr import const_expr
@flyc.kernel
def my_kernel(data: fx.Tensor, N: fx.Constexpr[int]):
tile_size = const_expr(N // 4)
for i in range_constexpr(tile_size):
...New kernels declare their LDS layout as a @fx.struct storage type and
allocate it with fx.SharedAllocator (from flydsl.expr.gpu, reached as
fx.SharedAllocator) inside the kernel body. Each field is an fx.Array[elem, count, align]; .allocate(StorageStruct).peek() returns a handle whose fields
expose typed views:
import flydsl.expr as fx
# Declare the LDS layout. Fields may be conditional on compile-time config.
@fx.struct
class SharedStorage:
s_red: fx.Array[fx.Float32, red_slots, 16] # red_slots elems, 16B aligned
s_red2: fx.Array[fx.Float32, red_slots, 16]
@flyc.kernel
def my_kernel(...):
# Allocate the storage struct in LDS (inside the @kernel body).
lds = fx.SharedAllocator().allocate(SharedStorage).peek()
# Get a logical view over each field and use it with the layout API.
s_red = lds.s_red.view(fx.make_layout(red_slots, 1))
s_red2 = lds.s_red2.view(fx.make_layout(red_slots, 1))By default SharedAllocator is static=True: each leaf emits a per-leaf static
LDS global that the compiler sizes, so launch(smem=...) is left unset. Use
static=False (dynamic) mode to have the launch wrapper auto-infer smem from
SharedAllocator.allocated_bytes when smem=None (an explicit smem must be
>= that size). See kernels/gemm/preshuffle_gemm.py and
kernels/norm/rmsnorm_kernel.py for real usage.
| Architecture | LDS per CU |
|---|---|
gfx942 (MI300X) |
64 KB |
gfx950 (MI350/MI355X) |
160 KB |
gfx1201 (Radeon AI PRO R9700) |
64 KB |
gfx1250 |
320 KB |
@flyc.jit
def launch(data: fx.Tensor, stream: fx.Stream = fx.Stream(None)):
my_kernel(data).launch(
grid=(num_blocks_x, num_blocks_y, num_blocks_z),
block=(threads_x, threads_y, threads_z),
smem=shared_mem_bytes, # dynamic shared memory
stream=stream, # CUDA/HIP stream
)Grid and block dimensions accept:
int— static valueir.Value— dynamic MLIR value- Tuple of 1–3 values — missing dimensions default to 1
@flyc.jit
def launch(data: fx.Tensor, M: fx.Int32, stream: fx.Stream = fx.Stream(None)):
grid_x = M // 256
my_kernel(data, M).launch(
grid=(grid_x, 1, 1),
block=(256, 1, 1),
stream=stream,
)from flydsl.expr import gpu
# Workgroup barrier (s_barrier)
gpu.barrier()JIT-compiled functions are cached automatically:
- In-memory cache — keyed by argument type signature
- Disk cache — stored in
~/.flydsl/cache/(configurable viaFLYDSL_RUNTIME_CACHE_DIR) - Cache key includes: source code hash, dependency sources, closure values, FlyDSL version, LLVM version
The cache is invalidated when:
- Source code of the function or its dependencies changes
- Argument types change (different tensor shapes/dtypes)
Constexprvalues change- FlyDSL or LLVM version changes
The JIT disk cache auto-invalidates when kernel source code or closure values change. Set FLYDSL_RUNTIME_ENABLE_CACHE=0 only when modifying C++ passes or non-closure helper functions:
FLYDSL_RUNTIME_ENABLE_CACHE=0 python my_script.py # or: rm -rf ~/.flydsl/cacheCOMPILE_ONLY=1 python my_script.pyFLYDSL_DUMP_IR=1 FLYDSL_DUMP_DIR=./my_dumps python my_script.pyUse the supported debug controls to print IR during compilation:
FLYDSL_DEBUG_PRINT_ORIGIN_IR=1 python my_script.py
FLYDSL_DEBUG_PRINT_AFTER_ALL=1 python my_script.pyThe internal CompiledArtifact object also stores compiled and source IR,
but it is not a public return value of @flyc.jit or flyc.compile.
FLYDSL_DEBUG_AST_DIFF=1 python my_script.pyShows the diff between original and rewritten AST for debugging control flow transformations.
From kernels/gemm/preshuffle_gemm.py:
import flydsl.compiler as flyc
import flydsl.expr as fx
from flydsl.expr import gpu, range_constexpr, rocdl
from flydsl.expr.typing import T
def compile_preshuffle_gemm(*, N, K, tile_m, tile_n, tile_k,
in_dtype="fp8", out_dtype="bf16",
epilogue="none", lds_stage=2, ...):
a_lds_elems = tile_m * tile_k
# Declare the LDS layout as a storage struct.
@fx.struct
class SharedStorage:
a0: fx.Array[layout_elem, a_lds_elems, 16]
if lds_stage == 2:
a1: fx.Array[layout_elem, a_lds_elems, 16]
@flyc.kernel
def kernel_gemm(
arg_c: fx.Tensor, arg_a: fx.Tensor, arg_b: fx.Tensor,
arg_scale_a: fx.Tensor, arg_scale_b: fx.Tensor, arg_bias: fx.Tensor,
i32_m: fx.Int32, i32_n: fx.Int32,
tiled_mma_arg: fx.TiledMma, tiled_copy_g2s: fx.TiledCopy,
):
tid = fx.thread_idx.x
bid_x, bid_y, _ = fx.block_idx
gA = fx.rocdl.make_buffer_tensor(arg_a, ...)
gB = fx.rocdl.make_buffer_tensor(arg_b)
gC = fx.rocdl.make_buffer_tensor(arg_c, ...)
# Allocate LDS and take typed views over the storage fields.
lds = fx.SharedAllocator().allocate(SharedStorage).peek()
# ... GEMM implementation using MFMA, LDS, tiling ...
@flyc.jit
def launch_fn(
arg_c: fx.Tensor, arg_a: fx.Tensor, arg_b: fx.Tensor,
arg_scale_a: fx.Tensor, arg_scale_b: fx.Tensor, arg_bias: fx.Tensor,
M_val: fx.Int32, N_val: fx.Int32,
stream: fx.Stream = fx.Stream(None),
):
kernel_gemm(arg_c, arg_a, arg_b, arg_scale_a, arg_scale_b, arg_bias,
M_val, N_val, tiled_mma, tiled_copy_g2s).launch(
grid=(grid_x, grid_y), block=(256,), stream=stream,
)
return launch_fnM is a runtime argument (i32_m / M_val), not a compile-time parameter, so
one compiled kernel serves any M. Select output post-processing with
epilogue= ("none", "bias", "bias_relu", "bias_silu", "bias_gelu").
Writing a new kernel?
│
├── Simple element-wise?
│ ├── Use @flyc.kernel + @flyc.jit
│ ├── fx.gpu.thread_idx.x for thread indexing
│ └── See tests/kernels/test_vec_add.py
│
├── Reduction (norm, softmax)?
│ ├── Reductions are inline (wave/block reduce helpers within the kernel)
│ └── See kernels/norm/rmsnorm_kernel.py, kernels/norm/softmax_kernel.py
│
├── Matrix multiply (GEMM)?
│ ├── Use @flyc.kernel + fx.SharedAllocator + MFMA
│ ├── B-preshuffle layout from kernels/common/mma/mfma_preshuffle_pipeline.py
│ └── See kernels/gemm/preshuffle_gemm.py
│
├── Need shared memory?
│ ├── Declare a @fx.struct storage layout
│ ├── Allocate with fx.SharedAllocator().allocate(Storage).peek()
│ └── Take .view(...) over each field inside @kernel
│
└── Need compile-time specialization?
├── Use Constexpr[T] parameters
└── Use range_constexpr() for unrolled loops
| File | Description |
|---|---|
python/flydsl/compiler/__init__.py |
Public API: jit, kernel, from_dlpack |
python/flydsl/compiler/jit_function.py |
@jit decorator, MlirCompiler, JitCacheManager |
python/flydsl/compiler/kernel_function.py |
@kernel decorator, KernelFunction, KernelLauncher |
python/flydsl/compiler/jit_executor.py |
CompiledArtifact (compiled IR and lazy ExecutionEngine wrapper) |
python/flydsl/compiler/jit_argument.py |
JitArgumentRegistry and tensor/pointer argument adapters |
python/flydsl/compiler/ast_rewriter.py |
ASTRewriter — Python AST → MLIR control flow |
python/flydsl/expr/typing.py |
Types (T), Tensor, Stream, Constexpr |
python/flydsl/expr/arith.py |
Arithmetic operations |
python/flydsl/expr/gpu.py |
GPU operations (thread indices, barriers, ...) |
python/flydsl/expr/rocdl/ |
ROCm dialect intrinsics (MFMA/WMMA, buffer, TDM, cluster) |
python/flydsl/expr/primitive.py |
Layout algebra primitives (make_shape, crd2idx, etc.) |
python/flydsl/expr/gpu.py |
SharedAllocator, GPU ops (thread indices, barriers, ...) |
kernels/gemm/preshuffle_gemm.py |
Preshuffle GEMM kernel example |
tests/kernels/test_vec_add.py |
Vector add kernel test |
tests/kernels/test_preshuffle_gemm.py |
Preshuffle GEMM test |