DOCSEARCHbyAlgolia
Submit Docs

FlyDSL Documentation

repository·main·Indexed 42 days ago

https://github.com/rocm/flydsl

A domain-specific language for defining and running GPU kernels on ROCm. FlyDSL utilizes a two-stage model with @flyc.kernel for device-side logic (traced to MLIR) and @flyc.jit for host-side launch wrappers. It provides tools for layout algebra, register tensors, copy atoms, and integration with PyTorch via dlpack. The project includes a strict API stability model for its expr (fx) and compiler modules.

Tokens
58.9K
Snippets
145
Records
220
Agent score
55%

What's inside FlyDSL

  1. Overview of FlyDSL Components

    main

    FlyDSL is a Python-based Domain Specific Language (DSL) designed for authoring high-performance GPU kernels. The project is composed of several key layers:

    • FlyDSL (Python Front-end): The primary interface for kernel authors, located in python/flydsl/.
    • FlyDSL Compiler: A trace-based JIT compiler using @flyc.jit and @flyc.kernel decorators, located in python/flydsl/compiler/.
    • FlyDSL Expr: The library of DSL expression operations (arithmetic, vector, GPU, buffer, and ROCm-specific rocdl), located in python/flydsl/expr/.
    • Fly Dialect: An MLIR dialect implementing layout algebra, located in include/flydsl/Dialect/Fly/.
    • MlirCompiler: The end-to-end MLIR pass pipeline that transforms the DSL into binary, located in python/flydsl/compiler/jit_function.py.
    • JITCFunction: A wrapper for the MLIR ExecutionEngine used for JIT execution, located in python/flydsl/compiler/jit_executor.py.

  2. Explore pre-built FlyDSL GPU kernels

    main

    FlyDSL provides a collection of pre-built GPU kernels located in the kernels/ directory. These kernels are organized into functional subpackages and can be used as ready-to-use components or as reference implementations for developing your own kernels.

    Subpackages include:

    • gemm/: General Matrix Multiplication
    • norm/: Normalization
    • attention/: Attention mechanisms
    • moe/: Mixture-of-Experts
    • mma/: Matrix Multiply-Accumulate
    • common/: Shared utilities
    • conv/: Convolutions
    • comm/: Communication primitives

  3. Control flow and compile-time constants

    main

    FlyDSL provides mechanisms to handle both runtime and compile-time control flow.

    Python Loops

    The ASTRewriter handles two types of loops:

    • Compile-time unrolled loops: Use for i in range_constexpr(N):. These are fully unrolled in the generated IR.
    • Runtime loops: Use standard Python for i in range(runtime_value):. These are lowered by the AST rewriter for execution at runtime.

    Compile-time Constants

    Use const_expr() to mark a value as a compile-time constant. This is useful for values derived from fx.Constexpr parameters.

    from flydsl.expr import const_expr

