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

Trace-based JIT compilation, specialization, and C object export

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

CompiledArtifact

Serializable compiled MLIR plus lazy ExecutionEngine state

python/flydsl/compiler/jit_executor.py

CompiledFunction

Callable specialization with export_to_c() support

python/flydsl/compiler/jit_function.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, compile, argument adapters
│   │   ├── _aot.py                   # Internal C object/header export implementation
│   │   ├── jit_function.py           # @jit decorator, MlirCompiler, JitCacheManager
│   │   ├── kernel_function.py        # @kernel decorator, KernelFunction, KernelLauncher
│   │   ├── jit_executor.py           # CompiledArtifact and 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                # Thread-level tiled copy/MMA views and helpers
│   │   ├── 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
│
├── 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_bf16_gfx1250.py      # GFX1250 BF16/FP16 GEMM
│   │   ├── gemm_a8w8_gfx1250.py      # GFX1250 FP8 GEMM (per-token/channel, blockscale)
│   │   ├── gemm_a8w4_mxscale_gfx1250.py # GFX1250 FP8 x MXFP4 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/          # 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
│   ├── common/mma/                   # Shared MMA pipeline helpers
│   │   └── mfma_preshuffle_pipeline.py # Preshuffle layout and XCD remapping
│   ├── conv/                         # Convolution kernels
│   │   ├── conv3d_implicit.py        # Implicit-GEMM 3D convolution
│   │   └── conv3d_implicit_fp8.py    # FP8 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 integration tests and examples
│   ├── 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-rocdl-fastmath-ops,              │
   │              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="..."}       │
   └────────────────────────────────────────────────────────┘
        │
        ▼
   CompiledArtifact (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)

The tables explain the stages conceptually. The following excerpt is included directly from the backend at documentation build time, so it is the exact pass definition for this checkout:

pre_binary_fragments = [
    "fly-rewrite-func-signature",
    "fly-canonicalize",
    "fly-rocdl-expand-ops",
    "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",
    f"gpu.module(convert-scf-to-cf,cse,"
    f"convert-rocdl-fastmath-ops,"
    f"convert-gpu-to-rocdl{{chipset={chip} index-bitwidth=0 runtime=HIP use-bare-ptr-memref-call-conv=true}},"
    f"fly-rocdl-cluster-attr)",
]
binary_prep_fragments = [
    f"rocdl-attach-target{{{self._format_pass_opts(rocdl_opts)}}}",
    "convert-scf-to-cf",
    "convert-cf-to-llvm",
    "gpu-to-llvm{use-bare-pointers-for-host=true use-bare-pointers-for-kernels=true}",
    "convert-vector-to-llvm",
    "convert-arith-to-llvm",
    "convert-func-to-llvm",
    "reconcile-unrealized-casts",
    *(
        ["ensure-debug-info-scope-on-llvm-func{emission-kind=LineTablesOnly}"]
        if env.debug.enable_debug_info
        else []
    ),
]
binary_fragment = f'gpu-module-to-binary{{format=fatbin opts="{" ".join(bin_cli_opts)}"}}'

#

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 operations, 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-rocdl-fastmath-ops, convert-gpu-to-rocdl{chipset=gfxNNN ...}, fly-rocdl-cluster-attr)

Inside the GPU module: SCF→CF, CSE, ROCDL fast-math ops lowering, 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: CompiledArtifact lazily creates an MLIR ExecutionEngine and invokes the compiled code.

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

3.4 Precompilation and C export#

flyc.compile(launcher, *specialization_args) uses the same tracing, lowering, and cache path as a normal @flyc.jit call, then returns a CompiledFunction that retains the selected CompiledArtifact. Live device arguments preserve the initial launch; null pointer wrappers and compatible non-device tensors are compile placeholders, so compilation and export can run without initializing a GPU runtime when the target architecture is specified explicitly.

CompiledFunction.export_to_c(file_path, file_name, function_prefix="") rewrites a fresh copy of the already-lowered module for standalone linkage. It:

  1. namespaces and internalizes the module’s definitions;

  2. selects the backend GPU object and attaches the AOT module handler;

  3. emits a packed int32_t entry(void **args) wrapper and ABI metadata;

  4. emits a position-independent host object and matching C header; and

  5. embeds the backend AOT runtime archive into the host object. Only backend system libraries remain as link-time dependencies.

This export path reuses the same GPU binary as the compiled specialization; it does not retrace the Python launcher. Backend-specific implementation hooks and the _AOTCompiledFunction helper remain internal APIs.

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.thread_idx.x
    bid = fx.block_idx.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 CompiledArtifact#

Stores the compiled MLIR and lazily creates MLIR’s ExecutionEngine for JIT execution:

  • Thread-safe with lazy engine initialization.

  • Serializable (pickle) for disk caching.

  • Supports packed calling convention via ctypes.

  • Provides .dump(), .ir, and .source_ir for inspecting compiled or original IR.

4.5 CompiledFunction#

The public result of flyc.compile(). It keeps the compiled artifact alive, provides a positional-only fast call path for runtime arguments, and exposes export_to_c() for producing a linkable host object and generated C header. Runtime initialization is deferred for placeholder-based compilations until the callable is invoked.

4.6 DslType / JitArgument protocols#

Extensible type system for mapping Python values to MLIR. The language-level contracts, including Storable, are in DSL protocols.

# DslType protocol — for values used inside kernel/jit functions
class DslType(Protocol):
    @classmethod
    def __construct_from_ir_values__(
        cls, values: List[ir.Value], exemplar: "DslType | None" = None
    ) -> "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 __cache_signature__(self) -> object: ...
    def __c_abi_spec__(self) -> List[Tuple[type, Callable]]: ...

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 __cache_signature__(self): ...
    def __c_abi_spec__(self): ...

__c_abi_spec__() returns ordered (ctype, fill) slots for the packed C interface; they need not be one-to-one with __get_ir_types__(). Do not implement __get_c_pointers__ — that hook has been removed.

4.7 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

CDNA 5, 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, flyc.compile, CompiledFunction, MlirCompiler, and JitCacheManager

python/flydsl/compiler/_aot.py

Internal C ABI construction and object/header export used by CompiledFunction.export_to_c

python/flydsl/compiler/kernel_function.py

@kernel decorator, KernelFunction, KernelLauncher, CompilationContext

python/flydsl/compiler/jit_executor.py

CompiledArtifact — serialized IR and lazy ExecutionEngine wrapper

python/flydsl/compiler/jit_argument.py

JitArgumentRegistry, tensor adapters, 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

ThrCopy, ThrMma, and tiled-operation helpers

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