Architecture & compilation pipeline guide

Architecture & compilation pipeline guide#

This guide covers the FlyDSL project structure, compilation stages, key abstractions, and environment configuration.

Quick Reference#

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

1. Project structure#

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
│   │   ├── gpu.py                    # GPU dialect ops (thread_idx, block_idx, barrier)
│   │   └── 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)
│   │   ├── 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
│   │   ├── 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)
│   │   ├── mxfp_moe/               # Fused a4w4/a8w4 MoE 2-stage (device fp4 re-quant)
│   │   ├── 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

2. Architecture#

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 and standard MLIR dialects (gpu, arith, scf, memref, vector, rocdl)

  • Compiles through the MlirCompiler pass pipeline (Fly → ROCDL → LLVM → HSACO)

  • Caches compiled kernels to disk for fast reuse

  • 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, and are then compiled through the MlirCompiler pipeline.

3. Compilation pipeline#

3.1 High-level flow#

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)

3.2 Pipeline stages in detail#

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.

3.3 JIT compilation flow#

When a @flyc.jit function is called:

  1. Cache check: Look up by argument type signature (in-memory → disk).

  2. AST rewriting: ASTRewriter.transform converts Python for/if to MLIR scf.for/scf.if.

  3. MLIR module creation: Sets up gpu.container_module with target.

  4. Argument conversion: convert_to_jit_arguments maps Python args to IR types.

  5. Function tracing: Execute the transformed function body to generate MLIR ops.

  6. GPU kernel emission: @kernel calls emit gpu.func into gpu.module.

  7. Pipeline compilation: MlirCompiler.compile() runs the full pass pipeline.

  8. Execution: JITCFunction wraps MLIR ExecutionEngine to invoke the compiled code.

  9. Cache store: Serialize the compiled function to disk for future runs.

4. Key abstractions#

4.1 @flyc.jit — Host launcher#

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:

  • The first call triggers compilation; subsequent calls with the same type signature use the cached binary.

  • Constexpr[T] parameters become compile-time constants and affect the cache key.

  • Tensor parameters map to memref descriptors via DLPack.

  • Stream parameters pass a CUDA/HIP stream to the GPU runtime.

  • When called inside an existing MLIR context, the function acts as a normal composable function.

4.2 @flyc.kernel — GPU kernel#

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:

  • You can only call this inside a @flyc.jit function.

  • Calling the kernel returns a KernelLauncher; you must call .launch() to emit the launch op.

  • Supports Constexpr[T] for compile-time specialization.

  • Emits a gpu.func with gpu.kernel attribute into the gpu.module.

4.3 KernelLauncher#

Calling a @kernel function returns a KernelLauncher. 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,
)

4.4 JITCFunction#

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 or original IR.

4.5 DslType / JitArgument protocols#

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, and Constexpr[T]

To 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): ...

4.6 ASTRewriter#

Transforms Python control flow to MLIR ops at the AST level:

  • for i in range(n) → scf.for

  • for i in range_constexpr(n) → compile-time unrolled loop

  • if condition → scf.if

  • const_expr(value) → compile-time constant

5. Environment variables#

5.1 Compilation options (FLYDSL_COMPILE_*)#

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).

5.2 Debug options (FLYDSL_DEBUG_*)#

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).

5.3 Runtime options (FLYDSL_RUNTIME_*)#

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).

5.4 Architecture detection priority#

get_rocm_arch() in runtime/device.py checks in the following order:

  1. FLYDSL_GPU_ARCH env var

  2. HSA_OVERRIDE_GFX_VERSION env var (supports 9.4.2 → gfx942 format)

  3. rocm_agent_enumerator system tool

  4. Default: gfx942

6. Target hardware#

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)

7. IR dump workflow#

Enable with FLYDSL_DUMP_IR=1:

FLYDSL_DUMP_IR=1 FLYDSL_DUMP_DIR=./dumps python test_my_kernel.py

This produces 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.

8. Source files#

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.)
Contents