@flyc.kernel
def my_kernel(data: fx.Tensor, N: fx.Constexpr[int]):
    tile_size = const_expr(N // 4)
    for i in range_constexpr(tile_size):
        ...
    from flydsl.expr import const_expr

@flyc.kernel
def my_kernel(data: fx.Tensor, N: fx.Constexpr[int]):
    # Compile-time unrolled loop
    for i in range_constexpr(N):
        ...

    # Runtime loop
    for i in range(runtime_value):
        ...

    # Using const_expr
    tile_size = const_expr(N // 4)
    for i in range_constexpr(tile_size):
        ...

  4. Core types in CuTe layout algebra

    main

    FlyDSL uses the CuTe layout algebra to describe multidimensional data layouts. The fundamental building blocks are:

    • Shape: A tuple of positive integers describing dimensions. Use fx.make_shape(M, N).
    • Stride: A tuple of integers describing step sizes per dimension. Use fx.make_stride(s0, s1).
    • Layout: A pair (Shape, Stride) that maps logical coordinates to physical offsets. Use fx.make_layout(shape, stride).
    • Coord: A tuple of integers identifying a position in logical space. Use fx.make_coord(i, j).
    shape = fx.make_shape(128, 64)
stride = fx.make_stride(1, 128)    # Column-major
layout = fx.make_layout(shape, stride)
coord = fx.make_coord(3, 5)

  5. Understand FlyDSL API stability and compatibility

    main

    FlyDSL uses a stability model to define which APIs are safe to use across minor releases (e.g., 0.3 → 0.4).

    • stable: Backward compatible across minor releases. Call paths, signatures, and semantics are guaranteed to work.
    • unstable: Everything else. These APIs may change or be removed in any minor release without notice.
    • deprecated: A lifecycle marker for a stable API. It remains protected by compatibility rules during its deprecation window and can only be removed after that window ends.

    The Returned-Object Rule: If an API is stable, the object it returns has a stable public interface. This means its non-underscore members (e.g., obj.member) and Python special methods (e.g., obj.__add__) are stable. However, the concrete implementation class, its constructor, or its import path may change as long as the public interface is preserved.

  6. Resolve LDS bank conflicts with XOR swizzling

    main

    LDS is banked (32 banks on gfx942, 64 banks on gfx950). Row-major tiles often cause bank conflicts because every row lands on the same banks.

    To fix this, use an XOR swizzle to fold the row index into the column address at 16-byte granularity. This adds minimal overhead (~1 SALU op) and requires zero extra LDS memory.

    Important: Apply the swizzle identically on both the write (global→LDS) and read (LDS→VGPR) paths. If only one side swizzles, the data will be read from the wrong location.

    Alternative: If swizzling is too complex to integrate, you can add 1–4 elements of padding per row to break stride alignment, though this consumes extra LDS memory.

    def swizzle_xor16(row, col_bytes, k_blocks16):
    return col_bytes ^ ((row % k_blocks16) * 16)   # k_blocks16 = tile_k_bytes // pack // 16

  7. Understand FlyDSL arithmetic types: Numeric and Vector

    main

    FlyDSL's arithmetic layer is built on two primary types: Numeric (scalars) and Vector (SIMD).

    • Numeric: The base scalar type wrapping an ir.Value. It includes:
      • Integers: Int4, Int8, Int16, Int32, Int64, Int128, and their unsigned counterparts (Uint8, etc.).
      • Floats: Float16, BFloat16, Float32, Float64.
      • Narrow Floats: Storage/transport types like Float8E5M2 (not intended for direct arithmetic).
      • Boolean: The i1 type used for comparisons and predicates.
    • Vector: A fixed-length sequence of N elements of a single Numeric type. Vectors inherit scalar operators, which are applied elementwise. A scalar operand used with a Vector is automatically broadcast across all lanes.

    Type classes act as constructors. For example, Float32(x) constructs or casts a value to Float32.

    import numpy as np
# Example of type construction and vector broadcasting
val = Int32(5)
vec = Float32x4(1.0)  # All four lanes are 1.0

  8. Use derived tiled operations for partitioning

    main

    For high-level tiled copy and MMA partitioning, use the classes in flydsl.expr.derived.

    Core Abstractions:

    • TiledCopy: Multi-thread tiled copy. Use get_slice(tid) to obtain a ThrCopy view.
    • TiledMma: Multi-thread tiled MMA. Use get_slice(tid) to obtain a ThrMma view.
    • ThrCopy: Per-thread copy view. Provides partition_S(src), partition_D(dst), and retile(t).
    • ThrMma: Per-thread MMA view. Provides partition_A(a), partition_B(b), and partition_C(c).

    Advanced Operations:

    • fx.make_layout_tv(thr, val): Build a thread-value layout.
    • fx.make_tiled_copy_A/B/C(copy_atom, tiled_mma): Create a TiledCopy matched to specific MMA operands.
    • fx.gather(copy_atom, base_iter, offset_tensor, dst_tensor, *, pred=None): Indexed load (dst = base[offset]) via a copy atom. The offset_tensor must be of shape (TV, Rest...).
    • fx.scatter(copy_atom, src_tensor, base_iter, offset_tensor, *, pred=None): Indexed store (base[offset] = src) via a copy atom.

  9. How the JIT pipeline handles linked externs

    main

    The FlyDSL JIT pipeline integrates external bitcode without requiring the compiler to import the framework itself. The flow is as follows:

    1. Registration: Calling a linked extern inside a @flyc.kernel body registers its metadata on the active CompilationContext.
    2. Compilation: JitFunction.__call__ snapshots link_libs and post_load_processors and passes them to MlirCompiler.compile and CompiledArtifact.
    3. Execution: The Python executor (jit_executor.py) uses the post_load_processors to initialize the module. If post_load_processors were registered but no modules were loaded, the executor raises a RuntimeError to prevent silent failures on uninitialized device globals.

  10. Understand occupancy and resource limits (VGPR, LDS, SGPR)

    main

    Occupancy (waves per SIMD) is determined by the minimum of three resource limiters, capped by the hardware maximum:

    occupancy (waves/SIMD) = min(vgpr_limit, lds_limit, sgpr_limit, HW_MAX)

    1. VGPR (Vector General Purpose Registers)

    On CDNA3 (gfx942) and CDNA4 (gfx950), arch_vgpr (used for VALU, VMEM, LDS, and prefetch) and accum_vgpr (used for MFMA writeback) share a combined 512-entry budget per SIMD.

    • If combined usage $\le 64$: 8 waves/SIMD (HW max).
    • If combined usage $\le 256$: 2 waves/SIMD.
    • If combined usage $> 512$: SPILL (severe regression).

    2. LDS (Local Data Share)

    LDS is a per-CU pool shared by all resident workgroups.

    • gfx942: 64 KB
    • gfx950: 160 KB
    • gfx1250: 320 KB

    3. SGPR (Scalar General Purpose Registers)

    SGPRs (kernel args, buffer descriptors) also gate occupancy. On CDNA, there are ~800 SGPRs per SIMD, allocated in granules of 16.

  11. Work with Copy Atoms and Tiled Copies

    main

    Copy atoms define the hardware-level instruction for moving data (e.g., 128-bit generic copies or AMD-specific buffer descriptor copies). Tiled copies use these atoms to perform large-scale data movement across layouts.

    Copy Atom Types

    • fx.UniversalCopy128b(), fx.UniversalCopy64b(), fx.UniversalCopy32b(): Generic bit-width copies.
    • fx.UniversalCopy(bits): Custom bit-width generic copy.
    • fx.rocdl.BufferCopy128b(), fx.rocdl.BufferCopy64b(), fx.rocdl.BufferCopy32b(): AMD buffer-descriptor copies.
    • fx.rocdl.make_tdm_atom(tensor, extents, ...): gfx1250 TDM async Global↔LDS whole-tile copy.

    Construction Workflow

    1. Create an atom: fx.make_copy_atom(atom_type, elem_type).
    2. Create a thread-value layout: fx.make_layout_tv(thr_layout, val_layout).
    3. Create a tiled copy: fx.make_tiled_copy(copy_atom, layout_tv, tile_mn).
    4. Match to MMA (Optional): Use fx.make_tiled_copy_A/B/C(copy_atom, tiled_mma) to align copies with Matrix Multiply-Accumulate partitioning.
    # Create copy atom (copy_op_type, elem_type)
copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), fx.Float32)

# Build thread-value layout from thread and value layouts
tiler_mn, layout_tv = fx.make_layout_tv(thr_layout, val_layout)

# Make tiled copy from copy atom + layout + tile
tiled_copy = fx.make_tiled_copy(copy_atom, layout_tv, tile_mn)

# Make tiled copy matched to a TiledMma's A/B/C partitioning
tiled_copy_a = fx.make_tiled_copy_A(copy_atom, tiled_mma)
tiled_copy_b = fx.make_tiled_copy_B(copy_atom, tiled_mma)
tiled_copy_c = fx.make_tiled_copy_C(copy_atom, tiled_mma)

  12. Create Nested/Hierarchical Layouts

    main

    FlyDSL supports nested shapes to represent multi-level tiling hierarchies (e.g., Block $\rightarrow$ Warp $\rightarrow$ Thread $\rightarrow$ Instruction). This is commonly used in GEMM kernels.

    Use fx.make_shape to define a nested shape where subsequent modes can be tuples.

    Example: A shape with 9 elements in the first mode and a $(4, 8)$ structure in the second mode (totaling 32 elements).

    # Nested shape: 9 elements in first mode, (4, 8) = 32 elements in second
shape = fx.make_shape(9, (4, 8))