Testing & benchmarking guide

This guide covers the FlyDSL test infrastructure, how to run tests and benchmarks, how to write new tests, and how to measure and compare performance.

Quick reference

Category	Location	Requires GPU	Description
MLIR lit tests	`tests/mlir/LayoutAlgebra/`, `tests/mlir/Conversion/`, `tests/mlir/Transforms/`	No	Verify Fly dialect lowering
Python tests	`tests/python/examples/`	Varies	Python integration and JIT-cache examples
GPU kernel tests	`tests/kernels/test_*.py`	Yes	Full compilation → GPU execution
Cache pre-warming examples	`tests/python/examples/`	Varies	Populate the normal JIT cache before workload execution

Run GEMM tests:

bash scripts/run_tests.sh

Run benchmarks:

bash scripts/run_benchmark.sh

1. Test categories

1.1 MLIR lit tests (`tests/mlir/`)

MLIR-based tests organized by category, verified using the fly-opt tool. These tests confirm that Fly dialect operations lower correctly to standard MLIR dialects without requiring a GPU.

Directories:

Directory	Tests	Description
`LayoutAlgebra/`	`coalesce.mlir`, `composition.mlir`, `construction.mlir`, `coordinate.mlir`, `divide.mlir`, `int_tuple.mlir`, `product.mlir`, `size_cosize.mlir`	Layout algebra operations
`Conversion/`	`gpu_ops.mlir`, `memref_alloca.mlir`, `memref_ops.mlir`, `mma_atom.mlir`, `pointer_ops.mlir`, `type_conversion.mlir`	Dialect conversion passes
`Transforms/`	`canonicalize.mlir`, `layout_lowering.mlir`	Transformation passes

Running individually:

# Build fly-opt first if needed
cmake --build build-fly --target fly-opt -j$(nproc)

# Run a single test
build-fly/bin/fly-opt --fly-canonicalize tests/mlir/LayoutAlgebra/construction.mlir

Editor support (MLIR LSP): build-fly/bin/flydsl-lsp-server is a thin MlirLspServerMain wrapper with Fly (+ backend) dialects registered. Point the editor's MLIR LSP client / mlir-lsp-server path at that binary for .mlir diagnostics, hover, go-to-definition, and completion. It is not a Python DSL / kernel language server.

cmake --build build-fly --target flydsl-lsp-server -j$(nproc)

1.2 Python tests (`tests/python/`)

Python-based integration tests, including an example that pre-warms the normal JIT cache. It does not create standalone linkable artifacts.

Running:

python tests/python/examples/aot_example.py

1.3 GPU kernel tests (`tests/kernels/`)

Full end-to-end tests that compile FlyDSL kernels, execute them on the GPU, and validate results against a PyTorch reference.

Files:

Test file	Kernel	Description
`test_vec_add.py`	VecAdd	Vector addition (C = A + B)
`test_softmax.py`	Softmax	Row-wise softmax
`test_layernorm.py`	LayerNorm	Layer normalization
`test_rmsnorm.py`	RMSNorm	RMS normalization
`test_preshuffle_gemm.py`	GEMM	Preshuffle MFMA GEMM (fp8/int8/fp16/bf16)
`test_moe_gemm.py`	MoE GEMM	Mixture-of-Experts GEMM
`test_moe_reduce.py`	MoE Reduce	MoE reduction kernel
`test_pa.py`	Paged Attn	Paged attention decode
`test_quant.py`	Quantization	Quantization ops
`test_ref.py`	Reference	Reference implementations

Running individually:

python tests/kernels/test_softmax.py
python tests/kernels/test_preshuffle_gemm.py --in_dtype fp8 -M 16 -N 5120 -K 8192

1.4 JIT cache pre-warming example (`tests/python/examples/`)

The script retains its historical filename, but it exercises compile-only JIT cache population:

tests/python/examples/
└── aot_example.py      # JIT cache pre-warming for preshuffle GEMM

2. Test runner scripts

2.1 `scripts/run_tests.sh`

Runs the full FlyDSL test suite:

bash scripts/run_tests.sh

Features:

Auto-discovers build directory (build-fly/)
Auto-selects the GPU with the most free VRAM when HIP_VISIBLE_DEVICES is unset
Sets up PYTHONPATH and LD_LIBRARY_PATH, and exports FLYDSL_RUN_QUANT=1
Runs pytest over tests/kernels/, tests/language/, tests/unit/, tests/system/, tests/extension/, and tests/python/examples/
Runs the standalone examples/ scripts and the MLIR FileCheck tests (tests/mlir/)
By default skips large_shape-marked tests; set RUN_TESTS_FULL=1 to include them. multi_gpu and benchmark tests remain separate.
Skips the MLIR/FileCheck stage when FileCheck is unavailable unless FLYDSL_REQUIRE_FILECHECK=1 is set
Fail-fast: exits on the first failure

