[Feat] Add an experimental cuda nvvm backend

.gitignore

@@ -7,9 +7,6 @@ cmake-build-*/
 # GPU core dumps
 gpucore.*
 
-# IR dumps
-my_ir_dumps*/
-
 # git 
 *.log
 *.diff
@@ -66,8 +63,6 @@ docs/_build/
 python/flydsl/_mlir
 
 # Benchmark/accuracy CSVs emitted by tests/kernels harnesses
-fmha_perf_*.csv
-run_pa_decode_ps_test.*.csv
 .humanize/
 
 # rocprofv3 raw counter/trace output

CLAUDE.md

@@ -58,7 +58,8 @@ FlyDSL/
 │   └── mlir_flydsl/               # MLIR Python binding package source
 ├── include/flydsl/                # C++ TableGen headers for Fly / FlyROCDL dialects and passes
 ├── lib/                           # C++ dialect implementation, conversions, runtime wrappers, Python bindings
-│   └── Dialect/FlyROCDL/{CDNA3,CDNA4,GFX11,GFX120X,GFX1250}/  # Per-subtarget atom lowering: MmaAtom (MFMA on CDNA3/4, WMMA on GFX11/120X/1250) + CopyAtom (Buffer/LDS, CDNA3/4 only; TDM on GFX1250)
+│   ├── Dialect/FlyROCDL/{CDNA3,CDNA4,...}/  # Per-subtarget atom lowering: MmaAtom (MFMA on CDNA3/4, WMMA on GFX11/120X/1250) + CopyAtom (Buffer/LDS, CDNA3/4 only; TDM on GFX1250)
+│   └── Dialect/FlyNVVM/{SM80, ...}/      # NVIDIA atom lowering: mma.sync.aligned, cp.async, ldmatrix (nvvm backend only)
 ├── tools/                         # fly-opt
 ├── kernels/                       # Production kernels, importable as kernels.*
 ├── tests/
@@ -67,7 +68,9 @@ FlyDSL/
 │   ├── system/                    # Cross-cutting compile/system tests
 │   ├── mlir/                      # FileCheck tests driven by scripts/run_tests.sh
 │   └── python/examples/           # AOT compile/cache pytest tests (aot_example.py)
-├── examples/                      # 01-vectorAdd, 02-tiledCopy, 03-tiledMma, 04-preshuffle_gemm
+├── examples/                      # Target-neutral, run on every backend
+│   ├── rocm/                      # AMD ROCm only
+│   └── cuda/                      # NVIDIA CUDA only
 ├── scripts/                       # build, test, benchmark, wheel, debug helper scripts
 ├── docs/                          # Sphinx documentation source
 ├── thirdparty/                    # Vendored dlpack and tvm-ffi
