From 4c53162e2a82a1ece0e69bad158a0224928aa091 Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Fri, 21 Aug 2026 10:55:22 -0400 Subject: [PATCH] Split benchmark into cuda/ and rocm/ platform directories Reorganize the flat source layout into two platform-specific trees so the CUDA and ROCm variants can diverge without conflicting: - rocm/ carries forward the existing sources, including the rocWMMA backend (mxm_rocwmma.h, transform_rocwmma.h) and the HIP entry points. - cuda/ adds the CUDA-specific variant: WMMA backend (mxm_wmma.h, transform_wmma.h), dmma.h, transform_cublasdx_mxm.h, and a validate_levels.cu counterpart to the HIP validator. Each directory gets its own CMakeLists.txt, README.md, CLAUDE.md, and LICENSE so it builds standalone. The root .gitignore is kept in place so build output stays untracked. Co-Authored-By: Claude Opus 5 --- cuda/.clangd | 12 + cuda/.gitignore | 42 ++ cuda/CLAUDE.md | 150 +++++++ cuda/CMakeLists.txt | 98 +++++ LICENSE => cuda/LICENSE | 0 cuda/README.md | 211 ++++++++++ cuda/dmma.h | 122 ++++++ mxm.h => cuda/mxm.h | 0 cuda/mxm_cublasdx.h | 391 ++++++++++++++++++ cuda/mxm_level2.h | 58 +++ cuda/mxm_level3.h | 86 ++++ cuda/mxm_level4.h | 169 ++++++++ cuda/mxm_level5.h | 192 +++++++++ cuda/mxm_level7.h | 249 +++++++++++ cuda/mxm_wmma.h | 220 ++++++++++ cuda/transform.h | 91 ++++ cuda/transform_cublasdx.h | 254 ++++++++++++ cuda/transform_cublasdx_mxm.h | 102 +++++ cuda/transform_kron.h | 136 ++++++ cuda/transform_level2.h | 65 +++ cuda/transform_level3.h | 83 ++++ cuda/transform_level4.h | 87 ++++ cuda/transform_level5.h | 87 ++++ cuda/transform_level7.h | 80 ++++ cuda/transform_wmma.h | 90 ++++ cuda/transformbench.cu | 244 +++++++++++ cuda/util.h | 182 ++++++++ cuda/validate_levels.cu | 184 +++++++++ rocm/.gitignore | 42 ++ CLAUDE.md => rocm/CLAUDE.md | 0 CMakeLists.txt => rocm/CMakeLists.txt | 0 rocm/LICENSE | 28 ++ README.md => rocm/README.md | 0 rocm/mxm.h | 373 +++++++++++++++++ mxm_cublasdx.h => rocm/mxm_cublasdx.h | 0 mxm_level2.h => rocm/mxm_level2.h | 0 mxm_level3.h => rocm/mxm_level3.h | 12 +- mxm_level4.h => rocm/mxm_level4.h | 0 mxm_level5.h => rocm/mxm_level5.h | 0 mxm_level7.h => rocm/mxm_level7.h | 0 mxm_rocwmma.h => rocm/mxm_rocwmma.h | 0 transform.h => rocm/transform.h | 0 .../transform_cublasdx.h | 0 transform_kron.h => rocm/transform_kron.h | 0 transform_level2.h => rocm/transform_level2.h | 0 transform_level3.h => rocm/transform_level3.h | 0 transform_level4.h => rocm/transform_level4.h | 0 transform_level5.h => rocm/transform_level5.h | 0 transform_level7.h => rocm/transform_level7.h | 0 .../transform_rocwmma.h | 0 transformbench.cu => rocm/transformbench.cu | 0 transformbench.hip => rocm/transformbench.hip | 0 util.h => rocm/util.h | 0 .../validate_levels.hip | 0 54 files changed, 4134 insertions(+), 6 deletions(-) create mode 100644 cuda/.clangd create mode 100644 cuda/.gitignore create mode 100644 cuda/CLAUDE.md create mode 100644 cuda/CMakeLists.txt rename LICENSE => cuda/LICENSE (100%) create mode 100644 cuda/README.md create mode 100644 cuda/dmma.h rename mxm.h => cuda/mxm.h (100%) create mode 100644 cuda/mxm_cublasdx.h create mode 100644 cuda/mxm_level2.h create mode 100644 cuda/mxm_level3.h create mode 100644 cuda/mxm_level4.h create mode 100644 cuda/mxm_level5.h create mode 100644 cuda/mxm_level7.h create mode 100644 cuda/mxm_wmma.h create mode 100644 cuda/transform.h create mode 100644 cuda/transform_cublasdx.h create mode 100644 cuda/transform_cublasdx_mxm.h create mode 100644 cuda/transform_kron.h create mode 100644 cuda/transform_level2.h create mode 100644 cuda/transform_level3.h create mode 100644 cuda/transform_level4.h create mode 100644 cuda/transform_level5.h create mode 100644 cuda/transform_level7.h create mode 100644 cuda/transform_wmma.h create mode 100644 cuda/transformbench.cu create mode 100644 cuda/util.h create mode 100644 cuda/validate_levels.cu create mode 100644 rocm/.gitignore rename CLAUDE.md => rocm/CLAUDE.md (100%) rename CMakeLists.txt => rocm/CMakeLists.txt (100%) create mode 100644 rocm/LICENSE rename README.md => rocm/README.md (100%) create mode 100644 rocm/mxm.h rename mxm_cublasdx.h => rocm/mxm_cublasdx.h (100%) rename mxm_level2.h => rocm/mxm_level2.h (100%) rename mxm_level3.h => rocm/mxm_level3.h (89%) rename mxm_level4.h => rocm/mxm_level4.h (100%) rename mxm_level5.h => rocm/mxm_level5.h (100%) rename mxm_level7.h => rocm/mxm_level7.h (100%) rename mxm_rocwmma.h => rocm/mxm_rocwmma.h (100%) rename transform.h => rocm/transform.h (100%) rename transform_cublasdx.h => rocm/transform_cublasdx.h (100%) rename transform_kron.h => rocm/transform_kron.h (100%) rename transform_level2.h => rocm/transform_level2.h (100%) rename transform_level3.h => rocm/transform_level3.h (100%) rename transform_level4.h => rocm/transform_level4.h (100%) rename transform_level5.h => rocm/transform_level5.h (100%) rename transform_level7.h => rocm/transform_level7.h (100%) rename transform_rocwmma.h => rocm/transform_rocwmma.h (100%) rename transformbench.cu => rocm/transformbench.cu (100%) rename transformbench.hip => rocm/transformbench.hip (100%) rename util.h => rocm/util.h (100%) rename validate_levels.hip => rocm/validate_levels.hip (100%) diff --git a/cuda/.clangd b/cuda/.clangd new file mode 100644 index 0000000..77e3d1a --- /dev/null +++ b/cuda/.clangd @@ -0,0 +1,12 @@ +CompileFlags: + # compile_commands.json already has all flags; nothing to add/remove + Remove: + - -O3 # clangd doesn't need optimization; speeds up indexing + +Index: + Background: Build + +Diagnostics: + Suppress: + - pp_including_mainfile_in_preamble + - unknown_builtin diff --git a/cuda/.gitignore b/cuda/.gitignore new file mode 100644 index 0000000..7f3e00b --- /dev/null +++ b/cuda/.gitignore @@ -0,0 +1,42 @@ +# Prerequisites +*.d + +# Compiled Object files +*.slo +*.lo +*.o +*.obj + +# Precompiled Headers +*.gch +*.pch + +# Compiled Dynamic libraries +*.so +*.dylib +*.dll + +# Fortran module files +*.mod +*.smod + +# Compiled Static libraries +*.lai +*.la +*.a +*.lib + +# Executables +*.exe +*.out +*.app + +# custom build environment +build_vscode* +build* +#clangd +clangd* +#compile_commands.json +compile_commands.json +rocroof/ + diff --git a/cuda/CLAUDE.md b/cuda/CLAUDE.md new file mode 100644 index 0000000..62a15ba --- /dev/null +++ b/cuda/CLAUDE.md @@ -0,0 +1,150 @@ +# transformbench-cuda + +CUDA-only port of the MRA (Multi-Resolution Analysis) transform benchmark — a +batched 3D tensor-times-matrix contraction. The parent directory holds the +dual-target (HIP + CUDA) original; this tree targets **NVIDIA only** and drops +every HIP branch. See [README.md](README.md) for the full porting rationale. + +## Mathematical core + +The transform applies a K×K matrix B (transposed) along each dimension of a +K×K×K tensor: + +``` +for d in {0, 1, 2}: + C <- B^T x A (in-place, cycling through a workspace) + +GEMM shape per pass: + A (input): K^2 x K col-major + B (matrix): K x K row-major + C (output): K^2 x K row-major + FLOPs: 2 * K^2 * K * K per pass -> 3 * 2 * K^4 per full transform +``` + +## Optimization levels + +| Level | File | Technique | Threads | Notes | +|---|---|---|---|---| +| L1 | `mxm.h` / `transform.h` | Global memory only | K×min(K,128/K) | Correctness reference for `validate_levels` | +| L2 | `mxm_level2.h` / `transform_level2.h` | B cached in shared memory | 128 | Eliminates B HBM redundancy | +| L3 | `mxm_level3.h` / `transform_level3.h` | Register blocking (K-templated) | 128 | `acc[K]` in registers | +| L4 | `mxm_level4.h` / `transform_level4.h` | FP64 tensor cores, one warp | 32 | sm_80+; L3 fallback | +| L5 | `mxm_level5.h` / `transform_level5.h` | Tensor cores + A staged in smem | 256 | sm_80+; 8 warps | +| L6 | `mxm_wmma.h` / `transform_wmma.h` | `nvcuda::wmma`, warp per tile | 64–1024 | sm_80+; pads N to 8 | +| L7 | `mxm_level7.h` / `transform_level7.h` | B resident in registers, 3 GEMMs | 256 | sm_80+; K ∈ {8,16} | +| L8 | `transform_kron.h` | Kronecker product GEMM (cuBLAS) | cuBLAS | K⁶·8 B, caps near K=16 | +| L9 | `transform_cublasdx.h` | cuBLASDx, 3 GEMMs fused | cuBLASDx | K ∈ {8,10,16,20} | +| L10 | `transform_cublasdx_mxm.h` | cuBLASDx per-pass block GEMM | 128 | Uses `mra::mTxmq_cublasdx` | + +**Default level**: L9 if cuBLASDx is available, else L3. + +### FP64 tensor cores (levels 4, 5, 7) + +NVIDIA's FP64 matrix instruction is `mma.sync.m8n8k4.f64` — an **8×8×4** tile +across a **32-lane warp**, against CDNA's 16×16×4 across a 64-lane wavefront. +All three levels go through `nvcuda::wmma` (see `dmma.h`) rather than raw PTX, +so `load_matrix_sync` owns the lane→element mapping. Constraints that shape the +code: + +- shape must be 8×8×4; no other FP64 fragment geometry exists +- requires `__CUDA_ARCH__ >= 800` +- `ldm` for a `double` must be a multiple of 2 elements — **all padded shared + memory strides must stay even** (L5 pads by +2, where the AMD source pads +1) +- every lane of the warp must reach `mma_sync` + +`MRA_DMMA_SUPPORTED` (host, from `MRA_CUDA_ARCH`) and `MRA_HAVE_DMMA` (device, +from `__CUDA_ARCH__`) gate the paths; they must stay in agreement. + +### Level 3 — register blocking (the portable workhorse) + +``` +for i in 0..K^2-1 (parallel over threads): + acc[K] = 0 // register array, K doubles per thread + for k in 0..K-1: + aki = A[k, i] // load A once per k + for j in 0..K-1: + acc[j] += aki * B[k, j] + for j in 0..K-1: + C[i, j] = acc[j] +``` + +K is a compile-time template parameter, so each K gets its own kernel binary and +register pressure stays proportional to K rather than max(K). On the +`transform-mem-alloc` branch the shared-memory staging of B in +`mTxmq_level3_k` is commented out — L3 reads B from global memory and relies on +L2 cache. `transform_level3_shmem_size` still reports K²·sizeof(T), so the +allocation is made but unused. + +## Key source files + +| File | Role | +|---|---| +| `transformbench.cu` | Benchmark driver — option parsing, timing loop, FLOPs reporting | +| `validate_levels.cu` | Correctness test: any level vs the L1 reference | +| `util.h` | Launch macros (`CALL_KERNEL`, `CONFIGURE_KERNEL`), memory macros, option parser | +| `dmma.h` | FP64 tensor-core support layer: fragment aliases, tile constants, availability macros | +| `mxm_cublasdx.h` | cuBLASDx `GEMMBuilder` + the `mTxmq_cublasdx` block GEMM | + +## Building + +```bash +mkdir build && cd build +cmake .. -DCMAKE_BUILD_TYPE=Release -DCMAKE_EXPORT_COMPILE_COMMANDS=ON +make -j +``` + +`Release` is required — `Debug` inflates register/shared-memory use enough to +break compilation. CMake fetches cuBLASDx v25.06 via `FetchContent` unless +`-DUSE_CUBLASDX=OFF`. Use `-DCMAKE_CUDA_ARCHITECTURES=90` for Hopper. + +Symlink `compile_commands.json` to this directory for clangd: + +```bash +ln -sf build/compile_commands.json compile_commands.json +``` + +## Running + +```bash +./transformbench_cuda [options] + -K transform order (default 16) + -N number of tensors in batch (default 2048) + -M max concurrent blocks (default 512) + -n task submissions per timing rep (default 500) + -r timing repetitions (default 5) + -l optimization level 1-10 (default: auto) + -s number of concurrent streams (default 4) + +# Sweep levels for K=16 +for L in 1 2 3 4 5 6 7 9 10; do ./transformbench_cuda -K 16 -N 2048 -n 100 -l $L; done + +# Correctness +./validate_levels -l 5 -K 16 +``` + +Output (one line per timing rep): + +``` +Transform;level=L3-regblk;nfuncs=2048;nblocks=512;K=16;tasks=100;threads={128,1,1};smem=2048;Time(us)=12345;GFlop=403.0;Gflop/s=32.6 +``` + +## FLOPs accounting + +- **L1–L7, L9, L10**: reported as `3 × 2 × K⁴ × nfuncs × ntasks` (mathematical + minimum — useful throughput) +- **L8**: reported as `2 × K⁶ × nfuncs × ntasks` (actual GEMM work — inflated + versus the rest because K⁶ ≫ 3K⁴) + +Do not compare L8 GFlop/s directly to the other levels. + +## Architecture notes + +- K-templated kernels (L3–L7): compile-time K avoids over-allocating registers + across K values +- One K³-sized workspace per block; workspace and output are ping-ponged across + the three passes — except L7, which keeps all three passes inside one kernel + call so B never leaves registers +- L7 reuses a single K³ shared buffer in place via the MADNESS pointer trick: + C written row-major as `buf[i*K+j]` is reread col-major as `buf[k*K²+i]`. + Never pad that buffer — the two views must alias exactly. +- Multiple streams (default 4) allow kernel overlap for throughput measurement diff --git a/cuda/CMakeLists.txt b/cuda/CMakeLists.txt new file mode 100644 index 0000000..9a2893d --- /dev/null +++ b/cuda/CMakeLists.txt @@ -0,0 +1,98 @@ +cmake_minimum_required(VERSION 3.18) + +project(transformbench_cuda LANGUAGES CXX CUDA) + +# sm_80 (A100) is the minimum for FP64 tensor cores, which levels 4-7 use. +# Override with -DCMAKE_CUDA_ARCHITECTURES=90 for Hopper. +if (NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES 80) +endif() + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 20) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) + +if (NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE Release CACHE STRING "" FORCE) +endif() + +option(USE_CUBLASDX "Fetch and build against cuBLASDx (levels 9 and 10)" ON) +option(USE_SUGGEST_LAYOUT "Use suggested layout instead of get_layout" ON) +option(DEBUG_TENSOR_TYPE "Compile-time print cute tensor types (breaks build)" OFF) +set(USE_CUBLASDX_VERSION "25.06" CACHE STRING "Version of cuBLASDx to use") + +# MRA_CUDA_ARCH gates the FP64 tensor-core paths at compile time; take the +# first entry when several architectures are requested. +list(GET CMAKE_CUDA_ARCHITECTURES 0 MRA_PRIMARY_ARCH) +string(REGEX REPLACE "[^0-9]" "" MRA_PRIMARY_ARCH "${MRA_PRIMARY_ARCH}") +message(STATUS "Building for CUDA architecture sm_${MRA_PRIMARY_ARCH}") +if (MRA_PRIMARY_ARCH LESS 80) + message(WARNING "sm_${MRA_PRIMARY_ARCH} has no FP64 tensor cores; " + "levels 4-7 will run their level-3 fallback") +endif() + +find_package(CUDAToolkit REQUIRED) + +# --------------------------------------------------------------------------- +# cuBLASDx (levels 9 and 10) +# --------------------------------------------------------------------------- +set(MRA_CUBLASDX_FOUND OFF) +if (USE_CUBLASDX) + include(FetchContent) + FetchContent_Declare( + cublasdx + URL https://developer.download.nvidia.com/compute/cublasdx/redist/cublasdx/nvidia-mathdx-${USE_CUBLASDX_VERSION}.0.tar.gz + ) + FetchContent_MakeAvailable(cublasdx) + FetchContent_GetProperties(cublasdx + SOURCE_DIR CUBLASDX_SOURCE_DIR + BINARY_DIR CUBLASDX_BINARY_DIR + ) + + find_package(mathdx REQUIRED COMPONENTS cublasdx + HINTS ${CUBLASDX_SOURCE_DIR}/nvidia/mathdx/${USE_CUBLASDX_VERSION}/) + if (TARGET mathdx::cublasdx) + message(STATUS "Found cuBLASDx at ${mathdx_CUBLASDX_DIR}") + set(MRA_CUBLASDX_FOUND ON) + else() + message(FATAL_ERROR "cuBLASDx not found") + endif() +endif() + +# --------------------------------------------------------------------------- +# Common settings +# --------------------------------------------------------------------------- +add_library(libmra INTERFACE) +target_compile_definitions(libmra INTERFACE + MRA_HAVE_CUDA=1 + MRA_CUDA_ARCH=${MRA_PRIMARY_ARCH}) +# constexpr host functions in device code, and extended lambdas +target_compile_options(libmra INTERFACE + $<$:--expt-relaxed-constexpr> + $<$:--extended-lambda>) +# cuBLAS backs the Kronecker level +target_link_libraries(libmra INTERFACE CUDA::cublas) + +if (MRA_CUBLASDX_FOUND) + target_link_libraries(libmra INTERFACE mathdx::cublasdx) +endif() + +function(mra_configure_target tgt) + target_link_libraries(${tgt} PUBLIC libmra) + set_target_properties(${tgt} PROPERTIES CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") + if (USE_SUGGEST_LAYOUT) + target_compile_definitions(${tgt} PUBLIC USE_SUGGEST_LAYOUT) + endif() + if (DEBUG_TENSOR_TYPE) + target_compile_definitions(${tgt} PUBLIC DEBUG_TENSOR_TYPE) + endif() +endfunction() + +# The benchmark driver +add_executable(transformbench_cuda transformbench.cu) +mra_configure_target(transformbench_cuda) + +# Correctness test: validate any optimization level against the L1 reference +add_executable(validate_levels validate_levels.cu) +mra_configure_target(validate_levels) diff --git a/LICENSE b/cuda/LICENSE similarity index 100% rename from LICENSE rename to cuda/LICENSE diff --git a/cuda/README.md b/cuda/README.md new file mode 100644 index 0000000..9d49c18 --- /dev/null +++ b/cuda/README.md @@ -0,0 +1,211 @@ +# transformbench — CUDA port + +A CUDA-only rewrite of the MRA transform benchmark that lives in the parent +directory. Same mathematical kernel, same measurement harness, same output +format; every HIP/ROCm construct has been replaced with its CUDA equivalent, +and the AMD matrix-core levels have been re-derived for NVIDIA FP64 tensor +cores rather than mechanically translated. + +## The kernel being measured + +A batched 3D tensor-times-matrix contraction: a K×K matrix B is applied +(transposed) along each dimension of a K×K×K tensor. + +``` +for d in {0, 1, 2}: + C <- B^T x A (cycling through a workspace) + +GEMM shape per pass: + A (input): K^2 x K col-major + B (matrix): K x K row-major + C (output): K^2 x K row-major + FLOPs: 2 * K^2 * K * K per pass -> 3 * 2 * K^4 per full transform +``` + +## Optimization levels + +| Level | Files | Technique | Threads | Needs | +|---|---|---|---|---| +| L1 | `mxm.h` / `transform.h` | Global memory only — correctness reference | K×min(K,128/K) | — | +| L2 | `mxm_level2.h` / `transform_level2.h` | B cached in shared memory | 128 | — | +| L3 | `mxm_level3.h` / `transform_level3.h` | Register blocking, `acc[K]` in registers | 128 | — | +| L4 | `mxm_level4.h` / `transform_level4.h` | FP64 tensor cores, one warp | 32 | sm_80 | +| L5 | `mxm_level5.h` / `transform_level5.h` | FP64 tensor cores, A staged in shared memory | 256 | sm_80 | +| L6 | `mxm_wmma.h` / `transform_wmma.h` | `nvcuda::wmma`, one warp per output tile | 64–1024 | sm_80 | +| L7 | `mxm_level7.h` / `transform_level7.h` | FP64 tensor cores, B resident in registers | 256 | sm_80 | +| L8 | `transform_kron.h` | Single K³×K³ DGEMM via Kronecker product | cuBLAS | — | +| L9 | `transform_cublasdx.h` | cuBLASDx, three GEMMs fused in shared memory | cuBLASDx-chosen | cuBLASDx | +| L10 | `transform_cublasdx_mxm.h` | cuBLASDx as a per-pass block GEMM | 128 | cuBLASDx | + +Default level: **L9** when cuBLASDx is available, otherwise **L3**. + +Levels 4–7 fall back to the L3 register-blocking kernel whenever the target is +older than sm_80 or K does not admit a whole tiling; the fallback is transparent +and the reported level name does not change. + +## How the AMD levels were ported + +The bulk of the tree is a direct translation — `hipMalloc`→`cudaMalloc`, +`hipStream_t`→`cudaStream_t`, `__HIP_DEVICE_COMPILE__`→`__CUDA_ARCH__`, +hipBLAS→cuBLAS. Levels 4–7 are not, because they are built on AMD matrix cores +whose geometry has no NVIDIA twin: + +| | CDNA (gfx90a/gfx940) | NVIDIA (sm_80+) | +|---|---|---| +| instruction | `v_mfma_f64_16x16x4f64` | `mma.sync.m8n8k4.f64` | +| tile (M×N×K) | 16 × 16 × 4 | 8 × 8 × 4 | +| lanes cooperating | 64 (wavefront) | 32 (warp) | +| accumulators per lane | 4 | 2 | + +The AMD sources index MFMA operands by hand — thread *t* supplies `A[t/4][t%4]`, +holds output rows `(t/16)*4 + 0..3`, and so on. That lane mapping is specific to +CDNA. The ports drive the tensor cores through `nvcuda::wmma` instead, so +`load_matrix_sync` owns the lane→element mapping; this keeps the translation +correct without hard-coding a register layout. What each level *demonstrates* is +preserved exactly — who stages A, where B lives, how many warps cooperate. See +[dmma.h](dmma.h) for the shared support layer. + +Consequences worth knowing: + +- **Tile size halves (16→8)**, so every tile-count constant is re-derived. L4/L5/L7 + need K to be a multiple of 8 (K = 8, 16, 32) where the AMD versions need 16. +- **L6 column padding drops from 16 to 8.** An 8-wide fragment expresses a K < 16 + GEMM directly, so rocWMMA's LDS-resident small-K special case (`transform_klt16`) + has no counterpart and is not carried over. L6 also stages output **per warp** + (8×8 scratch) instead of staging the whole M×N_PAD output as rocWMMA does — + the full-output variant would need ~75 KB of shared memory at K=20. +- **L5's shared-memory padding is +2, not +1.** The WMMA API requires a `double` + leading dimension to be a multiple of 16 bytes / 8 bytes = 2 elements, so the + AMD +1 pad would be rejected. Chunking is otherwise identical: K=16 → 4 strips + of 64 rows, K=32 → 8 strips of 128. +- **L7 drops the XOR swizzle.** `load_matrix_sync`/`store_matrix_sync` compute + their own lane addresses from a base pointer and a stride, leaving nowhere to + inject an address permutation, and padding would break the pointer trick that + the in-place buffer reuse depends on. The K²-stride bank conflicts the AMD + swizzle avoids are therefore accepted here. Correctness is unaffected; + L7 shared-memory throughput is the thing to watch when profiling. +- **L7 is limited to K ∈ {8, 16}.** A warp holds its whole A partition in + registers; at K=32 that is 128 fragments per lane and would spill. Other K + values are dispatched to L3 by the host-side submit function. + +## Deliberate differences from the parent tree + +These are fixes, not translation artifacts — each one is a place where the +dual-target sources had drifted. + +1. **The level map is complete.** `transformbench.cu` upstream includes + `transform_cublasdx.h` but never dispatches to it, and its `level_names[]` + still describes level 5 as cuBLASDx while case 5 calls the AMD MFMA kernel. + Here levels 1–8 keep their upstream meaning and cuBLASDx gets levels 9 and 10. +2. **L1 is unambiguously the reference.** Upstream, `transform.h` includes + `mxm_cublasdx.h`, which declares `mTxmq(long, long, long, …)`; every call site + passes `int`, which binds exactly to the `size_type` overload in `mxm.h`, so + the cuBLASDx path was silently unreachable. `mxm_cublasdx.h`'s entry point is + renamed `mTxmq_cublasdx` here and reached through level 10 on purpose, leaving + L1 a genuine global-memory reference for `validate_levels`. +3. **L7's fallback runs all three passes.** The AMD `mTxmq_level7_k` fallback + calls `mTxmq_level3_impl` once, but its caller expects the full three-GEMM + chain, so a non-MFMA target silently produced a one-pass result. The CUDA + fallback ping-pongs through the workspace like L3 does. +4. **`gemm7_pass` no longer marks aliasing pointers `__restrict__`.** GEMM 2 + passes the same buffer as source and destination. +5. **Streams are destroyed** at the end of `transform_bench`. + +`mxm_level3.h` is carried over verbatim, including the commented-out staging +copy on the `transform-mem-alloc` branch — L3 currently reads B straight from +global memory and relies on L2 cache. The comment in `transform_level3.h` was +corrected to say so. + +## Building + +```bash +mkdir build && cd build +cmake .. -DCMAKE_BUILD_TYPE=Release -DCMAKE_EXPORT_COMPILE_COMMANDS=ON +make -j +``` + +`Release` is the default and is required — `Debug` inflates register and +shared-memory use enough to break compilation. + +Useful options: + +| Option | Default | Meaning | +|---|---|---| +| `CMAKE_CUDA_ARCHITECTURES` | `80` | `90` for Hopper. Below 80 disables FP64 tensor cores; L4–L7 then run their L3 fallback. | +| `USE_CUBLASDX` | `ON` | Fetches cuBLASDx v25.06 via `FetchContent`. `OFF` drops levels 9 and 10. | +| `USE_SUGGEST_LAYOUT` | `ON` | cuBLASDx suggested layout instead of `get_layout`. | +| `DEBUG_TENSOR_TYPE` | `OFF` | Compile-time print of cute tensor types (breaks the build by design). | + +For clangd: + +```bash +ln -sf build/compile_commands.json compile_commands.json +``` + +## Running + +```bash +./transformbench_cuda [options] + -K transform order (default 16) + -N number of tensors in batch (default 2048) + -M max concurrent blocks (default 512) + -n task submissions per timing rep (default 500) + -r timing repetitions (default 5) + -l optimization level 1-10 (default: auto) + -s number of concurrent streams (default 4) + +# Sweep levels at K=16 +for L in 1 2 3 4 5 6 7 9 10; do ./transformbench_cuda -K 16 -N 2048 -n 100 -l $L; done + +# Sweep K at L3 +for K in 6 8 10 12 16 20 32; do ./transformbench_cuda -K $K -N 2048 -n 100 -l 3; done +``` + +Output, one line per timing rep: + +``` +Transform;level=L3-regblk;nfuncs=2048;nblocks=512;K=16;tasks=100;threads={128,1,1};smem=2048;Time(us)=12345;GFlop=403.0;Gflop/s=32.6 +``` + +## Correctness + +```bash +./validate_levels [-l ] [-K ] [-N ] +``` + +Compares any level against the L1 reference; with no `-K` it sweeps +K ∈ {6, 8, 10, 12, 16}. Passing means max relative error < 1e-10. +Levels with narrower dispatch tables (L7: K ∈ {8,16}; L9: K ∈ {8,10,16,20}) +report the K values they do not handle. + +## K support by level + +| Level | K values on the accelerated path | Other K | +|---|---|---| +| L1–L3 | 6, 8, 10, 12, 16, 20, 32 | — | +| L4, L5 | 8, 16, 32 | L3 fallback | +| L6 | 4, 8, 12, 16, 20, 32 (any K % 4 == 0) | L3 fallback | +| L7 | 8, 16 | dispatched to L3 | +| L8 | any, but K⁶·8 bytes caps it near K=16 | — | +| L9 | 8, 10, 16, 20 | throws | +| L10 | 6, 8, 10, 12, 16, 20, 32 | prints diagnostic | + +## FLOPs accounting + +- **L1–L7, L9, L10**: reported as `3 × 2 × K⁴ × nfuncs × ntasks` — the + mathematical minimum, i.e. useful throughput. +- **L8**: reported as `2 × K⁶ × nfuncs × ntasks` — the actual GEMM work, which is + inflated relative to the others because K⁶ ≫ 3K⁴. + +Do not compare L8 GFlop/s directly against the other levels; it counts far more +FLOPs for the same mathematical result. + +## Build verification status + +The port has been checked by parsing every translation unit and force-instantiating +every device entry point for all supported K, with both the tensor-core paths and +the fallback paths enabled. It has **not** been compiled with `nvcc` or run on a +GPU — no CUDA toolkit was available in the environment where it was written. The +`nvcuda::wmma` fragment/layout combination used here (`matrix_a` col_major, +`matrix_b` row_major, FP64 8×8×4) is the one the rocWMMA source uses and is +within the documented API, but it is the first thing to check if a build fails. diff --git a/cuda/dmma.h b/cuda/dmma.h new file mode 100644 index 0000000..1b3a2f4 --- /dev/null +++ b/cuda/dmma.h @@ -0,0 +1,122 @@ +#pragma once + +#include "util.h" + +/** + * FP64 tensor-core (DMMA) support layer. + * + * This header is the CUDA stand-in for the `__builtin_amdgcn_mfma_f64_16x16x4f64` + * intrinsic that levels 4, 5 and 7 use on CDNA. The two hardware units are not + * interchangeable, and the differences drive every shape constant downstream: + * + * CDNA (gfx90a/gfx940) NVIDIA (sm_80+) + * instruction v_mfma_f64_16x16x4f64 mma.sync.m8n8k4.f64 + * tile (M x N x K) 16 x 16 x 4 8 x 8 x 4 + * lanes cooperating 64 (wavefront) 32 (warp) + * accumulators per lane 4 2 + * + * The AMD sources index MFMA operands by hand (thread t supplies A[t/4][t%4] + * and so on). That lane mapping is specific to CDNA and does not carry over, + * so the ports below drive the tensor cores through `nvcuda::wmma` instead: + * `load_matrix_sync` owns the lane->element mapping, which keeps the port + * correct without hard-coding an undocumented register layout. The structural + * choices that distinguish the levels from one another - who stages A, where B + * lives, how many warps cooperate - are preserved exactly. + * + * Constraints inherited from the WMMA API for `double`: + * - shape must be 8 x 8 x 4; no other FP64 fragment geometry exists + * - requires __CUDA_ARCH__ >= 800 (A100 / H100); sm_70 has no FP64 tensor core + * - `ldm` must be a multiple of 16 bytes / sizeof(double) = 2 elements, so all + * padded shared-memory strides below are kept even + * - every lane of the warp must reach the mma_sync call + * + * Levels that cannot use DMMA - either because the GPU predates sm_80 or + * because K is not a multiple of the tile size - fall back to the level-3 + * register-blocking kernel, exactly as the AMD sources fall back for K values + * without a native MFMA shape. + */ + +/* Host-visible availability: MRA_CUDA_ARCH is supplied by CMake. */ +#if defined(MRA_CUDA_ARCH) && (MRA_CUDA_ARCH >= 80) +# define MRA_DMMA_SUPPORTED 1 +#else +# define MRA_DMMA_SUPPORTED 0 +#endif + +/* Device-side availability: only true while compiling for sm_80 or newer. */ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) +# define MRA_HAVE_DMMA 1 +#else +# define MRA_HAVE_DMMA 0 +#endif + +#if MRA_HAVE_DMMA +#include +#endif + +namespace mra { +namespace detail { + +/* FP64 tensor-core tile geometry (the only shape NVIDIA offers). */ +constexpr int DMMA_M = 8; +constexpr int DMMA_N = 8; +constexpr int DMMA_K = 4; + +/* Round n up to the next multiple of DMMA_N. */ +constexpr int dmma_pad_n(int n) { + return ((n + DMMA_N - 1) / DMMA_N) * DMMA_N; +} + +/** + * True when K admits the plain (unpadded) DMMA tiling used by levels 4, 5 and 7: + * K % DMMA_K == 0 so the contraction divides into whole 4-deep steps + * K % DMMA_N == 0 so the output columns divide into whole 8-wide tiles + * K % 8 == 0 implies both, and also K^2 % 8 == 0 for the row tiles. + */ +constexpr bool dmma_supports_k(int K) { + return (K % DMMA_N) == 0 && K > 0; +} + +#if MRA_HAVE_DMMA + +/** + * Fragment aliases for C[M x N] = A^T[M x K] * B[K x N]. + * + * A is stored col-major as a[k * M + i] (that is, A^T[i][k]), so the matrix_a + * fragment is declared col_major with ldm = M: element [i][k] is then read from + * ptr[k * ldm + i], which is precisely the source layout. B is row-major with + * ldm = N. This mirrors the operand handling in the rocWMMA source one-to-one. + */ +using FragA = nvcuda::wmma::fragment; +using FragB = nvcuda::wmma::fragment; +using FragC = nvcuda::wmma::fragment; + +/** Load one 8x4 tile of A^T starting at row `row`, contraction offset `k`. */ +__device__ __forceinline__ +void dmma_load_a(FragA& frag, const double* a, int k, int row, int ldm) { + nvcuda::wmma::load_matrix_sync(frag, a + (size_t)k * ldm + row, ldm); +} + +/** Load one 4x8 tile of B starting at contraction offset `k`, column `col`. */ +__device__ __forceinline__ +void dmma_load_b(FragB& frag, const double* b, int k, int col, int ldm) { + nvcuda::wmma::load_matrix_sync(frag, b + (size_t)k * ldm + col, ldm); +} + +/** Store one 8x8 accumulator tile to row-major C at [row][col]. */ +__device__ __forceinline__ +void dmma_store_c(double* c, const FragC& frag, int row, int col, int ldm) { + nvcuda::wmma::store_matrix_sync(c + (size_t)row * ldm + col, frag, ldm, + nvcuda::wmma::mem_row_major); +} + +#endif // MRA_HAVE_DMMA + +} // namespace detail +} // namespace mra diff --git a/mxm.h b/cuda/mxm.h similarity index 100% rename from mxm.h rename to cuda/mxm.h diff --git a/cuda/mxm_cublasdx.h b/cuda/mxm_cublasdx.h new file mode 100644 index 0000000..84fc114 --- /dev/null +++ b/cuda/mxm_cublasdx.h @@ -0,0 +1,391 @@ +#ifndef MRA_OPS_MXM_CUBLASDX_H +#define MRA_OPS_MXM_CUBLASDX_H + +#include "util.h" + +/** + * An implementation of A^T x B using cublasdx. + * We assume that A is tall-and-skinny (K^2 x K) and B is square (K x K). + * There is some code to cover the case where A is square and B is wide-and-skinny + * but that is not yet implemented and we don't use it yet. + */ + +#define MRA_CUBLASDX_BLOCK_C 0 + +#if __has_include() + +#define MRA_HAVE_CUBLASDX 1 + +#if !defined(MRA_CUDA_ARCH) || MRA_CUDA_ARCH < 70 +#error "MRA_CUDA_ARCH must be defined and >= 70 to use cublasdx" +#endif + +#include + +#if MRA_CUDA_ARCH == 70 +#define MRA_CUBLASDX_SM 700 +#define MRA_CUBLASDX_MAX_SHM (30*1024) +#elif MRA_CUDA_ARCH == 80 +#define MRA_CUBLASDX_SM 800 +#define MRA_CUBLASDX_MAX_SHM (40*1024) +#elif MRA_CUDA_ARCH == 90 +#define MRA_CUBLASDX_SM 900 +#define MRA_CUBLASDX_MAX_SHM (110*1024) +#else +#warning "Unknown MRA_CUDA_ARCH for cublasdx, using 80" +#define MRA_CUBLASDX_SM 800 +#endif + +#ifdef DEBUG_TENSOR_TYPE +#define PRINT_TENSOR_TYPE(t) cute::print_type(t) +#else // DEBUG_TENSOR_TYPE +#define PRINT_TENSOR_TYPE(t) +#endif // DEBUG_TENSOR_TYPE + +// get the layout for tensor t in GEMM +#ifdef USE_SUGGEST_LAYOUT +#define GET_SHARED_LAYOUT(op, t) op::suggest_layout_smem_##t() +#else // USE_SUGGEST_LAYOUT +#define GET_SHARED_LAYOUT(op, t) op::get_layout_smem_##t() +#endif // USE_SUGGEST_LAYOUT + + +namespace mra { + + namespace detail { + + constexpr int CUBLAS_MIN_MN = 16; + + template + constexpr int cublasdx_max_mn() { + // K^2 for square B/A, double buffering for A/B and C + auto max_nm = ((MRA_CUBLASDX_MAX_SHM / sizeof(T)) - K*K) / ((3+MRA_CUBLASDX_BLOCK_C)*K); + // round down to the nearest power of 2 + // TODO: std::log2 is constexpr only since C++26 + //int p = std::pow(2, (int)std::log2(max_nm)); + int l = 1; + while ((l<<1) <= max_nm) l <<= 1; + return std::min(l, K*K); + } + + template + struct GEMMBuilder { + + private: + using BaseGEMM = decltype(cublasdx::Precision() + + cublasdx::Type() + + cublasdx::Function() + + cublasdx::SM() // TODO + + cublasdx::Block() + + cublasdx::MaxAlignment()); + using GEMM_ = decltype(BaseGEMM() + cublasdx::Size() + + cublasdx::Arrangement()); + using GEMM_suggested_ld = cublasdx::suggested_leading_dimension_of_t; + public: + using GEMM = decltype(GEMM_() + GEMM_suggested_ld()); + }; + + template + __forceinline__ + __device__ void mTxmq_cublasdx_core(auto&& a_shared_tensor, auto&& b_shared_tensor, + auto&& c_tensor, + auto&& load = [](){}, auto&& prefetch = [](){}) { + + using alignment = cublasdx::alignment_of; + + /* load data to shared memory */ + load(); + /* wait for load to complete */ + cublasdx::copy_wait(); + + /* prefetch data for next iteration */ + prefetch(); + + // Execute using register API + auto [c_register_fragment, partitioner] = GEMM().execute(a_shared_tensor, b_shared_tensor); + + // Store back to global memory using cublasdx::copy_fragment API + + cublasdx::copy_fragment(c_register_fragment, c_tensor, partitioner); + } + + /** + * Compute the shared memory requirements for a given GEMM. + * Takes into account double buffering of A (block_a) and B (block_b) as well as + * staging of results through shared memory (block_c). + */ + template + constexpr int cublasdx_shmem_size_for(bool block_a, bool block_b, bool block_c) { + auto calc = cublasdx::make_shared_storage_calculator() + .add(cublasdx::alignment_of_v_a, sizeof(typename GEMM::a_value_type), GET_SHARED_LAYOUT(GEMM, a)) + .add(cublasdx::alignment_of_v_b, sizeof(typename GEMM::b_value_type), GET_SHARED_LAYOUT(GEMM, b)); + if (block_a) { + calc.add(cublasdx::alignment_of_v_a, sizeof(typename GEMM::a_value_type), GET_SHARED_LAYOUT(GEMM, a)); + } + if (block_b) { + calc.add(cublasdx::alignment_of_v_b, sizeof(typename GEMM::b_value_type), GET_SHARED_LAYOUT(GEMM, b)); + } + if (block_c) { + // double buffering of C + calc.add(cublasdx::alignment_of_v_c, sizeof(typename GEMM::c_value_type), GET_SHARED_LAYOUT(GEMM, c)); + calc.add(cublasdx::alignment_of_v_c, sizeof(typename GEMM::c_value_type), GET_SHARED_LAYOUT(GEMM, c)); + } + + int shared_memory_size = calc.get(); + return shared_memory_size; + } + + template + constexpr int cublasdx_shmem_size_k() { + constexpr auto blockdims = max_thread_dims(K); + using BaseGEMM = decltype(cublasdx::Precision() + + cublasdx::Type() + + cublasdx::Function() + + cublasdx::SM() // TODO + + cublasdx::Block() + + cublasdx::BlockDim() + + cublasdx::MaxAlignment()); + constexpr auto max_mn = cublasdx_max_mn(); + using GEMMBlockA = typename GEMMBuilder::GEMM; + auto size = cublasdx_shmem_size_for(true, false, MRA_CUBLASDX_BLOCK_C); + return size; + } + + template + __forceinline__ + __device__ void mTxmq_cublasdx_block(T* c, const T* a, const T* b) { + constexpr auto blockdims = max_thread_dims(K); + extern __shared__ __align__(16) char smem[]; + constexpr auto max_mn = cublasdx_max_mn(); + /* assuming aT = bT = cT for now */ + using GEMM = typename GEMMBuilder::GEMM; + + using alignment = cublasdx::alignment_of; + + + if constexpr (M == K*K) { + constexpr auto num_iter = M/max_mn; + //if (is_team_lead()) printf("mTxmq_cublasdx_block: max_mn %d, shared_memory %u, smem %p, M = %d, N = %d, K = %d iter %d\n", max_mn, cublasdx_shmem_size_for(true, false, true), smem, M, N, K, num_iter); + //__syncthreads(); + + if constexpr (num_iter > 0) { + auto [smem_a, smem_b, smem_a_n, smem_c, smem_c_n] = + cublasdx::shared_memory::slice_into_pointers( + smem, + cublasdx::alignment_of_v_a, cublasdx::cosize(GET_SHARED_LAYOUT(GEMM, a)), + cublasdx::alignment_of_v_b, cublasdx::cosize(GET_SHARED_LAYOUT(GEMM, b)), + cublasdx::alignment_of_v_a, cublasdx::cosize(GET_SHARED_LAYOUT(GEMM, a)), + cublasdx::alignment_of_v_c, cublasdx::cosize(GET_SHARED_LAYOUT(GEMM, c)), + cublasdx::alignment_of_v_c, cublasdx::cosize(GET_SHARED_LAYOUT(GEMM, c))); + + /* copy b tensor into shared memory and leave there */ + auto b_global_tensor = cublasdx::make_tensor(b, GEMM::get_layout_gmem_b()); + auto b_shared_tensor = cublasdx::make_tensor(smem_b, GET_SHARED_LAYOUT(GEMM, b)); + cublasdx::copy(b_global_tensor, b_shared_tensor); + PRINT_TENSOR_TYPE(b_global_tensor); + PRINT_TENSOR_TYPE(b_shared_tensor); + + auto a_shared_tensor = cublasdx::make_tensor(smem_a, GET_SHARED_LAYOUT(GEMM, a)); + auto a_shared_tensor_n = cublasdx::make_tensor(smem_a_n, GET_SHARED_LAYOUT(GEMM, a)); + + auto c_shared_tensor = cublasdx::make_tensor(smem_c, GET_SHARED_LAYOUT(GEMM, c)); + auto c_shared_tensor_n = cublasdx::make_tensor(smem_c_n, GET_SHARED_LAYOUT(GEMM, c)); + + int i; // used past the for loop below + + auto make_c_global_tensor = [&](int i){ + return cublasdx::make_tensor(c+((i*max_mn)*N), GEMM::get_layout_gmem_c()); + }; + + auto store_c = [&]() { +#if MRA_CUBLASDX_BLOCK_C + auto c_shared_tensor = cublasdx::make_tensor(smem_c, GET_SHARED_LAYOUT(GEMM, c)); + __syncthreads(); // make sure prior computations are done + auto c_global_tensor = make_c_global_tensor(i-1); + cublasdx::copy(c_shared_tensor, c_global_tensor); +#endif // MRA_CUBLASDX_BLOCK_C + }; + for (i = 0; i < num_iter; i++) { + // Make global memory tensors + auto a_global_tensor = cublasdx::make_tensor(a+(i*max_mn), GEMM::get_layout_gmem_a(cute::Int{})); + auto a_shared_tensor = cublasdx::make_tensor(smem_a, GET_SHARED_LAYOUT(GEMM, a)); + auto a_shared_tensor_n = cublasdx::make_tensor(smem_a_n, GET_SHARED_LAYOUT(GEMM, a)); + + auto c_shared_tensor = cublasdx::make_tensor(smem_c, GET_SHARED_LAYOUT(GEMM, c)); + auto c_shared_tensor_n = cublasdx::make_tensor(smem_c_n, GET_SHARED_LAYOUT(GEMM, c)); + + PRINT_TENSOR_TYPE(a_global_tensor); + PRINT_TENSOR_TYPE(a_shared_tensor); + PRINT_TENSOR_TYPE(make_c_global_tensor(i)); + PRINT_TENSOR_TYPE(c_shared_tensor); + //auto c_global_tensor = cublasdx::make_tensor(c+((i*max_mn)*N), GEMM::get_layout_gmem_c()); + mTxmq_cublasdx_core(a_shared_tensor, b_shared_tensor, +#if MRA_CUBLASDX_BLOCK_C + c_shared_tensor, +#else // MRA_CUBLASDX_BLOCK_C + /* global tensor */ + make_c_global_tensor(i), +#endif // MRA_CUBLASDX_BLOCK_C + [&](){ + /* load only on first iteration, all others are prefetched */ + if (i == 0) { + //if (is_team_lead()) printf("Loading initial block %d\n", i); + cublasdx::copy(a_global_tensor, a_shared_tensor); + } + }, + [&](){ + /* store prior iteration's result */ + if (i > 0) { + //if (is_team_lead()) printf("Storing block %d\n", i-1); + store_c(); + } + /* prefetch into shared memory */ + if ((i+1) < num_iter) { + //if (is_team_lead()) printf("Prefetching block %d\n", i); + auto a_global_tensor = cublasdx::make_tensor(a+((i+1)*max_mn), GEMM::get_layout_gmem_a(cute::Int{})); + cublasdx::copy(a_global_tensor, a_shared_tensor_n); + } + }); + auto tmp_a = smem_a; + smem_a = smem_a_n; + smem_a_n = tmp_a; + auto tmp_c = smem_c; + smem_c = smem_c_n; + smem_c_n = tmp_c; + +#if 0 + auto tmp_a = a_shared_tensor; + a_shared_tensor = a_shared_tensor_n; + a_shared_tensor_n = tmp_a; + auto tmp_b = c_shared_tensor; + c_shared_tensor = c_shared_tensor_n; + c_shared_tensor_n = tmp_b; +#else + //std::swap(a_shared_tensor, a_shared_tensor_n); + //std::swap(c_shared_tensor, c_shared_tensor_n); +#endif // 0 + } + /* store the last block of C */ + store_c(); + } + + /* handle remainder */ + constexpr const auto R = M%max_mn; + if constexpr (0 < R) { + // Make global memory tensors + using GEMM = typename GEMMBuilder::GEMM; + auto [smem_a, smem_b, smem_c] = cublasdx::slice_shared_memory(smem, GET_SHARED_LAYOUT(GEMM, a), + GET_SHARED_LAYOUT(GEMM, b), + GET_SHARED_LAYOUT(GEMM, c)); + auto a_shared_tensor = cublasdx::make_tensor(smem_a, GET_SHARED_LAYOUT(GEMM, a)); + auto a_global_tensor = cublasdx::make_tensor(a+((M/max_mn)*max_mn), GEMM::get_layout_gmem_a(cute::Int{})); + auto b_global_tensor = cublasdx::make_tensor(b, GEMM::get_layout_gmem_b()); + auto b_shared_tensor = cublasdx::make_tensor(smem_b, GET_SHARED_LAYOUT(GEMM, b)); + auto c_global_tensor = cublasdx::make_tensor(c+((M/max_mn)*max_mn*N), GEMM::get_layout_gmem_c()); + auto c_shared_tensor = cublasdx::make_tensor(smem_c, GET_SHARED_LAYOUT(GEMM, c)); + mTxmq_cublasdx_core(a_shared_tensor, b_shared_tensor, +#if MRA_CUBLASDX_BLOCK_C + c_shared_tensor, +#else // MRA_CUBLASDX_BLOCK_C + c_global_tensor, +#endif // MRA_CUBLASDX_BLOCK_C + [&](){ + cublasdx::copy(a_global_tensor, a_shared_tensor); + cublasdx::copy(b_global_tensor, b_shared_tensor); + }, + [](){}); + /* move the C block back to global memory */ + cublasdx::copy(c_shared_tensor, c_global_tensor); + } + } else { + // TODO: implement! + static_assert(M == K*K, "N equal to K*K currently not supported"); + } + /* final sync */ + cublasdx::copy_wait(); + } + + } // namespace detail + + template + __forceinline__ + __device__ void mTxmq_cublasdx(long dimi, long dimj, long dimk, + cT* c, const aT* a, const bT* b) { + int M = dimi; + int N = dimj; + int K = dimk; + if (M == K*K) { + // A is tall and skinny, B is square + if (K == 6) { + detail::mTxmq_cublasdx_block<36, 6, 6>(c, a, b); + } else if (K == 8) { + detail::mTxmq_cublasdx_block<64, 8, 8>(c, a, b); + } else if (K == 10) { + detail::mTxmq_cublasdx_block<100, 10, 10>(c, a, b); + } else if (K == 12) { + detail::mTxmq_cublasdx_block<12*12, 12, 12>(c, a, b); + } else if (K == 16) { + detail::mTxmq_cublasdx_block<16*16, 16, 16>(c, a, b); + } else if (K == 20) { + detail::mTxmq_cublasdx_block<400, 20, 20>(c, a, b); + } else if (K == 32) { + detail::mTxmq_cublasdx_block<32*32, 32, 32>(c, a, b); + } else { + if (is_team_lead()) printf("mTxmq_cublasdx: Unsupport K = %d\n", K); + } + } else { + printf("mTxmq_cublasdx: Unknown configuration with M = %d, N = %d, K = %d\n", M, N, K); + } + /* make sure all is done */ + __syncthreads(); + } + + template + constexpr int mTxmq_cublasdx_shmem_size(int K) { + switch (K) { + case 6: return detail::cublasdx_shmem_size_k(); + case 8: return detail::cublasdx_shmem_size_k(); + case 10: return detail::cublasdx_shmem_size_k(); + case 12: return detail::cublasdx_shmem_size_k(); + case 16: return detail::cublasdx_shmem_size_k(); + case 20: return detail::cublasdx_shmem_size_k(); + case 32: return detail::cublasdx_shmem_size_k(); + default: THROW("CUBLASdx: Unsupported K"); + } + } + + + namespace detail { + template + constexpr Dim3 cublasdx_blockdim_k() { + + return Dim3(MAX_THREADS_PER_BLOCK, 1, 1); + constexpr auto max_mn = cublasdx_max_mn(); + using GEMM = typename GEMMBuilder::GEMM; + return GEMM::suggested_block_dim; + } + + } // namespace detail + template + constexpr Dim3 mTxmq_cublasdx_blockdim(int K) { + switch (K) { + case 6: return detail::cublasdx_blockdim_k(); + case 8: return detail::cublasdx_blockdim_k(); + case 10: return detail::cublasdx_blockdim_k(); + case 12: return detail::cublasdx_blockdim_k(); + case 16: return detail::cublasdx_blockdim_k(); + case 20: return detail::cublasdx_blockdim_k(); + case 32: return detail::cublasdx_blockdim_k(); + default: THROW("CUBLASdx: Unsupported K"); + } + } + +} // namespace mra + +#endif // __has_include() + +#endif // MRA_OPS_MXM_CUBLASDX_H diff --git a/cuda/mxm_level2.h b/cuda/mxm_level2.h new file mode 100644 index 0000000..68e23b4 --- /dev/null +++ b/cuda/mxm_level2.h @@ -0,0 +1,58 @@ +#pragma once + +#include "util.h" + +/** + * Level 2: B matrix loaded into shared memory once per mTxmq call. + * A is streamed from global memory. Threads are distributed over + * rows (i) rather than columns (j), so all 128 threads stay busy + * even for small K. + * + * c(i,j) = sum_k a(k,i)*b(k,j) + * A: K^2 x K col-major a[k,i] = a[k*dimi + i] + * B: K x K row-major b[k,j] = b[k*dimj + j] + * C: K^2 x K row-major c[i,j] = c[i*dimj + j] + */ + +namespace mra { + +/* Public entry-point: always clears C (mTxmq semantics, equivalent to Q=true) */ +template +__device__ void mTxmq_level2(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + extern __shared__ char smem_level2[]; + bT* b_shmem = reinterpret_cast(smem_level2); + + /* Cooperatively load B (dimk * dimj elements) into shared memory */ + for (int idx = threadIdx.x; idx < dimk * dimj; idx += blockDim.x) { + b_shmem[idx] = b[idx]; + } + __syncthreads(); + + /* Each thread handles a stripe of rows; full j and k loops are sequential */ + for (size_type i = (size_type)threadIdx.x; i < dimi; i += (size_type)blockDim.x) { + const aT* a_col_i = a + i; /* pointer to a[0,i] in col-major layout */ + cT* ci = c + i * dimj; + for (size_type j = 0; j < dimj; ++j) { + cT sum = cT(0); /* always clear: mTxmq semantics */ + const aT* aik = a_col_i; + for (size_type k = 0; k < dimk; ++k, aik += dimi) { + sum += (*aik) * b_shmem[k * dimj + j]; + } + ci[j] = sum; + } + } + __syncthreads(); +} + +template +constexpr size_type mTxmq_level2_shmem_size(size_type K) { + return K * K * sizeof(T); +} + +template +constexpr Dim3 mTxmq_level2_blockdim(int /*K*/) { + return Dim3(MAX_THREADS_PER_BLOCK, 1, 1); +} + +} // namespace mra diff --git a/cuda/mxm_level3.h b/cuda/mxm_level3.h new file mode 100644 index 0000000..b260313 --- /dev/null +++ b/cuda/mxm_level3.h @@ -0,0 +1,86 @@ +#pragma once + +#include "util.h" + +/** + * Level 3: B in shared memory + register accumulation. + * Each thread owns a full row of the output tile held in a compile-time + * register array T acc[K]. The k-loop loads a[k,i] once and FMAs it + * against all K columns of B (from shared memory), eliminating redundant global + * loads and keeping the hot loop inside the register file. + * + * Use mTxmq_level3_k (K known at compile time) so that each K value + * gets its own kernel binary with isolated register pressure. + * + * c(i,j) = sum_k a(k,i)*b(k,j) + * A: K^2 x K col-major a[k,i] = a[k*dimi + i] + * B: K x K row-major b[k,j] = b[k*dimj + j] + * C: K^2 x K row-major c[i,j] = c[i*dimj + j] + */ + +namespace mra { + +namespace detail { + +/** + * Inner kernel: B is already in b_shmem, register array acc[K] accumulates + * the dot product. Compile-time K keeps acc[] in VGPRs. + */ +template +__device__ void mTxmq_level3_impl(T* __restrict__ c, const T* a, const T* b_shmem) { + constexpr int DIMI = K * K; + + for (int i = (int)threadIdx.x; i < DIMI; i += (int)blockDim.x) { + T acc[K]; + + if constexpr (Q) { + for (int j = 0; j < K; ++j) acc[j] = T(0); + } else { + for (int j = 0; j < K; ++j) acc[j] = c[i * K + j]; + } + + /* k-loop: load a[k,i] once, FMA with all K entries of row k of B */ + const T* aik = a + i; /* a[0,i] in col-major */ + for (int k = 0; k < K; ++k, aik += DIMI) { + T aki = *aik; + for (int j = 0; j < K; ++j) { + acc[j] += aki * b_shmem[k * K + j]; + } + } + + for (int j = 0; j < K; ++j) c[i * K + j] = acc[j]; + } +} + +} // namespace detail + + +/** + * K-templated entry point — one binary per K value. + * Each instantiation sees only acc[K] for its specific K, + * keeping register pressure proportional to K rather than max(K). + */ +template +__device__ void mTxmq_level3_k(T* __restrict__ c, const T* a, const T* b) { + // extern __shared__ char smem_level3[]; + // T* b_shmem = reinterpret_cast(smem_level3); + + // for (int idx = (int)threadIdx.x; idx < K * K; idx += (int)blockDim.x) + // b_shmem[idx] = b[idx]; + // __syncthreads(); + + detail::mTxmq_level3_impl(c, a, b); + __syncthreads(); +} + +template +constexpr size_type mTxmq_level3_shmem_size(size_type K) { + return K * K * sizeof(T); +} + +template +constexpr Dim3 mTxmq_level3_blockdim(int /*K*/) { + return Dim3(MAX_THREADS_PER_BLOCK, 1, 1); +} + +} // namespace mra diff --git a/cuda/mxm_level4.h b/cuda/mxm_level4.h new file mode 100644 index 0000000..406c7bb --- /dev/null +++ b/cuda/mxm_level4.h @@ -0,0 +1,169 @@ +#pragma once + +#include "util.h" +#include "dmma.h" +#include "mxm_level3.h" /* for the Level-3 fallback */ + +/** + * Level 4: NVIDIA FP64 tensor cores (DMMA), one warp per thread block. + * + * This is the CUDA counterpart of the CDNA level-4 kernel. There, one + * 64-lane wavefront issues v_mfma_f64_16x16x4f64 to produce a 16x16 output + * tile with a 4-deep contraction. Here, one 32-lane warp issues + * mma.sync.m8n8k4.f64 to produce an 8x8 output tile with the same 4-deep + * contraction, so the tile grid is finer but the shape of the kernel - a single + * warp walking every output tile in sequence, with B parked in shared memory + * and A streamed straight from global - is unchanged. + * + * Block dimension is 32 threads (one warp), matching the AMD kernel's + * one-wavefront block. Level 5 is the multi-warp, shared-memory-staged variant. + * + * K must be a multiple of 8 for the DMMA path (K = 8, 16, 32). Every other K + * (6, 10, 12, 20) transparently falls back to the level-3 register-blocking + * kernel, as does any GPU older than sm_80. + * + * c(i,j) = sum_k a(k,i)*b(k,j) + * A: K^2 x K col-major a[k,i] = a[k*dimi + i] + * B: K x K row-major b[k,j] = b[k*dimj + j] + * C: K^2 x K row-major c[i,j] = c[i*dimj + j] + */ + +namespace mra { + +namespace detail { + +#if MRA_HAVE_DMMA + +/** + * DMMA kernel for compile-time K. Requires blockDim.x == 32 (one warp). + * B must already be resident in b_smem. + * + * Leading dimensions: A uses ldm = K^2 and B/C use ldm = K. Both are even for + * every supported K, satisfying the WMMA 16-byte stride requirement for double. + */ +template +__device__ void mTxmq_level4_dmma(T* __restrict__ c, const T* a, const T* b_smem) { + static_assert(std::is_same_v, + "mTxmq_level4_dmma: FP64 tensor cores operate on double only"); + static_assert(K % DMMA_N == 0, + "mTxmq_level4_dmma: K must be a multiple of 8 for the 8x8x4 tile"); + + constexpr int DIMI = K * K; + constexpr int ROW_TILES = DIMI / DMMA_M; /* 8-row tiles of the output */ + constexpr int COL_TILES = K / DMMA_N; /* 8-column tiles of the output */ + + /* One warp walks the whole (ROW_TILES x COL_TILES) grid. */ + for (int r = 0; r < ROW_TILES; ++r) { + for (int ct = 0; ct < COL_TILES; ++ct) { + FragC acc; + nvcuda::wmma::fill_fragment(acc, 0.0); + + /* K/4 steps of 4-deep contraction */ + for (int k = 0; k < K; k += DMMA_K) { + FragA a_frag; + FragB b_frag; + /* A^T[r*8 .. r*8+8, k .. k+4] read from the col-major K^2 x K source */ + dmma_load_a(a_frag, a, k, r * DMMA_M, DIMI); + /* B[k .. k+4, ct*8 .. ct*8+8] read from row-major shared memory */ + dmma_load_b(b_frag, b_smem, k, ct * DMMA_N, K); + nvcuda::wmma::mma_sync(acc, a_frag, b_frag, acc); + } + + dmma_store_c(c, acc, r * DMMA_M, ct * DMMA_N, K); + } + } +} + +#endif /* MRA_HAVE_DMMA */ + +} // namespace detail + + +/* Public entry-point: always clears C (mTxmq semantics, Q=true). */ +template +__device__ void mTxmq_level4(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + extern __shared__ char smem_level4[]; + bT* b_smem = reinterpret_cast(smem_level4); + + /* Load B into shared memory */ + for (int idx = (int)threadIdx.x; idx < dimk * dimj; idx += (int)blockDim.x) { + b_smem[idx] = b[idx]; + } + __syncthreads(); + +#if MRA_HAVE_DMMA + /* DMMA path: only for K divisible by 8 */ + if constexpr (std::is_same_v) { + if (dimi == dimj * dimj) { + if (dimj == 8) { + detail::mTxmq_level4_dmma(c, a, b_smem); + __syncthreads(); + return; + } else if (dimj == 16) { + detail::mTxmq_level4_dmma(c, a, b_smem); + __syncthreads(); + return; + } else if (dimj == 32) { + detail::mTxmq_level4_dmma(c, a, b_smem); + __syncthreads(); + return; + } + } + } + /* Fall through to Level-3 register blocking for other K values */ +#endif + + /* Level-3 fallback (also the path on pre-sm_80 hardware) */ + if (dimi == dimj * dimj) { + if (dimj == 6) detail::mTxmq_level3_impl(c, a, b_smem); + else if (dimj == 8) detail::mTxmq_level3_impl(c, a, b_smem); + else if (dimj == 10) detail::mTxmq_level3_impl(c, a, b_smem); + else if (dimj == 12) detail::mTxmq_level3_impl(c, a, b_smem); + else if (dimj == 16) detail::mTxmq_level3_impl(c, a, b_smem); + else if (dimj == 20) detail::mTxmq_level3_impl(c, a, b_smem); + else if (dimj == 32) detail::mTxmq_level3_impl(c, a, b_smem); + else { + if (is_team_lead()) printf("mTxmq_level4: unsupported K=%d\n", (int)dimj); + } + } + __syncthreads(); +} + +/** + * K-templated entry point - one binary per K value. + * Loads B into shared memory, then dispatches to DMMA (if available) or the + * Level-3 fallback. Requires blockDim.x == 32 (one warp) on the DMMA path. + */ +template +__device__ void mTxmq_level4_k(T* __restrict__ c, const T* a, const T* b) { + extern __shared__ char smem_level4[]; + T* b_smem = reinterpret_cast(smem_level4); + + for (int idx = (int)threadIdx.x; idx < K * K; idx += (int)blockDim.x) + b_smem[idx] = b[idx]; + __syncthreads(); + +#if MRA_HAVE_DMMA + if constexpr (detail::dmma_supports_k(K) && std::is_same_v) { + detail::mTxmq_level4_dmma(c, a, b_smem); + __syncthreads(); + return; + } +#endif + /* Level-3 register-blocking fallback */ + detail::mTxmq_level3_impl(c, a, b_smem); + __syncthreads(); +} + +template +constexpr size_type mTxmq_level4_shmem_size(size_type K) { + return K * K * sizeof(T); +} + +template +constexpr Dim3 mTxmq_level4_blockdim(int /*K*/) { + return Dim3(MRA_WARP_SIZE, 1, 1); /* one warp */ +} + +} // namespace mra diff --git a/cuda/mxm_level5.h b/cuda/mxm_level5.h new file mode 100644 index 0000000..e76bbd2 --- /dev/null +++ b/cuda/mxm_level5.h @@ -0,0 +1,192 @@ +#pragma once + +#include "util.h" +#include "dmma.h" +#include "mxm_level3.h" + +/** + * Level 5: shared-memory-staged A with FP64 tensor cores, 8-warp block. + * + * CUDA counterpart of the CDNA level-5 kernel. The AMD version runs 256 + * threads as 4 wavefronts of 64 and hands each wavefront a 16-row MFMA tile; + * this version runs the same 256 threads as 8 warps of 32 and hands each warp + * an 8-row DMMA tile. The staging strategy - the point of the level - is + * identical: A is pulled through shared memory in strips so that each element + * crosses the memory bus once instead of once per output column tile. + * + * For K=16, A is 256x16 (col-major). A is loaded in NCHUNKS strips of + * CHUNK_ROWS rows each. All 256 threads cooperate to load one strip into + * shared memory, then each of the 8 warps takes a disjoint set of 8x8 subtiles + * of that strip and runs mma.sync.m8n8k4.f64 against B (always resident in + * shared memory). Once all warps finish, the next strip is loaded. + * + * For K=16: + * CHUNK_ROWS = 64 (K^2/4) - loads all of A in 4 strips + * NCHUNKS = 4 + * Each warp: 8 rows per chunk = 1 tile of 8x8 per column tile + * Total C coverage: 4 chunks x (8 warps x 8 rows) = 256 rows = K^2 + * + * CHUNK_ROWS is chosen at compile time as the largest fraction of DIMI + * (DIMI/4, DIMI/8, DIMI/16) whose A strip fits alongside B in 64 KB of shared + * memory and that is a multiple of NWARPS*8 = 64, so tiles divide evenly. + * For K=32: CHUNK_ROWS=128, NCHUNKS=8, TILES_PER_WARP=2 per chunk. + * For K=8 the whole of A is one 64-row strip. + * + * Shared memory layout: + * [0 ]: B K x K row-major (K*K doubles, loaded once) + * [K*K ]: A_strip K x A_STRIDE col-major, A_STRIDE = CHUNK_ROWS + 2 + * a_smem[k*A_STRIDE + row_local] = A^T[row_base+row_local][k] + * + * The +2 padding shifts each k-column off the 32-bank alignment that an + * unpadded CHUNK_ROWS stride would land on, and keeps the stride even, which + * the WMMA API requires of a double leading dimension (16 bytes / 8 bytes = 2). + * The AMD source pads by +1 for the same reason; +1 would be rejected here. + * + * Global memory load pattern (K=16, CHUNK_ROWS=64): + * 256 threads load 1024 elements in 4 passes; each pass reads 64 consecutive + * doubles from one k-column of A - a fully coalesced 512-byte burst. + * + * K must be a multiple of 8 (K = 8, 16, 32); other K values and pre-sm_80 + * hardware fall back to the level-3 register-blocking kernel. + * + * c(i,j) = sum_k a(k,i)*b(k,j) + * A: K^2 x K col-major a[k,i] = a[k*K^2+i] + * B: K x K row-major b[k,j] = b[k*K +j] + * C: K^2 x K row-major c[i,j] = c[i*K +j] + */ + +namespace mra { + +namespace detail { + +constexpr int LEVEL5_NWARPS = 8; /* 256 threads */ +constexpr int LEVEL5_NTHREAD = LEVEL5_NWARPS * MRA_WARP_SIZE; +constexpr int LEVEL5_BUDGET = 64 * 1024; /* shared memory ceiling */ + +/* Padded column stride of the staged A strip: even, as WMMA requires. */ +constexpr int level5_a_stride(int chunk_rows) { return chunk_rows + 2; } + +/* A candidate chunk size is usable when it splits DIMI evenly, splits evenly + * across the 8 warps' 8-row tiles, and its strip fits beside B. */ +constexpr bool level5_chunk_ok(int K, int elem_bytes, int cr) { + return cr >= LEVEL5_NWARPS * DMMA_M + && (cr % (LEVEL5_NWARPS * DMMA_M)) == 0 + && ((K * K) % cr) == 0 + && K * level5_a_stride(cr) * elem_bytes + <= LEVEL5_BUDGET - K * K * elem_bytes; +} + +constexpr int level5_chunk_rows(int K, int elem_bytes) { + return level5_chunk_ok(K, elem_bytes, (K * K) / 4) ? (K * K) / 4 + : level5_chunk_ok(K, elem_bytes, (K * K) / 8) ? (K * K) / 8 + : level5_chunk_ok(K, elem_bytes, (K * K) / 16) ? (K * K) / 16 + : LEVEL5_NWARPS * DMMA_M; /* K=8: DIMI is itself one 64-row strip */ +} + +#if MRA_HAVE_DMMA + +template +__device__ void mTxmq_level5_dmma(T* __restrict__ c, const T* a, T* b_smem) { + static_assert(std::is_same_v, + "mTxmq_level5_dmma: FP64 tensor cores operate on double only"); + static_assert(K % DMMA_N == 0, "mTxmq_level5_dmma: K must be a multiple of 8"); + + constexpr int DIMI = K * K; + + constexpr int CHUNK_ROWS = level5_chunk_rows(K, (int)sizeof(T)); + constexpr int NCHUNKS = DIMI / CHUNK_ROWS; + constexpr int A_STRIDE = level5_a_stride(CHUNK_ROWS); + constexpr int ROWS_PER_WARP = CHUNK_ROWS / LEVEL5_NWARPS; + constexpr int TILES_PER_WARP = ROWS_PER_WARP / DMMA_M; + constexpr int COL_TILES = K / DMMA_N; + + static_assert(TILES_PER_WARP >= 1, "level 5: chunk too small for 8 warps"); + + const int tid_block = (int)threadIdx.x; /* 0..255 */ + const int warp_id = tid_block / MRA_WARP_SIZE; /* 0..7 */ + + /* A strip buffer sits directly after B in shared memory */ + T* a_smem = b_smem + DIMI; + + for (int chunk = 0; chunk < NCHUNKS; ++chunk) { + const int row_base = chunk * CHUNK_ROWS; /* first global A^T row in strip */ + + /* --- Cooperative load of the A strip (all 256 threads) --------------- */ + for (int idx = tid_block; idx < K * CHUNK_ROWS; idx += LEVEL5_NTHREAD) { + const int row_local = idx % CHUNK_ROWS; /* row within strip */ + const int k = idx / CHUNK_ROWS; /* k-column of A */ + a_smem[k * A_STRIDE + row_local] = a[(size_t)k * DIMI + row_base + row_local]; + } + __syncthreads(); /* strip fully staged before any DMMA begins */ + + /* --- DMMA: each warp owns TILES_PER_WARP consecutive 8-row tiles ----- */ + const int warp_row_start = warp_id * ROWS_PER_WARP; + + for (int t = 0; t < TILES_PER_WARP; ++t) { + const int local_row = warp_row_start + t * DMMA_M; /* tile start in strip */ + + for (int ct = 0; ct < COL_TILES; ++ct) { + FragC acc; + nvcuda::wmma::fill_fragment(acc, 0.0); + + /* K/4 steps of 4-deep contraction */ + for (int kb = 0; kb < K; kb += DMMA_K) { + FragA a_frag; + FragB b_frag; + dmma_load_a(a_frag, a_smem, kb, local_row, A_STRIDE); + dmma_load_b(b_frag, b_smem, kb, ct * DMMA_N, K); + nvcuda::wmma::mma_sync(acc, a_frag, b_frag, acc); + } + + dmma_store_c(c, acc, row_base + local_row, ct * DMMA_N, K); + } + } + + __syncthreads(); /* all warps done before the strip is overwritten */ + } +} + +#endif /* MRA_HAVE_DMMA */ + +} // namespace detail + + +template +__device__ void mTxmq_level5_k(T* __restrict__ c, const T* a, const T* b) { + extern __shared__ char smem_level5[]; + T* b_smem = reinterpret_cast(smem_level5); + + /* All threads cooperate to load B once - it stays resident throughout */ + for (int idx = (int)threadIdx.x; idx < K * K; idx += (int)blockDim.x) + b_smem[idx] = b[idx]; + __syncthreads(); + +#if MRA_HAVE_DMMA + if constexpr (detail::dmma_supports_k(K) && std::is_same_v) { + detail::mTxmq_level5_dmma(c, a, b_smem); + __syncthreads(); + return; + } +#endif + detail::mTxmq_level3_impl(c, a, b_smem); + __syncthreads(); +} + +/* Host-side sizing - must mirror the constants used inside the kernel. */ +template +inline size_type mTxmq_level5_shmem_size(int K) { + const int DIMI = K * K; + if ((K % mra::detail::DMMA_N) != 0) { + return static_cast(DIMI * (int)sizeof(T)); /* level-3 fallback */ + } + const int chunk_rows = mra::detail::level5_chunk_rows(K, (int)sizeof(T)); + const int a_stride = mra::detail::level5_a_stride(chunk_rows); + return static_cast((DIMI + K * a_stride) * (int)sizeof(T)); +} + +template +constexpr Dim3 mTxmq_level5_blockdim(int /*K*/) { + return Dim3(detail::LEVEL5_NTHREAD, 1, 1); /* 8 warps */ +} + +} // namespace mra diff --git a/cuda/mxm_level7.h b/cuda/mxm_level7.h new file mode 100644 index 0000000..071aca3 --- /dev/null +++ b/cuda/mxm_level7.h @@ -0,0 +1,249 @@ +#pragma once + +#include "util.h" +#include "dmma.h" +#include "mxm_level3.h" + +/** + * Level 7: B resident in tensor-core registers across all three GEMMs. + * + * CUDA counterpart of the CDNA level-7 kernel. Both versions chase the same + * idea: the K x K matrix B is the one operand shared by all three passes of the + * transform, so load it into registers once and never touch memory for it + * again. The AMD source does this by hand, packing B into VGPRs according to + * the v_mfma_f64_16x16x4f64 lane mapping. Here a `wmma::fragment` + * *is* the register-resident operand, so the same trick falls out of holding + * the fragments live across the three passes. + * + * Block = 256 threads (8 warps). Warp w owns rows + * [w*K^2/8, (w+1)*K^2/8) of every output, as 8-row DMMA tiles. + * + * --- Pointer trick --- + * After each GEMM, C [K^2 x K] written row-major to shared memory is + * reinterpreted as A [K x K^2] col-major for the next GEMM via identical flat + * indices: + * Write: buf[i*K + j] (C row-major, i in [0,K^2), j in [0,K)) + * Read: buf[k*K^2 + i] (A col-major, k in [0,K), i in [0,K^2)) + * Since K*K^2 = K^3 = K^2*K, both index the same flat buffer - just different + * shapes. This is the standard MADNESS mTxmq pointer trick that makes the 3D + * separable transform work in place. + * + * --- Single-buffer reuse --- + * Within gemm7_pass every A fragment is pulled into registers before any write + * to dst, and a __syncthreads() separates the two phases, so the same shared + * buffer is safe to overwrite: + * GEMM 1: global A -> buf (shared, row-major) + * GEMM 2: buf -> buf (shared, in place) + * GEMM 3: buf -> C (global, row-major) + * + * Shared memory = K^3 * sizeof(T). For K=16: 16^3 * 8 = 32,768 bytes. + * + * --- A note on bank conflicts --- + * The AMD source applies an XOR swizzle to its LDS addresses, because a + * K^2-element stride aliases onto the same banks for every k. The same + * aliasing exists here (a 256-double stride is a multiple of the 32-bank + * cycle, so the four k-groups of a fragment contend 4 ways), but the swizzle + * cannot be carried over: `load_matrix_sync`/`store_matrix_sync` compute their + * own lane addresses from a base pointer and a stride, leaving no place to + * inject an address permutation. Padding is not an option either - it would + * break the pointer trick, which depends on the two views aliasing exactly. + * The conflicts are therefore accepted here; correctness is unaffected. + * + * Supported: K in {8, 16} on sm_80+. Other K values are dispatched to the + * level-3 kernel by the host-side submit function, and the device-side + * fallback below keeps the three-pass chain intact for pre-sm_80 builds. + */ + +namespace mra { + +namespace detail { + +constexpr int LEVEL7_NWARPS = 8; +constexpr int LEVEL7_NTHREAD = LEVEL7_NWARPS * MRA_WARP_SIZE; /* 256 */ + +/** + * K values that get the register-resident-B tensor-core path. + * + * K must be a multiple of 8 (tile geometry) and small enough that the whole A + * partition of a warp fits in registers: a warp holds + * (K^2/8/8) x (K/4) A fragments plus (K/4) x (K/8) B fragments. At K=16 that + * is 16 + 8 doubles per lane, plus 8 accumulator tiles; K=32 would need 128 A + * fragments per lane and spill. + */ +constexpr bool level7_supports_k(int K) { + return (K % DMMA_N) == 0 && K >= DMMA_N && K <= 16; +} + +#if MRA_HAVE_DMMA + +/** + * One GEMM pass: dst = src^T x B, with B already in b_frags. + * + * src is always addressed as A [K x K^2] col-major (ldm = K^2) and dst as + * C [K^2 x K] row-major (ldm = K) - the two halves of the pointer trick. That + * holds whether src/dst live in global or shared memory, so no layout template + * parameters are needed. + * + * The pass is split into load / sync / compute+store so that GEMM 2 can read + * and write the same buffer: no lane writes until every lane has read. + */ +template +__device__ __forceinline__ void gemm7_pass( + const T* src, + T* dst, + FragB b_frags[K / DMMA_K][K / DMMA_N], + int warp_row_offset) +{ + /* src and dst deliberately carry no __restrict__: GEMM 2 passes the same + * buffer for both, and promising the compiler they cannot alias would make + * that call ill-formed. */ + constexpr int K2 = K * K; + constexpr int ROWS_PER_WARP = K2 / LEVEL7_NWARPS; + constexpr int TILES_PER_WARP = ROWS_PER_WARP / DMMA_M; + constexpr int NSTEPS = K / DMMA_K; + constexpr int COL_TILES = K / DMMA_N; + + /* --- Pre-load this warp's whole A partition into registers -------------- */ + FragA a_frags[TILES_PER_WARP][NSTEPS]; + #pragma unroll + for (int t = 0; t < TILES_PER_WARP; ++t) { + #pragma unroll + for (int s = 0; s < NSTEPS; ++s) { + dmma_load_a(a_frags[t][s], src, + s * DMMA_K, /* contraction offset */ + warp_row_offset + t * DMMA_M, /* row in A^T */ + K2); /* col-major ldm */ + } + } + + /* Every lane has read; writes may now proceed even if dst aliases src. */ + __syncthreads(); + + /* --- Accumulate and store ---------------------------------------------- */ + #pragma unroll + for (int t = 0; t < TILES_PER_WARP; ++t) { + #pragma unroll + for (int ct = 0; ct < COL_TILES; ++ct) { + FragC acc; + nvcuda::wmma::fill_fragment(acc, 0.0); + #pragma unroll + for (int s = 0; s < NSTEPS; ++s) { + nvcuda::wmma::mma_sync(acc, a_frags[t][s], b_frags[s][ct], acc); + } + dmma_store_c(dst, acc, + warp_row_offset + t * DMMA_M, /* row in C */ + ct * DMMA_N, /* col in C */ + K); /* row-major ldm */ + } + } +} + +/** + * Three-GEMM chain for level 7. + * + * B is loaded into fragments once and stays in registers for all three passes. + * A single shared buffer (K^3 doubles) is reused in place. + */ +template +__device__ void mTxmq_level7_dmma( + T* __restrict__ c, /* output [K^2 x K] row-major, global */ + const T* __restrict__ a, /* input [K x K^2] col-major, global */ + const T* __restrict__ b, /* B [K x K] row-major, global */ + T* buf) /* shared scratch: K^3 doubles */ +{ + static_assert(std::is_same_v, + "mTxmq_level7_dmma: FP64 tensor cores operate on double only"); + static_assert(level7_supports_k(K), "mTxmq_level7_dmma: unsupported K"); + + constexpr int K2 = K * K; + constexpr int NSTEPS = K / DMMA_K; + constexpr int COL_TILES = K / DMMA_N; + + const int warp_id = (int)threadIdx.x / MRA_WARP_SIZE; + const int warp_row_offset = warp_id * (K2 / LEVEL7_NWARPS); + + /* Load B into registers once - resident for all three passes. */ + FragB b_frags[NSTEPS][COL_TILES]; + #pragma unroll + for (int s = 0; s < NSTEPS; ++s) { + #pragma unroll + for (int ct = 0; ct < COL_TILES; ++ct) { + dmma_load_b(b_frags[s][ct], b, s * DMMA_K, ct * DMMA_N, K); + } + } + + /* GEMM 1: A (global) -> buf (shared, row-major) */ + gemm7_pass(a, buf, b_frags, warp_row_offset); + __syncthreads(); + + /* GEMM 2: buf reread col-major via the pointer trick -> buf, in place */ + gemm7_pass(buf, buf, b_frags, warp_row_offset); + __syncthreads(); + + /* GEMM 3: buf reread col-major -> c (global, row-major) */ + gemm7_pass(buf, c, b_frags, warp_row_offset); +} + +#endif /* MRA_HAVE_DMMA */ + +} // namespace detail + + +/** + * Public interface: executes the full three-GEMM transform chain. + * + * `workspace` is used only by the non-tensor-core fallback, which ping-pongs + * between it and `c` exactly as the level-3 transform does. (The AMD source's + * fallback runs a single pass here rather than three, which silently produces a + * one-pass result; this version keeps the chain intact.) + */ +template +__device__ void mTxmq_level7_k( + T* __restrict__ c, + const T* __restrict__ a, + const T* __restrict__ b, + T* workspace) +{ + extern __shared__ char smem_level7[]; + T* buf = reinterpret_cast(smem_level7); + +#if MRA_HAVE_DMMA + if constexpr (detail::level7_supports_k(K) && std::is_same_v) { + (void)workspace; + detail::mTxmq_level7_dmma(c, a, b, buf); + return; + } +#endif + /* Fallback: level-3 register blocking, full three-pass chain. */ + for (int idx = (int)threadIdx.x; idx < K * K; idx += (int)blockDim.x) + buf[idx] = b[idx]; + __syncthreads(); + + T *t0 = workspace, *t1 = c; + { auto tmp = t0; t0 = t1; t1 = tmp; } + + detail::mTxmq_level3_impl(t0, a, buf); + __syncthreads(); + for (int n = 1; n < 3; ++n) { + detail::mTxmq_level3_impl(t1, t0, buf); + __syncthreads(); + auto tmp = t0; t0 = t1; t1 = tmp; + } +} + +template +inline size_type mTxmq_level7_shmem_size(int K) { + if (MRA_DMMA_SUPPORTED && detail::level7_supports_k(K)) { + /* Flat K^3 buffer: C [K^2 x K] row-major reinterpreted as A [K x K^2] + * col-major via the pointer trick. No padding - it would break aliasing. */ + return static_cast(K * K * K * (int)sizeof(T)); + } + return static_cast(K * K * (int)sizeof(T)); /* B only, fallback */ +} + +template +constexpr Dim3 mTxmq_level7_blockdim(int /*K*/) { + return Dim3(detail::LEVEL7_NTHREAD, 1, 1); +} + +} // namespace mra diff --git a/cuda/mxm_wmma.h b/cuda/mxm_wmma.h new file mode 100644 index 0000000..7cd4099 --- /dev/null +++ b/cuda/mxm_wmma.h @@ -0,0 +1,220 @@ +#pragma once + +#include "util.h" +#include "dmma.h" +#include "mxm_level3.h" + +/** + * Level 6: WMMA implementation of mTxmq, c(i,j) = sum_k a(k,i) * b(k,j). + * + * CUDA counterpart of the rocWMMA level. Both go through the vendor's + * warp-cooperative matrix API rather than raw intrinsics, but the fragment + * geometry differs and that changes the tiling arithmetic throughout: + * + * rocWMMA (CDNA) nvcuda::wmma (sm_80+) + * FP64 fragment 16 x 16 x 4 8 x 8 x 4 + * lanes per fragment 64 32 + * column padding N -> multiple of 16 N -> multiple of 8 + * + * Matrices: + * A : [K x K^2] col-major in the multiply (leading dimension = K^2) + * B : [K x K ] row-major (leading dimension = K) + * C : [K^2 x K] row-major (leading dimension = K) + * + * Supported K: any multiple of 4 (the FP64 contraction tile). That covers + * K = 4, 8, 12, 16, 20, 32; K = 6 and 10 fall back to level 3. + * + * Because the fragment is 8 wide rather than 16, K = 4 and 12 are the only + * supported values that still need column padding, and the smaller tile means + * the K < 16 special case that rocWMMA needs (its fragment cannot express a + * K < 16 GEMM at all) has no counterpart here - the WMMA path covers those K + * values directly. + * + * Work assignment: one warp per output tile, striding when there are more + * tiles than warps. Block size is capped at 32 warps (1024 threads). + * + * Shared memory: + * smem_b [K x N_PAD] zero-padded B + * smem_stage [NWARPS x 8 x 8] per-warp output tile, only when + * N is not a multiple of 8 + * + * When N is a multiple of 8 every column of a tile is valid and each warp + * stores straight to global C. Otherwise B's padding makes the surplus + * columns zero but they must not reach C, so each warp lands its tile in an + * 8x8 shared scratch and copies out only the N valid columns. (rocWMMA stages + * the entire M x N_PAD output instead; per-warp staging keeps shared memory + * bounded - the full-output variant would need 75 KB at K = 20.) + */ + +namespace mra { + +namespace detail { + +/* Zero-padded column count and tile counts for a given K. */ +constexpr int wmma_n_pad(int K) { return dmma_pad_n(K); } +constexpr int wmma_m_tiles(int K) { return (K * K) / DMMA_M; } +constexpr int wmma_n_tiles(int K) { return wmma_n_pad(K) / DMMA_N; } +constexpr int wmma_tiles(int K) { return wmma_m_tiles(K) * wmma_n_tiles(K); } + +constexpr int WMMA_MAX_WARPS = 32; /* 1024 threads, the block ceiling */ + +constexpr int wmma_warps(int K) { + return wmma_tiles(K) < WMMA_MAX_WARPS ? wmma_tiles(K) : WMMA_MAX_WARPS; +} + +/** K values the WMMA path handles: the contraction must tile by 4. */ +constexpr bool wmma_supports_k(int K) { + return K > 0 && (K % DMMA_K) == 0; +} + +/** True when B/C need column zero-padding. */ +constexpr bool wmma_needs_pad(int K) { return K != wmma_n_pad(K); } + +/** Shared memory in bytes: padded B, plus per-warp staging when padding. */ +template +constexpr size_type wmma_shmem_bytes(int K) { + return static_cast( + (K * wmma_n_pad(K) + + (wmma_needs_pad(K) ? wmma_warps(K) * DMMA_M * DMMA_N : 0)) + * (int)sizeof(T)); +} + +#if MRA_HAVE_DMMA + +/** + * Core device function: C[M x N] = A^T[M x K] x B[K x N] + * where M = K^2, N = K. + */ +template +__device__ void mTxmq_wmma_core(T* __restrict__ c, const T* a, const T* b, T* smem) { + static_assert(std::is_same_v, + "mTxmq_wmma_core: FP64 tensor cores operate on double only"); + static_assert(K % DMMA_K == 0, + "K must be divisible by the FP64 contraction tile (4)"); + + constexpr int M = K * K; + constexpr int N = K; + constexpr int N_PAD = wmma_n_pad(K); + constexpr int M_TILES = wmma_m_tiles(K); + constexpr int TOTAL_TILES = wmma_tiles(K); + constexpr int NWARPS = wmma_warps(K); + constexpr bool NEEDS_PAD = wmma_needs_pad(K); + + static_assert(M % DMMA_M == 0, + "K^2 must be 8-aligned; guaranteed whenever K % 4 == 0"); + + T* smem_b = smem; /* [K x N_PAD] */ + T* smem_stage = smem_b + K * N_PAD; /* [NWARPS x 8 x 8], padded case */ + + /* --- Phase 1: load B into smem_b with zero padding -------------------- */ + for (int idx = (int)threadIdx.x; idx < K * N_PAD; idx += (int)blockDim.x) { + const int ki = idx / N_PAD; + const int ni = idx % N_PAD; + smem_b[idx] = (ni < N) ? b[ki * N + ni] : T(0); + } + __syncthreads(); + + /* --- Phase 2: one warp per output tile, striding over tiles ----------- */ + const int warp_id = (int)threadIdx.x / MRA_WARP_SIZE; + const int lane = (int)threadIdx.x % MRA_WARP_SIZE; + + for (int tile = warp_id; tile < TOTAL_TILES; tile += NWARPS) { + const int tile_m = tile % M_TILES; + const int tile_n = tile / M_TILES; + const int m_start = tile_m * DMMA_M; + const int n_start = tile_n * DMMA_N; + + FragC c_frag; + nvcuda::wmma::fill_fragment(c_frag, 0.0); + + for (int k = 0; k < K; k += DMMA_K) { + /* A^T[m_start .. m_start+8, k .. k+4]. A is stored row-major as + * [K rows x M cols], i.e. a[k*M + i] = A[k][i] = A^T[i][k]. Declaring + * the fragment col_major with ldm = M makes element [m_i][k_j] resolve + * to ptr[k_j*M + m_i] = a[(k+k_j)*M + m_start+m_i] = A^T[m_start+m_i][k+k_j]. */ + FragA a_frag; + dmma_load_a(a_frag, a, k, m_start, M); + + /* B tile: smem_b[k .. k+4, n_start .. n_start+8], row-major, ldm = N_PAD */ + FragB b_frag; + dmma_load_b(b_frag, smem_b, k, n_start, N_PAD); + + nvcuda::wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); + } + + if constexpr (!NEEDS_PAD) { + /* Every column of the tile is valid: store straight to global C. + * Warps own disjoint row ranges, so there is no write conflict. */ + dmma_store_c(c, c_frag, m_start, n_start, N); + } else { + /* Surplus columns are zero but out of bounds for C. Land the tile in + * this warp's 8x8 scratch and copy out only the N valid columns. */ + T* stage = smem_stage + warp_id * DMMA_M * DMMA_N; + nvcuda::wmma::store_matrix_sync(stage, c_frag, DMMA_N, + nvcuda::wmma::mem_row_major); + __syncwarp(); + for (int idx = lane; idx < DMMA_M * DMMA_N; idx += MRA_WARP_SIZE) { + const int mi = idx / DMMA_N; + const int ni = idx % DMMA_N; + if (n_start + ni < N) { + c[(size_t)(m_start + mi) * N + n_start + ni] = stage[idx]; + } + } + __syncwarp(); /* scratch is reused on the next tile of this warp */ + } + } + __syncthreads(); +} + +#endif // MRA_HAVE_DMMA + +} // namespace detail + + +/** + * K-templated entry point - one binary per K value. + * Falls back to level-3 register blocking when K is not a multiple of 4 or the + * build targets a GPU without FP64 tensor cores. + */ +template +__device__ void mTxmq_wmma_k(T* __restrict__ c, const T* a, const T* b) { + extern __shared__ char smem_wmma[]; + T* smem = reinterpret_cast(smem_wmma); + +#if MRA_HAVE_DMMA + if constexpr (detail::wmma_supports_k(K) && std::is_same_v) { + detail::mTxmq_wmma_core(c, a, b, smem); + return; + } +#endif + for (int idx = (int)threadIdx.x; idx < K * K; idx += (int)blockDim.x) + smem[idx] = b[idx]; + __syncthreads(); + detail::mTxmq_level3_impl(c, a, smem); + __syncthreads(); +} + +template +inline size_type mTxmq_wmma_shmem_size(int K) { + if (MRA_DMMA_SUPPORTED && detail::wmma_supports_k(K)) { + return detail::wmma_shmem_bytes(K); + } + return static_cast(K * K * (int)sizeof(T)); /* B only, fallback */ +} + +/** + * Block dimension: one warp per output tile, capped at 32 warps. + * K= 4 : 2 warps = 64 threads + * K= 8 : 8 warps = 256 threads + * K=12 : 32 warps = 1024 threads (36 tiles, warps stride) + * K=16 : 32 warps = 1024 threads (64 tiles, warps stride) + */ +template +inline Dim3 mTxmq_wmma_blockdim(int K) { + if (MRA_DMMA_SUPPORTED && detail::wmma_supports_k(K)) { + return Dim3(detail::wmma_warps(K) * MRA_WARP_SIZE, 1, 1); + } + return Dim3(MAX_THREADS_PER_BLOCK, 1, 1); /* level-3 fallback */ +} + +} // namespace mra diff --git a/cuda/transform.h b/cuda/transform.h new file mode 100644 index 0000000..cd6164f --- /dev/null +++ b/cuda/transform.h @@ -0,0 +1,91 @@ +#ifndef HAVE_TRANSFORM_H +#define HAVE_TRANSFORM_H + +#include +#include "util.h" +#include "mxm.h" + +/***************************************** + * Level 1 - reference transform via mTxmq + * + * Everything stays in global memory: B is re-read from HBM for every row of A + * and A is re-read once per output column. This is the correctness reference + * that validate_levels compares every other level against, so it deliberately + * uses the plain mxm.h kernel and nothing else. + * + * (The dual-target sources also pulled in mxm_cublasdx.h here, where its + * `long`-typed mTxmq overload was silently shadowed by the exact-match + * `size_type` overload from mxm.h. In this tree the cuBLASDx block-GEMM path + * is reachable on a level of its own - see transform_cublasdx_mxm.h.) + *****************************************/ + +template +__device__ void transform( + int K, + const T* t, + const T* c, + T*& result, + T* workspace) +{ + constexpr const int ndim = 3; // fixed for benchmark + const T* pc = c; + T *t0=workspace, *t1=result; + { + auto tmp = t0; + t0 = t1; + t1 = tmp; + } + const int dimj = K; + int dimi = dimj*dimj; + mra::mTxmq(dimi, dimj, dimj, t0, t, pc); + for (int n=1; n +inline +LAUNCH_BOUNDS(MAX_THREADS_PER_BLOCK, 4) +__global__ void transform_kernel(int nfuncs, int K, const T* A, const T* B, T* C, T* workspace) { + + const T *a, *b; + T *c, *w; + int K2NDIM = K*K*K; + /* workspace is allocated for each thread-block */ + w = workspace + blockIdx.x * K2NDIM; + /* iterate over all tensors */ + for (int i = blockIdx.x; i < nfuncs; i += gridDim.x) { + a = A + i * K2NDIM; + b = B; + c = C + i * K2NDIM; + transform(K, a, b, c, w); + } +} + +template +inline int transform_shmem_size(int K) { + /* use whatever mTxm says we need */ + return mra::mTxmq_shmem_size(K); +} + +template +inline void submit_transform_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + Stream stream) +{ + Dim3 thread_dims = mra::mTxmq_blockdim(K); + assert(block_size(thread_dims) <= MAX_THREADS_PER_BLOCK); + size_type smem_size = mra::mTxmq_shmem_size(K); + size_type K2 = K*K; + if (smem_size < K2*(size_type)sizeof(T)) { + smem_size = K2*sizeof(T); + } + CONFIGURE_KERNEL(transform_kernel, smem_size); + CALL_KERNEL(transform_kernel, std::min(nfuncs, nblocks), thread_dims, smem_size, stream, (nfuncs, K, A, B, C, workspace)); +} + +#endif // HAVE_TRANSFORM_H diff --git a/cuda/transform_cublasdx.h b/cuda/transform_cublasdx.h new file mode 100644 index 0000000..f135d5b --- /dev/null +++ b/cuda/transform_cublasdx.h @@ -0,0 +1,254 @@ +#ifndef HAVE_TRANSFORM_CUBLASDX_H +#define HAVE_TRANSFORM_CUBLASDX_H + +#include "util.h" +#include "mxm_cublasdx.h" + + +/********************************************************************************** + * cublasDx implementation of transform + * + * The cublasDx implementation uses the cublasdx library directly to perform + * the tensor transformation. It relies on a single shared memory tensor + * to which the result of a GEMM is written in each iteration. + * The register fragment saves us the additional shared memory tensor + * that would be required to store the result of the GEMM. + **********************************************************************************/ + +#if __has_include() + +#define MRA_HAVE_CUBLASDX 1 + +template +__forceinline__ __device__ +void transform_cublasdx_k( + const T* t, // input tensor + const T* c, // input matrix + T* result) +{ + constexpr const int ndim = 3; // fixed for benchmark + using GEMM = typename mra::detail::GEMMBuilder::GEMM; + + using alignment = cublasdx::alignment_of; + + extern __shared__ __align__(16) char smem[]; + + auto [smem_a, smem_b] = + cublasdx::shared_memory::slice_into_pointers( + smem, + cublasdx::alignment_of_v_a, cublasdx::cosize(GET_SHARED_LAYOUT(GEMM, a)), + cublasdx::alignment_of_v_b, cublasdx::cosize(GET_SHARED_LAYOUT(GEMM, b))); + + + /* global memory tensors */ + auto a_global_tensor = cublasdx::make_tensor(t, GEMM::get_layout_gmem_a()); + auto b_global_tensor = cublasdx::make_tensor(c, GEMM::get_layout_gmem_b()); + auto c_global_tensor = cublasdx::make_tensor(result, GEMM::get_layout_gmem_c()); + + /* shared memory tensors */ + auto a_shared_tensor = cublasdx::make_tensor(smem_a, GET_SHARED_LAYOUT(GEMM, a)); + auto b_shared_tensor = cublasdx::make_tensor(smem_b, GET_SHARED_LAYOUT(GEMM, b)); + + cublasdx::copy(a_global_tensor, a_shared_tensor); + cublasdx::copy(b_global_tensor, b_shared_tensor); + + /* wait for loads to complete */ + cublasdx::copy_wait(); + + for (int n=0; n(c_register_fragment, a_shared_tensor, partitioner); + + /* wait for stores to complete */ + cublasdx::copy_wait(); + } + /* copy the result from shared memory to global memory */ + cublasdx::copy(a_shared_tensor, c_global_tensor); + + /* wait for the copy to complete */ + cublasdx::copy_wait(); +} + + +template +__forceinline__ __device__ +void transform_cublasdx( + const T* t, + const T* c, + T*& result, + T* workspace) +{ + return transform_cublasdx_k(t, c, result); + +#if 0 + (void)workspace; // unused in this implementation + switch (K) { + case 8 : transform_cublasdx_k(t, c, result); break; + case 10: transform_cublasdx_k(t, c, result); break; + case 16: transform_cublasdx_k(t, c, result); break; + case 20: transform_cublasdx_k(t, c, result); break; + default: + printf("Unsupported K value: %d\n", K); + return; + } +#endif // 0 + /* no need to synchronize here, cublasdx::copy_wait() synchronizes */ +} + +template +constexpr int transform_cublasdx_shmem_size_k() +{ + using GEMM = typename mra::detail::GEMMBuilder::GEMM; + + auto calc = cublasdx::make_shared_storage_calculator() + .add(cublasdx::alignment_of_v_a, sizeof(typename GEMM::a_value_type), GET_SHARED_LAYOUT(GEMM, a)) + .add(cublasdx::alignment_of_v_b, sizeof(typename GEMM::b_value_type), GET_SHARED_LAYOUT(GEMM, b)); + auto smem_size = calc.get(); + return smem_size; +} + +template +int transform_cublasdx_shmem_size(int K) +{ + switch (K) { + case 8 : return transform_cublasdx_shmem_size_k(); + case 10: return transform_cublasdx_shmem_size_k(); + case 16: return transform_cublasdx_shmem_size_k(); + case 20: return transform_cublasdx_shmem_size_k(); + default: + printf("Unsupported K value: %d\n", K); + return 0; + } + /* no need to synchronize here, cublasdx::copy_wait() synchronizes */ +} + + +template +constexpr auto transform_cublasdx_block_dim() +{ + using GEMM = typename mra::detail::GEMMBuilder::GEMM; + + return GEMM::suggested_block_dim; +} + + +template +constexpr auto transform_cublasdx_block_size() +{ + auto blockdims = transform_cublasdx_block_dim(); + return blockdims.x * blockdims.y * blockdims.z; +} + + +/* Runtime dispatcher over the compile-time block dims, for benchmark reporting. */ +template +inline Dim3 transform_cublasdx_blockdim(int K) +{ + switch (K) { + case 8 : { auto d = transform_cublasdx_block_dim(); return Dim3(d.x, d.y, d.z); } + case 10: { auto d = transform_cublasdx_block_dim(); return Dim3(d.x, d.y, d.z); } + case 16: { auto d = transform_cublasdx_block_dim(); return Dim3(d.x, d.y, d.z); } + case 20: { auto d = transform_cublasdx_block_dim(); return Dim3(d.x, d.y, d.z); } + default: + printf("transform_cublasdx_blockdim: unsupported K value: %d\n", K); + return Dim3(1, 1, 1); + } +} + + +template +LAUNCH_BOUNDS((transform_cublasdx_block_size()), 1) +__global__ void transform_cublasdx_kernel(int nfuncs, const T* A, const T* B, T* C, T* workspace) { + + const T *a, *b; + T *c, *w; + int K2NDIM = K*K*K; + /* workspace is allocated for each thread-block */ + w = workspace + blockIdx.x * K2NDIM; + /* iterate over all tensors */ + for (int i = blockIdx.x; i < nfuncs; i += gridDim.x) { + a = A + i * K2NDIM; + b = B; + c = C + i * K2NDIM; + transform_cublasdx(a, b, c, w); + } +} + +template +void submit_transform_cublasdx_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + cudaStream_t stream) +{ + auto smem_size = transform_cublasdx_shmem_size(K); + switch (K) { + case 8: { + CONFIGURE_KERNEL((transform_cublasdx_kernel), smem_size); + CALL_KERNEL((transform_cublasdx_kernel), std::min(nfuncs, nblocks), (transform_cublasdx_block_dim()), smem_size, stream, (nfuncs, A, B, C, workspace)); + break; + } + case 10: { + CONFIGURE_KERNEL((transform_cublasdx_kernel), smem_size); + CALL_KERNEL((transform_cublasdx_kernel), std::min(nfuncs, nblocks), (transform_cublasdx_block_dim()), smem_size, stream, (nfuncs, A, B, C, workspace)); + break; + } + case 16: { + CONFIGURE_KERNEL((transform_cublasdx_kernel), smem_size); + CALL_KERNEL((transform_cublasdx_kernel), std::min(nfuncs, nblocks), (transform_cublasdx_block_dim()), smem_size, stream, (nfuncs, A, B, C, workspace)); + break; + } + case 20: { + CONFIGURE_KERNEL((transform_cublasdx_kernel), smem_size); + CALL_KERNEL((transform_cublasdx_kernel), std::min(nfuncs, nblocks), (transform_cublasdx_block_dim()), smem_size, stream, (nfuncs, A, B, C, workspace)); + break; + } + default: + printf("Unsupported K value: %d\n", K); + throw std::runtime_error("Unsupported K value in transform_cublasdx_bench"); + } +} + +#else + +#define MRA_HAVE_CUBLASDX 0 + +template +void submit_transform_cublasdx_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + Stream stream) { + std::printf("CUBLASdx not available, cannot run benchmark\n"); +} + + +template +constexpr auto transform_cublasdx_block_size() { + return 1; +} + +template +int transform_cublasdx_shmem_size(int K) { + return 0; +} + +template +inline Dim3 transform_cublasdx_blockdim(int /*K*/) { + return Dim3(1, 1, 1); +} + +#endif // __has_include() + +#endif // HAVE_TRANSFORM_CUBLASDX_H \ No newline at end of file diff --git a/cuda/transform_cublasdx_mxm.h b/cuda/transform_cublasdx_mxm.h new file mode 100644 index 0000000..c33145e --- /dev/null +++ b/cuda/transform_cublasdx_mxm.h @@ -0,0 +1,102 @@ +#ifndef HAVE_TRANSFORM_CUBLASDX_MXM_H +#define HAVE_TRANSFORM_CUBLASDX_MXM_H + +#include +#include "util.h" +#include "mxm_cublasdx.h" + +/** + * Level 10 - cuBLASDx as a drop-in mTxmq. + * + * Where level 9 (transform_cublasdx.h) fuses all three passes and keeps the + * intermediate in shared memory, this level calls a cuBLASDx block GEMM once + * per pass, ping-ponging through the per-block workspace exactly like levels + * 1-3. The GEMM itself tiles A through shared memory with double buffering + * (mxm_cublasdx.h), so the two levels bracket the cost of the fusion: level 9 + * saves the round trip to global memory between passes, level 10 does not. + * + * In the dual-target sources this path was unreachable: mxm_cublasdx.h declared + * its entry point as `mTxmq(long, long, long, ...)`, and every call site passed + * ints, which bind exactly to the `size_type` reference overload in mxm.h. The + * entry point is named mTxmq_cublasdx here so the choice is explicit. + */ + +#if defined(MRA_HAVE_CUBLASDX) && MRA_HAVE_CUBLASDX + +template +__device__ void transform_cublasdx_mxm( + int K, + const T* t, + const T* c, + T*& result, + T* workspace) +{ + constexpr const int ndim = 3; + const T* pc = c; + T *t0 = workspace, *t1 = result; + { auto tmp = t0; t0 = t1; t1 = tmp; } + + const int dimj = K; + const int dimi = dimj * dimj; + mra::mTxmq_cublasdx(dimi, dimj, dimj, t0, t, pc); + for (int n = 1; n < ndim; ++n) { + mra::mTxmq_cublasdx(dimi, dimj, dimj, t1, t0, pc); + auto tmp = t0; t0 = t1; t1 = tmp; + } + /* mTxmq_cublasdx ends with __syncthreads() */ +} + +template +inline +LAUNCH_BOUNDS(MAX_THREADS_PER_BLOCK, 1) +__global__ void transform_kernel_cublasdx_mxm(int nfuncs, int K, + const T* A, const T* B, T* C, T* workspace) { + const int K2NDIM = K * K * K; + T* w = workspace + blockIdx.x * K2NDIM; + for (int i = blockIdx.x; i < nfuncs; i += gridDim.x) { + const T* a = A + i * K2NDIM; + T* c = C + i * K2NDIM; + transform_cublasdx_mxm(K, a, B, c, w); + } +} + +template +inline int transform_cublasdx_mxm_shmem_size(int K) { + return mra::mTxmq_cublasdx_shmem_size(K); +} + +template +inline Dim3 transform_cublasdx_mxm_blockdim(int K) { + return mra::mTxmq_cublasdx_blockdim(K); +} + +template +inline void submit_transform_cublasdx_mxm_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + Stream stream) +{ + Dim3 thread_dims = mra::mTxmq_cublasdx_blockdim(K); + int smem_size = mra::mTxmq_cublasdx_shmem_size(K); + CONFIGURE_KERNEL(transform_kernel_cublasdx_mxm, smem_size); + CALL_KERNEL(transform_kernel_cublasdx_mxm, std::min(nfuncs, nblocks), + thread_dims, smem_size, stream, + (nfuncs, K, A, B, C, workspace)); +} + +#else // MRA_HAVE_CUBLASDX + +template +inline int transform_cublasdx_mxm_shmem_size(int /*K*/) { return 0; } + +template +inline Dim3 transform_cublasdx_mxm_blockdim(int /*K*/) { return Dim3(1, 1, 1); } + +template +inline void submit_transform_cublasdx_mxm_bench(int, int, int, + const T*, const T*, T*, T*, Stream) { + std::printf("cuBLASDx not available, cannot run benchmark\n"); +} + +#endif // MRA_HAVE_CUBLASDX + +#endif // HAVE_TRANSFORM_CUBLASDX_MXM_H diff --git a/cuda/transform_kron.h b/cuda/transform_kron.h new file mode 100644 index 0000000..1809c10 --- /dev/null +++ b/cuda/transform_kron.h @@ -0,0 +1,136 @@ +#pragma once + +/** + * Level 8 - Kronecker product GEMM (cuBLAS). + * + * MATHEMATICAL BACKGROUND + * ----------------------- + * The standard 3-pass transform applies B^T along each mode of a K x K x K tensor: + * + * Pass 1: T1[j0,i1,i2] = SUM_{i0} A[i0,i1,i2] * B[i0,j0] (contract mode 0) + * Pass 2: T2[j0,j1,i2] = SUM_{i1} T1[j0,i1,i2] * B[i1,j1] (contract mode 1) + * Pass 3: C [j0,j1,j2] = SUM_{i2} T2[j0,j1,i2] * B[i2,j2] (contract mode 2) + * + * Vectorising the tensor (flattening to K^3 elements) turns this into a single + * matrix-vector product: + * + * vec(C) = KronMat * vec(A) + * + * where KronMat = B^T (x) B^T (x) B^T is the three-fold Kronecker product + * (a K^3 x K^3 matrix). Each entry is: + * + * KronMat[b, a] = B[a%K][b%K] * B[(a/K)%K][(b/K)%K] * B[a/K^2][b/K^2] + * + * with a = input linear index (i0 + K*i1 + K^2*i2) + * b = output linear index (j0 + K*j1 + K^2*j2) + * + * IMPLEMENTATION + * -------------- + * 1. build_kron_kernel - one GPU thread per (b, a) entry; called ONCE before + * the timing loop and cached for all subsequent batches. + * + * 2. submit_transform_kron_bench - calls cublasDgemm: + * + * C [K^3 x nfuncs] = KronMat [K^3 x K^3] x A [K^3 x nfuncs] + * + * Tensors are stored contiguously (tensor f occupies A[f*K^3 .. (f+1)*K^3-1]), + * so the batch dimension maps naturally to GEMM columns. + * + * TRADE-OFFS + * ---------- + * Pros + * - Single API call into a fully tuned library GEMM; one large DGEMM keeps + * the SMs busy where the per-tensor kernels of L1-L7 cannot. + * + * Cons + * - KronMat memory = K^6 x 8 bytes: 6 MB at K=10, 128 MB at K=16, + * 512 MB at K=20 - impractical beyond K ~ 16. + * - FLOPs reported are 2*K^6*N (actual GEMM work), not the 3*2*K^4*N + * mathematical minimum, so raw GFlop/s are not directly comparable + * to L1-L7 or L9/L10. + */ + +#include "util.h" +#include + +using blasHandle_t = cublasHandle_t; +#define BLAS_OP_N CUBLAS_OP_N +#define blasCreate cublasCreate +#define blasDestroy cublasDestroy +#define blasSetStream cublasSetStream +#define blasDgemm cublasDgemm + +// --------------------------------------------------------------------------- +// Kernel: build the K^3 x K^3 Kronecker product matrix (column-major). +// +// KronMat[I, J] = B^T[i0,j0] * B^T[i1,j1] * B^T[i2,j2] +// +// Index decomposition (first index fastest = column-major vector): +// I = i0 + K*i1 + K^2*i2 +// J = j0 + K*j1 + K^2*j2 +// +// B is row-major K x K, so B^T[i,j] = B[j*K + i]. +// --------------------------------------------------------------------------- +template +__global__ void build_kron_kernel(int K, const T* __restrict__ B, + T* __restrict__ KronMat) +{ + const int K3 = K * K * K; + const int I = blockIdx.x * blockDim.x + threadIdx.x; + const int J = blockIdx.y * blockDim.y + threadIdx.y; + if (I >= K3 || J >= K3) return; + + const int i0 = I % K, j0 = J % K; + const int i1 = (I / K) % K, j1 = (J / K) % K; + const int i2 = I / (K * K), j2 = J / (K * K); + + // B^T[i,j] = B[j*K + i] (B is row-major) + KronMat[(size_t)I + (size_t)J * K3] = B[j0*K + i0] * B[j1*K + i1] * B[j2*K + i2]; +} + +// --------------------------------------------------------------------------- +// Build the Kronecker matrix on the device (call once before timing). +// KronMat must already be allocated with K^3 x K^3 elements. +// --------------------------------------------------------------------------- +template +inline void build_kron_matrix(int K, const T* B_dev, T* KronMat_dev, + Stream stream) +{ + const int K3 = K * K * K; + dim3 block(16, 16); + dim3 grid((K3 + 15) / 16, (K3 + 15) / 16); + CALL_KERNEL(build_kron_kernel, grid, block, 0, stream, + (K, B_dev, KronMat_dev)); +} + +// --------------------------------------------------------------------------- +// Submit one round of the Kronecker GEMM (called inside the timing loop). +// +// C[K^3 x nfuncs] = KronMat[K^3 x K^3] x A[K^3 x nfuncs] +// +// A and C are treated as column-major (each contiguous K^3-block = one tensor). +// --------------------------------------------------------------------------- +template +inline void submit_transform_kron_bench(int nfuncs, int K, + const T* A, const T* KronMat, T* C, + blasHandle_t blas_handle, + Stream stream) +{ + const int K3 = K * K * K; + const double alpha = 1.0, beta = 0.0; + blasSetStream(blas_handle, stream); + blasDgemm(blas_handle, + BLAS_OP_N, BLAS_OP_N, + K3, nfuncs, K3, + &alpha, + KronMat, K3, + A, K3, + &beta, + C, K3); +} + +// Required by the benchmark dispatch (values are unused for the Kronecker level). +template +inline int kron_shmem_size(int /*K*/) { return 0; } + +inline Dim3 kron_blockdim(int /*K*/) { return {1, 1, 1}; } diff --git a/cuda/transform_level2.h b/cuda/transform_level2.h new file mode 100644 index 0000000..7256311 --- /dev/null +++ b/cuda/transform_level2.h @@ -0,0 +1,65 @@ +#pragma once + +#include "util.h" +#include "mxm_level2.h" + +/** + * Transform wrapper for Level-2 (B cached in shared memory). + * Follows the same structure as transform.h / transform_cublasdx.h. + */ + +template +__device__ void transform_level2( + int K, + const T* t, /* input tensor K^3 */ + const T* c, /* coefficient matrix K^2 */ + T*& result, /* output tensor K^3 (pointer updated on swap) */ + T* workspace) /* per-block scratch K^3 */ +{ + constexpr int ndim = 3; + const T* pc = c; + T *t0 = workspace, *t1 = result; + /* swap so t0 points at the output buffer first */ + auto tmp = t0; t0 = t1; t1 = tmp; + + const int dimj = K; + const int dimi = dimj * dimj; + + mra::mTxmq_level2(dimi, dimj, dimj, t0, t, pc); + for (int n = 1; n < ndim; ++n) { + mra::mTxmq_level2(dimi, dimj, dimj, t1, t0, pc); + auto tmp2 = t0; t0 = t1; t1 = tmp2; + } + /* mTxmq_level2 ends with __syncthreads(); no extra sync needed */ +} + +template +LAUNCH_BOUNDS(MAX_THREADS_PER_BLOCK, 4) +__global__ void transform_kernel_level2(int nfuncs, int K, + const T* A, const T* B, T* C, T* workspace) { + const int K2NDIM = K * K * K; + T* w = workspace + blockIdx.x * K2NDIM; + for (int i = blockIdx.x; i < nfuncs; i += gridDim.x) { + const T* a = A + i * K2NDIM; + T* c = C + i * K2NDIM; + transform_level2(K, a, B, c, w); + } +} + +template +inline int transform_level2_shmem_size(int K) { + return mra::mTxmq_level2_shmem_size(K); +} + +template +inline void submit_transform_level2_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + Stream stream) +{ + Dim3 thread_dims = mra::mTxmq_level2_blockdim(K); + auto smem_size = mra::mTxmq_level2_shmem_size(K); + CONFIGURE_KERNEL(transform_kernel_level2, smem_size); + CALL_KERNEL(transform_kernel_level2, std::min(nfuncs, nblocks), + thread_dims, smem_size, stream, + (nfuncs, K, A, B, C, workspace)); +} diff --git a/cuda/transform_level3.h b/cuda/transform_level3.h new file mode 100644 index 0000000..7b19f95 --- /dev/null +++ b/cuda/transform_level3.h @@ -0,0 +1,83 @@ +#pragma once + +#include "util.h" +#include "mxm_level3.h" + +/** + * Transform wrapper for Level-3 (B in shared memory + register accumulation). + * Each K value gets its own kernel binary via template, + * isolating register pressure to acc[K] for that specific K. + */ + +template +__device__ void transform_level3_k( + const T* t, + const T* c, + T*& result, + T* workspace) +{ + constexpr int ndim = 3; + constexpr int K2NDIM = K * K * K; + + T *t0 = workspace, *t1 = result; + auto tmp = t0; t0 = t1; t1 = tmp; + + /* B is passed straight through to mTxmq_level3_k, which currently reads it + * from global memory: the staging copy in mxm_level3.h is commented out. */ + mra::mTxmq_level3_k(t0, t, c); + for (int n = 1; n < ndim; ++n) { + mra::mTxmq_level3_k(t1, t0, c); + auto tmp2 = t0; t0 = t1; t1 = tmp2; + } +} + +/* One kernel binary per K — register pressure is proportional to K, not max(K). */ +template +LAUNCH_BOUNDS(MAX_THREADS_PER_BLOCK, 1) +__global__ void transform_kernel_level3_k(int nfuncs, + const T* A, const T* B, T* C, T* workspace) { + constexpr int K2NDIM = K * K * K; + T* w = workspace + blockIdx.x * K2NDIM; + for (int i = blockIdx.x; i < nfuncs; i += gridDim.x) { + const T* a = A + i * K2NDIM; + T* c = C + i * K2NDIM; + /* result pointer starts at c; workspace is w */ + T* result = c; + transform_level3_k(a, B, result, w); + } +} + +template +inline int transform_level3_shmem_size(int K) { + return mra::mTxmq_level3_shmem_size(K); +} + +template +inline void submit_transform_level3_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + Stream stream) +{ + Dim3 thread_dims = mra::mTxmq_level3_blockdim(K); + int smem_size = mra::mTxmq_level3_shmem_size(K); + +#define DISPATCH_L3(Kval) \ + case Kval: \ + CONFIGURE_KERNEL((transform_kernel_level3_k), smem_size); \ + CALL_KERNEL((transform_kernel_level3_k), std::min(nfuncs, nblocks), \ + thread_dims, smem_size, stream, \ + (nfuncs, A, B, C, workspace)); \ + break; + + switch (K) { + DISPATCH_L3( 6) + DISPATCH_L3( 8) + DISPATCH_L3(10) + DISPATCH_L3(12) + DISPATCH_L3(16) + DISPATCH_L3(20) + DISPATCH_L3(32) + default: + printf("submit_transform_level3_bench: unsupported K=%d\n", K); + } +#undef DISPATCH_L3 +} diff --git a/cuda/transform_level4.h b/cuda/transform_level4.h new file mode 100644 index 0000000..945d8fa --- /dev/null +++ b/cuda/transform_level4.h @@ -0,0 +1,87 @@ +#pragma once + +#include "util.h" +#include "mxm_level4.h" + +/** + * Transform wrapper for Level-4 (FP64 tensor cores on sm_80+, falling back to + * Level-3 on older targets or for K values without a whole 8x8 tiling). + * + * Each K value gets its own kernel binary via template, isolating + * register pressure to the specific K being compiled. Block dimension is 32 + * threads (one warp) - the CUDA analogue of the AMD kernel's single wavefront. + */ + +template +__device__ void transform_level4_k( + const T* t, + const T* c, + T*& result, + T* workspace) +{ + constexpr int ndim = 3; + + T *t0 = workspace, *t1 = result; + auto tmp = t0; t0 = t1; t1 = tmp; + + mra::mTxmq_level4_k(t0, t, c); + for (int n = 1; n < ndim; ++n) { + mra::mTxmq_level4_k(t1, t0, c); + auto tmp2 = t0; t0 = t1; t1 = tmp2; + } +} + +/* One kernel binary per K. */ +template +LAUNCH_BOUNDS(MRA_WARP_SIZE, 1) +__global__ void transform_kernel_level4_k(int nfuncs, + const T* A, const T* B, T* C, T* workspace) { + constexpr int K2NDIM = K * K * K; + T* w = workspace + blockIdx.x * K2NDIM; + for (int i = blockIdx.x; i < nfuncs; i += gridDim.x) { + const T* a = A + i * K2NDIM; + T* c = C + i * K2NDIM; + T* result = c; + transform_level4_k(a, B, result, w); + } +} + +template +inline int transform_level4_shmem_size(int K) { + return mra::mTxmq_level4_shmem_size(K); +} + +template +inline Dim3 transform_level4_blockdim(int K) { + return mra::mTxmq_level4_blockdim(K); +} + +template +inline void submit_transform_level4_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + Stream stream) +{ + Dim3 thread_dims = mra::mTxmq_level4_blockdim(K); + int smem_size = mra::mTxmq_level4_shmem_size(K); + +#define DISPATCH_L4(Kval) \ + case Kval: \ + CONFIGURE_KERNEL((transform_kernel_level4_k), smem_size); \ + CALL_KERNEL((transform_kernel_level4_k), std::min(nfuncs, nblocks), \ + thread_dims, smem_size, stream, \ + (nfuncs, A, B, C, workspace)); \ + break; + + switch (K) { + DISPATCH_L4( 6) + DISPATCH_L4( 8) + DISPATCH_L4(10) + DISPATCH_L4(12) + DISPATCH_L4(16) + DISPATCH_L4(20) + DISPATCH_L4(32) + default: + printf("submit_transform_level4_bench: unsupported K=%d\n", K); + } +#undef DISPATCH_L4 +} diff --git a/cuda/transform_level5.h b/cuda/transform_level5.h new file mode 100644 index 0000000..b91b3e3 --- /dev/null +++ b/cuda/transform_level5.h @@ -0,0 +1,87 @@ +#pragma once + +#include "util.h" +#include "mxm_level5.h" + +/** + * Transform wrapper for Level-5 (shared-memory-staged A, 8-warp DMMA block). + * + * Block = 256 threads (8 warps, all issuing DMMA). The 3-GEMM chain is + * preserved: each call to mTxmq_level5_k loads B once, then the tensor-core + * kernel iterates over A strips internally. __syncthreads() between steps + * ensures the full output of one GEMM is visible before the next starts. + */ + +template +__device__ void transform_level5_k( + const T* t, + const T* c, + T*& result, + T* workspace) +{ + constexpr int ndim = 3; + + T *t0 = workspace, *t1 = result; + auto tmp = t0; t0 = t1; t1 = tmp; + + mra::mTxmq_level5_k(t0, t, c); + for (int n = 1; n < ndim; ++n) { + mra::mTxmq_level5_k(t1, t0, c); + auto tmp2 = t0; t0 = t1; t1 = tmp2; + } +} + +/* One kernel binary per K. */ +template +LAUNCH_BOUNDS(mra::detail::LEVEL5_NTHREAD, 1) +__global__ void transform_kernel_level5_k(int nfuncs, + const T* A, const T* B, T* C, T* workspace) { + constexpr int K2NDIM = K * K * K; + T* w = workspace + blockIdx.x * K2NDIM; + for (int i = blockIdx.x; i < nfuncs; i += gridDim.x) { + const T* a = A + i * K2NDIM; + T* c = C + i * K2NDIM; + T* result = c; + transform_level5_k(a, B, result, w); + } +} + +template +inline int transform_level5_shmem_size(int K) { + return mra::mTxmq_level5_shmem_size(K); +} + +template +inline Dim3 transform_level5_blockdim(int K) { + return mra::mTxmq_level5_blockdim(K); +} + +template +inline void submit_transform_level5_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + Stream stream) +{ + Dim3 thread_dims = mra::mTxmq_level5_blockdim(K); + int smem_size = mra::mTxmq_level5_shmem_size(K); + +#define DISPATCH_L5(Kval) \ + case Kval: \ + CONFIGURE_KERNEL((transform_kernel_level5_k), smem_size); \ + CALL_KERNEL((transform_kernel_level5_k), std::min(nfuncs, nblocks), \ + thread_dims, smem_size, stream, \ + (nfuncs, A, B, C, workspace)); \ + break; + + switch (K) { + DISPATCH_L5( 6) + DISPATCH_L5( 8) + DISPATCH_L5(10) + DISPATCH_L5(12) + DISPATCH_L5(16) + DISPATCH_L5(20) + DISPATCH_L5(32) + default: + printf("submit_transform_level5_bench: unsupported K=%d\n", K); + } +#undef DISPATCH_L5 +} diff --git a/cuda/transform_level7.h b/cuda/transform_level7.h new file mode 100644 index 0000000..7ed9b48 --- /dev/null +++ b/cuda/transform_level7.h @@ -0,0 +1,80 @@ +#pragma once + +#include "util.h" +#include "mxm_level7.h" +#include "transform_level3.h" + +/** + * Transform wrapper for Level 7. + * + * Unlike levels 1-6 the three-GEMM chain runs inside a single call to + * mTxmq_level7_k, because B has to stay in tensor-core registers across all + * three GEMMs. One shared buffer of K^3 doubles is reused in place via the + * pointer trick; only the final output reaches global memory C. + * + * For K=16: shared memory = 16^3 * 8 = 32,768 bytes. occupancy=1 is retained + * to leave register headroom for the B fragments and the per-warp A partition. + * + * K values outside {8, 16} are dispatched to the level-3 kernel by the submit + * function below rather than being handled in-kernel, so the fallback keeps + * level-3's own block and shared-memory configuration. + */ + +template +LAUNCH_BOUNDS(mra::detail::LEVEL7_NTHREAD, 1) +__global__ void transform_kernel_level7_k(int nfuncs, + const T* A, const T* B, T* C, T* workspace) +{ + constexpr int K3 = K * K * K; + T* w = workspace + blockIdx.x * K3; + for (int i = blockIdx.x; i < nfuncs; i += gridDim.x) { + const T* a = A + i * K3; + T* c = C + i * K3; + mra::mTxmq_level7_k(c, a, B, w); + } +} + +template +inline int transform_level7_shmem_size(int K) { + return (int)mra::mTxmq_level7_shmem_size(K); +} + +template +inline Dim3 transform_level7_blockdim(int K) { + if (MRA_DMMA_SUPPORTED && mra::detail::level7_supports_k(K)) { + return mra::mTxmq_level7_blockdim(K); + } + return mra::mTxmq_level3_blockdim(K); +} + +template +inline void submit_transform_level7_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + Stream stream) +{ + /* K outside the register-resident-B range runs the level-3 kernel, which + * brings its own block dim and shared-memory size. */ + if (!MRA_DMMA_SUPPORTED || !mra::detail::level7_supports_k(K)) { + submit_transform_level3_bench(nfuncs, nblocks, K, A, B, C, workspace, stream); + return; + } + + Dim3 thread_dims = mra::mTxmq_level7_blockdim(K); + int smem_size = transform_level7_shmem_size(K); + +#define DISPATCH_L7(Kval) \ + case Kval: \ + CONFIGURE_KERNEL((transform_kernel_level7_k), smem_size); \ + CALL_KERNEL((transform_kernel_level7_k), std::min(nfuncs, nblocks), \ + thread_dims, smem_size, stream, \ + (nfuncs, A, B, C, workspace)); \ + break; + + switch (K) { + DISPATCH_L7( 8) + DISPATCH_L7(16) + default: + printf("submit_transform_level7_bench: unsupported K=%d\n", K); + } +#undef DISPATCH_L7 +} diff --git a/cuda/transform_wmma.h b/cuda/transform_wmma.h new file mode 100644 index 0000000..8a9e109 --- /dev/null +++ b/cuda/transform_wmma.h @@ -0,0 +1,90 @@ +#pragma once + +#include "util.h" +#include "mxm_wmma.h" + +/** + * Transform wrapper for Level-6 (nvcuda::wmma), the CUDA counterpart of the + * rocWMMA level. + * + * The three GEMM passes ping-pong between the per-block workspace and C, the + * same structure levels 2 and 3 use. The rocWMMA source keeps the whole + * tensor resident in LDS between passes; that is not reproduced here because + * an 8x8x4 fragment covers K < 16 directly, so the LDS-resident small-K + * special case it needs (transform_klt16) has no counterpart. + */ + +template +__device__ void transform_wmma_k( + const T* t, + const T* c, + T*& result, + T* workspace) +{ + constexpr int ndim = 3; + + T *t0 = workspace, *t1 = result; + auto tmp = t0; t0 = t1; t1 = tmp; + + mra::mTxmq_wmma_k(t0, t, c); + for (int n = 1; n < ndim; ++n) { + mra::mTxmq_wmma_k(t1, t0, c); + auto tmp2 = t0; t0 = t1; t1 = tmp2; + } +} + +/* One kernel binary per K. The launch bound is the block ceiling used by the + * widest K (32 warps); narrower K launch fewer threads. */ +template +LAUNCH_BOUNDS(mra::detail::WMMA_MAX_WARPS * MRA_WARP_SIZE, 1) +__global__ void transform_kernel_wmma_k(int nfuncs, + const T* A, const T* B, T* C, T* workspace) { + constexpr int K2NDIM = K * K * K; + T* w = workspace + blockIdx.x * K2NDIM; + for (int i = blockIdx.x; i < nfuncs; i += gridDim.x) { + const T* a = A + i * K2NDIM; + T* c = C + i * K2NDIM; + T* result = c; + transform_wmma_k(a, B, result, w); + } +} + +template +inline int transform_wmma_shmem_size(int K) { + return mra::mTxmq_wmma_shmem_size(K); +} + +template +inline Dim3 transform_wmma_blockdim(int K) { + return mra::mTxmq_wmma_blockdim(K); +} + +template +inline void submit_transform_wmma_bench(int nfuncs, int nblocks, int K, + const T* A, const T* B, T* C, T* workspace, + Stream stream) +{ + Dim3 thread_dims = mra::mTxmq_wmma_blockdim(K); + int smem_size = mra::mTxmq_wmma_shmem_size(K); + +#define DISPATCH_L6(Kval) \ + case Kval: \ + CONFIGURE_KERNEL((transform_kernel_wmma_k), smem_size); \ + CALL_KERNEL((transform_kernel_wmma_k), std::min(nfuncs, nblocks), \ + thread_dims, smem_size, stream, \ + (nfuncs, A, B, C, workspace)); \ + break; + + switch (K) { + DISPATCH_L6( 6) + DISPATCH_L6( 8) + DISPATCH_L6(10) + DISPATCH_L6(12) + DISPATCH_L6(16) + DISPATCH_L6(20) + DISPATCH_L6(32) + default: + printf("submit_transform_wmma_bench: unsupported K=%d\n", K); + } +#undef DISPATCH_L6 +} diff --git a/cuda/transformbench.cu b/cuda/transformbench.cu new file mode 100644 index 0000000..0cc5682 --- /dev/null +++ b/cuda/transformbench.cu @@ -0,0 +1,244 @@ +#include +#include +#include + +#include "util.h" +#include "transform.h" // L1 +#include "transform_level2.h" // L2 +#include "transform_level3.h" // L3 +#include "transform_level4.h" // L4 +#include "transform_level5.h" // L5 +#include "transform_wmma.h" // L6 +#include "transform_level7.h" // L7 +#include "transform_kron.h" // L8 +#include "transform_cublasdx.h" // L9 +#include "transform_cublasdx_mxm.h" // L10 + +/** + * Optimization levels (CUDA): + * 1 - L1: thread-parallel over j, serial k-loop, all global memory (reference) + * 2 - L2: B in shared memory, threads distributed over rows + * 3 - L3: B in shared memory + register accumulation (acc[K] in registers) + * 4 - L4: FP64 tensor cores (mma.sync.m8n8k4.f64), one warp; L3 fallback + * 5 - L5: FP64 tensor cores, A staged through shared memory, 8 warps + * 6 - L6: nvcuda::wmma, one warp per output tile + * 7 - L7: FP64 tensor cores with B resident in registers across all 3 GEMMs + * 8 - L8: single K^3 x K^3 DGEMM via the Kronecker product (cuBLAS) + * 9 - L9: cuBLASDx, all three GEMMs fused in shared memory + * 10 - L10: cuBLASDx as a per-pass block GEMM + * + * Levels 4-7 need sm_80 or newer; on older targets they fall back to L3. + * Levels 9 and 10 need cuBLASDx headers at build time. + */ + +template +void transform_bench(int nreps, int ntasks, int nfuncs, int nblocks, int K, int level, int num_streams) { + + std::vector streams(num_streams); // PaRSEC uses 4 streams by default + T* A, *B, *C, *workspace; + MALLOC(&A, (size_t)nfuncs * K * K * K * sizeof(T)); // N x KxKxK tensors + MALLOC(&B, (size_t)K * K * sizeof(T)); // KxK matrix + MALLOC(&C, (size_t)nfuncs * K * K * K * sizeof(T)); // N x KxKxK tensors + MALLOC(&workspace, (size_t)nblocks * K * K * K * sizeof(T)); // per-block scratch + + for (int i = 0; i < num_streams; ++i) { + CREATE_STREAM(&streams[i]); + } + + /* Warn early if a level is unavailable */ + if ((level == 9 || level == 10) && !MRA_HAVE_CUBLASDX) { + std::cerr << "Warning: level " << level << " (cuBLASDx) requested but cuBLASDx " + "was not found at build time; falling back to level 3\n"; + level = 3; + } + if (level >= 4 && level <= 7 && !MRA_DMMA_SUPPORTED) { + std::cerr << "Warning: level " << level << " needs FP64 tensor cores (sm_80+); " + "this build targets sm_" << MRA_CUDA_ARCH + << ", so the kernels will run their level-3 fallback\n"; + } + + /* Resolve default level */ + if (level <= 0) { + level = (MRA_HAVE_CUBLASDX) ? 9 : 3; + } + + const char* level_names[] = { + "", /* unused [0] */ + "L1-global", /* 1 */ + "L2-smem_b", /* 2 */ + "L3-regblk", /* 3 */ + "L4-dmma", /* 4 */ + "L5-dmma-staged", /* 5 */ + "L6-wmma", /* 6 */ + "L7-dmma-breg", /* 7 */ + "L8-kron", /* 8 */ + "L9-cublasdx", /* 9 */ + "L10-cublasdx-mxm" /* 10 */ + }; + + /* Print shmem and thread dims for this level */ + int smem_size = 0; + Dim3 thread_dims = {1, 1, 1}; + switch (level) { + case 1: + smem_size = mra::mTxmq_shmem_size(K); + thread_dims = mra::mTxmq_blockdim(K); + break; + case 2: + smem_size = transform_level2_shmem_size(K); + thread_dims = mra::mTxmq_level2_blockdim(K); + break; + case 3: + smem_size = transform_level3_shmem_size(K); + thread_dims = mra::mTxmq_level3_blockdim(K); + break; + case 4: + smem_size = transform_level4_shmem_size(K); + thread_dims = transform_level4_blockdim(K); + break; + case 5: + smem_size = transform_level5_shmem_size(K); + thread_dims = transform_level5_blockdim(K); + break; + case 6: + smem_size = transform_wmma_shmem_size(K); + thread_dims = transform_wmma_blockdim(K); + break; + case 7: + smem_size = transform_level7_shmem_size(K); + thread_dims = transform_level7_blockdim(K); + break; + case 8: + smem_size = kron_shmem_size(K); + thread_dims = kron_blockdim(K); + break; + case 9: + smem_size = transform_cublasdx_shmem_size(K); + thread_dims = transform_cublasdx_blockdim(K); + break; + case 10: + smem_size = transform_cublasdx_mxm_shmem_size(K); + thread_dims = transform_cublasdx_mxm_blockdim(K); + break; + } + + /* Level 8: build the Kronecker matrix once, before the timing loop */ + T* KronMat = nullptr; + blasHandle_t blas_handle{}; + if (level == 8) { + const int K3 = K * K * K; + const size_t kron_bytes = (size_t)K3 * K3 * sizeof(T); + std::cout << "L8-kron: allocating " << kron_bytes / (1024*1024.0) + << " MB for " << K3 << "x" << K3 << " Kronecker matrix\n"; + MALLOC(&KronMat, kron_bytes); + blasCreate(&blas_handle); + build_kron_matrix(K, B, KronMat, streams[0]); + SYNC_STREAM(streams[0]); + } + + std::chrono::time_point beg, end; + + for (int i = 0; i < nreps+1; ++i) { + beg = std::chrono::high_resolution_clock::now(); + for (int t = 0; t < ntasks; ++t) { + switch (level) { + case 1: + submit_transform_bench(nfuncs, nblocks, K, A, B, C, workspace, streams[t%num_streams]); + break; + case 2: + submit_transform_level2_bench(nfuncs, nblocks, K, A, B, C, workspace, streams[t%num_streams]); + break; + case 3: + submit_transform_level3_bench(nfuncs, nblocks, K, A, B, C, workspace, streams[t%num_streams]); + break; + case 4: + submit_transform_level4_bench(nfuncs, nblocks, K, A, B, C, workspace, streams[t%num_streams]); + break; + case 5: + submit_transform_level5_bench(nfuncs, nblocks, K, A, B, C, workspace, streams[t%num_streams]); + break; + case 6: + submit_transform_wmma_bench(nfuncs, nblocks, K, A, B, C, workspace, streams[t%num_streams]); + break; + case 7: + submit_transform_level7_bench(nfuncs, nblocks, K, A, B, C, workspace, streams[t%num_streams]); + break; + case 8: + submit_transform_kron_bench(nfuncs, K, A, KronMat, C, blas_handle, streams[t%num_streams]); + break; + case 9: + submit_transform_cublasdx_bench(nfuncs, nblocks, K, A, B, C, workspace, streams[t%num_streams]); + break; + case 10: + submit_transform_cublasdx_mxm_bench(nfuncs, nblocks, K, A, B, C, workspace, streams[t%num_streams]); + break; + } + } + for (int t = 0; t < num_streams; ++t) { + SYNC_STREAM(streams[t]); + } + end = std::chrono::high_resolution_clock::now(); + + /* skip warm-up */ + if (i > 0) { + auto us = (std::chrono::duration_cast(end - beg).count()); + /* L8 does one K^3 x K^3 GEMM per task (2*K^6 FLOPs); the rest do 3 passes (3*2*K^4 FLOPs) */ + uint64_t flops = (level == 8) + ? (uint64_t)ntasks * 2 * (uint64_t)K*K*K * (uint64_t)K*K*K * nfuncs + : (uint64_t)ntasks * K * K * K * K * 3 * 2 * nfuncs; + std::cout << "Transform" + << ";level=" << level_names[level] + << ";nfuncs=" << nfuncs + << ";nblocks=" << nblocks + << ";K=" << K + << ";tasks=" << ntasks + << ";threads={" << thread_dims.x << "," << thread_dims.y << "," << thread_dims.z << "}" + << ";smem=" << smem_size + << ";Time(us)=" << us + << ";GFlop=" << flops*1e-9 + << ";Gflop/s=" << (1e-3 * flops) / us + << std::endl; + } + } + + if (level == 8) { + blasDestroy(blas_handle); + FREE(KronMat); + } + + for (int i = 0; i < num_streams; ++i) { + DESTROY_STREAM(streams[i]); + } + + FREE(A); + FREE(B); + FREE(C); + FREE(workspace); +} + +int main(int argc, char **argv) { + + auto opt = OptionParser(argc, argv); + + int nreps = opt.parse("-r", 5); + int ntasks = opt.parse("-n", 500); + int N = opt.parse("-N", 2048); /* number of functions */ + int K = opt.parse("-K", 16); /* number of coefficients */ + int M = opt.parse("-M", 512); /* max number of blocks */ + int level = opt.parse("-l", 0); /* 0 = auto, 1-10 = explicit */ + int num_streams = opt.parse("-s", 4);/* number of concurrent streams to use */ + + /* Legacy -m flag: force level 1 */ + if (opt.exists("-m")) level = 1; + + std::cout << "Running benchmark" + << " nreps=" << nreps + << " ntasks=" << ntasks + << " N=" << N + << " K=" << K + << " M=" << M + << " level=" << (level <= 0 ? (MRA_HAVE_CUBLASDX ? 9 : 3) : level) + << std::endl; + + transform_bench(nreps, ntasks, N, M, K, level, num_streams); +} diff --git a/cuda/util.h b/cuda/util.h new file mode 100644 index 0000000..5807d11 --- /dev/null +++ b/cuda/util.h @@ -0,0 +1,182 @@ +#pragma once + +/** + * Cross-cutting helpers for the CUDA build: launch macros, error checking, + * thread-index helpers and the command-line option parser. + * + * This is the CUDA-only counterpart of the dual-target util.h in the parent + * directory; every HIP branch has been removed. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#define MAX_THREADS_PER_BLOCK 128 + +#define LAUNCH_BOUNDS(__NT, __NB) __launch_bounds__(__NT, __NB) + +typedef int32_t size_type; + +using Dim3 = dim3; + +/* NVIDIA warp size. The AMD sources assume a 64-lane wavefront; every place + * that constant appears has been re-derived for 32-lane warps. */ +#define MRA_WARP_SIZE 32 + +typedef cudaStream_t Stream; + +#define SYNC_STREAM(stream) (void)cudaStreamSynchronize(stream) +#define CREATE_STREAM(stream) (void)cudaStreamCreateWithFlags(stream, cudaStreamNonBlocking) +#define DESTROY_STREAM(stream) (void)cudaStreamDestroy(stream) + +#define MALLOC(ptr, size) (void)cudaMalloc(ptr, size) +#define FREE(ptr) (void)cudaFree(ptr) +#define MEMCPY_H2D(dst, src, size) (void)cudaMemcpy(dst, src, size, cudaMemcpyHostToDevice) +#define MEMCPY_D2H(dst, src, size) (void)cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost) + +#define CALL_KERNEL(name, block, thread, shared, stream, args) \ + do { \ + name<<>> args ; \ + { auto _err = cudaGetLastError(); \ + if (_err != cudaSuccess) { \ + std::cout << "kernel submission failed with " << shared << "B smem at " \ + << __FILE__ << ":" << __LINE__ << ": " \ + << cudaGetErrorString(_err) << std::endl; \ + throw std::runtime_error("kernel submission failed"); \ + } \ + } \ + } while (0) + +/** + * Opt in to more than the 48 KB of dynamic shared memory that a kernel gets by + * default. Tracked per instantiation so the (comparatively expensive) driver + * call happens only when the requirement grows. + */ +#define CONFIGURE_KERNEL(name, shared) \ + do { \ + static int smem_size_config = 0; \ + if (smem_size_config < (int)(shared)) { \ + cudaFuncSetAttribute(name, cudaFuncAttributeMaxDynamicSharedMemorySize, shared); \ + { auto _err = cudaGetLastError(); \ + if (_err != cudaSuccess) { \ + std::cout << "kernel configuration failed with " << shared << "B smem at " \ + << __FILE__ << ":" << __LINE__ << ": " \ + << cudaGetErrorString(_err) << std::endl; \ + throw std::runtime_error("kernel configuration failed"); \ + } \ + smem_size_config = (int)(shared); \ + } \ + } \ + } while (0) + + +#if defined(__CUDA_ARCH__) +#define HAVE_DEVICE_ARCH 1 +#define SCOPE __host__ __device__ inline +#define SYNCTHREADS() __syncthreads() +#define SHARED __shared__ +#define THROW(s) do { std::printf(s); __trap(); } while (0) +#else +#define SCOPE inline +#define SYNCTHREADS() +#define SHARED +#define THROW(s) do { throw std::runtime_error(s); } while (0) +#endif // __CUDA_ARCH__ + + +constexpr inline Dim3 max_thread_dims(int K) { + int x = K; + int y = std::min(K, MAX_THREADS_PER_BLOCK / x); + int z = 1; + return Dim3(x, y, z); +} + +constexpr inline int max_threads(int K) { + Dim3 thread_dims = max_thread_dims(K); + return thread_dims.x * thread_dims.y * thread_dims.z; +} + +__device__ inline int thread_id() { + return blockDim.x * ((blockDim.y * threadIdx.z) + threadIdx.y) + threadIdx.x; +} + +/* Two overloads rather than a `blockDim` default argument: nvcc rejects the + * builtin as a default argument of a __host__ __device__ function. */ +__host__ __device__ inline int block_size(Dim3 block) { + return block.x * block.y * block.z; +} + +__device__ inline int block_size() { + return blockDim.x * blockDim.y * blockDim.z; +} + +__device__ inline bool is_team_lead() { + return (0 == (threadIdx.x + threadIdx.y + threadIdx.z)); +} + + +struct OptionParser { + + private: + char **m_begin; + char **m_end; + + static inline const char *empty = ""; + + public: + OptionParser(int argc, char **argv) + : m_begin(argv), m_end(argv+argc) + { } + + std::string_view get(const std::string &option) { + char **itr = std::find(m_begin, m_end, option); + if (itr != m_end && ++itr != m_end) return std::string_view(*itr); + return std::string_view(empty); + } + + bool exists(const std::string &option) { + return std::find(m_begin, m_end, option) != m_end; + } + + int index(const std::string &option) { + char **itr = std::find(m_begin, m_end, option); + if (itr != m_end) return (int)(itr - m_end); + return -1; + } + + int parse(std::string_view option, int default_value) { + int N = default_value; + char **itr = std::find(m_begin, m_end, option); + if (++itr < m_end) { + N = std::stoi(*itr); + } + return N; + } + + long parse(std::string_view option, long default_value) { + long N = default_value; + char **itr = std::find(m_begin, m_end, option); + if (++itr < m_end) { + N = std::stol(*itr); + } + return N; + } + + double parse(std::string_view option, double default_value = 0.25) { + double N = default_value; + char **itr = std::find(m_begin, m_end, option); + if (++itr < m_end) { + N = std::stod(*itr); + } + return N; + } + + }; // struct OptionParser diff --git a/cuda/validate_levels.cu b/cuda/validate_levels.cu new file mode 100644 index 0000000..8d9b0f0 --- /dev/null +++ b/cuda/validate_levels.cu @@ -0,0 +1,184 @@ +/** + * Correctness test: compare any optimization level against the level-1 reference. + * + * Usage: + * ./validate_levels [-l ] [-K ] [-N ] + * + * -l level to validate (2-10, default 3) + * 2 L2: B cached in shared memory + * 3 L3: register blocking (K-templated) + * 4 L4: FP64 tensor cores, one warp (+ L3 fallback) + * 5 L5: FP64 tensor cores, A staged in shared memory, 8 warps + * 6 L6: nvcuda::wmma, one warp per output tile + * 7 L7: FP64 tensor cores, B resident in registers (K = 8, 16) + * 8 L8: Kronecker product GEMM (cuBLAS) + * 9 L9: cuBLASDx, three GEMMs fused (K = 8, 10, 16, 20) + * 10 L10: cuBLASDx per-pass block GEMM + * -K single K value; if omitted sweeps K in {6,8,10,12,16} + * -N batch size (default 16) + * + * The K sweep covers the values every level's dispatch table shares; levels + * with a narrower table report the K values they do not handle themselves. + */ + +#include +#include +#include +#include + +#include "util.h" +#include "transform.h" // L1 - reference +#include "transform_level2.h" // L2 +#include "transform_level3.h" // L3 +#include "transform_level4.h" // L4 +#include "transform_level5.h" // L5 +#include "transform_wmma.h" // L6 +#include "transform_level7.h" // L7 +#include "transform_kron.h" // L8 +#include "transform_cublasdx.h" // L9 +#include "transform_cublasdx_mxm.h" // L10 + +template +void test_level(int level, int K, int nfuncs) { + const int K3 = K * K * K; + const int nblocks = nfuncs; + + // Allocate and fill host arrays with random data + std::vector h_A(nfuncs * K3), h_B(K * K); + std::vector h_Cref(nfuncs * K3), h_Ctest(nfuncs * K3); + std::srand(42); + for (auto& v : h_A) v = (T)std::rand() / RAND_MAX; + for (auto& v : h_B) v = (T)std::rand() / RAND_MAX; + + // Device allocations + T *d_A, *d_B, *d_Cref, *d_Ctest, *d_workspace_ref, *d_workspace; + MALLOC(&d_A, (size_t)nfuncs * K3 * sizeof(T)); + MALLOC(&d_B, (size_t)K * K * sizeof(T)); + MALLOC(&d_Cref, (size_t)nfuncs * K3 * sizeof(T)); + MALLOC(&d_Ctest, (size_t)nfuncs * K3 * sizeof(T)); + MALLOC(&d_workspace_ref, (size_t)nfuncs * K3 * sizeof(T)); + MALLOC(&d_workspace, (size_t)nfuncs * K3 * sizeof(T)); + + T *d_KronMat = nullptr; + if (level == 8) { + MALLOC(&d_KronMat, (size_t)K3 * K3 * sizeof(T)); + } + + // Copy inputs to device + MEMCPY_H2D(d_A, h_A.data(), (size_t)nfuncs * K3 * sizeof(T)); + MEMCPY_H2D(d_B, h_B.data(), (size_t)K * K * sizeof(T)); + + Stream stream; + CREATE_STREAM(&stream); + + // --- Reference: level 1 --- + submit_transform_bench(nfuncs, nblocks, K, d_A, d_B, d_Cref, d_workspace_ref, stream); + SYNC_STREAM(stream); + + // --- Tested level --- + switch (level) { + case 2: + submit_transform_level2_bench(nfuncs, nblocks, K, d_A, d_B, d_Ctest, d_workspace, stream); + SYNC_STREAM(stream); + break; + case 3: + submit_transform_level3_bench(nfuncs, nblocks, K, d_A, d_B, d_Ctest, d_workspace, stream); + SYNC_STREAM(stream); + break; + case 4: + submit_transform_level4_bench(nfuncs, nblocks, K, d_A, d_B, d_Ctest, d_workspace, stream); + SYNC_STREAM(stream); + break; + case 5: + submit_transform_level5_bench(nfuncs, nblocks, K, d_A, d_B, d_Ctest, d_workspace, stream); + SYNC_STREAM(stream); + break; + case 6: + submit_transform_wmma_bench(nfuncs, nblocks, K, d_A, d_B, d_Ctest, d_workspace, stream); + SYNC_STREAM(stream); + break; + case 7: + submit_transform_level7_bench(nfuncs, nblocks, K, d_A, d_B, d_Ctest, d_workspace, stream); + SYNC_STREAM(stream); + break; + case 8: { + build_kron_matrix(K, d_B, d_KronMat, stream); + SYNC_STREAM(stream); + blasHandle_t blas_handle; + blasCreate(&blas_handle); + submit_transform_kron_bench(nfuncs, K, d_A, d_KronMat, d_Ctest, blas_handle, stream); + SYNC_STREAM(stream); + blasDestroy(blas_handle); + break; + } + case 9: + submit_transform_cublasdx_bench(nfuncs, nblocks, K, d_A, d_B, d_Ctest, d_workspace, stream); + SYNC_STREAM(stream); + break; + case 10: + submit_transform_cublasdx_mxm_bench(nfuncs, nblocks, K, d_A, d_B, d_Ctest, d_workspace, stream); + SYNC_STREAM(stream); + break; + default: + std::cerr << "Unknown level " << level << " (valid: 2-10)\n"; + FREE(d_A); FREE(d_B); FREE(d_Cref); FREE(d_Ctest); + FREE(d_workspace_ref); FREE(d_workspace); + if (d_KronMat) FREE(d_KronMat); + DESTROY_STREAM(stream); + return; + } + + // Copy results to host + MEMCPY_D2H(h_Cref.data(), d_Cref, (size_t)nfuncs * K3 * sizeof(T)); + MEMCPY_D2H(h_Ctest.data(), d_Ctest, (size_t)nfuncs * K3 * sizeof(T)); + + // Compare + T max_abs_err = 0, max_rel_err = 0; + for (int i = 0; i < nfuncs * K3; ++i) { + T abs_err = std::abs(h_Cref[i] - h_Ctest[i]); + T rel_err = abs_err / (std::abs(h_Cref[i]) + 1e-14); + max_abs_err = std::max(max_abs_err, abs_err); + max_rel_err = std::max(max_rel_err, rel_err); + } + + if (max_rel_err >= 1e-10) { + std::cout << "FAIL!\n"; + for (int i = 0; i < nfuncs * K3; ++i) { + if (i % K3 == 0) { + std::cout << " Function " << i / K3 << ":\n"; + } + int idx_k = (i % K3) / (K * K); + int idx_j = (i % (K * K)) / K; + int idx_i = i % K; + std::cout << " [" << idx_k << "][" << idx_j << "][" << idx_i << "] ref=" + << h_Cref[i] << " test=" << h_Ctest[i] << "\n"; + } + } + + std::cout << "K=" << K << " nfuncs=" << nfuncs << " level=" << level + << " max_abs_err=" << max_abs_err + << " max_rel_err=" << max_rel_err + << (max_rel_err < 1e-10 ? " PASS" : " FAIL") + << "\n"; + + FREE(d_A); FREE(d_B); FREE(d_Cref); FREE(d_Ctest); + FREE(d_workspace_ref); FREE(d_workspace); + if (d_KronMat) FREE(d_KronMat); + DESTROY_STREAM(stream); +} + +int main(int argc, char** argv) { + OptionParser opts(argc, argv); + int level = opts.parse(std::string("-l"), 3); + int nfuncs = opts.parse(std::string("-N"), 16); + + if (opts.exists(std::string("-K"))) { + int K = opts.parse(std::string("-K"), 8); + test_level(level, K, nfuncs); + } else { + for (int K : {6, 8, 10, 12, 16}) { + test_level(level, K, nfuncs); + } + } + return 0; +} diff --git a/rocm/.gitignore b/rocm/.gitignore new file mode 100644 index 0000000..7f3e00b --- /dev/null +++ b/rocm/.gitignore @@ -0,0 +1,42 @@ +# Prerequisites +*.d + +# Compiled Object files +*.slo +*.lo +*.o +*.obj + +# Precompiled Headers +*.gch +*.pch + +# Compiled Dynamic libraries +*.so +*.dylib +*.dll + +# Fortran module files +*.mod +*.smod + +# Compiled Static libraries +*.lai +*.la +*.a +*.lib + +# Executables +*.exe +*.out +*.app + +# custom build environment +build_vscode* +build* +#clangd +clangd* +#compile_commands.json +compile_commands.json +rocroof/ + diff --git a/CLAUDE.md b/rocm/CLAUDE.md similarity index 100% rename from CLAUDE.md rename to rocm/CLAUDE.md diff --git a/CMakeLists.txt b/rocm/CMakeLists.txt similarity index 100% rename from CMakeLists.txt rename to rocm/CMakeLists.txt diff --git a/rocm/LICENSE b/rocm/LICENSE new file mode 100644 index 0000000..c9017bf --- /dev/null +++ b/rocm/LICENSE @@ -0,0 +1,28 @@ +BSD 3-Clause License + +Copyright (c) 2025, Joseph Schuchart + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/rocm/README.md similarity index 100% rename from README.md rename to rocm/README.md diff --git a/rocm/mxm.h b/rocm/mxm.h new file mode 100644 index 0000000..8402972 --- /dev/null +++ b/rocm/mxm.h @@ -0,0 +1,373 @@ +#ifndef MRA_MXM_H +#define MRA_MXM_H + +#include "util.h" + + +#if __has_include() +#include +#define HAVE_BLASPP 1 +#endif // __has_include() + +namespace mra { + + +#ifndef MRA_HAVE_MTXM + +#if defined(HAVE_BLASPP) && !defined(HAVE_DEVICE_ARCH) + /** + * blaspp implementation of A^T * B + * c(i,j) += sum(k) a(k,i)*b(k,j) + */ + template + void mTxm(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + blas::gemm(blas::Layout::RowMajor, blas::Op::Trans, blas::Op::NoTrans, + dimi, dimj, dimk, + 1.0, a, dimi, b, dimj, + Q ? 0.0 : 1.0, c, dimj); + } +#else // HAVE_BLASPP + /** + * reference implementation, adapted from madness + * c(i,j) += sum(k) a(k,i)*b(k,j) + */ + template + SCOPE void mTxm(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + /* trivial 2D implementation for devices */ + if (threadIdx.z == 0) { + for (size_type i = threadIdx.y; i < dimi; i += blockDim.y) { + cT* ci = c + i*dimj; // the row of C all threads in dim x work on + const aT *aik_ptr = a + i; + if constexpr(Q) { + for (size_type j = threadIdx.x; j < dimj; j += blockDim.x) { + ci[j] = 0.0; + } + } + + for (long k=0; k + constexpr size_type mTxm_shmem_size(size_type K) { + return 0; + } + + template + constexpr Dim3 mTxm_blockdim(int K) { + return max_thread_dims(K); + } + + +#endif // MRA_HAVE_MTXM + +#ifndef MRA_HAVE_MTXMQ + + /** + * blaspp implementation of A^T * B + * c(i,j) = sum(k) a(k,i)*b(k,j) + */ + template + SCOPE void mTxmq(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + mTxm(dimi, dimj, dimk, c, a, b); + } + + template + constexpr size_type mTxmq_shmem_size(size_type K) { + return 0; + } + + template + constexpr Dim3 mTxmq_blockdim(int K) { + return mTxm_blockdim(K); + } + +#endif // MRA_HAVE_MTXMQ + + +#ifndef MRA_HAVE_MXM + +#if defined(HAVE_BLASPP) && !defined(HAVE_DEVICE_ARCH) + + /** + * blaspp implementation of A * B + * c(i,j) += sum(k) a(i,k)*b(k,j) + */ + template + void mxm(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, + dimi, dimj, dimk, + 1.0, a, dimk, b, dimj, + Q ? 0.0 : 1.0, c, dimj); + } +#else // defined(HAVE_BLASPP) && !defined(HAVE_DEVICE_ARCH) + /** + * reference implementation, adapted from madness + * + * c(i,j) += sum(k) a(i,k)*b(k,j) + */ + template + SCOPE void mxm(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + /* trivial 2D implementation for devices */ + if (threadIdx.z == 0) { + for (size_type i = threadIdx.y; i < dimi; i += blockDim.y) { + cT* ci = c + i*dimj; // the row of C all threads in dim x work on + const aT *ai_ptr = a + i*dimk; + if constexpr(Q) { + for (size_type j = threadIdx.x; j < dimj; j += blockDim.x) { + ci[j] = 0.0; + } + } + for (size_type j = threadIdx.x; j < dimj; j += blockDim.x) { + for (long k=0; k + constexpr size_type mxm_shmem_size(size_type K) { + return 0; + } + + template + constexpr Dim3 mxm_blockdim(int K) { + return max_thread_dims(K); + } + +#endif // MRA_HAVE_MXM + + +#ifndef MRA_HAVE_MXMQ + + /** + * reference implementation, adapted from madness + * + * c(i,j) = sum(k) a(i,k)*b(k,j) + */ + template + SCOPE void mxmq(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + mxm(dimi, dimj, dimk, c, a, b); + } + + template + constexpr size_type mxmq_shmem_size(size_type K) { + return 0; + } + + template + constexpr Dim3 mxmq_blockdim(int K) { + return mxm_blockdim(K); + } + +#endif // MRA_HAVE_MXMQ + + +#ifndef MRA_HAVE_MXMT + +#if defined(HAVE_BLASPP) && !defined(HAVE_DEVICE_ARCH) + + /** + * blaspp implementation of A * B^T + * c(i,j) += sum(k) a(i,k)*b(j,k) + */ + template + void mxmT(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + blas::gemm(blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::Trans, + dimi, dimj, dimk, + 1.0, a, dimk, b, dimk, + Q ? 0.0 : 1.0, c, dimj); + } + +#else // defined(HAVE_BLASPP) && !defined(HAVE_DEVICE_ARCH) + + /** + * reference implementation, adapted from madness + * + * c(i,j) += sum(k) a(i,k)*b(j,k) + */ + template + SCOPE void mxmT(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + /* trivial 2D implementation for devices */ + if (threadIdx.z == 0) { + for (size_type i = threadIdx.y; i < dimi; i += blockDim.y) { + cT* ci = c + i*dimj; // the row of C all threads in dim x work on + const aT *ai_ptr = a + i*dimk; + for (size_type j = threadIdx.x; j < dimj; j += blockDim.x) { + cT sum = 0.0; + for (size_type k=0; k + constexpr size_type mxmT_shmem_size(size_type K) { + return 0; + } + + template + constexpr Dim3 mxmT_blockdim(int K) { + return max_thread_dims(K); + } + +#endif // MRA_HAVE_MXMT + + +#ifndef MRA_HAVE_MXMTQ + + /** + * reference implementation, adapted from madness + * + * c(i,j) = sum(k) a(i,k)*b(j,k) + */ + template + SCOPE void mxmTq(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + mxmT(dimi, dimj, dimk, c, a, b); + } + + template + constexpr size_type mxmTq_shmem_size(size_type K) { + return 0; + } + + template + constexpr Dim3 mxmTq_blockdim(int K) { + return mxmT_blockdim(K); + } + + +#endif // MRA_HAVE_MXMTQ + + + +#ifndef MRA_HAVE_MTXMT + +#if defined(HAVE_BLASPP) && !defined(HAVE_DEVICE_ARCH) + + /** + * blaspp implementation of A^T * B^T + * c(i,j) += sum(k) a(k,i)*b(j,k) + */ + template + void mTxmT(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + blas::gemm(blas::Layout::RowMajor, blas::Op::Trans, blas::Op::Trans, + dimi, dimj, dimk, + 1.0, a, dimi, b, dimj, + Q ? 0.0 : 1.0, c, dimj); + } + +#else + + /** + * reference implementation, adapted from madness + * + * c(i,j) += sum(k) a(k,i)*b(j,k) + */ + template + SCOPE void mTxmT(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + /* trivial 2D implementation for devices */ + if (threadIdx.z == 0) { + for (size_type i = threadIdx.y; i < dimi; i += blockDim.y) { + cT* ci = c + i*dimj; // the row of C all threads in dim x work on + if constexpr(Q) { + for (size_type j = threadIdx.x; j < dimj; j += blockDim.x) { + ci[j] = 0.0; + } + } + for (size_type j = threadIdx.x; j < dimj; j += blockDim.x) { + const aT *aik_ptr = a + i; + for (long k=0; k + constexpr size_type mTxmT_shmem_size(size_type K) { + return 0; + } + + template + constexpr Dim3 mTxmT_blockdim(int K) { + return max_thread_dims(K); + } + + +#endif // MRA_HAVE_MTXMT + + + +#ifndef MRA_HAVE_MTXMTQ + + /** + * reference implementation, adapted from madness + * + * c(i,j) = sum(k) a(k,i)*b(j,k) + */ + template + SCOPE void mTxmTq(size_type dimi, size_type dimj, size_type dimk, + cT* __restrict__ c, const aT* a, const bT* b) { + mTxmT(dimi, dimj, dimk, c, a, b); + } + + template + constexpr size_type mTxmTq_shmem_size(size_type K) { + return 0; + } + + template + constexpr Dim3 mTxmTq_blockdim(int K) { + return mTxmT_blockdim(K); + } + +#endif // MRA_HAVE_MTXMTQ + +} // namespace mra + +#endif // MRA_MXM_H diff --git a/mxm_cublasdx.h b/rocm/mxm_cublasdx.h similarity index 100% rename from mxm_cublasdx.h rename to rocm/mxm_cublasdx.h diff --git a/mxm_level2.h b/rocm/mxm_level2.h similarity index 100% rename from mxm_level2.h rename to rocm/mxm_level2.h diff --git a/mxm_level3.h b/rocm/mxm_level3.h similarity index 89% rename from mxm_level3.h rename to rocm/mxm_level3.h index 86b5cc2..84ef805 100644 --- a/mxm_level3.h +++ b/rocm/mxm_level3.h @@ -62,14 +62,14 @@ __device__ void mTxmq_level3_impl(T* __restrict__ c, const T* a, const T* b_shme */ template __device__ void mTxmq_level3_k(T* __restrict__ c, const T* a, const T* b) { - extern __shared__ char smem_level3[]; - T* b_shmem = reinterpret_cast(smem_level3); + // extern __shared__ char smem_level3[]; + // T* b_shmem = reinterpret_cast(smem_level3); - for (int idx = (int)threadIdx.x; idx < K * K; idx += (int)blockDim.x) - b_shmem[idx] = b[idx]; - __syncthreads(); + // for (int idx = (int)threadIdx.x; idx < K * K; idx += (int)blockDim.x) + // b_shmem[idx] = b[idx]; + // __syncthreads(); - detail::mTxmq_level3_impl(c, a, b_shmem); + detail::mTxmq_level3_impl(c, a, b); __syncthreads(); } diff --git a/mxm_level4.h b/rocm/mxm_level4.h similarity index 100% rename from mxm_level4.h rename to rocm/mxm_level4.h diff --git a/mxm_level5.h b/rocm/mxm_level5.h similarity index 100% rename from mxm_level5.h rename to rocm/mxm_level5.h diff --git a/mxm_level7.h b/rocm/mxm_level7.h similarity index 100% rename from mxm_level7.h rename to rocm/mxm_level7.h diff --git a/mxm_rocwmma.h b/rocm/mxm_rocwmma.h similarity index 100% rename from mxm_rocwmma.h rename to rocm/mxm_rocwmma.h diff --git a/transform.h b/rocm/transform.h similarity index 100% rename from transform.h rename to rocm/transform.h diff --git a/transform_cublasdx.h b/rocm/transform_cublasdx.h similarity index 100% rename from transform_cublasdx.h rename to rocm/transform_cublasdx.h diff --git a/transform_kron.h b/rocm/transform_kron.h similarity index 100% rename from transform_kron.h rename to rocm/transform_kron.h diff --git a/transform_level2.h b/rocm/transform_level2.h similarity index 100% rename from transform_level2.h rename to rocm/transform_level2.h diff --git a/transform_level3.h b/rocm/transform_level3.h similarity index 100% rename from transform_level3.h rename to rocm/transform_level3.h diff --git a/transform_level4.h b/rocm/transform_level4.h similarity index 100% rename from transform_level4.h rename to rocm/transform_level4.h diff --git a/transform_level5.h b/rocm/transform_level5.h similarity index 100% rename from transform_level5.h rename to rocm/transform_level5.h diff --git a/transform_level7.h b/rocm/transform_level7.h similarity index 100% rename from transform_level7.h rename to rocm/transform_level7.h diff --git a/transform_rocwmma.h b/rocm/transform_rocwmma.h similarity index 100% rename from transform_rocwmma.h rename to rocm/transform_rocwmma.h diff --git a/transformbench.cu b/rocm/transformbench.cu similarity index 100% rename from transformbench.cu rename to rocm/transformbench.cu diff --git a/transformbench.hip b/rocm/transformbench.hip similarity index 100% rename from transformbench.hip rename to rocm/transformbench.hip diff --git a/util.h b/rocm/util.h similarity index 100% rename from util.h rename to rocm/util.h diff --git a/validate_levels.hip b/rocm/validate_levels.hip similarity index 100% rename from validate_levels.hip rename to rocm/validate_levels.hip