Environment setup:

PYTHONPATH="${BUILD_DIR}/python_packages:${REPO_ROOT}:${PYTHONPATH}"
LD_LIBRARY_PATH="${MLIR_LIBS_DIR}:${LD_LIBRARY_PATH}"

2.2 `scripts/run_benchmark.sh`

Specialized benchmarking harness for performance characterization.

Selected default configurations:

The script is the source of truth for the complete matrix; these excerpts show the configuration format.

# Softmax/RMSNorm: "M,N,dtype"
SOFTMAX_SHAPES='32768,8192,bf16'
RMSNORM_SHAPES='32768,8192,bf16'

# Preshuffle GEMM: "dtype,M,N,K,tile_m,tile_n,tile_k"
GEMM_SHAPES='
fp8,16,40960,5120,16,128,256
fp8,16,77824,5120,16,128,256
fp8,5120,5120,8320,64,256,128
fp8,9728,8192,8320,64,256,128
int8,9728,8192,8320,64,256,128
bf16,5120,5120,8320,64,256,128
'

# FP4 GEMM (gfx950 only): "M,N,K"
GEMM_FP4_SHAPES='8192,8192,8192'

Selective execution:

bash scripts/run_benchmark.sh                    # default: all benchmark groups
bash scripts/run_benchmark.sh softmax             # only softmax
bash scripts/run_benchmark.sh gemm moe            # GEMM and MoE
bash scripts/run_benchmark.sh --only softmax,rmsnorm
bash scripts/run_benchmark.sh --list              # list available ops

Output format: Tabular with TB/s and TFLOPS columns:

op             shape                              dtype       TB/s    TFLOPS
-------------- ---------------------------------- ---------- ---------- ----------
gemm           16x40960x5120                      fp8         1.234     56.789

Logs: Written to ${BENCH_LOG_DIR:-/tmp/flydsl_bench}/

3. Pytest configuration

3.1 `tests/conftest.py`

Pytest configuration with MLIR context fixtures for the Fly dialect.

Fixtures:

@pytest.fixture
def ctx():
    """Fresh MLIR context per test with dialects registered."""
    # Creates Context, yields object with: ctx.context, ctx.module, ctx.location

@pytest.fixture
def module(ctx):
    """Provides ctx.module."""

@pytest.fixture
def insert_point(ctx):
    """Sets insertion point to module body."""

Build discovery: The configuration supports multiple build layouts:

build-fly/python_packages (preferred)
build/python_packages/flydsl (fallback)

Session hook: Prevents pytest exit code 5 (no tests collected) from being treated as failure.

4. Performance measurement

4.1 `tests/test_common.py`

Core performance testing utilities (adapted from AIter).

perftest() decorator:

@perftest(num_iters=20, num_warmup=3, testGraph=False, num_rotate_args=0)
def my_kernel_test(Input, Output):
    # Kernel invocation
    ...

Features:

Device memory profiling to determine rotation count
Torch CUDA event timing
HIPGraph capture mode (testGraph=True)
Cache-aware iteration calculation

checkAllclose() function:

checkAllclose(output, reference, rtol=1e-2, atol=1e-2, tol_err_ratio=0.05)

Returns a mismatch ratio in [0, 1] (0 = pass).

verify_output() function:

verify_output(c_out, c_ref, atol=1e-2, rtol=1e-2, msg='')

High-level validation wrapper around checkAllclose.

4.2 `tests/kernels/benchmark_common.py`

Shared benchmark harness for performance comparison.

Key functions:

# Measure device time (torch CUDA events)
gpu_us = bench_gpu_us_torch(fn, warmup=20, iters=200)

5. Test utilities (`tests/utils.py`)

Weight utilities

from tests.utils import pertoken_quant, shuffle_weight

# Per-token quantization (handles NaN/Inf)
quantized, scales = pertoken_quant(tensor, dtype=torch.float8_e4m3fnuz)

# Weight preshuffle for MFMA (layout 16x16)
shuffled = shuffle_weight(weight, layout=(16, 16))

6. Writing new tests

6.1 PyIR test pattern (no GPU)

# tests/python/test_my_feature.py
import flydsl.expr as fx
from flydsl.expr.typing import T

def test_my_layout_op(ctx, insert_point):
    shape = fx.make_shape(4, 8)
    stride = fx.make_stride(8, 1)
    layout = fx.make_layout(shape, stride)
    result = fx.size(layout)
    ir_str = str(ctx.module)
    assert "fly.make_layout" in ir_str

