FlyDSL project structure, compilation stages, key abstractions, and configuration.
| Component | Description | Key File |
|---|---|---|
| FlyDSL | Python DSL front-end for authoring GPU kernels | python/flydsl/ |
| FlyDSL Compiler | @flyc.jit / @flyc.kernel — trace-based JIT compiler |
python/flydsl/compiler/ |
| FlyDSL Expr | DSL expression ops (arith, vector, gpu, buffer, rocdl) | python/flydsl/expr/ |
| Fly Dialect | Flexible Layout IR — MLIR dialect with layout algebra | include/flydsl/Dialect/Fly/ |
| MlirCompiler | End-to-end MLIR pass pipeline (DSL → binary) | python/flydsl/compiler/jit_function.py |
| JITCFunction | MLIR ExecutionEngine wrapper for JIT execution | python/flydsl/compiler/jit_executor.py |
FlyDSL/
├── include/flydsl/ # C++ dialect headers
│ └── Dialect/
│ ├── Fly/ # Fly layout dialect
│ │ ├── IR/
│ │ │ ├── FlyDialect.td # Dialect declaration (name = "fly")
│ │ │ ├── FlyOps.td # Layout ops (make_shape, crd2idx, composition, ...)
│ │ │ ├── FlyTypeDefs.td # Custom types (!fly.int_tuple, !fly.layout, ...)
│ │ │ ├── FlyAttrDefs.td # Attributes
│ │ │ └── FlyInterfaces.td # Op interfaces
│ │ └── Transforms/
│ │ ├── Passes.td # Pass declarations (fly-layout-lowering, etc.)
│ │ └── LayoutLowering.td # Layout lowering pass
│ └── FlyROCDL/ # FlyROCDL dialect (copy/MMA atoms)
│ └── IR/
│ ├── Dialect.td # FlyROCDL dialect declaration
│ ├── CopyAtom.td # Copy atom ops
│ └── MmaAtom.td # MMA atom ops
│
├── lib/ # C++ dialect implementation
│ ├── Dialect/Fly/ # Fly dialect ops, type inference, lowering
│ ├── Dialect/FlyROCDL/ # FlyROCDL dialect implementation
│ ├── Conversion/ # Dialect conversion passes
│ └── Transforms/ # Optimization passes
│
├── python/flydsl/ # Python DSL package
│ ├── __init__.py # Package version
│ ├── compiler/
│ │ ├── __init__.py # Public API: jit, kernel, from_dlpack
│ │ ├── jit_function.py # @jit decorator, MlirCompiler, JitCacheManager
│ │ ├── kernel_function.py # @kernel decorator, KernelFunction, KernelLauncher
│ │ ├── jit_executor.py # JITCFunction (ExecutionEngine wrapper)
│ │ ├── jit_argument.py # Argument conversion (Tensor, Stream, Int32)
│ │ ├── ast_rewriter.py # AST rewriting for Python control flow → MLIR
│ │ └── protocol.py # DslType / JitArgument protocols
│ ├── expr/
│ │ ├── __init__.py # Public expr API
│ │ ├── typing.py # Types (T.f32, Tensor, Stream, Constexpr)
│ │ ├── numeric.py # DSL numeric types (Float32, Int32, ...)
│ │ ├── primitive.py # Primitive operations (layout algebra, copy, gemm)
│ │ ├── derived.py # Derived types (CopyAtom, MmaAtom, TiledCopy)
│ │ ├── arith.py # Arithmetic dialect ops
│ │ ├── vector.py # Vector dialect ops
│ │ ├── gpu.py # GPU dialect ops (thread_idx, block_idx, barrier)
│ │ ├── buffer_ops.py # Buffer / memory operations
│ │ └── rocdl/ # ROCm-specific intrinsics (MFMA/WMMA, buffer, TDM, cluster)
│ ├── runtime/
│ │ └── device.py # get_rocm_arch() — GPU architecture detection
│ └── utils/
│ ├── env.py # EnvManager — typed environment config
│ ├── logger.py # Logging utilities
│ └── smem_allocator.py # SmemAllocator for LDS management
│
├── examples/ # Runnable examples
│ ├── 01-vectorAdd.py # Vector addition with layout algebra
│ ├── 02-tiledCopy.py # Tiled copy with partitioned tensors
│ ├── 03-tiledMma.py # Tiled MMA (GEMM) with MFMA atoms
│ └── 04-preshuffle_gemm.py # Preshuffle GEMM end-to-end example
│
├── kernels/ # Production GPU kernels (organized by domain)
│ ├── gemm/ # Dense GEMM kernels
│ │ ├── preshuffle_gemm.py # GEMM (preshuffle layout)
│ │ ├── blockscale_preshuffle_gemm.py # Blockscale GEMM
│ │ ├── mxfp4_preshuffle.py # MXFP4 preshuffle GEMM
│ │ ├── fp4_gemm_4wave.py # FP4 4-wave GEMM
│ │ ├── fp8_gemm_4wave.py # FP8 4-wave GEMM
│ │ ├── fp8_gemm_8wave.py # FP8 8-wave GEMM
│ │ ├── hgemm_splitk.py # FP16 GEMM split-K
│ │ ├── rdna_f16_gemm.py # RDNA FP16 GEMM
│ │ ├── rdna_fp8_preshuffle_gemm.py # RDNA FP8 GEMM
│ │ ├── gemm_common_gfx1250.py # GFX1250 GEMM common
│ │ ├── gemm_fp8fp4_gfx1250.py # GFX1250 FP8/FP4 GEMM
│ │ ├── wmma_gemm_gfx1250.py # GFX1250 WMMA GEMM
│ │ └── fp8_gemm_utils.py # FP8 GEMM helpers
│ ├── norm/ # Normalization kernels
│ │ ├── layernorm_kernel.py # LayerNorm (layout API)
│ │ ├── rmsnorm_kernel.py # RMSNorm (layout API)
│ │ └── softmax_kernel.py # Softmax (layout API)
│ ├── attention/ # Attention kernels
│ │ ├── pa_decode_fp8.py # Paged attention decode (FP8)
│ │ ├── pa_decode_swa.py # Paged attention sliding-window
│ │ ├── flash_attn_generic.py # FlashAttention generic fallback
│ │ ├── flash_attn_gfx950.py # FlashAttention gfx950 fast path
│ │ ├── mla_fwd_decode.py # MLA forward decode
│ │ └── fused_rope_cache_kernel.py # Fused RoPE + KV cache
│ ├── moe/ # Mixture-of-experts kernels
│ │ ├── moe_gemm_2stage.py # MoE GEMM (2-stage gate/up + reduce)
│ │ ├── moe_blockscale_2stage.py # MoE Blockscale GEMM
│ │ ├── mixed_moe_gemm_2stage.py # Mixed-precision MoE GEMM
│ │ ├── moe_sorting_kernel.py # MoE token sorting
│ │ └── topk_gating_softmax_kernel.py # Top-k gating softmax
│ ├── mma/ # Shared MMA pipeline helpers
│ │ ├── mfma_epilogues.py # MFMA epilogue helpers
│ │ ├── mfma_preshuffle_pipeline.py # Preshuffle helpers for MFMA kernels
│ │ └── pipeline_utils.py # Pipeline utility helpers
│ ├── conv/ # Convolution kernels
│ │ └── conv3d_implicit_8wave.py # Implicit-GEMM 3D convolution
│ ├── comm/ # Multi-GPU communication
│ │ └── custom_all_reduce.py # Multi-GPU all-reduce
│ └── common/ # Cross-domain kernel utilities
│ ├── kernels_common.py # Common kernel utilities
│ ├── layout_utils.py # Layout helpers
│ └── tensor_shim.py # GTensor/STensor abstraction
│
├── tests/
│ ├── mlir/ # MLIR-level tests (Conversion, LayoutAlgebra, Transforms)
│ ├── kernels/ # GPU kernel tests + benchmarks
│ ├── python/ # Python-based tests (examples, AOT)
│ ├── unit/ # Unit tests (streams, async, etc.)
│ ├── conftest.py # Pytest fixtures
│ ├── test_common.py # Shared test utilities
│ └── utils.py # Compilation helpers
│
└── scripts/ # Build and test helpers
├── build.sh # Build FlyDSL (CMake + ninja)
├── build_llvm.sh # Build MLIR from ROCm llvm-project
├── run_tests.sh # Run full test suite (pytest + examples + FileCheck)
├── run_benchmark.sh # Run benchmarks
└── dumpir.sh # Dump intermediate IR
The user-facing API lives in python/flydsl/. Kernel authors use @flyc.jit and @flyc.kernel decorators with expression operations from flydsl.expr:
- Traces Python functions via AST rewriting and execution
- Generates Fly dialect ops + standard MLIR dialects (gpu, arith, scf, memref, vector, rocdl)
- Compiles through the
MlirCompilerpass pipeline (Fly → ROCDL → LLVM → HSACO) - Caches compiled kernels to disk for fast re-use
- Executes via MLIR ExecutionEngine
The Fly dialect (include/flydsl/Dialect/Fly/) provides the MLIR-level layout algebra (composition, product, divide, coordinate mapping). Python DSL operations in flydsl.expr lower to Fly dialect ops during tracing, which are then compiled through the MlirCompiler pipeline.
Python Function (@flyc.kernel / @flyc.jit)
│
▼ AST Rewriting
Transformed Python Function
│
▼ Tracing (execution inside MLIR Context)
MLIR Module (fly, gpu, arith, scf, memref, vector dialects)
│
▼ MlirCompiler.compile()
┌────────────────────────────────────────────────────────┐
│ Stage A — pre_binary_fragments (Fly → ROCDL) │
│ fly-rewrite-func-signature │
│ fly-canonicalize │
│ fly-layout-lowering │
│ fly-int-swizzle-simplify │
│ canonicalize │
│ fly-convert-atom-call-to-ssa-form │
│ fly-promote-regmem-to-vectorssa │
│ convert-fly-to-rocdl │
│ canonicalize │
│ gpu.module(convert-scf-to-cf, cse, │
│ convert-gpu-to-rocdl{chipset=gfxNNN ...}, │
│ fly-rocdl-cluster-attr) │
├────────────────────────────────────────────────────────┤
│ Stage B — binary_prep_fragments (→ LLVM) │
│ rocdl-attach-target{chip=gfxNNN ...} │
│ convert-scf-to-cf │
│ convert-cf-to-llvm │
│ gpu-to-llvm{use-bare-pointers-...=true} │
│ convert-vector-to-llvm │
│ convert-arith-to-llvm │
│ convert-func-to-llvm │
│ reconcile-unrealized-casts │
│ ensure-debug-info-scope-on-llvm-func (optional) │
├────────────────────────────────────────────────────────┤
│ Stage C — binary_fragment │
│ gpu-module-to-binary{format=fatbin opts="..."} │
└────────────────────────────────────────────────────────┘
│
▼
JITCFunction (ExecutionEngine)
The pipeline is built by RocmBackend._pipeline_parts() in
python/flydsl/compiler/backends/rocm.py. The orchestrator
_pipeline_fragments_for_mode() in jit_function.py decides whether to run
the pipeline as a single combined pass list (pipeline_fragments()) or split
it for external LLVM codegen (external_binary_pipeline_fragments()). External
mode runs Stages A and B with the bundled MLIR runtime, then invokes the
external LLVM toolchain only for Stage C (gpu-module-to-binary).
Stage A — pre_binary_fragments (Fly dialect → ROCDL lowering)
| # | Pass | Description |
|---|---|---|
| 1 | fly-rewrite-func-signature |
Rewrite DSL types at function and SCF control-flow boundaries; lowers IntTuple / Layout / ComposedLayout / CoordTensor / MemRef to packed LLVM struct types and reconstructs them in the body via constructor ops. |
| 2 | fly-canonicalize |
FlyDSL-specific canonicalization (folds !fly.layout algebra when shapes are static). |
| 3 | fly-layout-lowering |
Lowers layout algebra (fly.crd2idx, partitions, divides) to concrete arith + vector ops. |
| 4 | fly-int-swizzle-simplify |
Algebraically simplifies the swizzle-shaped arith sequences emitted by applySwizzle. |
| 5 | canonicalize |
Standard MLIR canonicalization (constant folding, etc.). |
| 6 | fly-convert-atom-call-to-ssa-form |
Converts copy_atom_call / mma_atom_call to their SSA counterparts; promotes register tensors to vector SSA values. |
| 7 | fly-promote-regmem-to-vectorssa |
Promotes fly.make_ptr(register) memory semantics to vector SSA values (requires #6). |
| 8 | convert-fly-to-rocdl |
Lowers remaining Fly ops to MLIR upstream + ROCDL dialects (copy atoms → rocdl.buffer_load/store, or gfx1250 TDM → rocdl.tensor.load.to.lds / store.from.lds; MMA atoms → rocdl.mfma.* on CDNA, rocdl.wmma.* on gfx11/gfx1250). |
| 9 | canonicalize |
Second canonicalization round after ROCDL lowering. |
| 10 | gpu.module(convert-scf-to-cf, cse, convert-gpu-to-rocdl{chipset=gfxNNN ...}, fly-rocdl-cluster-attr) |
Inside the GPU module: SCF→CF, CSE, GPU intrinsics→ROCDL, then fly-rocdl-cluster-attr injects amdgpu-cluster-dims into the llvm.func passthrough. |
Stage B — binary_prep_fragments (LLVM lowering, host + kernel)
| # | Pass | Description |
|---|---|---|
| 11 | rocdl-attach-target{chip=gfxNNN ...} |
Attaches #rocdl.target<chip=gfxNNN> (plus fast/unsafe-math/wave64 options) to the GPU module for codegen. |
| 12 | convert-scf-to-cf |
Host-side SCF → ControlFlow. |
| 13 | convert-cf-to-llvm |
ControlFlow → LLVM dialect. |
| 14 | gpu-to-llvm{use-bare-pointers-for-host=true use-bare-pointers-for-kernels=true} |
GPU types and host launcher → LLVM. |
| 15 | convert-vector-to-llvm |
Vector → LLVM. |
| 16 | convert-arith-to-llvm |
Arith → LLVM. |
| 17 | convert-func-to-llvm |
Func → LLVM. |
| 18 | reconcile-unrealized-casts |
Final cast cleanup. |
When FLYDSL_DEBUG_ENABLE_DEBUG_INFO=1, Stage B appends
ensure-debug-info-scope-on-llvm-func{emission-kind=LineTablesOnly} after
reconcile-unrealized-casts and before Stage C.
Stage C — binary_fragment
| # | Pass | Description |
|---|---|---|
| 19 | gpu-module-to-binary{format=fatbin opts="..."} |
Invokes the LLVM AMDGPU backend and emits an HSA fatbin. |
gpu-kernel-outlining is no longer a pass in the runtime pipeline — kernel
outlining happens during Python tracing, when @flyc.kernel emits
gpu.func ops directly into a gpu.container_module.
When a @flyc.jit function is called:
- Cache check — look up by argument type signature (in-memory → disk)
- AST rewriting —
ASTRewriter.transformconverts Pythonfor/ifto MLIRscf.for/scf.if - MLIR module creation — sets up
gpu.container_modulewith target - Argument conversion —
convert_to_jit_argumentsmaps Python args to IR types - Function tracing — execute transformed function body to generate MLIR ops
- GPU kernel emission —
@kernelcalls emitgpu.funcintogpu.module - Pipeline compilation —
MlirCompiler.compile()runs the full pass pipeline - Execution —
JITCFunctionwraps MLIR ExecutionEngine for invoking the compiled code - Cache store — compiled function is serialized to disk for future runs
Decorates a Python function as a JIT-compiled host launcher:
import flydsl.compiler as flyc
import flydsl.expr as fx
@flyc.jit
def launch(a: fx.Tensor, b: fx.Tensor, n: fx.Constexpr[int],
stream: fx.Stream = fx.Stream(None)):
my_kernel(a, b, n).launch(grid=(n // 256,), block=(256,), stream=stream)Key behaviors:
- First call triggers compilation; subsequent calls with the same type signature use cached binary
Constexpr[T]parameters become compile-time constants (affect cache key)Tensorparameters map to memref descriptors via DLPackStreamparameters pass CUDA/HIP stream to the GPU runtime- When called inside an existing MLIR context, acts as a normal function (composable)
Decorates a Python function as a GPU kernel:
@flyc.kernel
def my_kernel(a: fx.Tensor, b: fx.Tensor, n: fx.Constexpr[int]):
tid = fx.gpu.thread_id("x")
bid = fx.gpu.block_id("x")
# ... kernel body ...Key behaviors:
- Can only be called inside a
@flyc.jitfunction - Calling returns a
KernelLauncher— you must call.launch()to emit the launch op - Supports
Constexpr[T]for compile-time specialization - Emits a
gpu.funcwithgpu.kernelattribute into thegpu.module
Returned by calling a @kernel function. Use .launch() to configure and emit the GPU launch:
launcher = my_kernel(a, b, 1024)
launcher.launch(
grid=(num_blocks, 1, 1),
block=(256, 1, 1),
smem=shared_mem_bytes,
stream=stream_value,
)Wraps MLIR's ExecutionEngine for JIT execution:
- Thread-safe with lazy engine initialization
- Serializable (pickle) for disk caching
- Supports packed calling convention via
ctypes - Provides
.print_ir()for debugging compiled/original IR
Extensible type system for mapping Python values to MLIR:
# DslType protocol — for values used inside kernel/jit functions
class DslType(Protocol):
@classmethod
def __construct_from_ir_values__(cls, values: List[ir.Value]) -> "DslType": ...
def __extract_to_ir_values__(self) -> List[ir.Value]: ...
# JitArgument protocol — for values passed at the host boundary
class JitArgument(Protocol):
def __get_ir_types__(self) -> List[ir.Type]: ...
def __get_c_pointers__(self) -> List[ctypes.c_void_p]: ...Built-in types: Tensor, Stream, Int32, Constexpr[T]
Register custom types:
from flydsl.compiler import JitArgumentRegistry
@JitArgumentRegistry.register(MyPythonType, dsl_type=MyDslType)
class MyJitArg:
def __get_ir_types__(self): ...
def __get_c_pointers__(self): ...Transforms Python control flow to MLIR ops at the AST level:
for i in range(n)→scf.forfor i in range_constexpr(n)→ compile-time unrolled loopif condition→scf.ifconst_expr(value)→ compile-time constant
| Variable | Default | Description |
|---|---|---|
FLYDSL_COMPILE_OPT_LEVEL |
2 |
Optimization level (0–3) |
COMPILE_ONLY |
0 |
If 1, compile without creating an executor. Returns None. |
ARCH |
auto-detect | Override target GPU architecture (e.g., gfx942, gfx950). |
| Variable | Default | Description |
|---|---|---|
FLYDSL_DUMP_IR |
false |
Dump intermediate IR at each pipeline stage. |
FLYDSL_DUMP_DIR |
~/.flydsl/debug |
Directory for IR dumps. |
FLYDSL_DEBUG_DUMP_ASM |
false |
Dump final AMD ISA assembly. |
FLYDSL_DEBUG_AST_DIFF |
false |
Print AST diff during rewrite. |
FLYDSL_DEBUG_PRINT_ORIGIN_IR |
false |
Print origin IR before compilation. |
FLYDSL_DEBUG_PRINT_AFTER_ALL |
false |
Print IR after each MLIR pass. |
FLYDSL_DEBUG_ENABLE_DEBUG_INFO |
false |
Generate debug info in compiled code. |
FLYDSL_DEBUG_ENABLE_VERIFIER |
true |
Verify IR module. |
FLYDSL_DEBUG_LOG_LEVEL |
WARNING |
Logging level (DEBUG, INFO, WARNING, ERROR). |
| Variable | Default | Description |
|---|---|---|
FLYDSL_RUNTIME_CACHE_DIR |
~/.flydsl/cache |
Directory for caching compiled kernels. |
FLYDSL_RUNTIME_ENABLE_CACHE |
true |
Enable kernel disk caching (in-memory cache is always active). |
get_rocm_arch() in runtime/device.py checks in order:
FLYDSL_GPU_ARCHenv varHSA_OVERRIDE_GFX_VERSIONenv var (supports9.4.2→gfx942format)rocm_agent_enumeratorsystem tool- Default:
gfx942
| Architecture | GPU | LDS per CU | Notes |
|---|---|---|---|
gfx942 |
MI300A / MI300X | 64 KB | CDNA 3, primary development target |
gfx950 |
MI350 / MI355X | 160 KB | CDNA 4, larger LDS |
gfx1201 |
Radeon AI PRO R9700 | 64 KB | RDNA 4 |
gfx1250 |
— | 320 KB | GFX12, wave32, WMMA, TDM ops |
gfx90a |
MI250X | 64 KB | CDNA 2 (verified platform) |
Enable with FLYDSL_DUMP_IR=1:
FLYDSL_DUMP_IR=1 FLYDSL_DUMP_DIR=./dumps python test_my_kernel.pyProduces numbered dump files (exact pass count tracks RocmBackend._pipeline_parts()):
dumps/my_func_name/
├── 00_origin.mlir
├── 01_fly_rewrite_func_signature.mlir
├── 02_fly_canonicalize.mlir
├── 03_fly_layout_lowering.mlir
├── 04_fly_int_swizzle_simplify.mlir
├── 05_canonicalize.mlir
├── 06_fly_convert_atom_call_to_ssa_form.mlir
├── 07_fly_promote_regmem_to_vectorssa.mlir
├── 08_convert_fly_to_rocdl.mlir
├── 09_canonicalize.mlir
├── 10_convert_scf_to_cf_cse_convert_gpu_to_rocdl.mlir
│ # also runs fly-rocdl-cluster-attr
├── 11_rocdl_attach_target.mlir
├── 12_convert_scf_to_cf.mlir
├── 13_convert_cf_to_llvm.mlir
├── 14_gpu_to_llvm.mlir
├── 15_convert_vector_to_llvm.mlir
├── 16_convert_arith_to_llvm.mlir
├── 17_convert_func_to_llvm.mlir
├── 18_reconcile_unrealized_casts.mlir
├── 19_gpu_module_to_binary.mlir
├── 20_llvm_ir.ll
└── 21_final_isa.s # AMD ISA assembly (best-effort)
If FLYDSL_DEBUG_ENABLE_DEBUG_INFO=1, the debug-info pass adds an extra
numbered dump before gpu_module_to_binary.
| File | Description |
|---|---|
python/flydsl/compiler/jit_function.py |
@jit decorator, MlirCompiler, JitCacheManager |
python/flydsl/compiler/kernel_function.py |
@kernel decorator, KernelFunction, KernelLauncher, CompilationContext |
python/flydsl/compiler/jit_executor.py |
JITCFunction — ExecutionEngine wrapper |
python/flydsl/compiler/jit_argument.py |
JitArgumentRegistry, TensorAdaptor, from_dlpack |
python/flydsl/compiler/ast_rewriter.py |
ASTRewriter — Python AST → MLIR control flow |
python/flydsl/compiler/protocol.py |
get_ir_types, extract_to_ir_values, construct_from_ir_values protocols |
python/flydsl/expr/typing.py |
Types (T), Tensor, Stream, Constexpr |
python/flydsl/expr/primitive.py |
Layout algebra primitives (make_shape, crd2idx, copy, gemm) |
python/flydsl/expr/derived.py |
Derived types (CopyAtom, MmaAtom, TiledCopy) |
python/flydsl/expr/numeric.py |
DSL numeric types (Float32, Int32, ...) |
python/flydsl/utils/env.py |
EnvManager — typed environment variable configuration |
python/flydsl/runtime/device.py |
get_rocm_arch() GPU detection |
include/flydsl/Dialect/Fly/IR/FlyOps.td |
Fly dialect op definitions |
include/flydsl/Dialect/Fly/Transforms/Passes.td |
Pass declarations (fly-layout-lowering, etc.) |