FlyDSL (Flexible layout python DSL)

A Python DSL and a MLIR stack for authoring high‑performance GPU kernels with explicit layouts and tiling.

FlyDSL is the Python front‑end of the project: a Flexible Layout Python DSL for expressing tiling, partitioning, data movement, and kernel structure at a high level.

FlyDSL: FlyDSL is powered by the Fly dialect: an end‑to‑end, MLIR‑native compiler stack for GPU kernels. Its core is the fly dialect—a first‑class layout IR with explicit algebra and coordinate mapping, plus a composable lowering pipeline to GPU/ROCDL.

Overview

• FlyDSL (Python DSL): author kernels in Python, compile them through the Fly dialect, and export compiled specializations as linkable C objects
  • Primary package: python/flydsl/
  • Kernel examples: kernels/ (importable as kernels.* from a source checkout; not installed in the wheel)
• Fly dialect: the layout IR and compiler foundation
  • Core abstractions: !fly.int_tuple, !fly.layout, !fly.coord_tensor, !fly.memref
  • Algebra ops: composition/product/divide/partition + coordinate mapping ops
• Embedded MLIR Python runtime (_mlir)
  • No external mlir python wheel is required: MLIR python bindings are included with the FlyDSL package/build artifacts

Repository layout

FlyDSL/
├── scripts/                   # build & test scripts
│   ├── build_llvm.sh          # build LLVM/MLIR from source (applies thirdparty/llvm-extensions)
│   ├── llvm_extension.sh      # create/rebase an LLVM extension
│   ├── build.sh               # build FlyDSL (C++ + Python bindings)
│   ├── run_tests.sh           # run tests
│   └── run_benchmark.sh       # run performance benchmarks
├── include/flydsl/            # C++ Fly/FlyROCDL dialect headers
├── lib/                       # C++ dialect implementation + Python bindings
├── python/
│   ├── flydsl/                # Python DSL sources
│   │   ├── expr/              # DSL expression API (typing, primitive, derived, arith, gpu, rocdl, math)
│   │   ├── compiler/          # JIT compilation pipeline (ast_rewriter, kernel_function, jit_function, backends/)
│   │   ├── runtime/           # Device runtime (device.py, device_runtime/)
│   │   ├── utils/             # Utilities (env, logger)
│   │   └── autotune.py        # Triton-style autotune module
│   └── mlir_flydsl/           # MLIR Python bindings (built, not edited)
├── 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/                   # Source-tree GPU kernels (not installed in the wheel)
├── tests/                     # All tests (kernels/, mlir/, unit/)
├── CMakeLists.txt             # top-level CMake
└── setup.py                   # Python packaging

Getting started

Prerequisites

• Python: Python 3.10+ with pip
• ROCm: required for GPU execution, tests, and benchmarks (tested on ROCm 6.x, 7.x)

Install FlyDSL

For most users, install the published package directly:

python -m pip install flydsl

Verify the install

python -c "import flydsl; print('FlyDSL installed')"

Build from source

Build from source only if you are developing FlyDSL itself or need a custom MLIR/LLVM build.

Prerequisites for source builds:

• Build tools: cmake (>=3.20), C++17 compiler, optionally ninja
• Python deps: nanobind, numpy, pybind11 (installed by scripts/build_llvm.sh; install them manually if you skip that step)

# Clone ROCm LLVM and build MLIR (takes ~30min with -j64)
bash scripts/build_llvm.sh -j64

# Build FlyDSL C++ dialects, compiler passes, and Python bindings
bash scripts/build.sh -j64

# Install in development mode
python -m pip install -e .

If you already have an MLIR build with Python bindings enabled, point to it instead:

python -m pip install nanobind numpy pybind11  # build.sh does not install these
export MLIR_PATH=/path/to/llvm-project/build-flydsl/mlir_install
MLIR_PATH=$MLIR_PATH bash scripts/build.sh -j64
python -m pip install -e .

Note: If MLIR_PATH is set in your environment pointing to a wrong LLVM build, unset MLIR_PATH first.

Run tests

Tests and examples require pytest, pandas, and a ROCm build of torch (not installed by pip install -e .):

python -m pip install pytest pandas
# torch must be a ROCm build matching your ROCm version (rocm7.2 shown):
python -m pip install torch --index-url https://download.pytorch.org/whl/rocm7.2

# Run GEMM correctness tests (fast, ~15s)
python -m pytest tests/kernels/test_preshuffle_gemm.py -m "not large_shape"

# Run performance benchmarks
bash scripts/run_benchmark.sh

Test layout, pytest markers, and environment variables used by the suite are documented in tests/README.md .

Quick reference

# Install from PyPI:
python -m pip install flydsl