@@ -95,9 +98,13 @@ Public docs are deployed from `.github/workflows/docs.yml` to
 
 ```bash
 bash scripts/build_llvm.sh -j64       # Build LLVM/MLIR once
-bash scripts/build.sh -j64            # Build FlyDSL C++ + Python bindings
+bash scripts/build.sh -j64            # Build FlyDSL C++ + Python bindings (rocdl backend)
 pip install -e .                      # Editable Python install
 
+# Backend selection (CMake cache var FLYDSL_BACKENDS; default "rocdl").
+# One backend per build for now; a combined "rocdl;nvvm" build is not supported yet.
+FLYDSL_BACKENDS="nvvm" bash scripts/build.sh -j64   # NVIDIA instead of AMD (needs a CUDA toolkit)
+
 # If not relying on editable install paths:
 export PYTHONPATH="${PWD}/build-fly/python_packages:${PWD}:${PYTHONPATH}"
 export LD_LIBRARY_PATH="${PWD}/build-fly/python_packages/flydsl/_mlir/_mlir_libs:${LD_LIBRARY_PATH}"
@@ -138,7 +145,7 @@ Use names from `python/flydsl/utils/env.py`; do not introduce alternate spelling
 
 | Purpose | Variable |
 |---|---|
-| Compile backend | `FLYDSL_COMPILE_BACKEND` (default `rocm`) |
+| Compile backend | `FLYDSL_COMPILE_BACKEND` (default `rocm`; `cuda` selects the NVVM backend) |
 | Override compile arch | `ARCH` |
 | Compile without execution | `COMPILE_ONLY` |
 | JIT cache directory | `FLYDSL_RUNTIME_CACHE_DIR` |
@@ -165,6 +172,8 @@ helper code that is not part of the traced closure.
 
 ## GPU Architecture Support
 
+AMD (`FLYDSL_COMPILE_BACKEND=rocm`, the default):
+
 | Arch | Chips | Wave size | MMA path | Notes |
 |---|---|---|---|---|
 | `gfx942` | MI300X / MI308X | 64 | MFMA | CDNA3 baseline; preshuffle GEMM, PA decode, CDNA BufferCopy |
@@ -183,6 +192,15 @@ RDNA and is wave32-true only for `gfx10*`/`gfx11*`/`gfx120*` prefixes; it does
 `tests/kernels/test_rdna_gemm.py` shows the gfx11* (v16 ABI) vs gfx120* (v8 ABI)
 kernel-selection pattern.
 
+NVIDIA (`FLYDSL_COMPILE_BACKEND=cuda`, requires a `FLYDSL_BACKENDS=nvvm` build):
+
+| Arch | Warp size | MMA path | Notes |
+|---|---|---|---|
+| `sm_80`+ | 32 | `mma.sync.aligned` | SM80 m16n8k16 f16->f32 MMA, SM80 `cp.async`, SM75+ `ldmatrix`. Arch string comes from `get_cuda_arch()`. |
+
+Target-specific NVIDIA atoms live in `python/flydsl/expr/nvvm/` (reached as
+`fx.nvvm`).
+
 ## Kernel Entry Points
 
 This is routing guidance, not a complete kernel inventory. Search the current `kernels/` tree before edits; keep user-facing catalogs in `docs/prebuilt_kernels_guide.md`.

README.md

@@ -49,10 +49,13 @@ FlyDSL/
 │   │   └── 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
+│   ├── 01-vectorAdd.py        # Vector addition
+│   ├── 02-gather_scatter.py   # Row gather/scatter
+│   ├── rocm/                  # AMD ROCm examples
+│   │   ├── 01-tiledCopy.py    # Tiled copy with partitioned tensors
+│   │   ├── 02-tiledMma.py     # Tiled MMA (GEMM) with MFMA atoms
+│   │   └── 03-preshuffle_gemm.py  # Preshuffle GEMM end-to-end example
+│   └── cuda/                  # NVIDIA CUDA examples
 ├── kernels/                   # Production GPU kernels (importable as `kernels.*`)
 ├── tests/                     # All tests (kernels/, mlir/, unit/)
 ├── CMakeLists.txt             # top-level CMake
@@ -358,7 +361,7 @@ 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`).
+See `examples/rocm/` for AMD examples including tiled copy (`02-tiledCopy.py`), tiled MMA (`03-tiledMma.py`), and preshuffle GEMM (`04-preshuffle_gemm.py`), and `examples/cuda/` for the NVIDIA NVVM examples.
 
 ## ✅ Testing Status
 

cmake/FlyDSLBackends.cmake

@@ -12,7 +12,7 @@
 
 set(FLYDSL_BACKENDS "rocdl"
     CACHE STRING "Enabled FlyDSL backend stacks (semicolon-separated)")
-set_property(CACHE FLYDSL_BACKENDS PROPERTY STRINGS rocdl)
+set_property(CACHE FLYDSL_BACKENDS PROPERTY STRINGS rocdl nvvm)
 
 # ---- Validate ----
 list(LENGTH FLYDSL_BACKENDS _n_backends)
@@ -23,7 +23,7 @@ if(_n_backends GREATER 5)
   message(FATAL_ERROR "FLYDSL_FOR_EACH_BACKEND supports at most 5 backends.")
 endif()
 
