ROCm / ROCm/FlyDSL

[Proposals Welcome] Grow `flydsl.extension` into a library of common GPU building blocks

Open
#1,016 0 comments 0 reactions 1 assignee View on GitHub

@sjfeng1999 is already working on this.

Since Aug 17, 2026.

enhancement FlyDSL/Customization
Dominant language
Python
Stars
282
Forks
120
Avg merge
1d 21h
Merged PRs (30d)
67

Description

Motivation

FlyDSL gives kernel authors two things today: a layout algebra, and atoms that move or multiply tiles. Both sit at the instruction/data-movement level. At the other end, kernels/ holds complete, workload-specific kernels.

The layer in between is missing -- the reusable algorithmic building blocks that every GPU stack eventually grows, and that CUB, rocPRIM, cuRAND and NVSHMEM occupy in the CUDA ecosystem. A block-wide reduction, a prefix scan, a counter-based RNG, an fp8 pack/unpack with per-block scales: none of these are specific to any one kernel, but today each kernel author has to rebuild them from the primitives.

That gap is costly in a way that is easy to underestimate:

  • Correctness risk concentrates in the hardest details. These blocks live or die on wave size, the right cross-lane instruction for the target, and cross-wave barrier placement -- exactly the things that differ across gfx942 / gfx950 / gfx11* / gfx120* / gfx1250 and are easiest to get subtly wrong.
  • They cannot be tested on their own. Written inside a kernel, a reduction is only ever exercised through an end-to-end kernel test, so a bug in it surfaces as "the kernel is slightly off" rather than as a failing unit test.
  • They cannot be shared. Not between kernels in this repo.

The long-term plan is to close that gap: port the building blocks that are common across the whole GPU ecosystem into the DSL itself, as first-class, independently tested, reusable libraries -- so that reaching for a block-wide scan is as ordinary as reaching for fx.copy(...).

The ground is already laid

This is not a request to build new infrastructure. #907 put the mechanism in place:

  • python/flydsl/extension/ -- home for libraries built on the expr primitives.
  • python/flydsl/extension/_dispatch.py -- Dispatcher gives every library a target-neutral implementation plus optional per-target overrides ({"rocm": "rocdl"}), resolved lazily at compile time. A library can be portable by default and still drop to target-specific instructions where it matters.
  • python/flydsl/expr/__init__.py:27 -- _LIBRARY_MODULES lazily exposes each library as fx.<name>, so adding one costs nothing at import time.
  • extension/random is the one library on it so far: a target-neutral RNG with a ROCDL override, plus tests under tests/extension/random/.

What is missing is scale and a public on-ramp: one library, no stated scope, and nowhere for anyone outside the core team to say "I need X".

Candidate areas

A starting point for discussion, explicitly not a roadmap or a priority order. Upstream analogues in parentheses:

Cooperative / collective -- warp and block level reduce, prefix scan, sort/merge, striped <-> blocked exchange, ballot/vote helpers
(CUB, rocPRIM)

Communication -- symmetric-heap addressing and one-sided put/get building blocks, so multi-GPU kernels stop being one bespoke all-reduce (NVSHMEM, ROCSHMEM, IRIS)

Things people actually miss but that aren't on this list are the most valuable replies to this issue.

Proposing a library -- open to everyone

We want proposals from users, not just from the core team. If you build kernels on top of FlyDSL and keep re-deriving the same helper, that helper is a candidate -- please open an issue for it. You do not need a working implementation or a design document. A useful proposal is roughly:

  1. What the primitive is, and the upstream library it mirrors (CUB / rocPRIM / cuRAND / NVSHMEM / ...), if any.
  2. Where you hit it -- the kernel or workload that made you write it by hand. One concrete site beats a general argument.
  3. A sketch of the call site you'd want -- a few lines of the API as you'd like to use it.
  4. Anything you already know about target divergence (wave64 vs wave32, DPP vs permlane, gfx1250 specifics), so we know whether it lands as target-neutral code or needs a Dispatcher override.

Design review, layout-algebra integration, and the target-specific paths are what maintainers can help with -- surfacing the need is the part only you can do. Prototypes as PRs are welcome too; extension/random is the shape to follow, with tests under tests/extension/<library>/ and a runnable example.

Non-goals

Fused kernels and workload-specific epilogues stay in kernels/. A library earns a place in extension/ when it is reusable across unrelated kernels, expressible on top of the expr primitives, and has semantics precise enough to test against a reference -- not merely because two call sites happen to share it today.

Comments, additions to the candidate list, and outright disagreement about the scope are all welcome on this issue.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.