From e6f051089d1b8da51b4869c7c472a51659df7956 Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Thu, 6 Aug 2026 18:53:44 +0800 Subject: [PATCH 1/2] [Feat] Add an experimental cuda nvvm backend --- .gitignore | 5 - CLAUDE.md | 26 +- README.md | 13 +- cmake/FlyDSLBackends.cmake | 4 +- cmake/backends/nvvm.cmake | 37 + ...gather_scatter.py => 02-gather_scatter.py} | 0 examples/cuda/01-MmaSync.py | 91 ++ examples/notebooks/README.md | 2 +- .../01-BufferCopy.py} | 0 examples/{03-tiledMma.py => rocm/02-MFMA.py} | 0 .../03-preshuffle_gemm.py} | 0 include/flydsl-c/FlyNVVMDialect.h | 27 + .../Conversion/FlyToNVVM/CMakeLists.txt | 6 + .../flydsl/Conversion/FlyToNVVM/FlyToNVVM.h | 14 + include/flydsl/Conversion/FlyToNVVM/Passes.td | 15 + include/flydsl/Dialect/FlyNVVM/CMakeLists.txt | 1 + include/flydsl/Dialect/FlyNVVM/IR/Atom.td | 11 + .../flydsl/Dialect/FlyNVVM/IR/CMakeLists.txt | 10 + include/flydsl/Dialect/FlyNVVM/IR/CopyAtom.td | 32 + include/flydsl/Dialect/FlyNVVM/IR/Dialect.h | 26 + include/flydsl/Dialect/FlyNVVM/IR/Dialect.td | 40 + include/flydsl/Dialect/FlyNVVM/IR/MmaAtom.td | 33 + lib/Bindings/Python/FlyNVVMExtension.cpp | 87 ++ lib/CAPI/Dialect/FlyNVVM/CMakeLists.txt | 10 + lib/CAPI/Dialect/FlyNVVM/FlyNVVMDialect.cpp | 26 + lib/Conversion/FlyToNVVM/CMakeLists.txt | 21 + lib/Conversion/FlyToNVVM/FlyToNVVM.cpp | 851 ++++++++++++++++++ lib/Dialect/FlyNVVM/CMakeLists.txt | 8 + lib/Dialect/FlyNVVM/Dialect.cpp | 25 + lib/Dialect/FlyNVVM/SM80/CopyAtom.cpp | 276 ++++++ lib/Dialect/FlyNVVM/SM80/MmaAtom.cpp | 184 ++++ lib/Runtime/CMakeLists.txt | 3 + lib/Runtime/CUDA/CMakeLists.txt | 60 ++ lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp | 304 +++++++ python/flydsl/compiler/backends/cuda.py | 192 ++++ python/flydsl/compiler/jit_function.py | 20 +- python/flydsl/expr/__init__.py | 1 + python/flydsl/expr/nvvm/__init__.py | 27 + python/flydsl/expr/nvvm/universal.py | 43 + .../flydsl/runtime/device_runtime/__init__.py | 4 + python/flydsl/runtime/device_runtime/cuda.py | 136 +++ python/mlir_flydsl/CMakeLists.txt | 40 + python/mlir_flydsl/dialects/FlyNVVM.td | 9 + python/mlir_flydsl/dialects/fly_nvvm.py | 9 + scripts/build.sh | 7 + scripts/run_tests.sh | 70 +- .../Conversion/fly-to-nvvm/copy_atom.mlir | 23 + .../mlir/Conversion/fly-to-nvvm/mma_atom.mlir | 28 + tests/unit/test_backend_cmake_defaults.py | 12 +- tests/unit/test_cuda_backend.py | 101 +++ tests/unit/test_gfx1250_atoms.py | 5 +- tests/unit/test_tdm_mcast_add_gfx1250.py | 29 +- tests/unit/test_universal_atomic.py | 34 +- 53 files changed, 2968 insertions(+), 70 deletions(-) create mode 100644 cmake/backends/nvvm.cmake rename examples/{05-gather_scatter.py => 02-gather_scatter.py} (100%) create mode 100644 examples/cuda/01-MmaSync.py rename examples/{02-tiledCopy.py => rocm/01-BufferCopy.py} (100%) rename examples/{03-tiledMma.py => rocm/02-MFMA.py} (100%) rename examples/{04-preshuffle_gemm.py => rocm/03-preshuffle_gemm.py} (100%) create mode 100644 include/flydsl-c/FlyNVVMDialect.h create mode 100644 include/flydsl/Conversion/FlyToNVVM/CMakeLists.txt create mode 100644 include/flydsl/Conversion/FlyToNVVM/FlyToNVVM.h create mode 100644 include/flydsl/Conversion/FlyToNVVM/Passes.td create mode 100644 include/flydsl/Dialect/FlyNVVM/CMakeLists.txt create mode 100644 include/flydsl/Dialect/FlyNVVM/IR/Atom.td create mode 100644 include/flydsl/Dialect/FlyNVVM/IR/CMakeLists.txt create mode 100644 include/flydsl/Dialect/FlyNVVM/IR/CopyAtom.td create mode 100644 include/flydsl/Dialect/FlyNVVM/IR/Dialect.h create mode 100644 include/flydsl/Dialect/FlyNVVM/IR/Dialect.td create mode 100644 include/flydsl/Dialect/FlyNVVM/IR/MmaAtom.td create mode 100644 lib/Bindings/Python/FlyNVVMExtension.cpp create mode 100644 lib/CAPI/Dialect/FlyNVVM/CMakeLists.txt create mode 100644 lib/CAPI/Dialect/FlyNVVM/FlyNVVMDialect.cpp create mode 100644 lib/Conversion/FlyToNVVM/CMakeLists.txt create mode 100644 lib/Conversion/FlyToNVVM/FlyToNVVM.cpp create mode 100644 lib/Dialect/FlyNVVM/CMakeLists.txt create mode 100644 lib/Dialect/FlyNVVM/Dialect.cpp create mode 100644 lib/Dialect/FlyNVVM/SM80/CopyAtom.cpp create mode 100644 lib/Dialect/FlyNVVM/SM80/MmaAtom.cpp create mode 100644 lib/Runtime/CUDA/CMakeLists.txt create mode 100644 lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp create mode 100644 python/flydsl/compiler/backends/cuda.py create mode 100644 python/flydsl/expr/nvvm/__init__.py create mode 100644 python/flydsl/expr/nvvm/universal.py create mode 100644 python/flydsl/runtime/device_runtime/cuda.py create mode 100644 python/mlir_flydsl/dialects/FlyNVVM.td create mode 100644 python/mlir_flydsl/dialects/fly_nvvm.py create mode 100644 tests/mlir/Conversion/fly-to-nvvm/copy_atom.mlir create mode 100644 tests/mlir/Conversion/fly-to-nvvm/mma_atom.mlir create mode 100644 tests/unit/test_cuda_backend.py diff --git a/.gitignore b/.gitignore index 0fb95bbef..199e1d3b7 100644 --- a/.gitignore +++ b/.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 diff --git a/CLAUDE.md b/CLAUDE.md index 1876f19bb..85caf61ed 100644 --- a/CLAUDE.md +++ b/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`. diff --git a/README.md b/README.md index 983b51bc6..761bf3e5d 100644 --- a/README.md +++ b/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 diff --git a/cmake/FlyDSLBackends.cmake b/cmake/FlyDSLBackends.cmake index 1f9b22dc8..5d05bf1cf 100644 --- a/cmake/FlyDSLBackends.cmake +++ b/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 diff --git a/cmake/backends/nvvm.cmake b/cmake/backends/nvvm.cmake new file mode 100644 index 000000000..3d00588bb --- /dev/null +++ b/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) diff --git a/examples/05-gather_scatter.py b/examples/02-gather_scatter.py similarity index 100% rename from examples/05-gather_scatter.py rename to examples/02-gather_scatter.py diff --git a/examples/cuda/01-MmaSync.py b/examples/cuda/01-MmaSync.py new file mode 100644 index 000000000..6f779de7c --- /dev/null +++ b/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) diff --git a/examples/notebooks/README.md b/examples/notebooks/README.md index b6219b8f0..b157a5832 100644 --- a/examples/notebooks/README.md +++ b/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) diff --git a/examples/02-tiledCopy.py b/examples/rocm/01-BufferCopy.py similarity index 100% rename from examples/02-tiledCopy.py rename to examples/rocm/01-BufferCopy.py diff --git a/examples/03-tiledMma.py b/examples/rocm/02-MFMA.py similarity index 100% rename from examples/03-tiledMma.py rename to examples/rocm/02-MFMA.py diff --git a/examples/04-preshuffle_gemm.py b/examples/rocm/03-preshuffle_gemm.py similarity index 100% rename from examples/04-preshuffle_gemm.py rename to examples/rocm/03-preshuffle_gemm.py diff --git a/include/flydsl-c/FlyNVVMDialect.h b/include/flydsl-c/FlyNVVMDialect.h new file mode 100644 index 000000000..cba3e001e --- /dev/null +++ b/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 diff --git a/include/flydsl/Conversion/FlyToNVVM/CMakeLists.txt b/include/flydsl/Conversion/FlyToNVVM/CMakeLists.txt new file mode 100644 index 000000000..40287e56c --- /dev/null +++ b/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) diff --git a/include/flydsl/Conversion/FlyToNVVM/FlyToNVVM.h b/include/flydsl/Conversion/FlyToNVVM/FlyToNVVM.h new file mode 100644 index 000000000..4650d25eb --- /dev/null +++ b/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 diff --git a/include/flydsl/Conversion/FlyToNVVM/Passes.td b/include/flydsl/Conversion/FlyToNVVM/Passes.td new file mode 100644 index 000000000..5597eadda --- /dev/null +++ b/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" + ]; +} diff --git a/include/flydsl/Dialect/FlyNVVM/CMakeLists.txt b/include/flydsl/Dialect/FlyNVVM/CMakeLists.txt new file mode 100644 index 000000000..f33061b2d --- /dev/null +++ b/include/flydsl/Dialect/FlyNVVM/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/include/flydsl/Dialect/FlyNVVM/IR/Atom.td b/include/flydsl/Dialect/FlyNVVM/IR/Atom.td new file mode 100644 index 000000000..896e98390 --- /dev/null +++ b/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 diff --git a/include/flydsl/Dialect/FlyNVVM/IR/CMakeLists.txt b/include/flydsl/Dialect/FlyNVVM/IR/CMakeLists.txt new file mode 100644 index 000000000..329cc4810 --- /dev/null +++ b/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) diff --git a/include/flydsl/Dialect/FlyNVVM/IR/CopyAtom.td b/include/flydsl/Dialect/FlyNVVM/IR/CopyAtom.td new file mode 100644 index 000000000..0410484e4 --- /dev/null +++ b/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 diff --git a/include/flydsl/Dialect/FlyNVVM/IR/Dialect.h b/include/flydsl/Dialect/FlyNVVM/IR/Dialect.h new file mode 100644 index 000000000..768dfb34f --- /dev/null +++ b/include/flydsl/Dialect/FlyNVVM/IR/Dialect.h @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#ifndef FLYDSL_DIALECT_FLYNVVM_IR_DIALECT_H +#define FLYDSL_DIALECT_FLYNVVM_IR_DIALECT_H + +#include "mlir/Bytecode/BytecodeOpInterface.h" +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/OpImplementation.h" +#include "mlir/IR/Types.h" + +#include "flydsl/Dialect/Fly/IR/FlyDialect.h" + +#include "flydsl/Dialect/FlyNVVM/IR/Dialect.h.inc" + +#define GET_TYPEDEF_CLASSES +#include "flydsl/Dialect/FlyNVVM/IR/Atom.h.inc" + +namespace mlir::fly_nvvm {} // namespace mlir::fly_nvvm + +#endif // FLYDSL_DIALECT_FLYNVVM_IR_DIALECT_H diff --git a/include/flydsl/Dialect/FlyNVVM/IR/Dialect.td b/include/flydsl/Dialect/FlyNVVM/IR/Dialect.td new file mode 100644 index 000000000..1d378f20e --- /dev/null +++ b/include/flydsl/Dialect/FlyNVVM/IR/Dialect.td @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#ifndef FLYNVVM_DIALECT +#define FLYNVVM_DIALECT + +include "mlir/IR/EnumAttr.td" +include "mlir/IR/AttrTypeBase.td" +include "mlir/IR/OpBase.td" + +include "flydsl/Dialect/Fly/IR/FlyInterfaces.td" + +def FlyNVVM_Dialect : Dialect { + let name = "fly_nvvm"; + let cppNamespace = "::mlir::fly_nvvm"; + + let dependentDialects = [ + "NVVM::NVVMDialect" + ]; + + let useDefaultTypePrinterParser = 1; +} + +class FlyNVVM_MmaOp traits = []> + : TypeDef, + DeclareTypeInterfaceMethods + ])> { + let mnemonic = typeMnemonic; +} + +class FlyNVVM_CopyOp traits = []> + : TypeDef, + DeclareTypeInterfaceMethods + ])> { + let mnemonic = typeMnemonic; +} + +#endif // FLYNVVM_DIALECT diff --git a/include/flydsl/Dialect/FlyNVVM/IR/MmaAtom.td b/include/flydsl/Dialect/FlyNVVM/IR/MmaAtom.td new file mode 100644 index 000000000..2c35a4f02 --- /dev/null +++ b/include/flydsl/Dialect/FlyNVVM/IR/MmaAtom.td @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#ifndef FLYNVVM_MMAATOM +#define FLYNVVM_MMAATOM + +include "flydsl/Dialect/FlyNVVM/IR/Dialect.td" + +//===----------------------------------------------------------------------===// +// MmaOp SM80 +// ptx: mma.sync.aligned +//===----------------------------------------------------------------------===// + +def FlyNVVM_MmaOpSM80_MmaSync : FlyNVVM_MmaOp<"MmaOpSM80_MmaSync", "sm80.mma.sync", []> { + let parameters = (ins + "int32_t":$m, + "int32_t":$n, + "int32_t":$k, + "Type":$elemTyA, + "Type":$elemTyB, + "Type":$elemTyAcc + ); + let assemblyFormat = "`<` custom($m, $n, $k) `,` `(` $elemTyA `,` $elemTyB `)` `->` $elemTyAcc `>`"; + + let builders = [ + TypeBuilderWithInferredContext<(ins "int32_t":$m, "int32_t":$n, "int32_t":$k, "Type":$elemTyA, "Type":$elemTyB, "Type":$elemTyAcc), [{ + return $_get(elemTyA.getContext(), m, n, k, elemTyA, elemTyB, elemTyAcc); + }]> + ]; + let genVerifyDecl = 1; +} + +#endif // FLYNVVM_MMAATOM diff --git a/lib/Bindings/Python/FlyNVVMExtension.cpp b/lib/Bindings/Python/FlyNVVMExtension.cpp new file mode 100644 index 000000000..0a0b8a033 --- /dev/null +++ b/lib/Bindings/Python/FlyNVVMExtension.cpp @@ -0,0 +1,87 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Value.h" + +#include "flydsl/Dialect/Fly/IR/FlyDialect.h" +#include "flydsl/Dialect/FlyNVVM/IR/Dialect.h" + +#include "BindingUtils.h" + +namespace nb = nanobind; +using namespace nb::literals; +using namespace ::mlir::fly; +using namespace ::mlir::fly_nvvm; + +namespace mlir { +namespace python { +namespace MLIR_BINDINGS_PYTHON_DOMAIN { +namespace fly_nvvm { + +struct PyCopyOpSM75_LdMatrixType : PyConcreteType { + FLYDSL_REGISTER_TYPE_BINDING(CopyOpSM75_LdMatrixType, "CopyOpSM75_LdMatrixType"); + + static void bindDerived(ClassTy &c) { + c.def_static( + "get", + [](int32_t num, bool trans, DefaultingPyMlirContext context) { + MLIRContext *ctx = unwrap(context.get()->get()); + return PyCopyOpSM75_LdMatrixType(context->getRef(), + wrap(CopyOpSM75_LdMatrixType::get(ctx, num, trans))); + }, + "num"_a, "trans"_a = false, nb::kw_only(), "context"_a = nb::none(), + "Create a CopyOpSM75_LdMatrixType (ldmatrix) with num tiles and transpose flag"); + } +}; + +struct PyMmaOpSM80_MmaSyncType : PyConcreteType { + FLYDSL_REGISTER_TYPE_BINDING(MmaOpSM80_MmaSyncType, "MmaOpSM80_MmaSyncType"); + + static void bindDerived(ClassTy &c) { + c.def_static( + "get", + [](int32_t m, int32_t n, int32_t k, PyType &elemTyA, PyType &elemTyB, PyType &elemTyAcc, + DefaultingPyMlirContext context) { + return PyMmaOpSM80_MmaSyncType( + context->getRef(), + wrap(MmaOpSM80_MmaSyncType::get(m, n, k, unwrap(elemTyA), unwrap(elemTyB), + unwrap(elemTyAcc)))); + }, + "m"_a, "n"_a, "k"_a, "elem_ty_a"_a, "elem_ty_b"_a, "elem_ty_acc"_a, nb::kw_only(), + "context"_a = nb::none(), + "Create a MmaOpSM80_MmaSyncType with m, n, k dimensions and element types"); + } +}; + +struct PyCopyOpSM80_CpAsyncType : PyConcreteType { + FLYDSL_REGISTER_TYPE_BINDING(CopyOpSM80_CpAsyncType, "CopyOpSM80_CpAsyncType"); + + static void bindDerived(ClassTy &c) { + c.def_static( + "get", + [](int32_t bit_size, DefaultingPyMlirContext context) { + MLIRContext *ctx = unwrap(context.get()->get()); + return PyCopyOpSM80_CpAsyncType(context->getRef(), + wrap(CopyOpSM80_CpAsyncType::get(ctx, bit_size))); + }, + "bit_size"_a, nb::kw_only(), "context"_a = nb::none(), + "Create a CopyOpSM80_CpAsyncType (cp.async.shared.global) with the given bit size"); + } +}; + +} // namespace fly_nvvm +} // namespace MLIR_BINDINGS_PYTHON_DOMAIN +} // namespace python +} // namespace mlir + +NB_MODULE(_mlirDialectsFlyNVVM, m) { + m.doc() = "MLIR Python FlyNVVM Extension"; + + // clang-format off + ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::fly_nvvm::PyCopyOpSM75_LdMatrixType::bind(m); + ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::fly_nvvm::PyMmaOpSM80_MmaSyncType::bind(m); + ::mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::fly_nvvm::PyCopyOpSM80_CpAsyncType::bind(m); + // clang-format on +} diff --git a/lib/CAPI/Dialect/FlyNVVM/CMakeLists.txt b/lib/CAPI/Dialect/FlyNVVM/CMakeLists.txt new file mode 100644 index 000000000..2787ed2f8 --- /dev/null +++ b/lib/CAPI/Dialect/FlyNVVM/CMakeLists.txt @@ -0,0 +1,10 @@ +add_mlir_public_c_api_library(MLIRCPIFlyNVVM + FlyNVVMDialect.cpp + + DEPENDS + FlyToNVVMPassIncGen + + LINK_LIBS PUBLIC + MLIRFlyNVVMDialect + MLIRFlyToNVVM +) diff --git a/lib/CAPI/Dialect/FlyNVVM/FlyNVVMDialect.cpp b/lib/CAPI/Dialect/FlyNVVM/FlyNVVMDialect.cpp new file mode 100644 index 000000000..643c264b4 --- /dev/null +++ b/lib/CAPI/Dialect/FlyNVVM/FlyNVVMDialect.cpp @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#include "flydsl-c/FlyNVVMDialect.h" + +#include "flydsl/Conversion/FlyToNVVM/FlyToNVVM.h" +#include "flydsl/Dialect/FlyNVVM/IR/Dialect.h" +#include "mlir/CAPI/IR.h" +#include "mlir/CAPI/Registration.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassRegistry.h" + +namespace mlir { +#define GEN_PASS_REGISTRATION +#include "flydsl/Conversion/FlyToNVVM/Passes.h.inc" +} // namespace mlir + +MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(FlyNVVM, fly_nvvm, mlir::fly_nvvm::FlyNVVMDialect) + +void mlirRegisterFlyToNVVMConversionPass(void) { mlir::registerFlyToNVVMConversionPass(); } + +void flydsl_register_nvvm_dialects(MlirDialectRegistry registry) { + unwrap(registry)->insert(); +} + +void flydsl_register_nvvm_passes(void) { mlirRegisterFlyToNVVMConversionPass(); } diff --git a/lib/Conversion/FlyToNVVM/CMakeLists.txt b/lib/Conversion/FlyToNVVM/CMakeLists.txt new file mode 100644 index 000000000..74166ddfe --- /dev/null +++ b/lib/Conversion/FlyToNVVM/CMakeLists.txt @@ -0,0 +1,21 @@ +add_mlir_conversion_library(MLIRFlyToNVVM + FlyToNVVM.cpp + + DEPENDS + MLIRFlyIncGen + MLIRFlyNVVMIncGen + FlyToNVVMPassIncGen + + LINK_LIBS PUBLIC + MLIRFlyDialect + MLIRFlyNVVMDialect + + MLIRArithDialect + MLIRIR + MLIRLLVMDialect + MLIRNVVMDialect + MLIRPass + MLIRSCFDialect + MLIRTransforms + MLIRVectorDialect +) diff --git a/lib/Conversion/FlyToNVVM/FlyToNVVM.cpp b/lib/Conversion/FlyToNVVM/FlyToNVVM.cpp new file mode 100644 index 000000000..5add7acd5 --- /dev/null +++ b/lib/Conversion/FlyToNVVM/FlyToNVVM.cpp @@ -0,0 +1,851 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Vector/IR/VectorOps.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/ADT/StringSet.h" + +#include "flydsl/Conversion/FlyToNVVM/FlyToNVVM.h" +#include "flydsl/Dialect/Fly/IR/FlyDialect.h" +#include "flydsl/Dialect/Fly/Utils/IntTupleUtils.h" +#include "flydsl/Dialect/Fly/Utils/LayoutUtils.h" +#include "flydsl/Dialect/Fly/Utils/PointerUtils.h" +#include "flydsl/Dialect/FlyNVVM/IR/Dialect.h" + +namespace mlir { +#define GEN_PASS_DEF_FLYTONVVMCONVERSIONPASS +#include "flydsl/Conversion/FlyToNVVM/Passes.h.inc" +} // namespace mlir + +using namespace mlir; +using namespace mlir::fly; +using namespace mlir::fly_nvvm; + +namespace { + +unsigned mapToLLVMAddressSpace(AddressSpace addrSpace) { + switch (addrSpace) { + case AddressSpace::Generic: + return 0; + case AddressSpace::Global: + return 1; + case AddressSpace::Shared: + return 3; + case AddressSpace::Register: + // NVVM has no dedicated register address space; register-backed scratch + // is modelled with an alloca in the generic (0) address space. + return 0; + } + llvm_unreachable("unsupported address space"); +} + +unsigned mapAttrToLLVMAddressSpace(Attribute attr) { + if (auto e = dyn_cast(attr)) + return mapToLLVMAddressSpace(e.getValue()); + return 0; // default to generic address space +} + +// Create a freshly named shared-memory global of `[nbytes x i8]` in `addrSpace`, +// inserted at the start of `moduleOp`. The symbol name is `prefix` followed by a +// counter chosen to avoid collisions with existing globals in the module. Shared +// by the static-shared `make_ptr` and dynamic-shared `get_dyn_shared` lowerings. +static LLVM::GlobalOp createSharedGlobal(ConversionPatternRewriter &rewriter, + gpu::GPUModuleOp moduleOp, Location loc, StringRef prefix, + int64_t nbytes, int64_t align, unsigned addrSpace) { + llvm::StringSet<> existingNames; + for (auto globalOp : moduleOp.getBody()->getOps()) + existingNames.insert(globalOp.getSymName()); + + unsigned counter = 0; + SmallString<128> symName = SymbolTable::generateSymbolName<128>( + prefix, [&](StringRef candidate) { return existingNames.contains(candidate); }, counter); + + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(moduleOp.getBody()); + + auto arrayTy = LLVM::LLVMArrayType::get(IntegerType::get(rewriter.getContext(), 8), nbytes); + auto globalOp = LLVM::GlobalOp::create(rewriter, loc, arrayTy, + /*isConstant=*/false, LLVM::Linkage::External, symName, + /*value=*/Attribute(), + /*alignment=*/align, addrSpace); + // Shared (addrspace 3) symbols are per-block hardware and can never resolve + // across DSO. + globalOp.setDsoLocal(true); + return globalOp; +} + +class MakePtrOpLowering : public OpConversionPattern { +public: + MakePtrOpLowering(const TypeConverter &typeConverter, MLIRContext *context) + : OpConversionPattern(typeConverter, context) {} + + LogicalResult matchAndRewrite(MakePtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto flyPtrTy = dyn_cast(op.getResult().getType()); + if (!flyPtrTy) + return failure(); + + Location loc = op.getLoc(); + Attribute addrSpaceAttr = flyPtrTy.getAddressSpace(); + + if (isGenericAddressSpace(addrSpaceAttr)) { + auto dictAttrs = op.getDictAttrs(); + if (!dictAttrs) + return rewriter.notifyMatchFailure(op, "register make_ptr requires dictAttrs"); + auto allocSize = dictAttrs->getAs("allocSize"); + if (!allocSize) + return rewriter.notifyMatchFailure(op, "register make_ptr requires allocSize in ptrAttrs"); + unsigned llvmAS = mapToLLVMAddressSpace(AddressSpace::Register); + auto llvmPtrTy = LLVM::LLVMPointerType::get(rewriter.getContext(), llvmAS); + Value nElems = arith::ConstantIntOp::create(rewriter, loc, allocSize.getInt(), 64); + Type elemTy = projectToLLVMCompatibleElemTy(flyPtrTy.getElemTy()); + Value ptr = LLVM::AllocaOp::create(rewriter, loc, llvmPtrTy, elemTy, nElems, 0); + rewriter.replaceOp(op, ptr); + return success(); + } else if (isGenericAddressSpace(addrSpaceAttr)) { + // Static shared sub-allocation. Each op lowers to a freshly named + // `@__shared_alloc_` shared-memory global. + auto dictAttrs = op.getDictAttrs(); + if (!dictAttrs) + return rewriter.notifyMatchFailure( + op, "shared make_ptr requires dictAttrs={allocBytes, allocAlign}"); + auto allocBytesAttr = dictAttrs->getAs("allocBytes"); + auto allocAlignAttr = dictAttrs->getAs("allocAlign"); + if (!allocBytesAttr || !allocAlignAttr) + return rewriter.notifyMatchFailure( + op, "shared make_ptr requires allocBytes, allocAlign in dictAttrs"); + + auto moduleOp = op->getParentOfType(); + if (!moduleOp) + return op->emitError("shared make_ptr must be inside a gpu.module"); + + LLVM::GlobalOp global = + createSharedGlobal(rewriter, moduleOp, loc, "__shared_alloc", allocBytesAttr.getInt(), + allocAlignAttr.getInt(), /*addrSpace=*/3); + rewriter.replaceOpWithNewOp(op, global); + return success(); + } + + return rewriter.notifyMatchFailure(op, "unsupported make_ptr address space"); + } +}; + +class GetDynSharedOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(GetDynSharedOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + auto flyPtrTy = cast(op.getResult().getType()); + unsigned addrSpace = mapAttrToLLVMAddressSpace(flyPtrTy.getAddressSpace()); + + auto moduleOp = op->getParentOfType(); + if (!moduleOp) + return op->emitError("get_dyn_shared must be inside a gpu.module"); + + // Dynamic shared memory has a single logical region per kernel; reuse an + // existing `[0 x i8]` global if one is already present so multiple + // `get_dyn_shared` ops resolve to the same base symbol. + LLVM::GlobalOp sharedGlobal; + for (auto globalOp : moduleOp.getBody()->getOps()) { + if (auto arrayType = dyn_cast(globalOp.getType())) { + if (globalOp.getAddrSpace() == addrSpace && arrayType.getNumElements() == 0 && + globalOp.getAlignment().value_or(0) == 1024) { + sharedGlobal = globalOp; + break; + } + } + } + if (!sharedGlobal) { + sharedGlobal = createSharedGlobal(rewriter, moduleOp, loc, "__dynamic_shared", + /*nbytes=*/0, /*align=*/1024, addrSpace); + } + + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(op); + + auto basePtr = LLVM::AddressOfOp::create(rewriter, loc, sharedGlobal); + Type ptrType = basePtr->getResultTypes()[0]; + + auto i8Ty = IntegerType::get(rewriter.getContext(), 8); + Value sharedPtr = + LLVM::GEPOp::create(rewriter, loc, ptrType, i8Ty, basePtr, ArrayRef{0}); + + rewriter.replaceOp(op, sharedPtr); + return success(); + } +}; + +class IntToPtrOpLowering : public OpConversionPattern { +public: + IntToPtrOpLowering(const TypeConverter &typeConverter, MLIRContext *context) + : OpConversionPattern(typeConverter, context) {} + + LogicalResult matchAndRewrite(IntToPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto flyPtrTy = dyn_cast(op.getResult().getType()); + if (!flyPtrTy) + return failure(); + + auto resultTy = dyn_cast(getTypeConverter()->convertType(flyPtrTy)); + if (!resultTy) + return failure(); + + rewriter.replaceOpWithNewOp(op, resultTy, adaptor.getSrc()); + return success(); + } +}; + +class PtrToIntOpLowering : public OpConversionPattern { +public: + PtrToIntOpLowering(const TypeConverter &typeConverter, MLIRContext *context) + : OpConversionPattern(typeConverter, context) {} + + LogicalResult matchAndRewrite(PtrToIntOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type resultTy = getTypeConverter()->convertType(op.getResult().getType()); + if (!resultTy) + return failure(); + + rewriter.replaceOpWithNewOp(op, resultTy, adaptor.getPtr()); + return success(); + } +}; + +class ToLLVMPtrOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(ToLLVMPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto convertedPtrTy = dyn_cast(adaptor.getPtr().getType()); + if (!convertedPtrTy) + return op.emitError("pointer address space does not lower to a bare NVVM LLVM pointer: ") + << cast(op.getPtr().getType()).getAddressSpace(); + + auto requestedPtrTy = dyn_cast(op.getResult().getType()); + if (!requestedPtrTy) + return op.emitError("expected an LLVM pointer result, got ") << op.getResult().getType(); + if (requestedPtrTy.getAddressSpace() != convertedPtrTy.getAddressSpace()) + return op.emitError("requested LLVM address space ") + << requestedPtrTy.getAddressSpace() << " but NVVM conversion requires " + << convertedPtrTy.getAddressSpace(); + + rewriter.replaceOp(op, adaptor.getPtr()); + return success(); + } +}; + +class ApplySwizzleOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(ApplySwizzleOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.replaceOp(op, adaptor.getPtr()); + return success(); + } +}; + +class RecastIterOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(RecastIterOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.replaceOp(op, adaptor.getSrc()); + return success(); + } +}; + +class AddOffsetOpLowering : public OpConversionPattern { +public: + AddOffsetOpLowering(const TypeConverter &typeConverter, MLIRContext *context) + : OpConversionPattern(typeConverter, context) {} + + LogicalResult matchAndRewrite(AddOffsetOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto loc = op.getLoc(); + Value base = adaptor.getPtr(); + Value offset = adaptor.getOffset(); + + auto flyPtrTy = dyn_cast(op.getPtr().getType()); + if (!flyPtrTy) + return failure(); + + auto offsetTy = dyn_cast(offset.getType()); + IntTupleAttr offsetAttr = offsetTy.getAttr(); + if (!offsetAttr.isLeaf()) + return rewriter.notifyMatchFailure(op, "offset must be a leaf int tuple"); + + Value offsetVal; + auto offsetInt = offsetAttr.extractIntFromLeaf(); + if (offsetInt.isStatic()) { + offsetVal = arith::ConstantIntOp::create(rewriter, loc, offsetInt.getValue(), 32); + } else { + Operation *defOp = offset.getDefiningOp(); + offsetVal = defOp->getOperand(0); + } + + auto ptrTy = dyn_cast(base.getType()); + if (!ptrTy) + return failure(); + + Type elemTy = projectToLLVMCompatibleElemTy(flyPtrTy.getElemTy()); + Value gep = LLVM::GEPOp::create(rewriter, loc, ptrTy, elemTy, base, ValueRange{offsetVal}); + rewriter.replaceOp(op, gep); + return success(); + } +}; + +class MakeViewOpLowering : public OpConversionPattern { +public: + MakeViewOpLowering(const TypeConverter &typeConverter, MLIRContext *context) + : OpConversionPattern(typeConverter, context) {} + + LogicalResult matchAndRewrite(MakeViewOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + if (isa(op.getResult().getType())) { + if (!op.getResult().use_empty()) + return rewriter.notifyMatchFailure(op, "coord_tensor result should have no uses"); + rewriter.eraseOp(op); + return success(); + } else { + Value base = adaptor.getIter(); + rewriter.replaceOp(op, base); + return success(); + } + } +}; + +class PtrLoadOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(PtrLoadOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + Value ptr = adaptor.getPtr(); + + auto flyPtrTy = dyn_cast(op.getPtr().getType()); + if (!flyPtrTy) + return failure(); + + Type loadTy = op.getResult().getType(); + + if (auto vecTy = dyn_cast(loadTy)) { + auto swizzle = flyPtrTy.getSwizzle(); + if (!swizzle.isTrivialSwizzle()) { + int64_t vecBytes = + vecTy.getNumElements() * vecTy.getElementType().getIntOrFloatBitWidth() / 8; + int64_t baseBytes = int64_t{1} << swizzle.getBase(); + if (baseBytes % vecBytes != 0) + return rewriter.notifyMatchFailure( + op, "vector ptr.load byte size must divide swizzle base granularity"); + } + } + + ptr = applySwizzleOnPtr(rewriter, loc, cast>(ptr), + flyPtrTy.getSwizzle()); + Value loaded = LLVM::LoadOp::create(rewriter, loc, loadTy, ptr); + rewriter.replaceOp(op, loaded); + return success(); + } +}; + +class PtrStoreOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(PtrStoreOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + Value ptr = adaptor.getPtr(); + Value value = adaptor.getValue(); + + auto flyPtrTy = dyn_cast(op.getPtr().getType()); + if (!flyPtrTy) + return failure(); + + if (auto vecTy = dyn_cast(value.getType())) { + auto swizzle = flyPtrTy.getSwizzle(); + if (!swizzle.isTrivialSwizzle()) { + int64_t vecBytes = + vecTy.getNumElements() * vecTy.getElementType().getIntOrFloatBitWidth() / 8; + int64_t baseBytes = int64_t{1} << swizzle.getBase(); + if (baseBytes % vecBytes != 0) + return rewriter.notifyMatchFailure( + op, "vector ptr.store byte size must divide swizzle base granularity"); + } + } + + ptr = applySwizzleOnPtr(rewriter, loc, cast>(ptr), + flyPtrTy.getSwizzle()); + LLVM::StoreOp::create(rewriter, loc, value, ptr); + rewriter.eraseOp(op); + return success(); + } +}; + +class MakeCopyAtomOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(MakeCopyAtomOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto copyAtomTy = dyn_cast(op.getResult().getType()); + if (!copyAtomTy) + return rewriter.notifyMatchFailure(op, "not a CopyAtomType"); + Type convertedTy = getTypeConverter()->convertType(copyAtomTy); + + auto statefulOp = dyn_cast(copyAtomTy.getCopyOp()); + if (statefulOp) { + Value state = statefulOp.getDefaultState(rewriter, op.getLoc()); + rewriter.replaceOp(op, state); + } else { + rewriter.replaceOpWithNewOp(op, convertedTy); + } + return success(); + } +}; + +class MakeMmaAtomOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(MakeMmaAtomOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto mmaAtomTy = dyn_cast(op.getResult().getType()); + if (!mmaAtomTy) + return rewriter.notifyMatchFailure(op, "not a MmaAtomType"); + Type convertedTy = getTypeConverter()->convertType(mmaAtomTy); + auto statefulOp = dyn_cast(mmaAtomTy.getMmaOp()); + if (statefulOp) { + Value state = statefulOp.getDefaultState(rewriter, op.getLoc()); + rewriter.replaceOp(op, state); + } else { + rewriter.replaceOpWithNewOp(op, convertedTy); + } + return success(); + } +}; + +class MakeTiledCopyOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(MakeTiledCopyOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.replaceOp(op, adaptor.getCopyAtom()); + return success(); + } +}; + +class MakeTiledMmaOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(MakeTiledMmaOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.replaceOp(op, adaptor.getMmaAtom()); + return success(); + } +}; + +class AtomSetValueOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(AtomSetValueOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type origAtomTy = op.getAtom().getType(); + StringAttr fieldAttr = op.getFieldAttr(); + Location loc = op.getLoc(); + + Value structVal = adaptor.getAtom(); + Value fieldVal = adaptor.getValue(); + Value result; + + if (auto copyAtomTy = dyn_cast(origAtomTy)) { + if (!copyAtomTy.isStateful()) + return rewriter.notifyMatchFailure(op, "CopyAtom is not stateful"); + result = copyAtomTy.setAtomState(rewriter, loc, structVal, fieldAttr, fieldVal); + } else if (auto mmaAtomTy = dyn_cast(origAtomTy)) { + if (!mmaAtomTy.isStateful()) + return rewriter.notifyMatchFailure(op, "MmaAtom is not stateful"); + result = mmaAtomTy.setAtomState(rewriter, loc, structVal, fieldAttr, fieldVal); + } else { + return rewriter.notifyMatchFailure(op, "atom is not CopyAtomType or MmaAtomType"); + } + + if (!result) + return rewriter.notifyMatchFailure(op, "setAtomState failed"); + + rewriter.replaceOp(op, result); + return success(); + } +}; + +class CopyAtomCallLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(CopyAtomCall op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type copyAtomType = op.getCopyAtom().getType(); + auto copyAtom = dyn_cast(copyAtomType); + if (!copyAtom) + return rewriter.notifyMatchFailure(op, "copyAtom is not CopyAtomType"); + + Value copyAtomVal = adaptor.getCopyAtom(); + Value src = adaptor.getSrc(); + Value dst = adaptor.getDst(); + Value pred = adaptor.getPred(); + + auto srcMemTy = dyn_cast(op.getSrc().getType()); + auto dstMemTy = dyn_cast(op.getDst().getType()); + + if (!srcMemTy || !dstMemTy) + return rewriter.notifyMatchFailure(op, "expected MemRef types on original op"); + if (srcMemTy.getElemTy() != dstMemTy.getElemTy()) + return rewriter.notifyMatchFailure(op, "src/dst element types mismatch"); + + Location loc = op.getLoc(); + + Type predMemTy = nullptr; + if (pred) { + predMemTy = dyn_cast(op.getPred().getType()); + if (!predMemTy) + return rewriter.notifyMatchFailure(op, "pred is not a MemRef type"); + } + + if (pred) { + if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcMemTy, dstMemTy, predMemTy, + copyAtomVal, src, dst, pred))) + return failure(); + } else { + if (failed(copyAtom.emitAtomCall(rewriter, loc, copyAtomType, srcMemTy, dstMemTy, copyAtomVal, + src, dst))) + return failure(); + } + rewriter.eraseOp(op); + return success(); + } +}; + +class CopyAtomCallSSALowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(CopyAtomCallSSA op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type copyAtomType = op.getCopyAtom().getType(); + auto copyAtom = dyn_cast(copyAtomType); + if (!copyAtom) + return rewriter.notifyMatchFailure(op, "copyAtom is not CopyAtomType"); + + Location loc = op.getLoc(); + bool hasResult = op.getResults().size() > 0; + Type srcTy = op.getSrc().getType(); + Value pred = adaptor.getPred(); + + Type resultTy = hasResult ? op.getResult(0).getType() : Type{}; + Type dstTy = op.getDst() ? op.getDst().getType() : Type{}; + + FailureOr result; + if (pred) { + result = copyAtom.emitAtomCallSSA(rewriter, loc, resultTy, copyAtomType, srcTy, dstTy, + op.getPred().getType(), adaptor.getCopyAtom(), + adaptor.getSrc(), adaptor.getDst(), pred); + } else { + result = copyAtom.emitAtomCallSSA(rewriter, loc, resultTy, copyAtomType, srcTy, dstTy, + adaptor.getCopyAtom(), adaptor.getSrc(), adaptor.getDst()); + } + if (failed(result)) + return failure(); + + if (hasResult) + rewriter.replaceOp(op, *result); + else + rewriter.eraseOp(op); + return success(); + } +}; + +class MmaAtomCallLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(MmaAtomCall op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto mmaAtomTy = dyn_cast(op.getMmaAtom().getType()); + if (!mmaAtomTy) + return rewriter.notifyMatchFailure(op, "expected MmaAtomType for mmaAtom operand"); + + Location loc = op.getLoc(); + + Value dPtr = adaptor.getD(); + Value aPtr = adaptor.getA(); + Value bPtr = adaptor.getB(); + Value cPtr = adaptor.getC(); + + if (!isa(dPtr.getType()) || + !isa(aPtr.getType()) || + !isa(bPtr.getType()) || !isa(cPtr.getType())) + return rewriter.notifyMatchFailure(op, "expected llvm.ptr operands after type conversion"); + + auto dMemTy = dyn_cast(op.getD().getType()); + auto aMemTy = dyn_cast(op.getA().getType()); + auto bMemTy = dyn_cast(op.getB().getType()); + auto cMemTy = dyn_cast(op.getC().getType()); + if (!dMemTy || !aMemTy || !bMemTy || !cMemTy) + return rewriter.notifyMatchFailure(op, "expected Fly memref types on original op"); + + if (failed(mmaAtomTy.emitAtomCall(rewriter, loc, mmaAtomTy, dMemTy, aMemTy, bMemTy, cMemTy, + adaptor.getMmaAtom(), dPtr, aPtr, bPtr, cPtr))) + return failure(); + + rewriter.eraseOp(op); + return success(); + } +}; + +class MmaAtomCallSSALowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(MmaAtomCallSSA op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto mmaAtomTy = dyn_cast(op.getMmaAtom().getType()); + if (!mmaAtomTy) + return rewriter.notifyMatchFailure(op, "expected MmaAtomType for mmaAtom operand"); + + Location loc = op.getLoc(); + bool hasResult = op.getResults().size() > 0; + + Type resultTy = hasResult ? op.getResult(0).getType() : Type{}; + Type dTy = op.getD() ? op.getD().getType() : Type{}; + Value dPtr = hasResult ? Value{} : adaptor.getD(); + + auto result = + mmaAtomTy.emitAtomCallSSA(rewriter, loc, resultTy, mmaAtomTy, dTy, op.getA().getType(), + op.getB().getType(), op.getC().getType(), adaptor.getMmaAtom(), + dPtr, adaptor.getA(), adaptor.getB(), adaptor.getC()); + if (failed(result)) + return failure(); + + if (hasResult) + rewriter.replaceOp(op, *result); + else + rewriter.eraseOp(op); + return success(); + } +}; + +/// Lower `gpu.launch_func` kernel operands so that any `!fly.memref` values are +/// replaced by their type-converted builtin `memref` values. This prevents +/// `unrealized_conversion_cast` materializations from remaining live after +/// partial conversion (e.g., when the surrounding `func.func` signature has +/// been converted to builtin memrefs). +class GpuLaunchFuncOpLowering : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(gpu::LaunchFuncOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto kernelRef = adaptor.getKernel(); + + auto grid = + gpu::KernelDim3{adaptor.getGridSizeX(), adaptor.getGridSizeY(), adaptor.getGridSizeZ()}; + auto block = + gpu::KernelDim3{adaptor.getBlockSizeX(), adaptor.getBlockSizeY(), adaptor.getBlockSizeZ()}; + + std::optional clusterSize = std::nullopt; + if (adaptor.getClusterSizeX() && adaptor.getClusterSizeY() && adaptor.getClusterSizeZ()) { + clusterSize = gpu::KernelDim3{adaptor.getClusterSizeX(), adaptor.getClusterSizeY(), + adaptor.getClusterSizeZ()}; + } + + // Preserve async token result type when present. + Type asyncTokenType = nullptr; + if (Value tok = op.getAsyncToken()) + asyncTokenType = tok.getType(); + + // There are two relevant builder signatures in this MLIR: + // - (kernel, ..., asyncTokenType, asyncDependencies, clusterSize) + // - (kernel, ..., asyncObject, clusterSize) + // Pick the one that matches the original op structure. + if (Value asyncObj = adaptor.getAsyncObject()) { + if (!adaptor.getAsyncDependencies().empty()) + return rewriter.notifyMatchFailure( + op, "launch_func has both asyncObject and asyncDependencies"); + + rewriter.replaceOpWithNewOp( + op, kernelRef, grid, block, adaptor.getDynamicSharedMemorySize(), + adaptor.getKernelOperands(), asyncObj, clusterSize); + return success(); + } + + rewriter.replaceOpWithNewOp( + op, kernelRef, grid, block, adaptor.getDynamicSharedMemorySize(), + adaptor.getKernelOperands(), asyncTokenType, adaptor.getAsyncDependencies(), clusterSize); + return success(); + } +}; + +class FlyTypeConverter : public TypeConverter { +public: + FlyTypeConverter() { + addConversion([](Type type) { return type; }); + + addConversion([&](FloatType floatTy) -> std::optional { + if (floatTy.getWidth() < 16) + return IntegerType::get(floatTy.getContext(), floatTy.getWidth()); + return std::nullopt; + }); + addConversion([&](VectorType vecTy) -> std::optional { + Type convertedElem = convertType(vecTy.getElementType()); + if (!convertedElem || convertedElem == vecTy.getElementType()) + return std::nullopt; + return VectorType::get(vecTy.getShape(), convertedElem, vecTy.getScalableDims()); + }); + addConversion([&](fly::MemRefType flyMemRefTy) -> Type { + unsigned as = mapAttrToLLVMAddressSpace(flyMemRefTy.getAddressSpace()); + return LLVM::LLVMPointerType::get(flyMemRefTy.getContext(), as); + }); + addConversion([&](fly::PointerType flyPtrTy) -> Type { + unsigned as = mapAttrToLLVMAddressSpace(flyPtrTy.getAddressSpace()); + return LLVM::LLVMPointerType::get(flyPtrTy.getContext(), as); + }); + addConversion([&](fly::CopyAtomType atomTy) -> Type { + if (atomTy.isStateful()) + return atomTy.getConvertedType(atomTy.getContext()); + return LLVM::LLVMStructType::getLiteral(atomTy.getContext(), {}); + }); + addConversion([&](fly::MmaAtomType atomTy) -> Type { + if (atomTy.isStateful()) + return atomTy.getConvertedType(atomTy.getContext()); + return LLVM::LLVMStructType::getLiteral(atomTy.getContext(), {}); + }); + addConversion( + [&](fly::TiledCopyType tiledTy) -> Type { return convertType(tiledTy.getCopyAtom()); }); + addConversion( + [&](fly::TiledMmaType tiledTy) -> Type { return convertType(tiledTy.getMmaAtom()); }); + } +}; + +class ExtractAlignedPointerAsIndexLowering + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult matchAndRewrite(ExtractAlignedPointerAsIndexOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + // fly.memref is a bare pointer; after type conversion the operand is llvm.ptr. + // Cast to the result type (e.g. llvm.ptr<0>) if address spaces differ. + Value src = adaptor.getSource(); + Type resultType = getTypeConverter()->convertType(op.getResult().getType()); + if (!resultType) + resultType = op.getResult().getType(); + if (src.getType() != resultType) + src = LLVM::AddrSpaceCastOp::create(rewriter, op.getLoc(), resultType, src); + rewriter.replaceOp(op, src); + return success(); + } +}; + +class FlyToNVVMConversionPass + : public mlir::impl::FlyToNVVMConversionPassBase { +public: + using mlir::impl::FlyToNVVMConversionPassBase< + FlyToNVVMConversionPass>::FlyToNVVMConversionPassBase; + + void runOnOperation() override { + MLIRContext *context = &getContext(); + RewritePatternSet patterns(context); + + ConversionTarget target(getContext()); + + target.addLegalDialect(); + target.addIllegalDialect(); + + // Constructors + target.addLegalOp(); + + FlyTypeConverter typeConverter; + + // Ensure function signatures are type-converted; otherwise conversions may rely on + // inserted unrealized casts that remain live. + target.addDynamicallyLegalOp( + [&](func::FuncOp op) { return typeConverter.isSignatureLegal(op.getFunctionType()); }); + target.addDynamicallyLegalOp( + [&](gpu::GPUFuncOp op) { return typeConverter.isSignatureLegal(op.getFunctionType()); }); + + // IMPORTANT: `gpu.launch_func` itself is in a legal dialect, but its kernel operands may + // still carry illegal `!fly.memref` types. If we don't mark it dynamically illegal in that + // case, partial conversion won't try to rewrite it, leaving `unrealized_conversion_cast` + // users alive and causing legalization failure. + target.addDynamicallyLegalOp([&](gpu::LaunchFuncOp op) { + auto isValueLegal = [&](Value v) { + if (!v) + return true; + return typeConverter.isLegal(v.getType()); + }; + + for (Value v : op.getKernelOperands()) + if (!isValueLegal(v)) + return false; + + if (!isValueLegal(op.getDynamicSharedMemorySize())) + return false; + + // Async operands are part of the operand list; keep them consistent as well. + for (Value dep : op.getAsyncDependencies()) + if (!isValueLegal(dep)) + return false; + if (!isValueLegal(op.getAsyncObject())) + return false; + + // Dimensions are typically index and already legal; no need to special-case. + return true; + }); + + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); + + // TODO: deprecated in the future + patterns.add(typeConverter, context); + + populateFunctionOpInterfaceTypeConversionPattern(patterns, typeConverter); + populateFunctionOpInterfaceTypeConversionPattern(patterns, typeConverter); + + if (failed(applyPartialConversion(getOperation(), target, std::move(patterns)))) + signalPassFailure(); + } +}; + +} // namespace diff --git a/lib/Dialect/FlyNVVM/CMakeLists.txt b/lib/Dialect/FlyNVVM/CMakeLists.txt new file mode 100644 index 000000000..752a50e35 --- /dev/null +++ b/lib/Dialect/FlyNVVM/CMakeLists.txt @@ -0,0 +1,8 @@ +add_mlir_dialect_library(MLIRFlyNVVMDialect + Dialect.cpp + SM80/MmaAtom.cpp + SM80/CopyAtom.cpp + + DEPENDS + MLIRFlyNVVMIncGen +) diff --git a/lib/Dialect/FlyNVVM/Dialect.cpp b/lib/Dialect/FlyNVVM/Dialect.cpp new file mode 100644 index 000000000..e2bf50e76 --- /dev/null +++ b/lib/Dialect/FlyNVVM/Dialect.cpp @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/DialectImplementation.h" +#include "llvm/ADT/StringExtras.h" +#include "llvm/ADT/TypeSwitch.h" + +#include "flydsl/Dialect/FlyNVVM/IR/Dialect.h" + +using namespace mlir; +using namespace mlir::fly; +using namespace mlir::fly_nvvm; + +#include "flydsl/Dialect/FlyNVVM/IR/Dialect.cpp.inc" + +#define GET_TYPEDEF_CLASSES +#include "flydsl/Dialect/FlyNVVM/IR/Atom.cpp.inc" + +void FlyNVVMDialect::initialize() { + addTypes< +#define GET_TYPEDEF_LIST +#include "flydsl/Dialect/FlyNVVM/IR/Atom.cpp.inc" + >(); +} diff --git a/lib/Dialect/FlyNVVM/SM80/CopyAtom.cpp b/lib/Dialect/FlyNVVM/SM80/CopyAtom.cpp new file mode 100644 index 000000000..b83ca13a3 --- /dev/null +++ b/lib/Dialect/FlyNVVM/SM80/CopyAtom.cpp @@ -0,0 +1,276 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors +// +// NVVM copy atoms used by the SM80 GEMM path: +// * CopyOpSM80_CpAsync — cp.async.{ca,cg}.shared.global (global -> shared, async) +// * CopyOpSM75_LdMatrix — ldmatrix.sync.aligned (shared -> register, sm_75+) +// +// Thread/bit layouts are derived from the PTX cp.async / ldmatrix operand ABI +// and FlyDSL copy-atom layout invariants, then cross-checked against independent +// SM80 GEMM references. + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Vector/IR/VectorOps.h" +#include "mlir/IR/BuiltinTypes.h" + +#include "flydsl/Dialect/Fly/IR/FlyDialect.h" +#include "flydsl/Dialect/Fly/Utils/PointerUtils.h" +#include "flydsl/Dialect/Fly/Utils/ThrValLayoutMacro.h.inc" +#include "flydsl/Dialect/FlyNVVM/IR/Dialect.h" + +using namespace mlir; +using namespace mlir::fly; + +namespace mlir::fly_nvvm { + +//===----------------------------------------------------------------------===// +// CopyOpSM80_CpAsync — cp.async.shared.global +// +// CuTe Copy_Traits>: one thread, per-thread bit +// layout (1, bits). dst is shared (addrspace 3), src is global (addrspace 1). +//===----------------------------------------------------------------------===// + +bool CopyOpSM80_CpAsyncType::isStatic() const { return true; } + +Value CopyOpSM80_CpAsyncType::rebuildStaticValue(OpBuilder &builder, Location loc, + Value currentValue) const { + if (currentValue && isa(currentValue.getDefiningOp())) + return nullptr; + return MakeCopyAtomOp::create(builder, loc, CopyAtomType::get(*this, getBitSize()), getBitSize()); +} + +Attribute CopyOpSM80_CpAsyncType::getThrLayout() const { return FxLayout(FxC(1), FxC(1)); } + +Attribute CopyOpSM80_CpAsyncType::getThrBitLayoutSrc() const { + return FxLayout(FxShape(FxC(1), FxC(getBitSize())), FxStride(FxC(1), FxC(1))); +} +Attribute CopyOpSM80_CpAsyncType::getThrBitLayoutDst() const { + return FxLayout(FxShape(FxC(1), FxC(getBitSize())), FxStride(FxC(1), FxC(1))); +} +Attribute CopyOpSM80_CpAsyncType::getThrBitLayoutRef() const { + return FxLayout(FxShape(FxC(1), FxC(getBitSize())), FxStride(FxC(1), FxC(1))); +} + +LogicalResult CopyOpSM80_CpAsyncType::verify(function_ref emitError, + int32_t bitSize) { + if (bitSize != 32 && bitSize != 64 && bitSize != 128) + return emitError() << "cp.async bitSize must be 32/64/128, got " << bitSize; + return success(); +} + +// cp.async has no SSA result (it is a void async DMA). Use the memref path. +FailureOr CopyOpSM80_CpAsyncType::emitAtomCallSSA(OpBuilder &, Location, Type, Type, Type, + Type, Value, Value, Value) const { + return failure(); +} + +FailureOr CopyOpSM80_CpAsyncType::emitAtomCallSSA(OpBuilder &, Location, Type, Type, Type, + Type, Type, Value, Value, Value, + Value) const { + return failure(); +} + +LogicalResult CopyOpSM80_CpAsyncType::emitAtomCall(OpBuilder &builder, Location loc, + Type copyAtomTy, Type srcMemTy, Type dstMemTy, + Value atomVal, Value src, Value dst) const { + MLIRContext *ctx = builder.getContext(); + int32_t sizeBytes = getBitSize() / 8; + + // cp.async is shared(3) <- global(1); cast either side if it arrived generic. + auto globalPtrTy = LLVM::LLVMPointerType::get(ctx, /*addrspace=*/1); + auto sharedPtrTy = LLVM::LLVMPointerType::get(ctx, /*addrspace=*/3); + Value srcCast = src; + if (srcCast.getType() != globalPtrTy) + srcCast = LLVM::AddrSpaceCastOp::create(builder, loc, globalPtrTy, srcCast); + Value dstCast = dst; + if (dstCast.getType() != sharedPtrTy) + dstCast = LLVM::AddrSpaceCastOp::create(builder, loc, sharedPtrTy, dstCast); + + // 16B copies use CG (bypass L1); 4/8B must use CA. + NVVM::LoadCacheModifierKind modifier = + sizeBytes == 16 ? NVVM::LoadCacheModifierKind::CG : NVVM::LoadCacheModifierKind::CA; + + NVVM::CpAsyncOp::create(builder, loc, dstCast, srcCast, builder.getI32IntegerAttr(sizeBytes), + NVVM::LoadCacheModifierKindAttr::get(ctx, modifier), + /*cpSize=*/Value{}); + return success(); +} + +LogicalResult CopyOpSM80_CpAsyncType::emitAtomCall(OpBuilder &builder, Location loc, + Type copyAtomTy, Type srcMemTy, Type dstMemTy, + Type predMemTy, Value atomVal, Value src, + Value dst, Value pred) const { + OpBuilder::InsertionGuard guard(builder); + auto predMemRefTy = cast(predMemTy); + Value predVal = LLVM::LoadOp::create(builder, loc, predMemRefTy.getElemTy(), pred); + auto ifOp = scf::IfOp::create(builder, loc, TypeRange{}, predVal, /*withElse=*/false); + builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); + return emitAtomCall(builder, loc, copyAtomTy, srcMemTy, dstMemTy, atomVal, src, dst); +} + +//===----------------------------------------------------------------------===// +// CopyOpSM75_LdMatrix — ldmatrix.sync.aligned.m8n8.x{1,2,4}[.trans].shared.b16 +// +// Bit layouts follow the PTX ldmatrix fragment mapping: +// num=4 non-trans -> four 32-bit registers per lane +// num=2 non-trans -> two 32-bit registers per lane +// num=1 non-trans -> one 32-bit register per lane +// num=4 trans -> eight packed 16-bit values per lane +// num=2 trans -> four packed 16-bit values per lane +// num=1 trans -> two packed 16-bit values per lane +//===----------------------------------------------------------------------===// + +bool CopyOpSM75_LdMatrixType::isStatic() const { return true; } + +Value CopyOpSM75_LdMatrixType::rebuildStaticValue(OpBuilder &builder, Location loc, + Value currentValue) const { + if (currentValue && isa(currentValue.getDefiningOp())) + return nullptr; + // val bits per the b16 element granularity; CopyAtom valBits is set by the + // Python make_copy_atom wrapper from the element type, so reuse 16 here. + return MakeCopyAtomOp::create(builder, loc, CopyAtomType::get(*this, 16), 16); +} + +Attribute CopyOpSM75_LdMatrixType::getThrLayout() const { return FxLayout(FxC(32), FxC(1)); } + +// Source: (src-thr, src-val) -> shared-memory bit. From CuTe SrcLayout. +Attribute CopyOpSM75_LdMatrixType::getThrBitLayoutSrc() const { + int32_t num = getNum(); + if (!getTrans()) { + // SM75_U32x{1,2,4}_LDSM_N + if (num == 1) // ((8,4),128):((128,0),1) + return FxLayout(FxShape(FxThr(8, 4), FxC(128)), FxStride(FxThr(128, 0), FxC(1))); + if (num == 2) // ((16,2),128):((128,0),1) + return FxLayout(FxShape(FxThr(16, 2), FxC(128)), FxStride(FxThr(128, 0), FxC(1))); + // num == 4: (32,128):(128,1) + return FxLayout(FxShape(FxC(32), FxC(128)), FxStride(FxC(128), FxC(1))); + } + // trans variants share the same Src shapes as the N variants. + if (num == 1) + return FxLayout(FxShape(FxThr(8, 4), FxC(128)), FxStride(FxThr(128, 0), FxC(1))); + if (num == 2) + return FxLayout(FxShape(FxThr(16, 2), FxC(128)), FxStride(FxThr(128, 0), FxC(1))); + return FxLayout(FxShape(FxC(32), FxC(128)), FxStride(FxC(128), FxC(1))); +} + +// Destination: (dst-thr, dst-val) -> register bit. From CuTe DstLayout. +Attribute CopyOpSM75_LdMatrixType::getThrBitLayoutDst() const { + int32_t num = getNum(); + if (!getTrans()) { + // SM75_U32x{1,2,4}_LDSM_N + if (num == 1) // (32,32):(32,1) + return FxLayout(FxShape(FxC(32), FxC(32)), FxStride(FxC(32), FxC(1))); + if (num == 2) // (32,(32,2)):(32,(1,1024)) + return FxLayout(FxShape(FxC(32), FxVal(32, 2)), FxStride(FxC(32), FxVal(1, 1024))); + // num == 4: (32,(32,4)):(32,(1,1024)) + return FxLayout(FxShape(FxC(32), FxVal(32, 4)), FxStride(FxC(32), FxVal(1, 1024))); + } + // SM75_U16x{2,4,8}_LDSM_T + if (num == 1) // ((4,8),(16,2)):((256,16),(1,128)) + return FxLayout(FxShape(FxThr(4, 8), FxVal(16, 2)), FxStride(FxThr(256, 16), FxVal(1, 128))); + if (num == 2) // ((4,8),(16,2,2)):((256,16),(1,128,1024)) + return FxLayout(FxShape(FxThr(4, 8), FxVal(16, 2, 2)), + FxStride(FxThr(256, 16), FxVal(1, 128, 1024))); + // num == 4: ((4,8),(16,2,4)):((256,16),(1,128,1024)) + return FxLayout(FxShape(FxThr(4, 8), FxVal(16, 2, 4)), + FxStride(FxThr(256, 16), FxVal(1, 128, 1024))); +} + +Attribute CopyOpSM75_LdMatrixType::getThrBitLayoutRef() const { return getThrBitLayoutDst(); } + +LogicalResult CopyOpSM75_LdMatrixType::verify(function_ref emitError, + int32_t num, bool trans) { + if (num != 1 && num != 2 && num != 4) + return emitError() << "ldmatrix num must be 1/2/4, got " << num; + return success(); +} + +FailureOr CopyOpSM75_LdMatrixType::emitAtomCallSSA(OpBuilder &builder, Location loc, + Type resultTy, Type copyAtomTyArg, + Type srcTyArg, Type dstTyArg, + Value atomVal, Value src, + Value dst) const { + MLIRContext *ctx = builder.getContext(); + int32_t num = getNum(); + Type i32Ty = builder.getI32Type(); + + Type ldResTy = + num > 1 ? cast(LLVM::LLVMStructType::getLiteral(ctx, SmallVector(num, i32Ty))) + : i32Ty; + + auto shape = NVVM::LdStMatrixShapeAttr::get(ctx, /*m=*/8, /*n=*/8); + Value loaded = NVVM::LdMatrixOp::create(builder, loc, ldResTy, src, /*num=*/num, + getTrans() ? NVVM::MMALayout::col : NVVM::MMALayout::row, + shape, NVVM::LdStMatrixEltType::B16); + + // Repack the num i32 registers into the result vector type. + if (num == 1) { + if (resultTy && loaded.getType() != resultTy) + loaded = LLVM::BitcastOp::create(builder, loc, resultTy, loaded); + return loaded; + } + + // Build vector from the struct, then bitcast to resultTy. + auto i32VecTy = VectorType::get({num}, i32Ty); + Value vec = LLVM::PoisonOp::create(builder, loc, i32VecTy); + for (int i = 0; i < num; ++i) { + Value el = LLVM::ExtractValueOp::create(builder, loc, loaded, ArrayRef{i}); + Value idx = arith::ConstantIntOp::create(builder, loc, i, 32); + vec = LLVM::InsertElementOp::create(builder, loc, i32VecTy, vec, el, idx); + } + Value res = vec; + + if (resultTy && res.getType() != resultTy) + res = LLVM::BitcastOp::create(builder, loc, resultTy, res); + return res; +} + +FailureOr CopyOpSM75_LdMatrixType::emitAtomCallSSA(OpBuilder &builder, Location loc, + Type resultTy, Type copyAtomTyArg, + Type srcTyArg, Type dstTyArg, + Type predTyArg, Value atomVal, Value src, + Value dst, Value pred) const { + assert(resultTy && "resultTy must be SSA Type"); + OpBuilder::InsertionGuard guard(builder); + auto ifOp = scf::IfOp::create(builder, loc, resultTy, pred, /*withElseRegion=*/true); + builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); + auto result = + emitAtomCallSSA(builder, loc, resultTy, copyAtomTyArg, srcTyArg, dstTyArg, atomVal, src, dst); + if (failed(result)) + return failure(); + scf::YieldOp::create(builder, loc, *result); + builder.setInsertionPointToStart(&ifOp.getElseRegion().front()); + scf::YieldOp::create(builder, loc, dst); + return ifOp.getResult(0); +} + +LogicalResult CopyOpSM75_LdMatrixType::emitAtomCall(OpBuilder &builder, Location loc, + Type copyAtomTyArg, Type srcMemTyArg, + Type dstMemTyArg, Value atomVal, Value src, + Value dst) const { + auto dstSSATy = fly::RegMem2SSAType(cast(dstMemTyArg), true); + auto res = emitAtomCallSSA(builder, loc, dstSSATy, copyAtomTyArg, srcMemTyArg, Type{}, atomVal, + src, Value{}); + if (failed(res)) + return failure(); + LLVM::StoreOp::create(builder, loc, *res, dst); + return success(); +} + +LogicalResult CopyOpSM75_LdMatrixType::emitAtomCall(OpBuilder &builder, Location loc, + Type copyAtomTyArg, Type srcMemTyArg, + Type dstMemTyArg, Type predMemTyArg, + Value atomVal, Value src, Value dst, + Value pred) const { + OpBuilder::InsertionGuard guard(builder); + auto predMemTy = cast(predMemTyArg); + Value predVal = LLVM::LoadOp::create(builder, loc, predMemTy.getElemTy(), pred); + auto ifOp = scf::IfOp::create(builder, loc, TypeRange{}, predVal, /*withElse=*/false); + builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); + return emitAtomCall(builder, loc, copyAtomTyArg, srcMemTyArg, dstMemTyArg, atomVal, src, dst); +} + +} // namespace mlir::fly_nvvm diff --git a/lib/Dialect/FlyNVVM/SM80/MmaAtom.cpp b/lib/Dialect/FlyNVVM/SM80/MmaAtom.cpp new file mode 100644 index 000000000..1b93303fc --- /dev/null +++ b/lib/Dialect/FlyNVVM/SM80/MmaAtom.cpp @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors +// +// ThrVal layouts + lowering for PTX Multiply-and-Accumulate Instruction: mma +// mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 +// (SM80 Ampere 16x8x16 f16->f32 tensor-core instruction). Operand fragment +// ABI (per NVIDIA PTX ISA): +// A: 4 x vector<2xf16> B: 2 x vector<2xf16> C/D: 4 x f32 +// +// Layouts are derived from the PTX operand ABI and FlyDSL's column-major atom +// coordinate convention, then cross-checked against independent SM80 GEMM +// references. The thread axis decomposes colexicographically as lane = T + 4*G, +// so the first thread sub-mode is T = lane & 3 (size 4) and the second is +// G = lane >> 2 (size 8). +// T = lane & 3 in [0, 4) threadID_in_group +// G = lane >> 2 in [0, 8) groupID + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/IR/BuiltinTypes.h" + +#include "flydsl/Dialect/Fly/IR/FlyDialect.h" +#include "flydsl/Dialect/Fly/Utils/ThrValLayoutMacro.h.inc" +#include "flydsl/Dialect/FlyNVVM/IR/Dialect.h" + +using namespace mlir; +using namespace mlir::fly; + +namespace sm80 { + +static LayoutAttr getThrValLayoutA(MLIRContext *ctx) { + auto getContext = [&]() { return ctx; }; + return FxLayout(FxShape(FxThr(4, 8), FxVal(2, 2, 2)), FxStride(FxThr(32, 1), FxVal(16, 8, 128))); +} + +static LayoutAttr getThrValLayoutB(MLIRContext *ctx) { + auto getContext = [&]() { return ctx; }; + return FxLayout(FxShape(FxThr(4, 8), FxVal(2, 2)), FxStride(FxThr(16, 1), FxVal(8, 64))); +} + +static LayoutAttr getThrValLayoutC(MLIRContext *ctx) { + auto getContext = [&]() { return ctx; }; + return FxLayout(FxShape(FxThr(4, 8), FxVal(2, 2)), FxStride(FxThr(32, 1), FxVal(16, 8))); +} + +} // namespace sm80 + +namespace mlir::fly_nvvm { + +bool MmaOpSM80_MmaSyncType::isStatic() const { return true; } + +Value MmaOpSM80_MmaSyncType::rebuildStaticValue(OpBuilder &builder, Location loc, + Value currentValue) const { + if (currentValue && isa(currentValue.getDefiningOp())) + return nullptr; + return MakeMmaAtomOp::create(builder, loc, MmaAtomType::get(*this)); +} + +Attribute MmaOpSM80_MmaSyncType::getThrLayout() const { return FxLayout(FxC(32), FxC(1)); } + +Attribute MmaOpSM80_MmaSyncType::getShapeMNK() const { + return IntTupleAttr::get(ArrayAttr::get(getContext(), {FxC(getM()), FxC(getN()), FxC(getK())})); +} + +Type MmaOpSM80_MmaSyncType::getValTypeA() const { return getElemTyA(); } +Type MmaOpSM80_MmaSyncType::getValTypeB() const { return getElemTyB(); } +Type MmaOpSM80_MmaSyncType::getValTypeC() const { return getElemTyAcc(); } +Type MmaOpSM80_MmaSyncType::getValTypeD() const { return getElemTyAcc(); } + +Attribute MmaOpSM80_MmaSyncType::getThrValLayoutA() const { + return sm80::getThrValLayoutA(getContext()); +} +Attribute MmaOpSM80_MmaSyncType::getThrValLayoutB() const { + return sm80::getThrValLayoutB(getContext()); +} +Attribute MmaOpSM80_MmaSyncType::getThrValLayoutC() const { + return sm80::getThrValLayoutC(getContext()); +} + +LogicalResult MmaOpSM80_MmaSyncType::verify(function_ref emitError, int32_t m, + int32_t n, int32_t k, Type elemTyA, Type elemTyB, + Type elemTyAcc) { + if (m != 16 || n != 8 || k != 16) + return emitError() << "unsupported SM80 mma.sync.aligned shape " << m << "x" << n << "x" << k + << ", only 16x8x16 is supported"; + if (!elemTyA.isF16() || !elemTyB.isF16()) + return emitError() << "SM80 mma.sync.aligned requires f16 inputs, got " << elemTyA << ", " + << elemTyB; + if (!elemTyAcc.isF32()) + return emitError() << "SM80 mma.sync.aligned requires f32 accumulator, got " << elemTyAcc; + return success(); +} + +// SSA form: operands arrive as packed register vectors +// a: vector<8xf16> b: vector<4xf16> c: vector<4xf32> +// and we return the result as vector<4xf32>. +FailureOr MmaOpSM80_MmaSyncType::emitAtomCallSSA(OpBuilder &builder, Location loc, + Type resultTy, Type mmaAtomTyArg, + Type dTyArg, Type aTyArg, Type bTyArg, + Type cTyArg, Value atomVal, Value d, + Value a, Value b, Value c) const { + MLIRContext *ctx = builder.getContext(); + Type f16Ty = Float16Type::get(ctx); + Type f32Ty = Float32Type::get(ctx); + auto f16x2Ty = VectorType::get({2}, f16Ty); + auto aPackTy = VectorType::get({8}, f16Ty); + auto bPackTy = VectorType::get({4}, f16Ty); + auto cPackTy = VectorType::get({4}, f32Ty); + + if (a.getType() != aPackTy) + a = LLVM::BitcastOp::create(builder, loc, aPackTy, a); + if (b.getType() != bPackTy) + b = LLVM::BitcastOp::create(builder, loc, bPackTy, b); + if (c.getType() != cPackTy) + c = LLVM::BitcastOp::create(builder, loc, cPackTy, c); + + // A: vector<8xf16> -> 4 x vector<2xf16> + SmallVector matA; + for (int i = 0; i < 4; ++i) + matA.push_back(LLVM::ShuffleVectorOp::create(builder, loc, f16x2Ty, a, a, + ArrayRef{2 * i, 2 * i + 1})); + // B: vector<4xf16> -> 2 x vector<2xf16> + SmallVector matB; + for (int i = 0; i < 2; ++i) + matB.push_back(LLVM::ShuffleVectorOp::create(builder, loc, f16x2Ty, b, b, + ArrayRef{2 * i, 2 * i + 1})); + // C: vector<4xf32> -> 4 x f32 + SmallVector matC; + for (int i = 0; i < 4; ++i) { + Value idx = arith::ConstantIntOp::create(builder, loc, i, 32); + matC.push_back(LLVM::ExtractElementOp::create(builder, loc, c, idx)); + } + + // Result struct of the intrinsic: !llvm.struct<(f32, f32, f32, f32)>. + auto resStructTy = LLVM::LLVMStructType::getLiteral(ctx, {f32Ty, f32Ty, f32Ty, f32Ty}); + + Value mma = NVVM::MmaOp::create( + builder, loc, resStructTy, matA, matB, matC, + /*shape=*/ArrayRef{16, 8, 16}, + /*b1Op=*/std::nullopt, + /*intOverflow=*/std::nullopt, + /*multiplicandPtxTypes=*/ + std::array{NVVM::MMATypes::f16, NVVM::MMATypes::f16}, + /*multiplicandLayouts=*/ + std::array{NVVM::MMALayout::row, NVVM::MMALayout::col}); + + // Repack the 4 scalar f32 results into vector<4xf32>. + Value res = LLVM::PoisonOp::create(builder, loc, cPackTy); + for (int i = 0; i < 4; ++i) { + Value el = LLVM::ExtractValueOp::create(builder, loc, mma, ArrayRef{i}); + Value idx = arith::ConstantIntOp::create(builder, loc, i, 32); + res = LLVM::InsertElementOp::create(builder, loc, cPackTy, res, el, idx); + } + // The accumulator fragment is always 4 x f32 for this instruction; bitcast if + // the caller asked for an equally sized but differently spelled type. + if (resultTy && res.getType() != resultTy) + res = LLVM::BitcastOp::create(builder, loc, resultTy, res); + return res; +} + +LogicalResult MmaOpSM80_MmaSyncType::emitAtomCall(OpBuilder &builder, Location loc, Type mmaAtomTy, + Type dMemTy, Type aMemTy, Type bMemTy, + Type cMemTy, Value atomVal, Value dPtr, + Value aPtr, Value bPtr, Value cPtr) const { + MLIRContext *ctx = builder.getContext(); + Type f16Ty = Float16Type::get(ctx); + Type f32Ty = Float32Type::get(ctx); + auto aPackTy = VectorType::get({8}, f16Ty); + auto bPackTy = VectorType::get({4}, f16Ty); + auto cPackTy = VectorType::get({4}, f32Ty); + + Value a = LLVM::LoadOp::create(builder, loc, aPackTy, aPtr); + Value b = LLVM::LoadOp::create(builder, loc, bPackTy, bPtr); + Value c = LLVM::LoadOp::create(builder, loc, cPackTy, cPtr); + auto res = emitAtomCallSSA(builder, loc, cPackTy, mmaAtomTy, Type{}, aPackTy, bPackTy, cPackTy, + atomVal, Value{}, a, b, c); + if (failed(res)) + return failure(); + LLVM::StoreOp::create(builder, loc, *res, dPtr); + return success(); +} + +} // namespace mlir::fly_nvvm diff --git a/lib/Runtime/CMakeLists.txt b/lib/Runtime/CMakeLists.txt index 5f2b26c2f..829257e1b 100644 --- a/lib/Runtime/CMakeLists.txt +++ b/lib/Runtime/CMakeLists.txt @@ -5,3 +5,6 @@ if("rocdl" IN_LIST FLYDSL_BACKENDS) add_subdirectory(ROCm) endif() +if("nvvm" IN_LIST FLYDSL_BACKENDS) + add_subdirectory(CUDA) +endif() diff --git a/lib/Runtime/CUDA/CMakeLists.txt b/lib/Runtime/CUDA/CMakeLists.txt new file mode 100644 index 000000000..7ae74a761 --- /dev/null +++ b/lib/Runtime/CUDA/CMakeLists.txt @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +find_package(CUDAToolkit QUIET) +if(NOT CUDAToolkit_FOUND) + message(FATAL_ERROR + "The 'nvvm' backend requires a CUDA toolkit to build the CUDA JIT runtime, " + "but CUDAToolkit was not found. Install CUDA 12+ or set " + "CUDAToolkit_ROOT, or drop 'nvvm' from FLYDSL_BACKENDS.") +endif() + +# Pick the cuda.h to build against explicitly instead of taking whatever +# FindCUDAToolkit resolved. Distros ship an old header in /usr/include (e.g. +# 11.5) that CMake prefers, which would both disagree with the ptxas the Python +# backend selects (_detect_cuda_toolkit() takes the newest /usr/local/cuda-*) +# and hide the CUDA 12 library-management API this runtime needs. Mirror that +# same "newest /usr/local/cuda-* first" order here. +function(_flydsl_pick_cuda_include out_var) + set(_candidates "") + if(DEFINED CUDAToolkit_ROOT) + list(APPEND _candidates "${CUDAToolkit_ROOT}") + endif() + file(GLOB _globbed "/usr/local/cuda-*") + list(SORT _globbed COMPARE NATURAL ORDER DESCENDING) + list(APPEND _candidates ${_globbed} "/usr/local/cuda" ${CUDAToolkit_INCLUDE_DIRS}) + foreach(_root ${_candidates}) + foreach(_inc "${_root}/include" "${_root}") + if(EXISTS "${_inc}/cuda.h") + file(STRINGS "${_inc}/cuda.h" _line REGEX "^#define CUDA_VERSION ") + string(REGEX MATCH "[0-9]+" _ver "${_line}") + if(_ver AND _ver GREATER_EQUAL 12000) + set(${out_var} "${_inc}" PARENT_SCOPE) + return() + endif() + endif() + endforeach() + endforeach() + set(${out_var} "" PARENT_SCOPE) +endfunction() + +_flydsl_pick_cuda_include(FLYDSL_CUDA_INCLUDE_DIR) +if(NOT FLYDSL_CUDA_INCLUDE_DIR) + message(FATAL_ERROR + "The 'nvvm' backend needs CUDA 12+ headers: the runtime uses the CUDA 12 " + "library-management API (cuLibraryLoadData) so one JIT artifact can run on " + "several devices. No cuda.h with CUDA_VERSION >= 12000 was found under " + "CUDAToolkit_ROOT or /usr/local/cuda*. Set CUDAToolkit_ROOT to a CUDA 12+ install.") +endif() +message(STATUS "FlyCudaJitRuntime cuda.h: ${FLYDSL_CUDA_INCLUDE_DIR}") + +add_library(FlyCudaJitRuntime SHARED FlyCudaRuntimeWrappers.cpp) +# BEFORE so the chosen cuda.h wins over any older one on the default path. +target_include_directories(FlyCudaJitRuntime BEFORE PRIVATE "${FLYDSL_CUDA_INCLUDE_DIR}") +target_include_directories(FlyCudaJitRuntime PRIVATE + ${LLVM_INCLUDE_DIRS} + ${MLIR_INCLUDE_DIRS} +) +target_compile_features(FlyCudaJitRuntime PRIVATE cxx_std_17) +target_link_libraries(FlyCudaJitRuntime PRIVATE CUDA::cuda_driver ${CMAKE_DL_LIBS}) +set_target_properties(FlyCudaJitRuntime PROPERTIES OUTPUT_NAME "fly_cuda_runtime") diff --git a/lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp b/lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp new file mode 100644 index 000000000..6fb568e18 --- /dev/null +++ b/lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp @@ -0,0 +1,304 @@ +//===- FlyCudaRuntimeWrappers.cpp - CUDA runtime for MLIR JIT -------------===// +// +// Derived from LLVM Project: mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Thin CUDA driver-API runtime wrappers for the MLIR ExecutionEngine JIT. +// Exposes the same vendor-neutral `mgpu*` symbol set as the ROCm fork +// (FlyRocmRuntimeWrappers.cpp) so the host-side launch / explicit-module +// offloading IR (which emits calls to mgpuModuleLoad / mgpuLaunchKernel / +// mgpuModuleUnload) links unchanged against either backend. +// +//===----------------------------------------------------------------------===// + +#include +#include +#include +#include + +#include "cuda.h" +#include "mlir/ExecutionEngine/CRunnerUtils.h" + +#define CUDA_REPORT_IF_ERROR(expr) \ + [](CUresult result) { \ + if (!result) \ + return; \ + const char *name = nullptr; \ + cuGetErrorName(result, &name); \ + if (!name) \ + name = ""; \ + fprintf(stderr, "'%s' failed with '%s'\n", #expr, name); \ + }(expr) + +thread_local static int32_t defaultDevice = 0; + +// Ensure a CUDA context is current for the duration of the instance. +// +// Unlike upstream's CudaRuntimeWrappers.cpp we do NOT unconditionally push the +// primary context of `defaultDevice`. FlyDSL is embedded in a host framework +// (PyTorch) that owns device selection and hands us its streams, so forcing +// device 0's context makes every launch on `cuda:N` (N > 0) fail with +// CUDA_ERROR_INVALID_HANDLE -- the stream belongs to a different context. The +// ROCm fork has no such problem because HIP simply inherits the calling +// thread's current device; this mirrors that behaviour. +// +// When a context is already current we use it as-is. Only in the standalone +// case (no host framework, nothing bound) do we retain and push the primary +// context of `defaultDevice`, which is also what makes mgpuSetDefaultDevice() +// meaningful. Retained contexts are cached and never released -- the process +// keeps its device contexts alive for its whole lifetime, as upstream does. +namespace { +class ScopedContext { +public: + ScopedContext() { + // Raw call: CUDA_ERROR_NOT_INITIALIZED here just means "no host framework + // has set anything up", which is the fallback path below, not an error. + CUcontext current = nullptr; + if (cuCtxGetCurrent(¤t) == CUDA_SUCCESS && current != nullptr) + return; + CUDA_REPORT_IF_ERROR(cuCtxPushCurrent(getPrimaryContext(defaultDevice))); + pushed = true; + } + ~ScopedContext() { + if (pushed) + CUDA_REPORT_IF_ERROR(cuCtxPopCurrent(nullptr)); + } + + // Retain (once per device) and return the primary context of `ordinal`. + static CUcontext getPrimaryContext(int32_t ordinal) { + static std::once_flag initFlag; + std::call_once(initFlag, [] { CUDA_REPORT_IF_ERROR(cuInit(/*flags=*/0)); }); + + static std::mutex mutex; + static std::map contexts; + std::lock_guard lock(mutex); + auto it = contexts.find(ordinal); + if (it != contexts.end()) + return it->second; + + CUdevice device; + CUDA_REPORT_IF_ERROR(cuDeviceGet(&device, /*ordinal=*/ordinal)); + CUcontext ctx = nullptr; + CUDA_REPORT_IF_ERROR(cuDevicePrimaryCtxRetain(&ctx, device)); + contexts[ordinal] = ctx; + return ctx; + } + +private: + bool pushed = false; +}; +} // namespace + +// Opt the kernel into `smem` bytes of dynamic shared memory. Anything above +// 48KB needs this explicit opt-in; warn up front when the request exceeds what +// the device allows, since the driver would otherwise only report a bare +// CUDA_ERROR_INVALID_VALUE. cuKernelSetAttribute takes the device explicitly. +static void setDynamicSharedMemory(CUkernel kernel, int32_t smem) { + if (smem <= 0) + return; + int32_t maxShmem = 0; + CUdevice device; + CUDA_REPORT_IF_ERROR(cuCtxGetDevice(&device)); + CUDA_REPORT_IF_ERROR(cuDeviceGetAttribute( + &maxShmem, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, device)); + if (maxShmem < smem) + fprintf(stderr, + "Requested shared memory (%dB) is larger than the maximum allowed " + "shared memory (%dB) for this device\n", + smem, maxShmem); + CUDA_REPORT_IF_ERROR( + cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem, kernel, device)); +} + +// MLIR's offloading handler loads each GPU binary exactly once per process (a +// module constructor calls mgpuModuleLoad), but a CUmodule/CUfunction pair is +// bound to the CUDA context it was created in -- reusing a device-0 CUfunction +// on device 1 fails with CUDA_ERROR_INVALID_HANDLE and silently wrong results. +// +// CUDA 12's library management API exists for exactly this: a CUlibrary is +// context-independent and the driver instantiates it per context on demand. So +// we hand MLIR a CUlibrary/CUkernel behind its CUmodule/CUfunction handles and +// resolve the context-bound CUfunction at launch. Requires a CUDA 12 / r525+ driver. + +extern "C" CUmodule mgpuModuleLoad(void *data, size_t /*gpuBlobSize*/) { + ScopedContext scopedContext; + CUlibrary library = nullptr; + CUDA_REPORT_IF_ERROR(cuLibraryLoadData(&library, data, nullptr, nullptr, 0, nullptr, nullptr, 0)); + return reinterpret_cast(library); +} + +extern "C" CUmodule mgpuModuleLoadJIT(void *data, int optLevel) { + ScopedContext scopedContext; + CUlibrary library = nullptr; + char jitErrorBuffer[4096] = {0}; + CUjit_option jitOptions[] = {CU_JIT_ERROR_LOG_BUFFER, CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES, + CU_JIT_OPTIMIZATION_LEVEL}; + void *jitOptionsVals[] = {jitErrorBuffer, reinterpret_cast(sizeof(jitErrorBuffer)), + reinterpret_cast(optLevel)}; + CUresult result = + cuLibraryLoadData(&library, data, jitOptions, jitOptionsVals, 3, nullptr, nullptr, 0); + if (result) { + fprintf(stderr, "JIT compilation failed with: '%s'\n", jitErrorBuffer); + CUDA_REPORT_IF_ERROR(result); + } + return reinterpret_cast(library); +} + +extern "C" void mgpuModuleUnload(CUmodule module) { + CUDA_REPORT_IF_ERROR(cuLibraryUnload(reinterpret_cast(module))); +} + +extern "C" CUfunction mgpuModuleGetFunction(CUmodule module, const char *name) { + CUkernel kernel = nullptr; + CUDA_REPORT_IF_ERROR(cuLibraryGetKernel(&kernel, reinterpret_cast(module), name)); + return reinterpret_cast(kernel); +} + +// The wrapper uses intptr_t instead of CUDA's unsigned int to match MLIR's +// index type, avoiding casts in the generated code. +extern "C" void mgpuLaunchKernel(CUfunction function, intptr_t gridX, intptr_t gridY, + intptr_t gridZ, intptr_t blockX, intptr_t blockY, intptr_t blockZ, + int32_t smem, CUstream stream, void **params, void **extra, + size_t /*paramsCount*/) { + ScopedContext scopedContext; + auto kernel = reinterpret_cast(function); + CUfunction fn = nullptr; + CUDA_REPORT_IF_ERROR(cuKernelGetFunction(&fn, kernel)); + if (!fn) + return; + setDynamicSharedMemory(kernel, smem); + CUDA_REPORT_IF_ERROR( + cuLaunchKernel(fn, gridX, gridY, gridZ, blockX, blockY, blockZ, smem, stream, params, extra)); +} + +// Stage-one CUDA backend does not support thread-block clusters (Hopper+). +// The cluster launch path (cuLaunchKernelEx / CUlaunchConfig) is gated behind +// CUDA-version macros and may be unavailable at compile time, so we degrade to +// a plain launch and only warn if a non-trivial cluster was actually requested. +// Real cluster support belongs to a later stage. +extern "C" void mgpuLaunchClusterKernel(CUfunction function, intptr_t clusterX, intptr_t clusterY, + intptr_t clusterZ, intptr_t gridX, intptr_t gridY, + intptr_t gridZ, intptr_t blockX, intptr_t blockY, + intptr_t blockZ, int32_t smem, CUstream stream, + void **params, void **extra, size_t /*paramsCount*/) { + ScopedContext scopedContext; + if ((clusterX > 1) || (clusterY > 1) || (clusterZ > 1)) { + fprintf(stderr, + "[mgpuLaunchClusterKernel] cluster=(%ld,%ld,%ld) requested but the " + "stage-one CUDA backend has no cluster support; falling back to a " + "plain launch.\n", + static_cast(clusterX), static_cast(clusterY), static_cast(clusterZ)); + } + auto kernel = reinterpret_cast(function); + CUfunction fn = nullptr; + CUDA_REPORT_IF_ERROR(cuKernelGetFunction(&fn, kernel)); + if (!fn) + return; + setDynamicSharedMemory(kernel, smem); + CUDA_REPORT_IF_ERROR( + cuLaunchKernel(fn, gridX, gridY, gridZ, blockX, blockY, blockZ, smem, stream, params, extra)); +} + +extern "C" CUstream mgpuStreamCreate() { + ScopedContext scopedContext; + CUstream stream = nullptr; + CUDA_REPORT_IF_ERROR(cuStreamCreate(&stream, CU_STREAM_NON_BLOCKING)); + return stream; +} + +extern "C" void mgpuStreamDestroy(CUstream stream) { + CUDA_REPORT_IF_ERROR(cuStreamDestroy(stream)); +} + +extern "C" void mgpuStreamSynchronize(CUstream stream) { + CUDA_REPORT_IF_ERROR(cuStreamSynchronize(stream)); +} + +extern "C" void mgpuStreamWaitEvent(CUstream stream, CUevent event) { + CUDA_REPORT_IF_ERROR(cuStreamWaitEvent(stream, event, /*flags=*/0)); +} + +extern "C" CUevent mgpuEventCreate() { + ScopedContext scopedContext; + CUevent event = nullptr; + CUDA_REPORT_IF_ERROR(cuEventCreate(&event, CU_EVENT_DISABLE_TIMING)); + return event; +} + +extern "C" void mgpuEventDestroy(CUevent event) { CUDA_REPORT_IF_ERROR(cuEventDestroy(event)); } + +extern "C" void mgpuEventSynchronize(CUevent event) { + CUDA_REPORT_IF_ERROR(cuEventSynchronize(event)); +} + +extern "C" void mgpuEventRecord(CUevent event, CUstream stream) { + CUDA_REPORT_IF_ERROR(cuEventRecord(event, stream)); +} + +extern "C" void *mgpuMemAlloc(uint64_t sizeBytes, CUstream /*stream*/, bool /*isHostShared*/) { + ScopedContext scopedContext; + CUdeviceptr ptr = 0; + if (sizeBytes != 0) + CUDA_REPORT_IF_ERROR(cuMemAlloc(&ptr, sizeBytes)); + return reinterpret_cast(ptr); +} + +extern "C" void mgpuMemFree(void *ptr, CUstream /*stream*/) { + CUDA_REPORT_IF_ERROR(cuMemFree(reinterpret_cast(ptr))); +} + +extern "C" void mgpuMemcpy(void *dst, void *src, size_t sizeBytes, CUstream stream) { + CUDA_REPORT_IF_ERROR(cuMemcpyAsync(reinterpret_cast(dst), + reinterpret_cast(src), sizeBytes, stream)); +} + +extern "C" void mgpuMemset32(void *dst, int value, size_t count, CUstream stream) { + CUDA_REPORT_IF_ERROR(cuMemsetD32Async(reinterpret_cast(dst), value, count, stream)); +} + +extern "C" void mgpuMemset16(void *dst, int shortValue, size_t count, CUstream stream) { + CUDA_REPORT_IF_ERROR( + cuMemsetD16Async(reinterpret_cast(dst), shortValue, count, stream)); +} + +extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) { + ScopedContext scopedContext; + CUDA_REPORT_IF_ERROR(cuMemHostRegister(ptr, sizeBytes, /*flags=*/0)); +} + +extern "C" void mgpuMemHostRegisterMemRef(int64_t rank, StridedMemRefType *descriptor, + int64_t elementSizeBytes) { + int64_t *sizes = descriptor->sizes; + int64_t *strides = sizes + rank; + + int64_t denseStride = 1; + for (int64_t i = rank - 1; i >= 0; --i) { + (void)strides; + denseStride *= sizes[i]; + } + auto sizeBytes = denseStride * elementSizeBytes; + auto *ptr = descriptor->data + descriptor->offset * elementSizeBytes; + mgpuMemHostRegister(ptr, sizeBytes); +} + +extern "C" void mgpuMemHostUnregister(void *ptr) { CUDA_REPORT_IF_ERROR(cuMemHostUnregister(ptr)); } + +extern "C" void mgpuMemHostUnregisterMemRef(int64_t /*rank*/, + StridedMemRefType *descriptor, + int64_t elementSizeBytes) { + auto *ptr = descriptor->data + descriptor->offset * elementSizeBytes; + mgpuMemHostUnregister(ptr); +} + +extern "C" void mgpuSetDefaultDevice(int32_t device) { + defaultDevice = device; + // Retain the new device's primary context eagerly so a bad ordinal is + // reported here rather than at the next launch (mirrors the ROCm fork's + // hipSetDevice). + (void)ScopedContext::getPrimaryContext(device); +} diff --git a/python/flydsl/compiler/backends/cuda.py b/python/flydsl/compiler/backends/cuda.py new file mode 100644 index 000000000..99434285c --- /dev/null +++ b/python/flydsl/compiler/backends/cuda.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +import glob +import os +import re +import shutil +from typing import List, Optional, Tuple + +from ...runtime.device_runtime.cuda import get_cuda_arch +from ...utils import env +from ...utils.logger import log +from .base import BaseBackend, GPUTarget + + +def _cuda_version_key(path: str): + """Sort key extracting the numeric version from a /usr/local/cuda-X.Y path.""" + m = re.search(r"cuda-(\d+)(?:\.(\d+))?", path) + if not m: + return (0, 0) + return (int(m.group(1)), int(m.group(2) or 0)) + + +def _detect_cuda_toolkit() -> Optional[str]: + """Locate a CUDA toolkit (the directory containing ``bin/ptxas``). + + The default ``/usr/bin/ptxas`` may be too old for newer SM targets (e.g. + sm_90 needs CUDA 12+). Resolution order: + + 1. ``CUDA_HOME`` / ``CUDA_PATH`` / ``CUDAToolkit_ROOT`` env vars. + 2. The newest ``/usr/local/cuda-*`` install, then ``/usr/local/cuda``. + 3. The toolkit root inferred from ``ptxas`` on ``PATH``. + + Returns ``None`` when no toolkit is found; callers must handle that + (the PTX dump is skipped, and device-binary codegen falls back to the + serializer's default toolkit, which may be too old). + """ + for var in ("CUDA_HOME", "CUDA_PATH", "CUDAToolkit_ROOT"): + p = os.environ.get(var) + if p and os.path.exists(os.path.join(p, "bin", "ptxas")): + return p + candidates = sorted(glob.glob("/usr/local/cuda-*"), key=_cuda_version_key, reverse=True) + candidates.append("/usr/local/cuda") + for p in candidates: + if os.path.exists(os.path.join(p, "bin", "ptxas")): + return p + + ptxas = shutil.which("ptxas") + if ptxas: + toolkit = os.path.dirname(os.path.dirname(os.path.realpath(ptxas))) + if os.path.exists(os.path.join(toolkit, "bin", "ptxas")): + return toolkit + + log().warning( + "No CUDA toolkit found (checked CUDA_HOME, CUDA_PATH, CUDAToolkit_ROOT, " + "/usr/local/cuda-*, /usr/local/cuda, and PATH). " + "Device-binary codegen will fall back to the default toolkit, which may be " + "too old for the target arch; set CUDA_HOME to a CUDA 12+ install." + ) + return None + + +class CudaBackend(BaseBackend): + """CUDA / NVIDIA compile backend (CUDA driver runtime, NVVM lowering). + + Stage-one backend: lowers Fly to NVVM/PTX via target-neutral Universal + atoms plus initial SM80 mma.sync.aligned / cp.async / ldmatrix atoms. + """ + + @staticmethod + def supports_target(target: GPUTarget) -> bool: + return target.backend == "cuda" + + @staticmethod + def detect_target() -> GPUTarget: + arch = env.compile.arch or get_cuda_arch() + if not arch: + raise RuntimeError( + "No CUDA device found, so the target architecture could not be detected. " + "Set FLYDSL_GPU_ARCH to an sm_* target to compile without a device." + ) + return GPUTarget(backend="cuda", arch=arch, warp_size=32) + + @classmethod + def make_target(cls, arch: str) -> GPUTarget: + return GPUTarget(backend="cuda", arch=arch, warp_size=32) + + @classmethod + def llvm_address_space(cls, address_space) -> int: + """Map an address space to its NVPTX LLVM representation. + + Must stay in sync with ``mapToLLVMAddressSpace`` in + ``lib/Conversion/FlyToNVVM/FlyToNVVM.cpp``. NVVM has no dedicated + register address space; register-backed scratch is an alloca in the + generic (0) address space, unlike AMDGPU's private (5). + """ + from ..._mlir.dialects.fly import AddressSpace + + mapping = { + AddressSpace.Generic: 0, + AddressSpace.Global: 1, + AddressSpace.Shared: 3, + AddressSpace.Register: 0, + } + try: + return mapping[address_space] + except KeyError: + raise ValueError(f"CUDA address space {address_space} does not lower to a bare LLVM pointer") from None + + # -- compile pipeline ------------------------------------------------ + + @staticmethod + def _format_pass_opts(opts: dict) -> str: + return " ".join(f"{k}={v}" for k, v in opts.items()) + + def _pipeline_parts(self, *, compile_hints: dict) -> Tuple[List[str], str]: + chip = self.target.arch # e.g. "sm_90" + + bin_cli_opts = [] + if env.debug.enable_debug_info: + bin_cli_opts.append("-g") + + nvvm_target_opts = { + "O": 3, + "chip": chip, + "fast": "true" if compile_hints.get("fast_fp_math") else "false", + "triple": "nvptx64-nvidia-cuda", + } + + pre_binary_fragments = [ + "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-nvvm", + "canonicalize", + "gpu.module(convert-scf-to-cf,cse," + "convert-gpu-to-nvvm{index-bitwidth=0 use-bare-ptr-memref-call-conv=true})", + ] + binary_prep_fragments = [ + f"nvvm-attach-target{{{self._format_pass_opts(nvvm_target_opts)}}}", + "convert-scf-to-cf", + "convert-cf-to-llvm", + "gpu-to-llvm{use-bare-pointers-for-host=true use-bare-pointers-for-kernels=true}", + "convert-vector-to-llvm", + "convert-arith-to-llvm", + "convert-func-to-llvm", + "reconcile-unrealized-casts", + *( + ["ensure-debug-info-scope-on-llvm-func{emission-kind=LineTablesOnly}"] + if env.debug.enable_debug_info + else [] + ), + ] + toolkit = _detect_cuda_toolkit() + toolkit_opt = f"toolkit={toolkit} " if toolkit else "" + binary_fragment = f'gpu-module-to-binary{{{toolkit_opt}format=fatbin opts="{" ".join(bin_cli_opts)}"}}' + return [*pre_binary_fragments, *binary_prep_fragments], binary_fragment + + def pipeline_fragments(self, *, compile_hints: dict) -> List[str]: + pre_binary_fragments, binary_fragment = self._pipeline_parts(compile_hints=compile_hints) + return [*pre_binary_fragments, binary_fragment] + + def external_binary_pipeline_fragments(self, *, compile_hints: dict) -> Tuple[List[str], str]: + return self._pipeline_parts(compile_hints=compile_hints) + + def gpu_module_targets(self) -> List[str]: + # The NVVM target attribute is injected by the `nvvm-attach-target` pass + # in the pipeline. Returning it here too would attach the target twice + # (the gpu.module would carry two `#nvvm.target` attrs), making + # gpu-module-to-binary serialize the module once per target. Leave it to + # the pass and return empty. + return [] + + # -- cache / fingerprint --------------------------------------------- + + def native_lib_patterns(self) -> List[str]: + return [ + "_mlirDialectsFly*.so", + "libFly*.so", + "libfly_cuda_runtime.so", + "_mlirRegisterEverything*.so", + ] + + def jit_runtime_lib_basenames(self) -> List[str]: + return [ + "libfly_cuda_runtime.so", + "libmlir_c_runner_utils.so", + ] diff --git a/python/flydsl/compiler/jit_function.py b/python/flydsl/compiler/jit_function.py index 1935ad114..b5cfa0ab8 100644 --- a/python/flydsl/compiler/jit_function.py +++ b/python/flydsl/compiler/jit_function.py @@ -7,6 +7,7 @@ import inspect import os import pickle +import re import threading import time import types @@ -659,20 +660,24 @@ def _extract_isa_text(mlir_asm: str) -> str: return "".join(chars) -def _dump_isa(*, dump_dir: Path, ctx: ir.Context, asm: str, verify: bool, stage_name: str = "15_final_isa"): +def _dump_isa( + *, dump_dir: Path, ctx: ir.Context, asm: str, verify: bool, stage_name: str = "15_final_isa", binary_fragment: str +): """Best-effort dump of final GPU ISA/assembly (.s). - Runs ``gpu-module-to-binary{format=isa}`` on a *cloned* module so the - main compilation is not affected. The raw ISA text is extracted from the - MLIR ``assembly = "..."`` attribute and written as a clean ``.s`` file. + Runs the backend's own ``gpu-module-to-binary`` fragment with ``format=isa`` + on a *cloned* module so the main compilation is not affected. The raw ISA + text is extracted from the MLIR ``assembly = "..."`` attribute and written + as a clean ``.s`` file. """ try: mod = ir.Module.parse(asm, context=ctx) di_pass = ( "ensure-debug-info-scope-on-llvm-func{emission-kind=LineTablesOnly}," if env.debug.enable_debug_info else "" ) + isa_fragment = re.sub(r"\bformat=\w+", "format=isa", binary_fragment, count=1) pm = PassManager.parse( - f'builtin.module({di_pass}gpu-module-to-binary{{format=isa opts="{"-g" if env.debug.enable_debug_info else ""}" section= toolkit=}})', + f"builtin.module({di_pass}{isa_fragment})", context=ctx, ) pm.enable_verifier(bool(verify)) @@ -830,12 +835,14 @@ def compile( print(f"[flydsl.compile] dump 00_origin -> {out}") asm_for_isa = None + isa_binary_fragment = None llir = None stage_num_base = 1 dump_fragments = pre_binary_fragments if external_binary else fragments for idx, frag in enumerate(dump_fragments): if frag.strip().startswith("gpu-module-to-binary"): llir = _extract_llvm_ir(module) + isa_binary_fragment = frag.strip() stage_num = stage_num_base + idx stage_name = f"{stage_num:02d}_{_stage_label_from_fragment(frag)}" @@ -882,7 +889,7 @@ def compile( print(f"[flydsl.compile] dump {ll_name} -> {dump_dir / f'{ll_name}.ll'}") next_stage += 1 - if asm_for_isa is not None: + if asm_for_isa is not None and isa_binary_fragment is not None: if not external_binary: isa_stage = f"{next_stage:02d}_final_isa" isa_out = _dump_isa( @@ -891,6 +898,7 @@ def compile( asm=asm_for_isa, verify=env.debug.enable_verifier, stage_name=isa_stage, + binary_fragment=isa_binary_fragment, ) if isa_out is not None: print(f"[flydsl.compile] dump {isa_stage} -> {isa_out}") diff --git a/python/flydsl/expr/__init__.py b/python/flydsl/expr/__init__.py index cb63ee208..6b38c604f 100644 --- a/python/flydsl/expr/__init__.py +++ b/python/flydsl/expr/__init__.py @@ -20,6 +20,7 @@ _BACKEND_MODULES = { "rocdl": ".rocdl", "tdm_ops": ".rocdl.tdm_ops", # deprecated, use .rocdl.tdm_ops instead + "nvvm": ".nvvm", } _LIBRARY_MODULES = { diff --git a/python/flydsl/expr/nvvm/__init__.py b/python/flydsl/expr/nvvm/__init__.py new file mode 100644 index 000000000..242db0bb7 --- /dev/null +++ b/python/flydsl/expr/nvvm/__init__.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""NVVM dialect extension for NVIDIA GPU programming. + +Target-specific tensor-core MMA atom and async/ldmatrix copy atom types for the +NVVM backend, plus the cp.async commit/wait-group barriers. Lazy-loaded from +``flydsl.expr`` so that ``import flydsl.expr`` does not require the FlyNVVM +Python bindings to be present. +""" + +from ..._mlir.dialects.nvvm import cp_async_commit_group as cp_async_commit_group +from ..._mlir.dialects.nvvm import cp_async_wait_group as cp_async_wait_group +from ..meta import dsl_loc_tracing +from .universal import * + + +@dsl_loc_tracing +def commit_group(): + """Commit pending cp.async operations into a group (cp.async.commit_group).""" + return cp_async_commit_group() + + +@dsl_loc_tracing +def wait_group(n): + """Wait until at most n cp.async groups remain in flight (cp.async.wait_group).""" + return cp_async_wait_group(n) diff --git a/python/flydsl/expr/nvvm/universal.py b/python/flydsl/expr/nvvm/universal.py new file mode 100644 index 000000000..2dc491458 --- /dev/null +++ b/python/flydsl/expr/nvvm/universal.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +from ..._mlir.dialects.fly_nvvm import ( + CopyOpSM75_LdMatrixType, + CopyOpSM80_CpAsyncType, + MmaOpSM80_MmaSyncType, +) +from ..._mlir.extras import types as T + + +def MmaSync(m, n, k, elem_ty_ab, elem_ty_acc=None): + """Create an SM80 ``mma.sync.aligned`` MMA op type. + + Maps to PTX ``mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32``. Only the + 16x8x16 f16->f32 shape is currently supported. + """ + ty_ab = elem_ty_ab.ir_type if hasattr(elem_ty_ab, "ir_type") else elem_ty_ab + if elem_ty_acc is None: + ty_acc = T.f32() + else: + ty_acc = elem_ty_acc.ir_type if hasattr(elem_ty_acc, "ir_type") else elem_ty_acc + return MmaOpSM80_MmaSyncType.get(m, n, k, ty_ab, ty_ab, ty_acc) + + +def CpAsync(bit_size): + """Create an SM80 cp.async.shared.global copy atom (global -> shared, async).""" + return CopyOpSM80_CpAsyncType.get(bit_size) + + +CpAsync32b = lambda: CopyOpSM80_CpAsyncType.get(32) +CpAsync64b = lambda: CopyOpSM80_CpAsyncType.get(64) +CpAsync128b = lambda: CopyOpSM80_CpAsyncType.get(128) + + +def LdMatrix(num, trans=False): + """Create an SM75+ ldmatrix copy atom (shared -> register), 8x8 b16 tiles.""" + return CopyOpSM75_LdMatrixType.get(num, trans) + + +LdMatrixX1 = lambda trans=False: CopyOpSM75_LdMatrixType.get(1, trans) +LdMatrixX2 = lambda trans=False: CopyOpSM75_LdMatrixType.get(2, trans) +LdMatrixX4 = lambda trans=False: CopyOpSM75_LdMatrixType.get(4, trans) diff --git a/python/flydsl/runtime/device_runtime/__init__.py b/python/flydsl/runtime/device_runtime/__init__.py index 9babdac9c..3281c3b7b 100644 --- a/python/flydsl/runtime/device_runtime/__init__.py +++ b/python/flydsl/runtime/device_runtime/__init__.py @@ -20,17 +20,20 @@ from ...utils import env from .base import DeviceRuntime +from .cuda import CudaDeviceRuntime from .rocm import RocmDeviceRuntime # Compile-backend id -> device-runtime kind (single string namespace). COMPILE_BACKEND_TO_RUNTIME_KIND: Dict[str, str] = { "rocm": "rocm", + "cuda": "cuda", } _EXTRA_MAPPINGS: Dict[str, str] = {} _builtin_runtimes: Dict[str, Type[DeviceRuntime]] = { "rocm": RocmDeviceRuntime, + "cuda": CudaDeviceRuntime, } _runtime_cls_override: Optional[Type[DeviceRuntime]] = None @@ -163,6 +166,7 @@ def get_device_runtime() -> DeviceRuntime: __all__ = [ "COMPILE_BACKEND_TO_RUNTIME_KIND", + "CudaDeviceRuntime", "DeviceRuntime", "RocmDeviceRuntime", "ensure_compile_runtime_compatible", diff --git a/python/flydsl/runtime/device_runtime/cuda.py b/python/flydsl/runtime/device_runtime/cuda.py new file mode 100644 index 000000000..9c1c28c91 --- /dev/null +++ b/python/flydsl/runtime/device_runtime/cuda.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""CUDA driver-API device runtime (NVIDIA GPU stack).""" + +from __future__ import annotations + +import ctypes +import functools +import os +from typing import ClassVar, Optional + +from .base import DeviceRuntime + +# Cached CUDA driver handle (``libcuda``); cached once. +_CUDA_LIB = None +_CUDA_LIB_TRIED = False + +CUDA_SUCCESS = 0 + + +def _load_cuda(): + global _CUDA_LIB, _CUDA_LIB_TRIED + if not _CUDA_LIB_TRIED: + _CUDA_LIB_TRIED = True + for soname in ("libcuda.so", "libcuda.so.1"): + try: + _CUDA_LIB = ctypes.CDLL(soname) + break + except OSError: + continue + return _CUDA_LIB + + +@functools.lru_cache(maxsize=1) +def _driver() -> Optional[ctypes.CDLL]: + """Return an initialized ``libcuda`` handle, or None when unavailable. + + Everything here goes through the driver rather than ``nvidia-smi`` on + purpose: the driver enumerates *visible* devices, so ``CUDA_VISIBLE_DEVICES`` + (and its ordering / UUID forms) is honored automatically and agrees with + what PyTorch sees. ``nvidia-smi`` reports physical devices and would + disagree with the runtime. + """ + lib = _load_cuda() + if lib is None: + return None + try: + if lib.cuInit(0) != CUDA_SUCCESS: + return None + except Exception: + return None + return lib + + +@functools.lru_cache(maxsize=1) +def get_cuda_device_count() -> int: + """Number of CUDA devices visible to this process. 0 when unavailable.""" + lib = _driver() + if lib is None: + return 0 + count = ctypes.c_int(0) + try: + if lib.cuDeviceGetCount(ctypes.byref(count)) != CUDA_SUCCESS: + return 0 + except Exception: + return 0 + return int(count.value) + + +def _cuda_current_device() -> int: + """Active CUDA device index via ``cuCtxGetDevice`` (falls back to 0).""" + lib = _driver() + if lib is None: + return 0 + try: + dev = ctypes.c_int(0) + if lib.cuCtxGetDevice(ctypes.byref(dev)) == CUDA_SUCCESS: + return int(dev.value) + except Exception: + pass + return 0 + + +def _compute_capability(ordinal: int) -> Optional[str]: + """Compute capability of visible device ``ordinal`` as ``sm_XX``.""" + lib = _driver() + if lib is None: + return None + try: + dev = ctypes.c_int(0) + if lib.cuDeviceGet(ctypes.byref(dev), ordinal) != CUDA_SUCCESS: + return None + major, minor = ctypes.c_int(0), ctypes.c_int(0) + # CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_{MAJOR,MINOR} + if lib.cuDeviceGetAttribute(ctypes.byref(major), 75, dev) != CUDA_SUCCESS: + return None + if lib.cuDeviceGetAttribute(ctypes.byref(minor), 76, dev) != CUDA_SUCCESS: + return None + except Exception: + return None + return f"sm_{major.value}{minor.value}" + + +@functools.lru_cache(maxsize=None) +def _cuda_arch_from_hardware(ordinal: int) -> str: + """Cached compute capability of a visible device.""" + return _compute_capability(ordinal) + + +def get_cuda_arch() -> str: + """Best-effort CUDA GPU arch string (e.g. ``'sm_90'``). + + Honors ``ARCH`` / ``FLYDSL_GPU_ARCH`` when they name an ``sm_*`` target, + otherwise reports the capability of the *currently selected* device. + """ + env = os.environ.get("ARCH") or os.environ.get("FLYDSL_GPU_ARCH") + if env and env.startswith("sm_"): + return env + return _cuda_arch_from_hardware(_cuda_current_device()) + + +class CudaDeviceRuntime(DeviceRuntime): + """CUDA driver-API runtime; matches compile backend ``cuda``. + + Both ``device_count()`` and ``current_device_id()`` query the CUDA driver, + so they agree with ``CUDA_VISIBLE_DEVICES`` and with PyTorch. + """ + + kind: ClassVar[str] = "cuda" + + def device_count(self) -> int: + return get_cuda_device_count() + + def current_device_id(self) -> int: + return _cuda_current_device() diff --git a/python/mlir_flydsl/CMakeLists.txt b/python/mlir_flydsl/CMakeLists.txt index 3d6456f37..6c7198c41 100644 --- a/python/mlir_flydsl/CMakeLists.txt +++ b/python/mlir_flydsl/CMakeLists.txt @@ -27,6 +27,17 @@ if(FLYDSL_HAS_ROCDL) ) endif() +if(FLYDSL_HAS_NVVM) + declare_mlir_dialect_python_bindings( + ADD_TO_PARENT FlyPythonSources + ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/" + TD_FILE dialects/FlyNVVM.td + SOURCES + dialects/fly_nvvm.py + DIALECT_NAME fly_nvvm + ) +endif() + # NOTE: Do NOT link MLIRFlyDialect/MLIRFlyROCDLDialect here via PRIVATE_LINK_LIBS. # These symbols are already provided by FlyPythonCAPI.so (via MLIRCPIFly's transitive # dependencies in EMBED_CAPI_LINK_LIBS). Statically linking them here creates DUPLICATE @@ -59,6 +70,19 @@ if(FLYDSL_HAS_ROCDL) ) endif() +if(FLYDSL_HAS_NVVM) + declare_mlir_python_extension(FlyPythonSources.Core.fly_nvvm + MODULE_NAME _mlirDialectsFlyNVVM + ADD_TO_PARENT FlyPythonSources + ROOT_DIR "${PROJECT_SOURCE_DIR}/lib/Bindings/Python" + PYTHON_BINDINGS_LIBRARY nanobind + SOURCES + FlyNVVMExtension.cpp + PRIVATE_LINK_LIBS + LLVMSupport + ) +endif() + # NOTE: Do NOT link MLIRFlyToROCDL or other C++ libs via PRIVATE_LINK_LIBS. # Doing so statically embeds a copy of MLIRPass (and its pass registry) into # _mlirRegisterEverything.so. Passes registered via registerFlyPasses() then @@ -236,6 +260,16 @@ else() set(_FLY_COPY_ROCDL_TABLEGEN "") endif() +if(FLYDSL_HAS_NVVM) + set(_FLY_COPY_NVVM_TABLEGEN + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${MLIR_BINARY_DIR}/python/mlir_flydsl/dialects/_fly_nvvm_ops_gen.py" + "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/dialects/_fly_nvvm_ops_gen.py" + ) +else() + set(_FLY_COPY_NVVM_TABLEGEN "") +endif() + add_custom_target(CopyFlyPythonSources ALL COMMAND ${CMAKE_COMMAND} -E copy_directory "${PROJECT_SOURCE_DIR}/python/flydsl" @@ -248,6 +282,7 @@ add_custom_target(CopyFlyPythonSources ALL "${MLIR_BINARY_DIR}/python/mlir_flydsl/dialects/_fly_enum_gen.py" "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/dialects/_fly_enum_gen.py" ${_FLY_COPY_ROCDL_TABLEGEN} + ${_FLY_COPY_NVVM_TABLEGEN} COMMAND ${CMAKE_COMMAND} -E copy_if_different "$" "${_MLIR_LIBS_DIR}/libmlir_c_runner_utils.so" @@ -270,3 +305,8 @@ if(TARGET FlyJitRuntime) LIBRARY_OUTPUT_DIRECTORY "${_MLIR_LIBS_DIR}") add_dependencies(FlyPythonCAPI FlyJitRuntime) endif() +if(TARGET FlyCudaJitRuntime) + set_target_properties(FlyCudaJitRuntime PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${_MLIR_LIBS_DIR}") + add_dependencies(FlyPythonCAPI FlyCudaJitRuntime) +endif() diff --git a/python/mlir_flydsl/dialects/FlyNVVM.td b/python/mlir_flydsl/dialects/FlyNVVM.td new file mode 100644 index 000000000..cad225ea0 --- /dev/null +++ b/python/mlir_flydsl/dialects/FlyNVVM.td @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#ifndef PYTHON_BINDINGS_FLYNVVM_OPS +#define PYTHON_BINDINGS_FLYNVVM_OPS + +include "flydsl/Dialect/FlyNVVM/IR/Atom.td" + +#endif // PYTHON_BINDINGS_FLYNVVM_OPS diff --git a/python/mlir_flydsl/dialects/fly_nvvm.py b/python/mlir_flydsl/dialects/fly_nvvm.py new file mode 100644 index 000000000..3fb337afc --- /dev/null +++ b/python/mlir_flydsl/dialects/fly_nvvm.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +# isort: skip_file +# ruff: noqa: F401,F403 +from ._fly_nvvm_ops_gen import * +from ._fly_nvvm_ops_gen import _Dialect + +from .._mlir_libs._mlirDialectsFlyNVVM import * diff --git a/scripts/build.sh b/scripts/build.sh index 5ab8921e3..b1a49c2e2 100755 --- a/scripts/build.sh +++ b/scripts/build.sh @@ -116,6 +116,13 @@ cmake_args=( -DPython3_EXECUTABLE="$(which python3)" -DHIP_PLATFORM="${HIP_PLATFORM}" ) +# Enabled backend stacks (semicolon-separated), e.g. "rocdl", "rocdl;nvvm", +# "nvvm". Defaults to the CMake default (rocdl) when unset. Override with: +# FLYDSL_BACKENDS="rocdl;nvvm" bash scripts/build.sh +if [ -n "${FLYDSL_BACKENDS:-}" ]; then + cmake_args+=(-DFLYDSL_BACKENDS="${FLYDSL_BACKENDS}") + echo " FLYDSL_BACKENDS: ${FLYDSL_BACKENDS}" +fi if [ -n "${NANOBIND_DIR}" ]; then cmake_args+=(-Dnanobind_DIR="${NANOBIND_DIR}") fi diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh index 1a417826f..7741c6032 100644 --- a/scripts/run_tests.sh +++ b/scripts/run_tests.sh @@ -39,41 +39,66 @@ if [[ ":${LD_LIBRARY_PATH:-}:" != *":${MLIR_LIBS_DIR}:"* ]]; then export LD_LIBRARY_PATH="${MLIR_LIBS_DIR}:${LD_LIBRARY_PATH:-}" fi +# Compile backend selects which example directory and pytest suites apply. +_compile_backend="${FLYDSL_COMPILE_BACKEND:-rocm}" +_compile_backend="${_compile_backend,,}" + pytest_args=(-v --no-header --tb=short) +_marker_expr="" if [ "${RUN_TESTS_FULL:-0}" != "1" ]; then - pytest_args+=(-m "not large_shape") + _marker_expr="not large_shape" +fi +if [ "${_compile_backend}" == "cuda" ]; then + # tests/kernels, tests/system and the AOT suite are ROCm kernels and + # rocm_lower tests; only the backend-agnostic tier is meaningful here. + pytest_paths=(tests/language/ tests/unit/) + _marker_expr="l0_backend_agnostic${_marker_expr:+ and ${_marker_expr}}" +else + pytest_paths=(tests/kernels/ tests/language/ tests/unit/ tests/system/ tests/python/examples/) +fi +if [ -n "${_marker_expr}" ]; then + pytest_args+=(-m "${_marker_expr}") fi # --------------------------------------------------------------------------- # 1. All pytest-based tests (kernels + language + unit + system + examples) # --------------------------------------------------------------------------- echo "========================================================================" -echo "Pytest: kernels + language + unit + system + examples" +echo "Pytest: ${pytest_paths[*]}" echo "========================================================================" -python3 -m pytest \ - tests/kernels/ \ - tests/language/ \ - tests/unit/ \ - tests/system/ \ - tests/python/examples/ \ - "${pytest_args[@]}" +python3 -m pytest "${pytest_paths[@]}" "${pytest_args[@]}" # --------------------------------------------------------------------------- # 2. Standalone example scripts (not pytest) +# +# examples/*.py -> target-neutral examples, run on every compile backend +# examples/rocm/*.py -> ROCm/HIP examples, run on the rocm compile backend +# examples/cuda/*.py -> NVVM/CUDA examples, run on the cuda compile backend +# examples/cuda/bench -> developer benchmark harnesses, never run here (they +# need an external CUTLASS checkout and take minutes) # --------------------------------------------------------------------------- +if [ "${_compile_backend}" == "cuda" ]; then + _example_subdir="cuda" +else + _example_subdir="rocm" +fi +_example_dirs=("${REPO_ROOT}/examples" "${REPO_ROOT}/examples/${_example_subdir}") + echo "" echo "========================================================================" -echo "Examples (examples/)" +echo "Examples (examples/ + examples/${_example_subdir}/)" echo "========================================================================" # Whitelist from tests/arch_compat.py (single source of truth for arch compat). _RDNA_EXAMPLE_WHITELIST=$(python3 -c "from tests.arch_compat import RDNA_COMPATIBLE_EXAMPLES; print(' '.join(RDNA_COMPATIBLE_EXAMPLES))" 2>/dev/null || echo "") _gpu_arch=$(python3 -c "from flydsl.runtime.device import get_rocm_arch; print(get_rocm_arch())" 2>/dev/null || echo "unknown") -for example in "${REPO_ROOT}"/examples/*.py; do + +for _dir in "${_example_dirs[@]}"; do +for example in "${_dir}"/*.py; do [ -f "${example}" ] || continue - name="$(basename "${example}")" - if [[ "${_gpu_arch}" != gfx9* ]] && ! echo "${_RDNA_EXAMPLE_WHITELIST}" | grep -qw "${name}"; then + name="${example#${REPO_ROOT}/examples/}" + if [[ "${_compile_backend}" != "cuda" && "${_gpu_arch}" != gfx9* ]] && ! echo "${_RDNA_EXAMPLE_WHITELIST}" | grep -qw "$(basename "${example}")"; then echo " SKIP ${name} (not in RDNA whitelist, arch: ${_gpu_arch})" continue fi @@ -85,6 +110,7 @@ for example in "${REPO_ROOT}"/examples/*.py; do fi echo " PASS ${name}" done +done # --------------------------------------------------------------------------- # 3. MLIR FileCheck tests @@ -96,9 +122,12 @@ echo "========================================================================" FLY_OPT="${BUILD_DIR}/bin/fly-opt" FILECHECK="" +_enabled_backends="${FLYDSL_BACKENDS:-rocdl}" if [ -f "${BUILD_DIR}/CMakeCache.txt" ]; then _mlir_dir=$(grep '^MLIR_DIR:' "${BUILD_DIR}/CMakeCache.txt" | sed 's|^MLIR_DIR:[A-Z]*=||') [ -n "${_mlir_dir}" ] && FILECHECK="${_mlir_dir}/../../../bin/FileCheck" + _cache_backends=$(grep '^FLYDSL_BACKENDS:' "${BUILD_DIR}/CMakeCache.txt" | sed 's|^FLYDSL_BACKENDS:[A-Z]*=||' || true) + [ -n "${_cache_backends}" ] && _enabled_backends="${_cache_backends}" fi [ -z "${FILECHECK}" ] || [ ! -x "${FILECHECK}" ] && FILECHECK="$(which FileCheck 2>/dev/null || true)" @@ -106,9 +135,24 @@ if [ -z "${FILECHECK}" ] || [ ! -x "${FILECHECK}" ]; then echo " SKIP FileCheck not found; skipping MLIR lit tests." else +backend_enabled() { + [[ ";${_enabled_backends};" == *";$1;"* ]] +} + +# A test needs a backend if it runs that backend's conversion pass OR mentions +# its dialect anywhere: target atom types (!fly_rocdl.*, !fly_nvvm.*) fail to +# parse under target-neutral passes too when the dialect is not registered. for f in $(find "${REPO_ROOT}/tests/mlir" -name "*.mlir" -type f 2>/dev/null | sort); do run_line=$(grep '^// RUN:' "$f" | head -1 | sed 's|^// RUN: *||') [ -z "$run_line" ] && continue + if grep -q 'fly_nvvm\|convert-fly-to-nvvm' "$f" && ! backend_enabled nvvm; then + echo " SKIP ${f#${REPO_ROOT}/tests/mlir/} (nvvm backend not enabled)" + continue + fi + if grep -q 'fly_rocdl\|convert-fly-to-rocdl' "$f" && ! backend_enabled rocdl; then + echo " SKIP ${f#${REPO_ROOT}/tests/mlir/} (rocdl backend not enabled)" + continue + fi cmd=$(echo "$run_line" | sed "s|%fly-opt|${FLY_OPT}|g; s|%FileCheck|${FILECHECK}|g; s|%s|${f}|g; s|FileCheck|${FILECHECK}|g") if eval "$cmd" > /tmp/filecheck_out.log 2>&1; then echo " PASS ${f#${REPO_ROOT}/tests/mlir/}" diff --git a/tests/mlir/Conversion/fly-to-nvvm/copy_atom.mlir b/tests/mlir/Conversion/fly-to-nvvm/copy_atom.mlir new file mode 100644 index 000000000..3d9c8f590 --- /dev/null +++ b/tests/mlir/Conversion/fly-to-nvvm/copy_atom.mlir @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors +// RUN: %fly-opt %s --convert-fly-to-nvvm | FileCheck %s + +// NVVM SM80 copy atom lowering: +// cp.async (global -> shared) -> nvvm.cp.async.shared.global (cg, 16B) +// ldmatrix (shared -> register) -> nvvm.ldmatrix m8n8 x4 b16 + +// CHECK-LABEL: @test_cp_async +func.func @test_cp_async(%s: !fly.memref, %d: !fly.memref) { + %atom = fly.make_copy_atom {valBits = 16 : i32} : !fly.copy_atom, 16> + // CHECK: nvvm.cp.async.shared.global {{.*}}, {{.*}}, 16, cache = cg : !llvm.ptr<3>, !llvm.ptr<1> + fly.copy_atom_call(%atom, %s, %d) : (!fly.copy_atom, 16>, !fly.memref, !fly.memref) -> () + return +} + +// CHECK-LABEL: @test_ldmatrix_x4 +func.func @test_ldmatrix_x4(%s: !fly.memref, %d: !fly.memref) { + %atom = fly.make_copy_atom {valBits = 16 : i32} : !fly.copy_atom, 16> + // CHECK: nvvm.ldmatrix {{.*}} {eltType = #nvvm.ld_st_matrix_elt_type, layout = #nvvm.mma_layout, num = 4 : i32, shape = #nvvm.ld_st_matrix_shape} : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)> + fly.copy_atom_call(%atom, %s, %d) : (!fly.copy_atom, 16>, !fly.memref, !fly.memref) -> () + return +} diff --git a/tests/mlir/Conversion/fly-to-nvvm/mma_atom.mlir b/tests/mlir/Conversion/fly-to-nvvm/mma_atom.mlir new file mode 100644 index 000000000..1f308de13 --- /dev/null +++ b/tests/mlir/Conversion/fly-to-nvvm/mma_atom.mlir @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors +// RUN: %fly-opt %s --convert-fly-to-nvvm | FileCheck %s + +// NVVM SM80 mma.sync.aligned atom call lowering: +// fly.mma_atom_call_ssa -> nvvm.mma.sync (m16n8k16, f16 -> f32) +// A is unpacked into 4 x vector<2xf16>, B into 2 x vector<2xf16>, C into 4 x f32, +// and the result struct is repacked into vector<4xf32>. + +// CHECK-LABEL: @test_mma_sync_aligned_ssa +// CHECK-SAME: (%[[A:.*]]: vector<8xf16>, %[[B:.*]]: vector<4xf16>, %[[C:.*]]: vector<4xf32>) +func.func @test_mma_sync_aligned_ssa( + %a: vector<8xf16>, + %b: vector<4xf16>, + %c: vector<4xf32>) -> vector<4xf32> { + %atom = fly.make_mma_atom : !fly.mma_atom f32>> + // CHECK: llvm.shufflevector %[[A]], %[[A]] [0, 1] + // CHECK: llvm.shufflevector %[[A]], %[[A]] [2, 3] + // CHECK: llvm.shufflevector %[[A]], %[[A]] [4, 5] + // CHECK: llvm.shufflevector %[[A]], %[[A]] [6, 7] + // CHECK: llvm.shufflevector %[[B]], %[[B]] [0, 1] + // CHECK: llvm.shufflevector %[[B]], %[[B]] [2, 3] + // CHECK: nvvm.mma.sync + // CHECK-SAME: shape = #nvvm.shape + // CHECK-SAME: -> !llvm.struct<(f32, f32, f32, f32)> + %res = fly.mma_atom_call_ssa(%atom, %a, %b, %c) : (!fly.mma_atom f32>>, vector<8xf16>, vector<4xf16>, vector<4xf32>) -> vector<4xf32> + return %res : vector<4xf32> +} diff --git a/tests/unit/test_backend_cmake_defaults.py b/tests/unit/test_backend_cmake_defaults.py index b3fd9d0d1..932767852 100644 --- a/tests/unit/test_backend_cmake_defaults.py +++ b/tests/unit/test_backend_cmake_defaults.py @@ -18,8 +18,8 @@ def test_cmake_default_backend_stays_rocdl(): text = (_REPO_ROOT / "cmake" / "FlyDSLBackends.cmake").read_text() assert 'set(FLYDSL_BACKENDS "rocdl"' in text - assert "set_property(CACHE FLYDSL_BACKENDS PROPERTY STRINGS rocdl)" in text - assert "set(_FLYDSL_BACKENDS_ALLOWED rocdl)" in text + assert "set_property(CACHE FLYDSL_BACKENDS PROPERTY STRINGS rocdl nvvm)" in text + assert "set(_FLYDSL_BACKENDS_ALLOWED rocdl nvvm)" in text def test_rocm_runtime_is_only_added_for_rocdl_backend(): @@ -50,12 +50,12 @@ def test_future_backend_descriptor_is_opt_in(tmp_path): text = (_REPO_ROOT / "cmake" / "FlyDSLBackends.cmake").read_text() text = text.replace( - "set_property(CACHE FLYDSL_BACKENDS PROPERTY STRINGS rocdl)", - "set_property(CACHE FLYDSL_BACKENDS PROPERTY STRINGS rocdl dummy)", + "set_property(CACHE FLYDSL_BACKENDS PROPERTY STRINGS rocdl nvvm)", + "set_property(CACHE FLYDSL_BACKENDS PROPERTY STRINGS rocdl nvvm dummy)", ) text = text.replace( - "set(_FLYDSL_BACKENDS_ALLOWED rocdl)", - "set(_FLYDSL_BACKENDS_ALLOWED rocdl dummy)", + "set(_FLYDSL_BACKENDS_ALLOWED rocdl nvvm)", + "set(_FLYDSL_BACKENDS_ALLOWED rocdl nvvm dummy)", ) (cmake_dir / "FlyDSLBackends.cmake").write_text(text) diff --git a/tests/unit/test_cuda_backend.py b/tests/unit/test_cuda_backend.py new file mode 100644 index 000000000..e425145c5 --- /dev/null +++ b/tests/unit/test_cuda_backend.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""Backend-agnostic checks for the CUDA/NVVM compile backend. + +These exercise the pure-Python surface (registry, target detection, pipeline +shape, address-space mapping) so the NVVM backend keeps regression coverage on +CI runners that have no NVIDIA GPU and no ``nvvm`` in ``FLYDSL_BACKENDS``. +Anything needing the FlyNVVM bindings or a device belongs elsewhere. +""" + +import pytest + +from flydsl.compiler.backends import get_backend +from flydsl.compiler.backends.cuda import CudaBackend +from flydsl.runtime.device_runtime.cuda import get_cuda_arch + +pytestmark = pytest.mark.l0_backend_agnostic + + +def test_backend_is_discovered_under_the_cuda_name(): + backend = get_backend("cuda", arch="sm_80") + assert isinstance(backend, CudaBackend) + assert backend.target.backend == "cuda" + assert backend.target.arch == "sm_80" + assert backend.target.warp_size == 32 + + +def test_supports_target_is_exclusive(): + assert CudaBackend.supports_target(CudaBackend.make_target("sm_90")) + assert not CudaBackend.supports_target(get_backend("rocm", arch="gfx942").target) + + +@pytest.mark.parametrize("arch", ["sm_80", "sm_90"]) +def test_pipeline_carries_the_nvvm_stages_and_chip(arch): + fragments = CudaBackend(CudaBackend.make_target(arch)).pipeline_fragments(compile_hints={}) + joined = "\n".join(fragments) + assert "convert-fly-to-nvvm" in joined + assert "convert-gpu-to-nvvm" in joined + assert f"chip={arch}" in joined + # The target attribute is attached by nvvm-attach-target, never twice. + assert "nvvm-attach-target" in joined + assert CudaBackend(CudaBackend.make_target(arch)).gpu_module_targets() == [] + # The device binary must be the last fragment. + assert fragments[-1].startswith("gpu-module-to-binary") + + +def test_external_binary_split_matches_the_full_pipeline(): + backend = CudaBackend(CudaBackend.make_target("sm_90")) + pre, binary = backend.external_binary_pipeline_fragments(compile_hints={}) + assert [*pre, binary] == backend.pipeline_fragments(compile_hints={}) + + +def test_fast_fp_math_hint_reaches_the_nvvm_target(): + backend = CudaBackend(CudaBackend.make_target("sm_90")) + assert "fast=true" in "\n".join(backend.pipeline_fragments(compile_hints={"fast_fp_math": True})) + assert "fast=false" in "\n".join(backend.pipeline_fragments(compile_hints={})) + + +def test_llvm_address_space_matches_the_nvvm_conversion(): + # Must stay in sync with mapToLLVMAddressSpace in FlyToNVVM.cpp. NVVM has no + # dedicated register address space, so Register lowers to generic (0). + from flydsl._mlir.dialects.fly import AddressSpace + + assert CudaBackend.llvm_address_space(AddressSpace.Generic) == 0 + assert CudaBackend.llvm_address_space(AddressSpace.Global) == 1 + assert CudaBackend.llvm_address_space(AddressSpace.Shared) == 3 + assert CudaBackend.llvm_address_space(AddressSpace.Register) == 0 + + +def test_jit_runtime_libs_name_the_cuda_wrapper(): + backend = CudaBackend(CudaBackend.make_target("sm_90")) + assert "libfly_cuda_runtime.so" in backend.jit_runtime_lib_basenames() + assert "libfly_cuda_runtime.so" in backend.native_lib_patterns() + + +@pytest.mark.parametrize("value", ["sm_80", "sm_90", "sm_100"]) +def test_get_cuda_arch_honors_the_arch_override(monkeypatch, value): + monkeypatch.setenv("ARCH", value) + monkeypatch.delenv("FLYDSL_GPU_ARCH", raising=False) + assert get_cuda_arch() == value + + +def test_get_cuda_arch_ignores_a_non_sm_override(monkeypatch): + """A gfx arch left over from a ROCm session must not leak into sm_* land. + + Detection falls through to the driver, which reports ``sm_XX`` when a device + is visible and ``None`` when there is none (as on a ROCm CI runner) -- never + the gfx value. + """ + monkeypatch.setenv("ARCH", "gfx942") + monkeypatch.delenv("FLYDSL_GPU_ARCH", raising=False) + arch = get_cuda_arch() + assert arch is None or arch.startswith("sm_") + + +def test_cuda_runtime_kind_is_registered(): + from flydsl.runtime.device_runtime import COMPILE_BACKEND_TO_RUNTIME_KIND, CudaDeviceRuntime + + assert COMPILE_BACKEND_TO_RUNTIME_KIND["cuda"] == "cuda" + assert CudaDeviceRuntime.kind == "cuda" diff --git a/tests/unit/test_gfx1250_atoms.py b/tests/unit/test_gfx1250_atoms.py index 2016065df..ca01e1a9b 100644 --- a/tests/unit/test_gfx1250_atoms.py +++ b/tests/unit/test_gfx1250_atoms.py @@ -13,7 +13,10 @@ import pytest -pytestmark = [pytest.mark.l0_backend_agnostic] +# No GPU needed, but these do construct vendor target-dialect types, so this is +# l1b rather than l0: the FlyROCDL bindings only exist in a build that includes +# the ROCDL backend. +pytestmark = [pytest.mark.l1b_target_dialect, pytest.mark.rocm_lower] from flydsl._mlir import ir # noqa: E402 diff --git a/tests/unit/test_tdm_mcast_add_gfx1250.py b/tests/unit/test_tdm_mcast_add_gfx1250.py index f4fbd4359..42ebda519 100644 --- a/tests/unit/test_tdm_mcast_add_gfx1250.py +++ b/tests/unit/test_tdm_mcast_add_gfx1250.py @@ -16,16 +16,25 @@ import pytest -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl._mlir import ir -from flydsl._mlir.dialects import vector -from flydsl.compiler.kernel_function import CompilationContext -from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, tdm_ops -from flydsl.expr.rocdl import cluster -from flydsl.expr.typing import T -from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr, get_op_result_or_value -from kernels.common.gfx1250_cluster import compute_mcast_masks +# TDM / cluster multicast is gfx1250-only, so this module imports the ROCDL +# backend package unconditionally. Skip the whole module on a build without the +# ROCDL bindings (e.g. FLYDSL_BACKENDS=nvvm), otherwise the imports below fail +# at collection time, before the arch guard further down can run. +pytest.importorskip( + "flydsl._mlir.dialects.rocdl", + reason="requires a build with the ROCDL backend (FLYDSL_BACKENDS=rocdl)", +) + +import flydsl.compiler as flyc # noqa: E402 +import flydsl.expr as fx # noqa: E402 +from flydsl._mlir import ir # noqa: E402 +from flydsl._mlir.dialects import vector # noqa: E402 +from flydsl.compiler.kernel_function import CompilationContext # noqa: E402 +from flydsl.expr import arith, as_ir_value, const_expr, gpu, range_constexpr, tdm_ops # noqa: E402 +from flydsl.expr.rocdl import cluster # noqa: E402 +from flydsl.expr.typing import T # noqa: E402 +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr, get_op_result_or_value # noqa: E402 +from kernels.common.gfx1250_cluster import compute_mcast_masks # noqa: E402 try: import torch diff --git a/tests/unit/test_universal_atomic.py b/tests/unit/test_universal_atomic.py index b612a7d8c..d1c53548a 100644 --- a/tests/unit/test_universal_atomic.py +++ b/tests/unit/test_universal_atomic.py @@ -20,6 +20,21 @@ if torch is None or not torch.cuda.is_available(): pytest.skip("CUDA/ROCm not available", allow_module_level=True) +# UniversalAtomic itself is target-neutral, but the Agent/Workgroup/Wavefront +# sync scopes are AMD-specific and ``fx.rocdl`` only imports in a build that +# includes the ROCDL backend. These feed ``parametrize``, which is evaluated at +# collection time, so gate them here rather than skipping inside the test. +try: + _ROCDL_SYNC_SCOPES = [ + fx.rocdl.SyncScope.Agent, + fx.rocdl.SyncScope.Workgroup, + fx.rocdl.SyncScope.Wavefront, + ] +except ImportError: + _ROCDL_SYNC_SCOPES = [] + +SYNC_SCOPES = [fx.SyncScope.System, fx.SyncScope.SingleThread, *_ROCDL_SYNC_SCOPES] + @flyc.kernel def reduce_add_kernel( @@ -68,16 +83,7 @@ def reduce_add( ) -@pytest.mark.parametrize( - "syncscope", - [ - fx.SyncScope.System, - fx.SyncScope.SingleThread, - fx.rocdl.SyncScope.Agent, - fx.rocdl.SyncScope.Workgroup, - fx.rocdl.SyncScope.Wavefront, - ], -) +@pytest.mark.parametrize("syncscope", SYNC_SCOPES) def test_reduce_add_atomic(syncscope): BLOCK_DIM = 64 N = BLOCK_DIM * 4 @@ -157,13 +163,7 @@ def test_reduce_max_atomic(): if __name__ == "__main__": - for _scope in [ - fx.SyncScope.System, - fx.SyncScope.SingleThread, - fx.rocdl.SyncScope.Agent, - fx.rocdl.SyncScope.Workgroup, - fx.rocdl.SyncScope.Wavefront, - ]: + for _scope in SYNC_SCOPES: test_reduce_add_atomic(_scope) test_reduce_max_atomic() print("ALL PASSED") From feb45ce0cd3a2d45d48fec6a7f09776d2e51a695 Mon Sep 17 00:00:00 2001 From: Feng Shijie Date: Fri, 7 Aug 2026 15:55:15 +0800 Subject: [PATCH 2/2] fix --- README.md | 5 +++-- lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp | 24 +++++++++++++++------ python/flydsl/expr/__init__.py | 2 +- 3 files changed, 22 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 761bf3e5d..ab7129054 100644 --- a/README.md +++ b/README.md @@ -52,10 +52,11 @@ FlyDSL/ │ ├── 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 +│ │ ├── 01-BufferCopy.py # Tiled copy with partitioned tensors +│ │ ├── 02-MFMA.py # Tiled MMA (GEMM) with MFMA atoms │ │ └── 03-preshuffle_gemm.py # Preshuffle GEMM end-to-end example │ └── cuda/ # NVIDIA CUDA examples +│ └── 01-MmaSync.py # Tiled MMA (GEMM) with mma.sync atoms ├── kernels/ # Production GPU kernels (importable as `kernels.*`) ├── tests/ # All tests (kernels/, mlir/, unit/) ├── CMakeLists.txt # top-level CMake diff --git a/lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp b/lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp index 6fb568e18..c47f3a5ac 100644 --- a/lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp +++ b/lib/Runtime/CUDA/FlyCudaRuntimeWrappers.cpp @@ -16,10 +16,12 @@ // //===----------------------------------------------------------------------===// +#include #include #include #include #include +#include #include "cuda.h" #include "mlir/ExecutionEngine/CRunnerUtils.h" @@ -274,14 +276,24 @@ extern "C" void mgpuMemHostRegister(void *ptr, uint64_t sizeBytes) { extern "C" void mgpuMemHostRegisterMemRef(int64_t rank, StridedMemRefType *descriptor, int64_t elementSizeBytes) { int64_t *sizes = descriptor->sizes; - int64_t *strides = sizes + rank; + [[maybe_unused]] int64_t *strides = sizes + rank; - int64_t denseStride = 1; - for (int64_t i = rank - 1; i >= 0; --i) { - (void)strides; - denseStride *= sizes[i]; + std::vector denseStrides(static_cast(rank)); + if (rank > 0) { + denseStrides[static_cast(rank - 1)] = sizes[rank - 1]; + for (int64_t i = rank - 2; i >= 0; --i) + denseStrides[static_cast(i)] = sizes[i] * denseStrides[static_cast(i + 1)]; } - auto sizeBytes = denseStride * elementSizeBytes; + auto sizeBytes = (rank > 0 ? denseStrides[0] : 1) * elementSizeBytes; + + for (int64_t i = 0; i < rank - 1; ++i) + denseStrides[static_cast(i)] = denseStrides[static_cast(i + 1)]; + if (rank > 0) + denseStrides[static_cast(rank - 1)] = 1; + + for (int64_t i = 0; i < rank; ++i) + assert(strides[i] == denseStrides[static_cast(i)]); + auto *ptr = descriptor->data + descriptor->offset * elementSizeBytes; mgpuMemHostRegister(ptr, sizeBytes); } diff --git a/python/flydsl/expr/__init__.py b/python/flydsl/expr/__init__.py index 6b38c604f..20993a557 100644 --- a/python/flydsl/expr/__init__.py +++ b/python/flydsl/expr/__init__.py @@ -19,7 +19,7 @@ _BACKEND_MODULES = { "rocdl": ".rocdl", - "tdm_ops": ".rocdl.tdm_ops", # deprecated, use .rocdl.tdm_ops instead + "tdm_ops": ".rocdl.tdm_ops", # deprecated, use fx.rocdl.tdm_ops instead "nvvm": ".nvvm", }