-set(_FLYDSL_BACKENDS_ALLOWED rocdl)
+set(_FLYDSL_BACKENDS_ALLOWED rocdl nvvm)
 foreach(_b ${FLYDSL_BACKENDS})
   if(NOT _b IN_LIST _FLYDSL_BACKENDS_ALLOWED)
     message(FATAL_ERROR

cmake/backends/nvvm.cmake

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+#
+# NVVM backend descriptor.
+# Self-registers into global properties consumed by downstream CMakeLists.txt.
+#
+# Stage one ships FlyNVVM SM80 atom types, FlyToNVVM conversion, Python
+# bindings, and CUDA runtime support. The Python-side properties below keep the
+# generated dialect bindings and stubs in sync with enabled backends.
+
+# TableGen / header subdirectories under include/flydsl/
+set_property(GLOBAL APPEND PROPERTY FLYDSL_BACKEND_INCLUDE_DIALECT_SUBDIRS "FlyNVVM")
+set_property(GLOBAL APPEND PROPERTY FLYDSL_BACKEND_INCLUDE_CONVERSION_SUBDIRS "FlyToNVVM")
+
+# C++ library subdirectories under lib/
+set_property(GLOBAL APPEND PROPERTY FLYDSL_BACKEND_LIB_DIALECT_SUBDIRS "FlyNVVM")
+set_property(GLOBAL APPEND PROPERTY FLYDSL_BACKEND_LIB_CONVERSION_SUBDIRS "FlyToNVVM")
+
+# CAPI wrapper subdirectory under lib/CAPI/Dialect/
+set_property(GLOBAL APPEND PROPERTY FLYDSL_BACKEND_CAPI_SUBDIRS "FlyNVVM")
+
+# CAPI link targets for _mlirRegisterEverything (EMBED_CAPI_LINK_LIBS)
+set_property(GLOBAL APPEND PROPERTY FLYDSL_BACKEND_EMBED_CAPI_LIBS "MLIRCPIFlyNVVM")
+
+# Link targets for fly-opt
+set_property(GLOBAL APPEND PROPERTY FLYDSL_BACKEND_FLYOPT_LINK_LIBS "MLIRCPIFlyNVVM")
+
+# Upstream MLIR dialect sources needed by this backend's Python bindings
+set_property(GLOBAL APPEND PROPERTY FLYDSL_BACKEND_UPSTREAM_DIALECT_SOURCES
+  "MLIRPythonSources.Dialects.nvvm")
+
+# Stubgen modules for this backend
+set_property(GLOBAL APPEND PROPERTY FLYDSL_BACKEND_STUBGEN_MODULES
+  "flydsl._mlir._mlir_libs._mlirDialectsFlyNVVM")
+
+# Convenience boolean for Python CMakeLists gating of NVVM-specific bindings.
+set(FLYDSL_HAS_NVVM ON)

examples/cuda/01-MmaSync.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+#
+# Run:
+#   FLYDSL_COMPILE_BACKEND=cuda FLYDSL_RUNTIME_KIND=cuda \
+#     python3 examples/cuda/01-MmaSync.py
+
+import torch
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+
+# One mma.sync.aligned instruction tile: M=16, N=8, K=16.
+INST_M = 16
+INST_N = 8
+INST_K = 16
+
+
+@flyc.kernel
+def gemm_kernel(
+    A: fx.Tensor,  # (M, K) row-major
+    B: fx.Tensor,  # (N, K) row-major  (so C = A @ B^T)
+    C: fx.Tensor,  # (M, N) row-major
+):
+    tid = fx.thread_idx.x
+    bid = fx.block_idx.x
+
+    bA = fx.zipped_divide(A, (INST_M, INST_K))
+    bB = fx.zipped_divide(B, (INST_N, INST_K))
+    bC = fx.zipped_divide(C, (INST_M, INST_N))
+
+    bA = fx.slice(bA, (None, bid))
+    bB = fx.slice(bB, (None, bid))
+    bC = fx.slice(bC, (None, bid))
+
+    mma_atom = fx.make_mma_atom(fx.nvvm.MmaSync(16, 8, 16, fx.Float16))
+    tiled_mma = fx.make_tiled_mma(mma_atom, fx.make_layout((1, 1, 1), (0, 0, 0)))
+    thr_mma = tiled_mma.thr_slice(tid)
+
+    copy_atom_f16 = fx.make_copy_atom(fx.UniversalCopy16b(), fx.Float16)
+    copy_atom_f32 = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Float32)
+    tiled_copy_A = fx.make_tiled_copy_A(copy_atom_f16, tiled_mma)
+    tiled_copy_B = fx.make_tiled_copy_B(copy_atom_f16, tiled_mma)
+    tiled_copy_C = fx.make_tiled_copy_C(copy_atom_f32, tiled_mma)
+
+    thr_copy_A = tiled_copy_A.get_slice(tid)
+    thr_copy_B = tiled_copy_B.get_slice(tid)
+    thr_copy_C = tiled_copy_C.get_slice(tid)
+
+    copy_src_A = thr_copy_A.partition_S(bA)
+    copy_src_B = thr_copy_B.partition_S(bB)
+    copy_dst_C = thr_copy_C.partition_S(bC)
+
+    frag_A = thr_mma.make_fragment_A(bA)
+    frag_B = thr_mma.make_fragment_B(bB)
+    frag_C = thr_mma.make_fragment_C(bC)
+
+    copy_frag_A = thr_copy_A.retile(frag_A)
+    copy_frag_B = thr_copy_B.retile(frag_B)
+    copy_frag_C = thr_copy_C.retile(frag_C)
+
+    fx.copy(copy_atom_f16, copy_src_A, copy_frag_A, pred=None)
+    fx.copy(copy_atom_f16, copy_src_B, copy_frag_B, pred=None)
+
+    frag_C.fill(0)
+    fx.gemm(mma_atom, frag_C, frag_A, frag_B, frag_C)
+
+    fx.copy(copy_atom_f32, copy_frag_C, copy_dst_C, pred=None)
+
+
+@flyc.jit
+def nvvm_gemm(
+    A: fx.Tensor,
+    B: fx.Tensor,
+    C: fx.Tensor,
+    stream: fx.Stream = fx.Stream(None),
+):
+    gemm_kernel(A, B, C).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream)
+
+
+M, N, K = INST_M, INST_N, INST_K
+A = torch.randn(M, K, dtype=torch.float16).cuda()
+B = torch.randn(N, K, dtype=torch.float16).cuda()
+C = torch.zeros(M, N, dtype=torch.float32).cuda()
+
+nvvm_gemm(A, B, C, stream=torch.cuda.Stream())
+torch.cuda.synchronize()
+
+expected = A.float() @ B.float().T
+is_correct = torch.allclose(C, expected, atol=1e-2, rtol=1e-2)
+print("Result correct:", is_correct)

