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