6.2 GPU kernel test pattern (new API)

# tests/kernels/test_my_kernel.py
import torch
import flydsl.compiler as flyc
import flydsl.expr as fx
from flydsl.expr import gpu
from tests.test_common import checkAllclose

@flyc.kernel
def my_kernel(A: fx.Tensor, B: fx.Tensor, N: fx.Constexpr[int]):
    tid = gpu.thread_idx.x
    bid = gpu.block_idx.x
    # ... kernel body ...

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

def test_my_kernel():
    N = 1024
    A = torch.randn(N, device="cuda", dtype=torch.float32)
    B = torch.empty(N, device="cuda", dtype=torch.float32)

    launch(A, B, N)

    # Reference
    ref = A  # or some computation

    # Validate
    err = checkAllclose(B, ref, rtol=1e-2, atol=1e-2)
    assert err == 0, f"Mismatch: {err * 100:.2f}%"

6.3 Benchmark test pattern

from tests.kernels.benchmark_common import bench_gpu_us_torch

def benchmark_my_kernel():
    # Setup
    launch_fn = compile_my_kernel(...)

    def run():
        launch_fn(input_tensor, output_tensor)

    # Measure
    gpu_us = bench_gpu_us_torch(run, warmup=20, iters=200)

    # Compute metrics
    total_bytes = 2 * M * N * elem_size
    bandwidth_tbs = total_bytes / (gpu_us * 1e-6) / 1e12
    print(f"Time: {gpu_us:.1f} us, Bandwidth: {bandwidth_tbs:.2f} TB/s")

7. GEMM test CLI arguments

The test_preshuffle_gemm.py test supports extensive CLI configuration:

python tests/kernels/test_preshuffle_gemm.py \
    --in_dtype fp8 \
    -M 16 -N 5120 -K 8192 \
    --tile_m 16 --tile_n 128 --tile_k 256 \
    --lds_stage 2 \
    --num_iters 20 \
    --num_warmup 3 \
    --no_aiter_bench \
    --test_graph        # or -tg for HIPGraph mode

# FP4 weight path (gfx950 only)
python tests/kernels/test_preshuffle_gemm.py \
    --in_dtype fp4 --wfp4 \
    -M 8192 -N 8192 -K 8192

8. Test configuration via environment variables

Variable	Used by	Description
`ROCDSL_SOFTMAX_SHAPES`	`test_softmax.py`	Override softmax test shapes (`"M,N,dtype;..."`)
`ROCDSL_LAYERNORM_SHAPES`	`test_layernorm.py`	Override layernorm test shapes
`FLYDSL_DUMP_IR`	Compiler	Dump intermediate IR at each pipeline stage
`FLYDSL_DUMP_DIR`	Compiler	IR dump directory (default: `~/.flydsl/debug`)
`FLYDSL_RUNTIME_CACHE_DIR`	Compiler	Cache directory (default: `~/.flydsl/cache`)
`RUN_TESTS_FULL`	`run_tests.sh`	Set to `1` to include `large_shape` cases; multi-GPU and benchmark tests remain separate
`BENCH_LOG_DIR`	`run_benchmark.sh`	Benchmark log directory (default: `/tmp/flydsl_bench`)

9. IR dump workflow

Via `MlirCompiler`

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

Produces numbered .mlir files per pipeline stage plus final_isa.s.

Dedicated IR dump script

bash scripts/dumpir.sh python my_kernel.py

Running a hand-edited `.s`

To try "what if these few instructions were different" without rebuilding the compiler, edit the dumped ISA and run the unmodified test against it with tools/inject_asm/inject_asm.py:

python3 tools/inject_asm/inject_asm.py prepare ./ws -- python -m pytest tests/kernels/test_my_kernel.py -k one_case
vim ./ws/edited.s
python3 tools/inject_asm/inject_asm.py run ./ws                # the edited assembly
python3 tools/inject_asm/inject_asm.py run ./ws --baseline     # the unedited dump, under the same conditions

prepare runs the command twice into a fresh workspace: once with FLYDSL_DUMP_IR for the ISA, and once with a scratch FLYDSL_RUNTIME_CACHE_DIR for the compiled kernel, since a dump run leaves the cache alone. It copies the dump to original.s (read-only) and edited.s. run assembles edited.s, swaps it into the cache entry's gpu.binary, and reruns the command with FLYDSL_RUNTIME_RUN_ONLY=1, so a cache miss is an error rather than a silent recompile. This works because a cache entry stores the compiled module as MLIR text and rebuilds its ExecutionEngine from it on load.

The command must compile exactly one JIT function in one specialization; narrow it with -k otherwise. Constexpr specializations of a kernel share its symbol, so the dump keeps only the last one compiled and their cache entries cannot be told apart. Keep every kernel symbol unchanged in edited.s, since the host launches kernels by name.