examples/notebooks/README.md

@@ -21,7 +21,7 @@ last.
 
 The whole API these notebooks cover, in one place — enough to write a kernel without
 reading the source. The MMA atoms (`make_mma_atom`, `make_tiled_mma`, `gemm`) are the
-one piece left for later; `examples/03-tiledMma.py` is the worked reference.
+one piece left for later.
 
 ```python
 # Kernel + launch (00)

include/flydsl-c/FlyNVVMDialect.h

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 FlyDSL Project Contributors
+
+#ifndef FLYDSL_C_FLYNVVMDIALECT_H
+#define FLYDSL_C_FLYNVVMDIALECT_H
+
+#include "mlir-c/IR.h"
+#include "mlir-c/Support.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(FlyNVVM, fly_nvvm);
+
+MLIR_CAPI_EXPORTED void mlirRegisterFlyToNVVMConversionPass(void);
+
+/// Backend plugin registration: insert all NVVM dialects into \p registry.
+MLIR_CAPI_EXPORTED void flydsl_register_nvvm_dialects(MlirDialectRegistry registry);
+/// Backend plugin registration: register all NVVM passes.
+MLIR_CAPI_EXPORTED void flydsl_register_nvvm_passes(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // FLYDSL_C_FLYNVVMDIALECT_H

include/flydsl/Conversion/FlyToNVVM/CMakeLists.txt

@@ -0,0 +1,6 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name FlyToNVVM)
+mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix FlyToNVVM)
+mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix FlyToNVVM)
+
+add_mlir_generic_tablegen_target(FlyToNVVMPassIncGen)

include/flydsl/Conversion/FlyToNVVM/FlyToNVVM.h

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 FlyDSL Project Contributors
+
+#ifndef CONVERSION_FLYTONVVM_FLYTONVVM_H
+#define CONVERSION_FLYTONVVM_FLYTONVVM_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+#define GEN_PASS_DECL_FLYTONVVMCONVERSIONPASS
+#include "flydsl/Conversion/FlyToNVVM/Passes.h.inc"
+} // namespace mlir
+
+#endif // CONVERSION_FLYTONVVM_FLYTONVVM_H

include/flydsl/Conversion/FlyToNVVM/Passes.td

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 FlyDSL Project Contributors
+
+include "mlir/Pass/PassBase.td"
+
+def FlyToNVVMConversionPass : Pass<"convert-fly-to-nvvm"> {
+  let summary = "Lower Fly to MLIR upstream and nvvm dialects ";
+  let dependentDialects = [
+    "arith::ArithDialect",
+    "scf::SCFDialect",
+    "vector::VectorDialect",
+    "LLVM::LLVMDialect",
+    "NVVM::NVVMDialect"
+  ];
+}

include/flydsl/Dialect/FlyNVVM/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(IR)

include/flydsl/Dialect/FlyNVVM/IR/Atom.td

@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 FlyDSL Project Contributors
+
+#ifndef FLYNVVM_ATOM
+#define FLYNVVM_ATOM
+
+include "flydsl/Dialect/FlyNVVM/IR/Dialect.td"
+include "flydsl/Dialect/FlyNVVM/IR/MmaAtom.td"
+include "flydsl/Dialect/FlyNVVM/IR/CopyAtom.td"
+
+#endif // FLYNVVM_ATOM

include/flydsl/Dialect/FlyNVVM/IR/CMakeLists.txt

@@ -0,0 +1,10 @@
+set(LLVM_TARGET_DEFINITIONS Dialect.td)
+
+mlir_tablegen(Dialect.h.inc -gen-dialect-decls)
+mlir_tablegen(Dialect.cpp.inc -gen-dialect-defs)
+
+set(LLVM_TARGET_DEFINITIONS Atom.td)
+mlir_tablegen(Atom.h.inc -gen-typedef-decls -typedefs-dialect=fly_nvvm)
+mlir_tablegen(Atom.cpp.inc -gen-typedef-defs -typedefs-dialect=fly_nvvm)
+
+add_public_tablegen_target(MLIRFlyNVVMIncGen)

include/flydsl/Dialect/FlyNVVM/IR/CopyAtom.td

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 FlyDSL Project Contributors
+
+#ifndef FLYNVVM_COPYATOM
+#define FLYNVVM_COPYATOM
+
+include "flydsl/Dialect/FlyNVVM/IR/Dialect.td"
+
+//===----------------------------------------------------------------------===//
+// CopyOp SM75 — PTX Warp-level Matrix Load Instruction: ldmatrix
+//   ldmatrix.sync.aligned.m8n8.x{1,2,4}[.trans].shared.b16
+//===----------------------------------------------------------------------===//
+
+def FlyNVVM_CopyOpSM75_LdMatrix : FlyNVVM_CopyOp<"CopyOpSM75_LdMatrix", "sm75.ldmatrix", []> {
+  let parameters = (ins "int32_t":$num, "bool":$trans);
+  let assemblyFormat = "`<` `num` `=` $num `,` `trans` `=` $trans `>`";
+  let genVerifyDecl = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// CopyOp SM80 — PTX Data Movement and Conversion Instruction: cp.async
+//   cp.async.{ca,cg}.shared.global (global -> shared, asynchronous)
+//===----------------------------------------------------------------------===//
+
+def FlyNVVM_CopyOpSM80_CpAsync : FlyNVVM_CopyOp<"CopyOpSM80_CpAsync", "sm80.cp.async", []> {
+  let parameters = (ins "int32_t":$bitSize);
+  let assemblyFormat = "`<` $bitSize `>`"; // TODO: cache modifiers
+  let genVerifyDecl = 1;
+}
+
+
+#endif // FLYNVVM_COPYATOM