# Full source build from scratch:
bash scripts/build_llvm.sh -j64   # one-time: build LLVM/MLIR
bash scripts/build.sh -j64        # build FlyDSL
python -m pip install -e .        # install in dev mode
bash scripts/run_tests.sh         # verify

# Rebuild after code changes (C++ only):
bash scripts/build.sh -j64

# Rebuild after Python-only changes:
# No rebuild needed — editable install picks up changes automatically.

Troubleshooting

• Wrong LLVM picked up (std::gcd not found, redeclaration errors)

  • unset MLIR_PATH and let build.sh auto-detect, or set it to the correct path.

• No module named flydsl

  • Run python -m pip install flydsl. For source checkouts, run python -m pip install -e . after building.

• MLIR .so load errors

  • Add MLIR build lib dir to the loader path:

    export LD_LIBRARY_PATH=$(pwd)/build-fly/python_packages/flydsl/_mlir/_mlir_libs:$LD_LIBRARY_PATH
    

Documentation

Full documentation: rocm.github.io/FlyDSL

• Kernel cache issues (stale results after code changes)
  • The JIT disk cache auto-invalidates on source/closure changes; only needed for C++ pass or non-closure helper changes
  • Clear manually: rm -rf ~/.flydsl/cache or export FLYDSL_RUNTIME_ENABLE_CACHE=0

📐 Layout System

FlyDSL introduces a layout system to express complex data mapping patterns on GPUs (tiling, swizzling, vectorization).

Core Abstractions

1. Shape: The extent of dimensions (e.g., (M, N)).
2. Stride: The distance between elements in memory (e.g., (1, M) for column-major).
3. Layout: A pair of (Shape, Stride) that maps a logical Coordinate to a physical linear Index.

Formula: Index = dot(Coord, Stride) = sum(c_i * s_i)

Operations

• Construction: make_shape, make_stride, make_layout, make_coord
• Mapping:
  • crd2idx(coord, layout) -> index: Convert logical coordinate to physical index.
  • idx2crd(index, layout) -> coord: Convert physical index to logical coordinate.
• Inspection: size, cosize, rank
• Algebra:
  • composition(A, B): Compose layouts (A ∘ B).
  • product(A, B): Combine layouts (Logical, Tiled, Blocked, etc.).
  • divide(A, B): Partition layout A by B (Logical, Tiled, etc.).

🐍 Python API (flydsl)

@flyc.kernel / @flyc.jit API

import flydsl.compiler as flyc
import flydsl.expr as fx
from flydsl.expr import arith, gpu

@flyc.kernel
def my_kernel(arg_a: fx.Tensor, arg_b: fx.Tensor, n: fx.Constexpr[int]):
    tid = gpu.thread_idx.x
    bid = gpu.block_idx.x
    # ... kernel body using layout ops ...

@flyc.jit
def launch(arg_a: fx.Tensor, arg_b: fx.Tensor, n: fx.Constexpr[int],
           stream: fx.Stream = fx.Stream(None)):
    grid_x = (n + 255) // 256
    my_kernel(arg_a, arg_b, n).launch(
        grid=(grid_x, 1, 1),
        block=(256, 1, 1),
        stream=stream,
    )

Compilation Pipeline

On first call, @flyc.jit traces the Python function into an MLIR module, then compiles it through MlirCompiler. The pass list is built by RocmBackend._pipeline_parts() in three stages — see docs/architecture_guide.md for the per-pass table.

Python Function (@flyc.kernel / @flyc.jit)
        │
        ▼  AST Rewriting + Tracing
   MLIR Module (fly, gpu, arith, scf, memref, vector dialects)
        │
        ▼  MlirCompiler.compile()
   ┌──────────────────────────────────────────────────────────┐
   │ 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{...}, fly-rocdl-cluster-attr) │
   ├──────────────────────────────────────────────────────────┤
   │ B. binary_prep_fragments  (→ LLVM)                       │
   │    rocdl-attach-target{chip=gfxNNN} →                    │
   │    convert-scf-to-cf → convert-cf-to-llvm →              │
   │    gpu-to-llvm → convert-vector/arith/func-to-llvm →     │
   │    reconcile-unrealized-casts                            │
   ├──────────────────────────────────────────────────────────┤
   │ C. binary_fragment                                       │
   │    gpu-module-to-binary{format=fatbin}                   │
   └──────────────────────────────────────────────────────────┘
        │
        ▼
   Cached Compiled Artifact (ExecutionEngine)

Compiled kernels are cached to disk (~/.flydsl/cache/) and reused on subsequent calls with the same type signature.

⚙️ Hierarchical Kernel Control

FlyDSL keeps the tiling hierarchy explicit across block, warp, thread, and instruction scopes using layout algebra:

import flydsl.expr as fx

# Define thread and value layouts for tiled copy
thr_layout = fx.make_layout((THR_M, THR_N), (1, THR_M))
val_layout = fx.make_layout((VAL_M, VAL_N), (1, VAL_M))