FlyDSL assembles with the LLVM it was built with, while the tool defaults to the ROCm clang, which is usually older; both link with the ROCm ld.lld. prepare checks that original.s reassembles to the compiler's code object byte for byte, which holds on gfx942 and gfx950. Dumps for gfx11 and later use directives the ROCm clang cannot read, so for those pass prepare --clang with the clang of the LLVM FlyDSL was built with; run reuses it.

inject_asm.py inject CACHE_DIR FILE.s is the single patching step on its own, for a cache you populated yourself. --check compares instead of patching and --save-hsaco also writes the code object next to the .s. It refuses any directory that contains or lies inside the default cache ~/.flydsl/cache, where a patched entry would be served to every later run.

10. Source files

File	Description
`scripts/run_tests.sh`	Full test runner (pytest + examples + FileCheck)
`scripts/run_benchmark.sh`	Benchmark harness with configurable shapes
`scripts/dumpir.sh`	IR dump helper script
`tools/inject_asm/inject_asm.py`	Run a test against a hand-edited ISA dump (debug only)
`tests/conftest.py`	Pytest fixtures (MLIR context, module, insert point)
`tests/test_common.py`	`perftest()`, `checkAllclose()`, `verify_output()`
`tests/utils.py`	`pertoken_quant()`, `shuffle_weight()`
`tests/kernels/benchmark_common.py`	`bench_gpu_us_torch()`, benchmark harness
`tests/mlir/LayoutAlgebra/`, `tests/mlir/Conversion/`, `tests/mlir/Transforms/`	MLIR lit tests
`tests/python/examples/`	Python integration and cache pre-warming examples
`tests/kernels/test_*.py`	GPU kernel test modules
`tests/python/examples/aot_example.py`	Compile-only JIT cache pre-warming example

Navigation Menu

Repository navigation

testing_benchmarking_guide.md

testing_benchmarking_guide.md

Testing & benchmarking guide

Quick reference

1. Test categories

1.1 MLIR lit tests (`tests/mlir/`)

1.2 Python tests (`tests/python/`)

1.3 GPU kernel tests (`tests/kernels/`)

1.4 JIT cache pre-warming example (`tests/python/examples/`)

2. Test runner scripts

2.1 `scripts/run_tests.sh`

2.2 `scripts/run_benchmark.sh`

3. Pytest configuration

3.1 `tests/conftest.py`

4. Performance measurement

4.1 `tests/test_common.py`

4.2 `tests/kernels/benchmark_common.py`

5. Test utilities (`tests/utils.py`)

Weight utilities

6. Writing new tests

6.1 PyIR test pattern (no GPU)

6.2 GPU kernel test pattern (new API)

6.3 Benchmark test pattern

7. GEMM test CLI arguments

8. Test configuration via environment variables

9. IR dump workflow

Via `MlirCompiler`

Dedicated IR dump script

Running a hand-edited `.s`

10. Source files

FilesExpand file tree

testing_benchmarking_guide.md

Latest commit

History

testing_benchmarking_guide.md

File metadata and controls

Testing & benchmarking guide

Quick reference

1. Test categories

1.1 MLIR lit tests (tests/mlir/)

1.2 Python tests (tests/python/)

1.3 GPU kernel tests (tests/kernels/)

1.4 JIT cache pre-warming example (tests/python/examples/)

2. Test runner scripts

2.1 scripts/run_tests.sh

2.2 scripts/run_benchmark.sh

3. Pytest configuration

3.1 tests/conftest.py

4. Performance measurement

4.1 tests/test_common.py

4.2 tests/kernels/benchmark_common.py

5. Test utilities (tests/utils.py)

Weight utilities

6. Writing new tests

6.1 PyIR test pattern (no GPU)

6.2 GPU kernel test pattern (new API)

6.3 Benchmark test pattern

7. GEMM test CLI arguments

8. Test configuration via environment variables

9. IR dump workflow

Via MlirCompiler

Dedicated IR dump script

Running a hand-edited .s

10. Source files

1.1 MLIR lit tests (`tests/mlir/`)

1.2 Python tests (`tests/python/`)

1.3 GPU kernel tests (`tests/kernels/`)

1.4 JIT cache pre-warming example (`tests/python/examples/`)

2.1 `scripts/run_tests.sh`

2.2 `scripts/run_benchmark.sh`

3.1 `tests/conftest.py`

4.1 `tests/test_common.py`

4.2 `tests/kernels/benchmark_common.py`

5. Test utilities (`tests/utils.py`)

Via `MlirCompiler`

Running a hand-edited `.s`