# Create tiled copy with vectorized atoms
copy_atom = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32)
layout_thr_val = fx.raked_product(thr_layout, val_layout)
tile_mn = fx.make_tile(fx.make_layout(THR_M, 1), fx.make_layout(VAL_M, 1))
tiled_copy = fx.make_tiled_copy(copy_atom, layout_thr_val, tile_mn)

# Partition tensor across blocks and threads
thr_copy = tiled_copy.get_slice(tid)
partition_src = thr_copy.partition_S(block_tile_A)
partition_dst = thr_copy.partition_D(register_fragment)

# Execute copy
fx.copy(copy_atom, partition_src, partition_dst)

With per-level partitions, you can allocate register fragments, emit predicate masks, and schedule MFMA/vector instructions while retaining full knowledge of the execution hierarchy.

🧮 Minimal VecAdd Example

This condensed snippet mirrors examples/01-vectorAdd.py, showing how to define GPU kernels with layout algebra and tiled copies:

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

@flyc.kernel
def vectorAddKernel(
    A: fx.Tensor, B: fx.Tensor, C: fx.Tensor,
    block_dim: fx.Constexpr[int],
):
    bid = fx.block_idx.x
    tid = fx.thread_idx.x

    # Partition tensors by block
    tA = fx.logical_divide(A, fx.make_layout(block_dim, 1))
    tB = fx.logical_divide(B, fx.make_layout(block_dim, 1))
    tC = fx.logical_divide(C, fx.make_layout(block_dim, 1))

    tA = fx.slice(tA, (None, bid))
    tB = fx.slice(tB, (None, bid))
    tC = fx.slice(tC, (None, bid))

    tA = fx.logical_divide(tA, fx.make_layout(1, 1))
    tB = fx.logical_divide(tB, fx.make_layout(1, 1))
    tC = fx.logical_divide(tC, fx.make_layout(1, 1))

    # Load to registers, compute, store via copy atoms
    copyAtom = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32)
    rA = fx.make_rmem_tensor(1, fx.Float32)
    rB = fx.make_rmem_tensor(1, fx.Float32)
    rC = fx.make_rmem_tensor(1, fx.Float32)

    fx.copy_atom_call(copyAtom, fx.slice(tA, (None, tid)), rA)
    fx.copy_atom_call(copyAtom, fx.slice(tB, (None, tid)), rB)

    vC = fx.arith.addf(fx.memref_load_vec(rA), fx.memref_load_vec(rB))
    fx.memref_store_vec(vC, rC)
    fx.copy_atom_call(copyAtom, rC, fx.slice(tC, (None, tid)))

@flyc.jit
def vectorAdd(
    A: fx.Tensor, B: fx.Tensor, C: fx.Tensor,
    n: fx.Int32,  # dynamic int32
    const_n: fx.Constexpr[int],  # static int32, affects JIT cache-key
    stream: fx.Stream = fx.Stream(None),
):
    block_dim = 64
    grid_x = (n + block_dim - 1) // block_dim
    vectorAddKernel(A, B, C, block_dim).launch(
        grid=(grid_x, 1, 1), block=[block_dim, 1, 1], stream=stream,
    )

# Usage
n = 128
A = torch.randint(0, 10, (n,), dtype=torch.float32).cuda()
B = torch.randint(0, 10, (n,), dtype=torch.float32).cuda()
C = torch.zeros(n, dtype=torch.float32).cuda()
vectorAdd(A, B, C, n, n + 1, stream=torch.cuda.Stream())

torch.cuda.synchronize()
print("Result correct:", torch.allclose(C, A + B))

See examples/ for more examples including tiled copy (02-tiledCopy.py), tiled MMA (03-tiledMma.py), and preshuffle GEMM (04-preshuffle_gemm.py).

✅ Testing Status

Targeted platforms:

• AMD MI300X/MI308X (gfx942), AMD MI350/MI355X (gfx950), gfx1250, Radeon AI PRO R9700 (gfx1201)
• Linux / ROCm 6.x and 7.x. Hardware coverage varies by CI runner and revision; consult the current workflow results for validation status.

🙏 Acknowledgements

FlyDSL's design is inspired by ideas from several projects:

• Categorical Foundations for CuTe Layouts — mathematical framework for layout algebra (companion code)
• NVIDIA CUTLASS — CuTe layout algebra concepts (BSD-3-Clause parts only; no EULA-licensed code was referenced)
• ROCm Composable Kernel — tile-based kernel design patterns for AMD GPUs
• ROCm AIter — test infrastructure and performance comparison baselines (MIT)
• Triton — Python DSL for GPU kernel authoring
• HipKittens — minimal, opinionated C++ embedded primitives for fast AMD AI kernels (part of the ThunderKittens family)

📄 License

Apache License 2.0