From 1b51a1a7c779282c003c7ee1e427072684c446bf Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Fri, 24 Jul 2026 15:55:02 +0530 Subject: [PATCH 01/13] fix: support runtime GEMM dimensions for dynamic shapes --- .../backends/cuda/sofieBLAS_cublas.hpp | 97 +++++++++---------- 1 file changed, 44 insertions(+), 53 deletions(-) diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index f0b9e57..7063448 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -31,22 +31,6 @@ } while (0) -struct PairHash { - std::size_t - operator()(const std::pair &p) const noexcept { - std::size_t h1 = std::hash{}(p.first); - std::size_t h2 = std::hash{}(p.second); - return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2)); - } -}; - -struct PairEq { - bool operator()(const std::pair &a, - const std::pair &b) const noexcept { - return a.first == b.first && a.second == b.second; - } -}; - struct DescKey { int transA; // CUBLAS_OP_N / CUBLAS_OP_T encoded as int int transB; @@ -69,8 +53,8 @@ struct DescKeyHash { struct AlgoKey { DescKey dk; - std::size_t rowsA, colsA; // physical dimensions of A in layoutStore - std::size_t rowsB, colsB; // physical dimensions of B in layoutStore + std::size_t rowsA, colsA; // physical dimensions of A + std::size_t rowsB, colsB; // physical dimensions of B bool operator==(const AlgoKey &o) const noexcept { return dk == o.dk && rowsA == o.rowsA && colsA == o.colsA @@ -98,9 +82,12 @@ class BlasCuda { size_t workspaceSize = 1u << 25; // 32 MB (was 4 MB) cudaStream_t stream = nullptr; - std::unordered_map, - cublasLtMatrixLayout_t, PairHash, PairEq> - layoutStore; + // One persistent layout descriptor per matrix role, re-stamped with the + // runtime dimensions before each matmul. The descriptor is host-side metadata + // consumed by cublasLtMatmul at the call, so a single object can be reused + // across shapes - this is what lets one Session serve dynamic (runtime) sizes. + enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 }; + cublasLtMatrixLayout_t roleLayout[3] = {}; std::unordered_map descStore; @@ -129,8 +116,8 @@ class BlasCuda { } ~BlasCuda() { - for (auto &[key, layout] : layoutStore) - if (layout) cublasLtMatrixLayoutDestroy(layout); + for (auto L : roleLayout) + if (L) cublasLtMatrixLayoutDestroy(L); for (auto &[key, desc] : descStore) if (desc) cublasLtMatmulDescDestroy(desc); if (preference) cublasLtMatmulPreferenceDestroy(preference); @@ -149,22 +136,9 @@ class BlasCuda { } } - void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, - std::size_t lda, std::size_t ldb, std::size_t ldc, - char transa, char transb) { - // Physical A: (m×k) if NoTrans, (k×m) if Trans - if (transa == 'N' || transa == 'n') - checkAndAddLayout(m, k, lda); - else - checkAndAddLayout(k, m, lda); - // Physical B: (k×n) if NoTrans, (n×k) if Trans - if (transb == 'N' || transb == 'n') - checkAndAddLayout(k, n, ldb); - else - checkAndAddLayout(n, k, ldb); - // C is always (m×n) - checkAndAddLayout(m, n, ldc); - } + // No-op kept for the generated ctor's API; layouts are created lazily now. + void addLayoutConfig(std::size_t, std::size_t, std::size_t, + std::size_t, std::size_t, std::size_t, char, char) {} template inline void @@ -370,14 +344,25 @@ class BlasCuda { : std::make_pair(n, k); } - void checkAndAddLayout(std::size_t rows, std::size_t cols, std::size_t ld) { - auto key = std::make_pair(rows, cols); - if (layoutStore.find(key) == layoutStore.end()) { - cublasLtMatrixLayout_t layout = nullptr; - CHECK_CUBLAS( - cublasLtMatrixLayoutCreate(&layout, CUDA_R_32F, rows, cols, ld)); - layoutStore.emplace(key, layout); + // Resolve a matrix role's layout at the runtime dims: create the descriptor + // once, then overwrite its dims in place on later calls. ld = rows (dense, + // column-major, as the generated calls produce). + cublasLtMatrixLayout_t stampLayout(LayoutRole role, + const std::pair &key) { + const uint64_t rows = key.first, cols = key.second; + const int64_t ld = static_cast(key.first); + cublasLtMatrixLayout_t &L = roleLayout[role]; + if (!L) { + CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&L, CUDA_R_32F, rows, cols, ld)); + } else { + CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( + L, CUBLASLT_MATRIX_LAYOUT_ROWS, &rows, sizeof(rows))); + CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( + L, CUBLASLT_MATRIX_LAYOUT_COLS, &cols, sizeof(cols))); + CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( + L, CUBLASLT_MATRIX_LAYOUT_LD, &ld, sizeof(ld))); } + return L; } cublasLtMatmulDesc_t &getOrCreateDesc(cublasOperation_t transA, @@ -421,12 +406,13 @@ class BlasCuda { return it->second; auto &desc = getOrCreateDesc(transA, transB, epilogue); + auto lA = stampLayout(ROLE_A, kA); + auto lB = stampLayout(ROLE_B, kB); + auto lC = stampLayout(ROLE_C, kC); // C and D share the same layout cublasLtMatmulHeuristicResult_t h{}; int returnedResults = 0; CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic( - ltHandle, desc, - layoutStore.at(kA), layoutStore.at(kB), - layoutStore.at(kC), layoutStore.at(kC), + ltHandle, desc, lA, lB, lC, lC, preference, 1, &h, &returnedResults)); if (returnedResults == 0) { std::cerr << "[sofieBLAS] No suitable cuBLASLt algorithm found for " @@ -459,12 +445,17 @@ class BlasCuda { &bias_ptr, sizeof(bias_ptr))); } + // Re-stamp the shared role layouts to this call's shape right before the + // matmul (the algo-cache hit path in getOrComputeAlgo skips stamping). + auto lA = stampLayout(ROLE_A, kA); + auto lB = stampLayout(ROLE_B, kB); + auto lC = stampLayout(ROLE_C, kC); CHECK_CUBLAS(cublasLtMatmul( ltHandle, desc, - &alpha, A, layoutStore.at(kA), - B, layoutStore.at(kB), - &beta, D_in, layoutStore.at(kC), - C_out, layoutStore.at(kC), + &alpha, A, lA, + B, lB, + &beta, D_in, lC, + C_out, lC, &h.algo, d_workspace, workspaceSize, stream)); } }; From b4fc25e5eed2a0177617f4933360f04e525dd9a6 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Tue, 4 Aug 2026 16:26:59 +0530 Subject: [PATCH 02/13] feat: resolve GEMM algorithms per call site instead of per shape --- .../backends/cuda/sofieBLAS_cublas.hpp | 134 ++++++++++++++++-- 1 file changed, 119 insertions(+), 15 deletions(-) diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index 7063448..324668c 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -5,9 +5,11 @@ #include #include #include +#include #include #include #include +#include #include "sofieBLAS/core.hpp" #include @@ -74,6 +76,17 @@ struct AlgoKeyHash { } }; +// A call site's maximum shape, as declared by addLayoutConfig from the +// generated Session constructor. +struct ShapeEnvelope { + std::size_t rowsA, colsA, rowsB, colsB, rowsC, colsC; +}; + +struct LayoutStats { + std::size_t heuristicQueries = 0; + std::size_t envelopeRejects = 0; // envelope algorithm unusable at the call +}; + class BlasCuda { cublasLtHandle_t ltHandle = nullptr; cublasHandle_t handle = nullptr; // legacy cuBLAS for batched ops @@ -94,7 +107,15 @@ class BlasCuda { std::unordered_map algoCache; + // call-site envelopes declared by addLayoutConfig + std::vector envelopes; + + LayoutStats stats; + public: + const LayoutStats &layoutStats() const { return stats; } + std::size_t algoCacheSize() const { return algoCache.size(); } + BlasCuda(const BlasCuda &) = delete; BlasCuda &operator=(const BlasCuda &) = delete; BlasCuda(BlasCuda &&) = delete; @@ -136,9 +157,31 @@ class BlasCuda { } } - // No-op kept for the generated ctor's API; layouts are created lazily now. - void addLayoutConfig(std::size_t, std::size_t, std::size_t, - std::size_t, std::size_t, std::size_t, char, char) {} + // Records the call site's envelope. The generated constructor evaluates its + // shape expressions with its own parameters, so for a dynamic model these are + // the largest dims the call site will ever use. Layouts are created lazily, + // so nothing is registered here beyond the envelope and, with warmup on, the + // algorithm resolved for it. + void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, + std::size_t, std::size_t, std::size_t, + char transa, char transb) { + const auto kA = layoutKeyA(transa, m, k); + const auto kB = layoutKeyB(transb, k, n); + const std::pair kC{m, n}; + envelopes.push_back({kA.first, kA.second, kB.first, kB.second, m, n}); + + // The constructor does not know which epilogue this call site uses, so + // resolve all three. Unused ones cost one heuristic query each, off the + // inference path. + const cublasOperation_t tA = charToCuBlasTranspose(transa); + const cublasOperation_t tB = charToCuBlasTranspose(transb); + const cublasLtEpilogue_t eps[] = {CUBLASLT_EPILOGUE_DEFAULT, + CUBLASLT_EPILOGUE_BIAS, + CUBLASLT_EPILOGUE_RELU_BIAS}; + for (cublasLtEpilogue_t ep : eps) { + getOrComputeAlgo(tA, tB, ep, kA, kB, kC, /*required=*/false); + } + } template inline void @@ -365,6 +408,32 @@ class BlasCuda { return L; } + // Tightest declared envelope covering this call, or null if none does. + // Tightest matters: several envelopes may cover a small shape, but only the + // call site's own matches its weight dims exactly and so has zero excess + // on those axes. + const ShapeEnvelope * + findEnvelope(const std::pair &kA, + const std::pair &kB, + const std::pair &kC) const { + const ShapeEnvelope *best = nullptr; + std::size_t bestExcess = std::numeric_limits::max(); + for (const auto &e : envelopes) { + // colsA and rowsB are both the contraction dimension k, which comes from + // the weight tensor and never varies at runtime. Requiring an exact match + // on it stops one call site's envelope from serving another's shapes. + if (e.colsA != kA.second || e.rowsB != kB.first) + continue; + if (e.rowsA < kA.first || e.colsB < kB.second || + e.rowsC < kC.first || e.colsC < kC.second) + continue; + const std::size_t ex = (e.rowsA - kA.first) + (e.colsA - kA.second) + + (e.rowsB - kB.first) + (e.colsB - kB.second); + if (ex < bestExcess) { bestExcess = ex; best = &e; } + } + return best; + } + cublasLtMatmulDesc_t &getOrCreateDesc(cublasOperation_t transA, cublasOperation_t transB, cublasLtEpilogue_t epilogue) { @@ -393,17 +462,37 @@ class BlasCuda { return descStore.at(key); } - cublasLtMatmulHeuristicResult_t & + // Whether an algorithm can actually run this shape. cuBLASLt rejects some + // combinations, so an algorithm resolved at a call site's envelope is not + // guaranteed to work at every smaller shape it serves. + bool algoUsable(cublasLtMatmulDesc_t desc, const cublasLtMatmulAlgo_t &algo, + const std::pair &kA, + const std::pair &kB, + const std::pair &kC) { + auto lA = stampLayout(ROLE_A, kA); + auto lB = stampLayout(ROLE_B, kB); + auto lC = stampLayout(ROLE_C, kC); + cublasLtMatmulHeuristicResult_t chk{}; + return cublasLtMatmulAlgoCheck(ltHandle, desc, lA, lB, lC, lC, &algo, + &chk) == CUBLAS_STATUS_SUCCESS && + chk.workspaceSize <= workspaceSize; + } + + // required=false is used by constructor warmup, which speculatively resolves + // epilogues the call site may never use: those may legitimately have no + // algorithm and must not abort. + cublasLtMatmulHeuristicResult_t * getOrComputeAlgo(cublasOperation_t transA, cublasOperation_t transB, cublasLtEpilogue_t epilogue, const std::pair &kA, const std::pair &kB, - const std::pair &kC) { + const std::pair &kC, + bool required = true) { AlgoKey key{{(int)transA, (int)transB, (int)epilogue}, kA.first, kA.second, kB.first, kB.second}; auto it = algoCache.find(key); if (it != algoCache.end()) - return it->second; + return &it->second; auto &desc = getOrCreateDesc(transA, transB, epilogue); auto lA = stampLayout(ROLE_A, kA); @@ -414,7 +503,10 @@ class BlasCuda { CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic( ltHandle, desc, lA, lB, lC, lC, preference, 1, &h, &returnedResults)); + ++stats.heuristicQueries; if (returnedResults == 0) { + if (!required) + return nullptr; std::cerr << "[sofieBLAS] No suitable cuBLASLt algorithm found for " << "transA=" << transA << " transB=" << transB << " epilogue=" << epilogue @@ -422,8 +514,7 @@ class BlasCuda { << " B=[" << kB.first << "x" << kB.second << "]\n"; exit(EXIT_FAILURE); } - algoCache.emplace(key, h); - return algoCache.at(key); + return &algoCache.emplace(key, h).first->second; } void executeMatmul(cublasOperation_t transA, cublasOperation_t transB, @@ -434,10 +525,6 @@ class BlasCuda { const std::pair &kA, const std::pair &kB, const std::pair &kC) { - // Retrieve (or lazily compute) the cached algorithm for this shape - auto &h = getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); - - // Retrieve the cached descriptor and patch the real bias pointer in-place auto &desc = getOrCreateDesc(transA, transB, epilogue); if (bias_ptr) { CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( @@ -445,8 +532,25 @@ class BlasCuda { &bias_ptr, sizeof(bias_ptr))); } - // Re-stamp the shared role layouts to this call's shape right before the - // matmul (the algo-cache hit path in getOrComputeAlgo skips stamping). + // Resolve at this call site's declared envelope, so every runtime size it + // produces shares one cache entry. + const ShapeEnvelope *env = findEnvelope(kA, kB, kC); + const std::pair + aA = env ? std::make_pair(env->rowsA, env->colsA) : kA, + aB = env ? std::make_pair(env->rowsB, env->colsB) : kB, + aC = env ? std::make_pair(env->rowsC, env->colsC) : kC; + auto *h = getOrComputeAlgo(transA, transB, epilogue, aA, aB, aC); + + // Fall back to the exact shape when the envelope's algorithm cannot run + // it. cuBLASLt returns CUBLAS_STATUS_NOT_SUPPORTED for at least some + // shape/algorithm combinations; m=1 was the first observed. + if (env && !algoUsable(desc, h->algo, kA, kB, kC)) { + ++stats.envelopeRejects; + h = getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); + } + + // Re-stamp the shared role layouts to this call's exact shape. Resolving + // the algorithm leaves them at the envelope size, so this happens last. auto lA = stampLayout(ROLE_A, kA); auto lB = stampLayout(ROLE_B, kB); auto lC = stampLayout(ROLE_C, kC); @@ -456,7 +560,7 @@ class BlasCuda { B, lB, &beta, D_in, lC, C_out, lC, - &h.algo, d_workspace, workspaceSize, stream)); + &h->algo, d_workspace, workspaceSize, stream)); } }; From 21f938da4eca43e2388a4a9f9e93137fd2b19ad0 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Tue, 4 Aug 2026 16:26:59 +0530 Subject: [PATCH 03/13] test: cover multi-size use of one instance and bounded algorithm cache --- tests/test.cc | 90 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/tests/test.cc b/tests/test.cc index e106275..0f29561 100644 --- a/tests/test.cc +++ b/tests/test.cc @@ -567,6 +567,95 @@ static void runCudaTests() { } } +// One instance used at many sizes, which is what a dynamic-shape model does. +// earlier, any size other than the one registered at construction threw std::out_of_range from the layout lookup +// so this is the case the rest of the suite never covered. +static void runDynamicShapeTests() { + std::cout << "\n=== CUDA Dynamic-Shape Tests ===\n"; + + alpaka::PlatformCudaRt platform{}; + auto dev = alpaka::getDevByIdx(platform, 0u); + alpaka::Queue queue{dev}; + sofieBLAS blas(queue); + + alpaka::PlatformCpu hostPlatform{}; + auto hostDev = alpaka::getDevByIdx(hostPlatform, 0u); + + // Buffers hold MCAP rows while the call site declares MENV, so sizes above + // MENV stay in bounds and exercise the path where no envelope covers them. + constexpr int MCAP = 96, MENV = 64, N = 3, K = 5; + + auto hA = alpaka::allocBuf(hostDev, static_cast(MCAP * K)); + auto hB = alpaka::allocBuf(hostDev, static_cast(K * N)); + auto hC = alpaka::allocBuf(hostDev, static_cast(MCAP * N)); + float *A = alpaka::getPtrNative(hA); + float *B = alpaka::getPtrNative(hB); + float *C = alpaka::getPtrNative(hC); + fillSeq(A, MCAP * K, 0.5f, 0.25f); + fillSeq(B, K * N, 1.f, 0.5f); + + auto dA = alpaka::allocAsyncBuf(queue, static_cast(MCAP * K)); + auto dB = alpaka::allocAsyncBuf(queue, static_cast(K * N)); + auto dC = alpaka::allocAsyncBuf(queue, static_cast(MCAP * N)); + alpaka::memcpy(queue, dA, hA); + alpaka::memcpy(queue, dB, hB); + alpaka::wait(queue); + + // Declare the call site's largest shape, as a generated Session constructor + // does with its own arguments. + blas.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), + MENV, 'N', 'N'); + + std::vector ref; + auto runAt = [&](int m, const std::string &name) { + ref.assign(static_cast(m) * N, 0.f); + refMatmul(ref.data(), A, B, m, N, K, 1.f, 0.f, false, false); + blas.matmul('N', 'N', static_cast(m), static_cast(N), + static_cast(K), 1.f, dA, dB, 0.f, dC); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + checkClose(C, ref.data(), m * N, name); + }; + + // At, below and above the declared envelope, all on one instance. + for (int m : {MENV, 37, 8, 51, 1, MENV, MCAP}) + runAt(m, "cuda::dynamic m=" + std::to_string(m)); + + // The cache must not grow one entry per size. Entries are added only where + // the envelope's algorithm is rejected. Keying by shape fails only here. + const int nSizes = MENV - 1; + const std::size_t cacheBefore = blas.algoCacheSize(); + const std::size_t rejBefore = blas.layoutStats().envelopeRejects; + const std::size_t searchBefore = blas.layoutStats().heuristicQueries; + for (int m = 2; m <= MENV; ++m) + blas.matmul('N', 'N', static_cast(m), static_cast(N), + static_cast(K), 1.f, dA, dB, 0.f, dC); + alpaka::wait(queue); + const std::size_t added = blas.algoCacheSize() - cacheBefore; + const std::size_t rejected = blas.layoutStats().envelopeRejects - rejBefore; + + std::cout << " " << nSizes << " sizes added " << added + << " cache entries, " << rejected << " rejected\n"; + if (added < static_cast(nSizes)) { + std::cout << " PASS cuda::cache bounded\n"; + } else { + std::cerr << " FAIL [cuda::cache bounded] one entry per size\n"; + ++gFailures; + } + + // The constructor resolves every declared envelope, so a size it covers must + // not trigger a search. Only a rejected envelope algorithm may. + const std::size_t searched = + blas.layoutStats().heuristicQueries - searchBefore; + if (searched <= rejected) { + std::cout << " PASS cuda::no search during inference\n"; + } else { + std::cerr << " FAIL [cuda::no search during inference] " << searched + << " searches, only " << rejected << " explained by rejects\n"; + ++gFailures; + } +} + #endif // ALPAKA_ACC_GPU_CUDA_ENABLED // --------------------------------------------------------------------------- @@ -579,6 +668,7 @@ int main() { #endif #ifdef ALPAKA_ACC_GPU_CUDA_ENABLED runCudaTests(); + runDynamicShapeTests(); #endif std::cout << "\n"; From 5c4565192a78550078e98ff5181bdd9af3b1908f Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Wed, 5 Aug 2026 10:40:08 +0530 Subject: [PATCH 04/13] feat: optional LRU limit on the algorithm cache --- .../backends/cuda/sofieBLAS_cublas.hpp | 51 +++++++++++++++---- tests/test.cc | 32 ++++++++++++ 2 files changed, 73 insertions(+), 10 deletions(-) diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index 4227f66..80f1679 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -84,6 +85,7 @@ struct ShapeEnvelope { struct LayoutStats { std::size_t heuristicQueries = 0; std::size_t envelopeRejects = 0; // envelope algorithm unusable at the call + std::size_t evictions = 0; // entries dropped to stay under the limit }; class BlasCuda { @@ -103,8 +105,20 @@ class BlasCuda { std::unordered_map descStore; - std::unordered_map - algoCache; + // Recency handle is only maintained when a limit is set; with no limit the + // list stays empty and the iterator is never read. + struct CacheEntry { + cublasLtMatmulHeuristicResult_t h{}; + std::list::iterator lru{}; + }; + std::unordered_map algoCache; + std::list lruOrder; + // 0 = unbounded. Per-call-site resolution already bounds the cache by the + // number of Conv/Gemm in the model, so a limit only matters when infer() + // runs above the constructed size: those shapes have no envelope and are + // resolved individually. Set it below the number of call sites and + // construction will evict its own warmup. + std::size_t algoCacheLimit = 0; // call-site envelopes declared by addLayoutConfig std::vector envelopes; @@ -114,13 +128,16 @@ class BlasCuda { public: const LayoutStats &layoutStats() const { return stats; } std::size_t algoCacheSize() const { return algoCache.size(); } + void setAlgoCacheLimit(std::size_t n) { algoCacheLimit = n; } BlasCuda(const BlasCuda &) = delete; BlasCuda &operator=(const BlasCuda &) = delete; BlasCuda(BlasCuda &&) = delete; BlasCuda &operator=(BlasCuda &&) = delete; - BlasCuda(alpaka::QueueCudaRtNonBlocking &queue) : m_queue{queue} { + BlasCuda(alpaka::QueueCudaRtNonBlocking &queue, + std::size_t algoCacheLimit_ = 0) + : algoCacheLimit{algoCacheLimit_}, m_queue{queue} { stream = static_cast(m_queue.getNativeHandle()); CHECK_CUBLAS(cublasLtCreate(<Handle)); @@ -505,8 +522,11 @@ class BlasCuda { kB.first, kB.second}; auto it = algoCache.find(key); - if (it != algoCache.end()) - return &it->second; + if (it != algoCache.end()) { + if (algoCacheLimit) + lruOrder.splice(lruOrder.begin(), lruOrder, it->second.lru); + return &it->second.h; + } auto &desc = getOrCreateDesc(transA, transB, epilogue); auto lA = stampLayout(ROLE_A, kA); @@ -527,7 +547,17 @@ class BlasCuda { << " B=[" << kB.first << "x" << kB.second << "]\n"; exit(EXIT_FAILURE); } - return &algoCache.emplace(key, h).first->second; + auto ins = algoCache.emplace(key, CacheEntry{h, {}}).first; + if (algoCacheLimit) { + lruOrder.push_front(key); + ins->second.lru = lruOrder.begin(); + while (algoCache.size() > algoCacheLimit) { + algoCache.erase(lruOrder.back()); + lruOrder.pop_back(); + ++stats.evictions; + } + } + return &ins->second.h; } void executeMatmul(cublasOperation_t transA, cublasOperation_t transB, @@ -551,14 +581,15 @@ class BlasCuda { aA = env ? std::make_pair(env->rowsA, env->colsA) : kA, aB = env ? std::make_pair(env->rowsB, env->colsB) : kB, aC = env ? std::make_pair(env->rowsC, env->colsC) : kC; - auto *h = getOrComputeAlgo(transA, transB, epilogue, aA, aB, aC); + cublasLtMatmulHeuristicResult_t h = + *getOrComputeAlgo(transA, transB, epilogue, aA, aB, aC); // Fall back to the exact shape when the envelope's algorithm cannot run // it. cuBLASLt returns CUBLAS_STATUS_NOT_SUPPORTED for at least some // shape/algorithm combinations; m=1 was the first observed. - if (env && !algoUsable(desc, h->algo, kA, kB, kC)) { + if (env && !algoUsable(desc, h.algo, kA, kB, kC)) { ++stats.envelopeRejects; - h = getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); + h = *getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); } // Re-stamp the shared role layouts to this call's exact shape. Resolving @@ -567,7 +598,7 @@ class BlasCuda { auto lB = stampLayout(ROLE_B, kB); auto lC = stampLayout(ROLE_C, kC); CHECK_CUBLAS(cublasLtMatmul(ltHandle, desc, &alpha, A, lA, B, lB, &beta, - D_in, lC, C_out, lC, &h->algo, d_workspace, + D_in, lC, C_out, lC, &h.algo, d_workspace, workspaceSize, stream)); } }; diff --git a/tests/test.cc b/tests/test.cc index de65c72..10b5d29 100644 --- a/tests/test.cc +++ b/tests/test.cc @@ -635,6 +635,38 @@ static void runDynamicShapeTests() { ++gFailures; } + // With a limit set, shapes above the envelope are resolved individually and + // the cache must stay at the limit rather than growing per shape. + { + sofieBLAS capped(queue, 8); + capped.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), + MENV, 'N', 'N'); + std::vector cref; + float worst = 0.f; + for (int m = MENV + 1; m <= MCAP; ++m) { + cref.assign(static_cast(m) * N, 0.f); + refMatmul(cref.data(), A, B, m, N, K, 1.f, 0.f, false, false); + capped.matmul('N', 'N', static_cast(m), static_cast(N), + static_cast(K), 1.f, dA, dB, 0.f, dC); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + for (std::size_t i = 0; i < cref.size(); ++i) + worst = std::max(worst, std::abs(C[i] - cref[i])); + } + std::cout << " limit=8: " << capped.algoCacheSize() << " entries, " + << capped.layoutStats().evictions << " evictions over " + << (MCAP - MENV) << " above-envelope sizes, worst err " << worst + << "\n"; + if (capped.algoCacheSize() <= 8 && worst < 1e-3f) { + std::cout << " PASS cuda::cache limit honoured\n"; + } else { + std::cerr << " FAIL [cuda::cache limit honoured] " + << capped.algoCacheSize() << " entries, worst err " << worst + << "\n"; + ++gFailures; + } + } + // The constructor resolves every declared envelope, so a size it covers must // not trigger a search. Only a rejected envelope algorithm may. const std::size_t searched = From 09d312399a8855f18d2b7e3482f1206a72a3cd4c Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Mon, 17 Aug 2026 19:13:42 +0530 Subject: [PATCH 05/13] style: clang-format --- .../backends/cuda/sofieBLAS_cublas.hpp | 33 ++++++++++--------- tests/test.cc | 16 +++++---- 2 files changed, 28 insertions(+), 21 deletions(-) diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index 80f1679..08ac147 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -84,8 +84,8 @@ struct ShapeEnvelope { struct LayoutStats { std::size_t heuristicQueries = 0; - std::size_t envelopeRejects = 0; // envelope algorithm unusable at the call - std::size_t evictions = 0; // entries dropped to stay under the limit + std::size_t envelopeRejects = 0; // envelope algorithm unusable at the call + std::size_t evictions = 0; // entries dropped to stay under the limit }; class BlasCuda { @@ -99,7 +99,8 @@ class BlasCuda { // One persistent layout descriptor per matrix role, re-stamped with the // runtime dimensions before each matmul. The descriptor is host-side metadata // consumed by cublasLtMatmul at the call, so a single object can be reused - // across shapes - this is what lets one Session serve dynamic (runtime) sizes. + // across shapes - this is what lets one Session serve dynamic (runtime) + // sizes. enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 }; cublasLtMatrixLayout_t roleLayout[3] = {}; @@ -109,7 +110,7 @@ class BlasCuda { // list stays empty and the iterator is never read. struct CacheEntry { cublasLtMatmulHeuristicResult_t h{}; - std::list::iterator lru{}; + std::list::iterator lru{}; }; std::unordered_map algoCache; std::list lruOrder; @@ -127,7 +128,7 @@ class BlasCuda { public: const LayoutStats &layoutStats() const { return stats; } - std::size_t algoCacheSize() const { return algoCache.size(); } + std::size_t algoCacheSize() const { return algoCache.size(); } void setAlgoCacheLimit(std::size_t n) { algoCacheLimit = n; } BlasCuda(const BlasCuda &) = delete; @@ -190,9 +191,8 @@ class BlasCuda { // the largest dims the call site will ever use. Layouts are created lazily, // so nothing is registered here beyond the envelope and, with warmup on, the // algorithm resolved for it. - void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, - std::size_t, std::size_t, std::size_t, - char transa, char transb) { + void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, std::size_t, + std::size_t, std::size_t, char transa, char transb) { const auto kA = layoutKeyA(transa, m, k); const auto kB = layoutKeyB(transb, k, n); const std::pair kC{m, n}; @@ -418,10 +418,10 @@ class BlasCuda { // Resolve a matrix role's layout at the runtime dims: create the descriptor // once, then overwrite its dims in place on later calls. ld = rows (dense, // column-major, as the generated calls produce). - cublasLtMatrixLayout_t stampLayout(LayoutRole role, - const std::pair &key) { + cublasLtMatrixLayout_t + stampLayout(LayoutRole role, const std::pair &key) { const uint64_t rows = key.first, cols = key.second; - const int64_t ld = static_cast(key.first); + const int64_t ld = static_cast(key.first); cublasLtMatrixLayout_t &L = roleLayout[role]; if (!L) { CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&L, CUDA_R_32F, rows, cols, ld)); @@ -452,12 +452,15 @@ class BlasCuda { // on it stops one call site's envelope from serving another's shapes. if (e.colsA != kA.second || e.rowsB != kB.first) continue; - if (e.rowsA < kA.first || e.colsB < kB.second || - e.rowsC < kC.first || e.colsC < kC.second) + if (e.rowsA < kA.first || e.colsB < kB.second || e.rowsC < kC.first || + e.colsC < kC.second) continue; const std::size_t ex = (e.rowsA - kA.first) + (e.colsA - kA.second) + (e.rowsB - kB.first) + (e.colsB - kB.second); - if (ex < bestExcess) { bestExcess = ex; best = &e; } + if (ex < bestExcess) { + bestExcess = ex; + best = &e; + } } return best; } @@ -531,7 +534,7 @@ class BlasCuda { auto &desc = getOrCreateDesc(transA, transB, epilogue); auto lA = stampLayout(ROLE_A, kA); auto lB = stampLayout(ROLE_B, kB); - auto lC = stampLayout(ROLE_C, kC); // C and D share the same layout + auto lC = stampLayout(ROLE_C, kC); // C and D share the same layout cublasLtMatmulHeuristicResult_t h{}; int returnedResults = 0; CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic( diff --git a/tests/test.cc b/tests/test.cc index 10b5d29..83bd46e 100644 --- a/tests/test.cc +++ b/tests/test.cc @@ -560,8 +560,9 @@ static void runCudaTests() { } // One instance used at many sizes, which is what a dynamic-shape model does. -// earlier, any size other than the one registered at construction threw std::out_of_range from the layout lookup -// so this is the case the rest of the suite never covered. +// earlier, any size other than the one registered at construction threw +// std::out_of_range from the layout lookup so this is the case the rest of the +// suite never covered. static void runDynamicShapeTests() { std::cout << "\n=== CUDA Dynamic-Shape Tests ===\n"; @@ -586,9 +587,11 @@ static void runDynamicShapeTests() { fillSeq(A, MCAP * K, 0.5f, 0.25f); fillSeq(B, K * N, 1.f, 0.5f); - auto dA = alpaka::allocAsyncBuf(queue, static_cast(MCAP * K)); + auto dA = + alpaka::allocAsyncBuf(queue, static_cast(MCAP * K)); auto dB = alpaka::allocAsyncBuf(queue, static_cast(K * N)); - auto dC = alpaka::allocAsyncBuf(queue, static_cast(MCAP * N)); + auto dC = + alpaka::allocAsyncBuf(queue, static_cast(MCAP * N)); alpaka::memcpy(queue, dA, hA); alpaka::memcpy(queue, dB, hB); alpaka::wait(queue); @@ -646,8 +649,9 @@ static void runDynamicShapeTests() { for (int m = MENV + 1; m <= MCAP; ++m) { cref.assign(static_cast(m) * N, 0.f); refMatmul(cref.data(), A, B, m, N, K, 1.f, 0.f, false, false); - capped.matmul('N', 'N', static_cast(m), static_cast(N), - static_cast(K), 1.f, dA, dB, 0.f, dC); + capped.matmul('N', 'N', static_cast(m), + static_cast(N), static_cast(K), 1.f, dA, + dB, 0.f, dC); alpaka::memcpy(queue, hC, dC); alpaka::wait(queue); for (std::size_t i = 0; i < cref.size(); ++i) From e5b362c7246b04c1971249c0f8c3b78d7bc43751 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Mon, 17 Aug 2026 23:39:15 +0530 Subject: [PATCH 06/13] chore: cleanup, cache limit via constructor --- .../backends/cuda/sofieBLAS_cublas.hpp | 55 ++----------------- tests/test.cc | 15 ----- 2 files changed, 5 insertions(+), 65 deletions(-) diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index 08ac147..aa75a8d 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -76,16 +76,14 @@ struct AlgoKeyHash { } }; -// A call site's maximum shape, as declared by addLayoutConfig from the -// generated Session constructor. struct ShapeEnvelope { std::size_t rowsA, colsA, rowsB, colsB, rowsC, colsC; }; struct LayoutStats { std::size_t heuristicQueries = 0; - std::size_t envelopeRejects = 0; // envelope algorithm unusable at the call - std::size_t evictions = 0; // entries dropped to stay under the limit + std::size_t envelopeRejects = 0; + std::size_t evictions = 0; }; class BlasCuda { @@ -96,32 +94,21 @@ class BlasCuda { size_t workspaceSize = 1u << 25; // 32 MB cudaStream_t stream = nullptr; - // One persistent layout descriptor per matrix role, re-stamped with the - // runtime dimensions before each matmul. The descriptor is host-side metadata - // consumed by cublasLtMatmul at the call, so a single object can be reused - // across shapes - this is what lets one Session serve dynamic (runtime) - // sizes. enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 }; cublasLtMatrixLayout_t roleLayout[3] = {}; std::unordered_map descStore; - // Recency handle is only maintained when a limit is set; with no limit the - // list stays empty and the iterator is never read. + // algo cache entry struct CacheEntry { cublasLtMatmulHeuristicResult_t h{}; std::list::iterator lru{}; }; std::unordered_map algoCache; std::list lruOrder; - // 0 = unbounded. Per-call-site resolution already bounds the cache by the - // number of Conv/Gemm in the model, so a limit only matters when infer() - // runs above the constructed size: those shapes have no envelope and are - // resolved individually. Set it below the number of call sites and - // construction will evict its own warmup. + // 0 = unbounded std::size_t algoCacheLimit = 0; - // call-site envelopes declared by addLayoutConfig std::vector envelopes; LayoutStats stats; @@ -129,7 +116,6 @@ class BlasCuda { public: const LayoutStats &layoutStats() const { return stats; } std::size_t algoCacheSize() const { return algoCache.size(); } - void setAlgoCacheLimit(std::size_t n) { algoCacheLimit = n; } BlasCuda(const BlasCuda &) = delete; BlasCuda &operator=(const BlasCuda &) = delete; @@ -186,11 +172,6 @@ class BlasCuda { } } - // Records the call site's envelope. The generated constructor evaluates its - // shape expressions with its own parameters, so for a dynamic model these are - // the largest dims the call site will ever use. Layouts are created lazily, - // so nothing is registered here beyond the envelope and, with warmup on, the - // algorithm resolved for it. void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, std::size_t, std::size_t, std::size_t, char transa, char transb) { const auto kA = layoutKeyA(transa, m, k); @@ -198,9 +179,6 @@ class BlasCuda { const std::pair kC{m, n}; envelopes.push_back({kA.first, kA.second, kB.first, kB.second, m, n}); - // The constructor does not know which epilogue this call site uses, so - // resolve all three. Unused ones cost one heuristic query each, off the - // inference path. const cublasOperation_t tA = charToCuBlasTranspose(transa); const cublasOperation_t tB = charToCuBlasTranspose(transb); const cublasLtEpilogue_t eps[] = {CUBLASLT_EPILOGUE_DEFAULT, @@ -415,9 +393,6 @@ class BlasCuda { : std::make_pair(n, k); } - // Resolve a matrix role's layout at the runtime dims: create the descriptor - // once, then overwrite its dims in place on later calls. ld = rows (dense, - // column-major, as the generated calls produce). cublasLtMatrixLayout_t stampLayout(LayoutRole role, const std::pair &key) { const uint64_t rows = key.first, cols = key.second; @@ -436,10 +411,6 @@ class BlasCuda { return L; } - // Tightest declared envelope covering this call, or null if none does. - // Tightest matters: several envelopes may cover a small shape, but only the - // call site's own matches its weight dims exactly and so has zero excess - // on those axes. const ShapeEnvelope * findEnvelope(const std::pair &kA, const std::pair &kB, @@ -447,9 +418,6 @@ class BlasCuda { const ShapeEnvelope *best = nullptr; std::size_t bestExcess = std::numeric_limits::max(); for (const auto &e : envelopes) { - // colsA and rowsB are both the contraction dimension k, which comes from - // the weight tensor and never varies at runtime. Requiring an exact match - // on it stops one call site's envelope from serving another's shapes. if (e.colsA != kA.second || e.rowsB != kB.first) continue; if (e.rowsA < kA.first || e.colsB < kB.second || e.rowsC < kC.first || @@ -493,9 +461,6 @@ class BlasCuda { return descStore.at(key); } - // Whether an algorithm can actually run this shape. cuBLASLt rejects some - // combinations, so an algorithm resolved at a call site's envelope is not - // guaranteed to work at every smaller shape it serves. bool algoUsable(cublasLtMatmulDesc_t desc, const cublasLtMatmulAlgo_t &algo, const std::pair &kA, const std::pair &kB, @@ -509,9 +474,6 @@ class BlasCuda { chk.workspaceSize <= workspaceSize; } - // required=false is used by constructor warmup, which speculatively resolves - // epilogues the call site may never use: those may legitimately have no - // algorithm and must not abort. cublasLtMatmulHeuristicResult_t * getOrComputeAlgo(cublasOperation_t transA, cublasOperation_t transB, cublasLtEpilogue_t epilogue, @@ -534,7 +496,7 @@ class BlasCuda { auto &desc = getOrCreateDesc(transA, transB, epilogue); auto lA = stampLayout(ROLE_A, kA); auto lB = stampLayout(ROLE_B, kB); - auto lC = stampLayout(ROLE_C, kC); // C and D share the same layout + auto lC = stampLayout(ROLE_C, kC); cublasLtMatmulHeuristicResult_t h{}; int returnedResults = 0; CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic( @@ -577,8 +539,6 @@ class BlasCuda { sizeof(bias_ptr))); } - // Resolve at this call site's declared envelope, so every runtime size it - // produces shares one cache entry. const ShapeEnvelope *env = findEnvelope(kA, kB, kC); const std::pair aA = env ? std::make_pair(env->rowsA, env->colsA) : kA, @@ -587,16 +547,11 @@ class BlasCuda { cublasLtMatmulHeuristicResult_t h = *getOrComputeAlgo(transA, transB, epilogue, aA, aB, aC); - // Fall back to the exact shape when the envelope's algorithm cannot run - // it. cuBLASLt returns CUBLAS_STATUS_NOT_SUPPORTED for at least some - // shape/algorithm combinations; m=1 was the first observed. if (env && !algoUsable(desc, h.algo, kA, kB, kC)) { ++stats.envelopeRejects; h = *getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); } - // Re-stamp the shared role layouts to this call's exact shape. Resolving - // the algorithm leaves them at the envelope size, so this happens last. auto lA = stampLayout(ROLE_A, kA); auto lB = stampLayout(ROLE_B, kB); auto lC = stampLayout(ROLE_C, kC); diff --git a/tests/test.cc b/tests/test.cc index 83bd46e..fcb15b8 100644 --- a/tests/test.cc +++ b/tests/test.cc @@ -559,10 +559,6 @@ static void runCudaTests() { } } -// One instance used at many sizes, which is what a dynamic-shape model does. -// earlier, any size other than the one registered at construction threw -// std::out_of_range from the layout lookup so this is the case the rest of the -// suite never covered. static void runDynamicShapeTests() { std::cout << "\n=== CUDA Dynamic-Shape Tests ===\n"; @@ -574,8 +570,6 @@ static void runDynamicShapeTests() { alpaka::PlatformCpu hostPlatform{}; auto hostDev = alpaka::getDevByIdx(hostPlatform, 0u); - // Buffers hold MCAP rows while the call site declares MENV, so sizes above - // MENV stay in bounds and exercise the path where no envelope covers them. constexpr int MCAP = 96, MENV = 64, N = 3, K = 5; auto hA = alpaka::allocBuf(hostDev, static_cast(MCAP * K)); @@ -596,8 +590,6 @@ static void runDynamicShapeTests() { alpaka::memcpy(queue, dB, hB); alpaka::wait(queue); - // Declare the call site's largest shape, as a generated Session constructor - // does with its own arguments. blas.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), MENV, 'N', 'N'); @@ -612,12 +604,9 @@ static void runDynamicShapeTests() { checkClose(C, ref.data(), m * N, name); }; - // At, below and above the declared envelope, all on one instance. for (int m : {MENV, 37, 8, 51, 1, MENV, MCAP}) runAt(m, "cuda::dynamic m=" + std::to_string(m)); - // The cache must not grow one entry per size. Entries are added only where - // the envelope's algorithm is rejected. Keying by shape fails only here. const int nSizes = MENV - 1; const std::size_t cacheBefore = blas.algoCacheSize(); const std::size_t rejBefore = blas.layoutStats().envelopeRejects; @@ -638,8 +627,6 @@ static void runDynamicShapeTests() { ++gFailures; } - // With a limit set, shapes above the envelope are resolved individually and - // the cache must stay at the limit rather than growing per shape. { sofieBLAS capped(queue, 8); capped.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), @@ -671,8 +658,6 @@ static void runDynamicShapeTests() { } } - // The constructor resolves every declared envelope, so a size it covers must - // not trigger a search. Only a rejected envelope algorithm may. const std::size_t searched = blas.layoutStats().heuristicQueries - searchBefore; if (searched <= rejected) { From d71ddf54469151df2572f270259b49abe2498677 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Tue, 18 Aug 2026 17:04:37 +0530 Subject: [PATCH 07/13] docs: layout and caching behavior --- README.md | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 191fe2a..8c39bf7 100644 --- a/README.md +++ b/README.md @@ -96,7 +96,22 @@ sofieBLAS blas(queue); blas.matmul('N', 'N', size, size, size, 1.0f, dA, dB, 0.0f, dC); ``` -The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addLayoutConfig` (used to pre-register cuBLASLt/hipBLASLt matrix layouts for a given shape before the first `matmul`/`gemm` call on that shape). +The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addLayoutConfig` (declares a call site's shape ahead of the first call; on CUDA this feeds the algorithm cache described below, on HIP it pre-registers the matrix layouts for that shape). + +## Dynamic GEMM shapes and the algorithm cache (CUDA) + +One `BlasCuda` instance serves GEMM calls at sizes that vary at runtime. Matrix layouts are not tied to a shape: each call stamps its dimensions into a shared per-role descriptor right before the multiply. + +cuBLASLt algorithm selection is cached. `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)` declares the largest shape a call site will use (its envelope) and resolves the algorithm for it once, up front. Any later call at a size covered by an envelope reuses that entry, so sweeping sizes does not re-query the heuristic or grow the cache. A call with no covering envelope is resolved at its exact shape and cached per shape, and if an envelope's algorithm cannot run a particular size the call falls back to exact-shape resolution. + +The cache is unbounded by default. Pass a limit at construction to cap it with LRU eviction: + +```cpp +sofieBLAS blas(queue); // unbounded cache (default) +sofieBLAS capped(queue, 32); // at most 32 entries, LRU eviction +``` + +`algoCacheSize()` returns the current entry count and `layoutStats()` returns counters (`heuristicQueries`, `envelopeRejects`, `evictions`) for inspecting cache behaviour. ## Contributing From 1b0ea970521854e8aa0b3d82b264fa241e6b4d0e Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Wed, 19 Aug 2026 14:29:44 +0530 Subject: [PATCH 08/13] feat: dynamic gemm shapes and algorithm cache for hipBLASLt --- README.md | 11 +- .../backends/hip/sofieBLAS_hipblaslt.hpp | 215 +++++++++++++----- tests/test.cc | 111 +++++++++ 3 files changed, 272 insertions(+), 65 deletions(-) diff --git a/README.md b/README.md index 8c39bf7..5ea25f3 100644 --- a/README.md +++ b/README.md @@ -96,23 +96,26 @@ sofieBLAS blas(queue); blas.matmul('N', 'N', size, size, size, 1.0f, dA, dB, 0.0f, dC); ``` -The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addLayoutConfig` (declares a call site's shape ahead of the first call; on CUDA this feeds the algorithm cache described below, on HIP it pre-registers the matrix layouts for that shape). +The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addLayoutConfig` (declares a call site's shape ahead of the first call, which feeds the algorithm cache described below). -## Dynamic GEMM shapes and the algorithm cache (CUDA) +## Dynamic GEMM shapes and the algorithm cache -One `BlasCuda` instance serves GEMM calls at sizes that vary at runtime. Matrix layouts are not tied to a shape: each call stamps its dimensions into a shared per-role descriptor right before the multiply. +Both GPU backends behave the same way here. One instance serves GEMM calls at sizes that vary at runtime: matrix layouts are not tied to a shape, and each call stamps its dimensions into a shared per-role descriptor right before the multiply. -cuBLASLt algorithm selection is cached. `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)` declares the largest shape a call site will use (its envelope) and resolves the algorithm for it once, up front. Any later call at a size covered by an envelope reuses that entry, so sweeping sizes does not re-query the heuristic or grow the cache. A call with no covering envelope is resolved at its exact shape and cached per shape, and if an envelope's algorithm cannot run a particular size the call falls back to exact-shape resolution. +Algorithm selection is cached. `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)` declares the largest shape a call site will use (its envelope) and resolves the algorithm for it once, up front. Any later call at a size covered by an envelope reuses that entry, so sweeping sizes does not re-query the heuristic or grow the cache. A call with no covering envelope is resolved at its exact shape and cached per shape, and if an envelope's algorithm cannot run a particular size the call falls back to exact-shape resolution. The cache is unbounded by default. Pass a limit at construction to cap it with LRU eviction: ```cpp sofieBLAS blas(queue); // unbounded cache (default) sofieBLAS capped(queue, 32); // at most 32 entries, LRU eviction +sofieBLAS hipCapped(queue, 32); ``` `algoCacheSize()` returns the current entry count and `layoutStats()` returns counters (`heuristicQueries`, `envelopeRejects`, `evictions`) for inspecting cache behaviour. +The exact-shape fallback uses `cublasLtMatmulAlgoCheck` on CUDA. hipBLASLt has no equivalent in its core API, so the HIP backend uses `hipblaslt_ext::matmulIsAlgoSupported` from `hipblaslt-ext.hpp` for the same check. + ## Contributing diff --git a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp index 0e41e70..5ce5fa5 100644 --- a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp +++ b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp @@ -5,13 +5,17 @@ #include #include #include +#include +#include #include #include #include +#include #include "sofieBLAS/core.hpp" #include #include +#include #include #define CHECK_HIP(err) \ @@ -30,22 +34,6 @@ } \ } while (0) -struct PairHash { - std::size_t - operator()(const std::pair &p) const noexcept { - std::size_t h1 = std::hash{}(p.first); - std::size_t h2 = std::hash{}(p.second); - return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2)); - } -}; - -struct PairEq { - bool operator()(const std::pair &a, - const std::pair &b) const noexcept { - return a.first == b.first && a.second == b.second; - } -}; - struct DescKey { int transA; // HIPBLAS_OP_N / HIPBLAS_OP_T encoded as int int transB; @@ -66,8 +54,8 @@ struct DescKeyHash { struct AlgoKey { DescKey dk; - std::size_t rowsA, colsA; // physical dimensions of A in layoutStore - std::size_t rowsB, colsB; // physical dimensions of B in layoutStore + std::size_t rowsA, colsA; + std::size_t rowsB, colsB; bool operator==(const AlgoKey &o) const noexcept { return dk == o.dk && rowsA == o.rowsA && colsA == o.colsA && rowsB == o.rowsB && colsB == o.colsB; @@ -89,6 +77,16 @@ struct AlgoKeyHash { } }; +struct ShapeEnvelope { + std::size_t rowsA, colsA, rowsB, colsB, rowsC, colsC; +}; + +struct LayoutStats { + std::size_t heuristicQueries = 0; + std::size_t envelopeRejects = 0; + std::size_t evictions = 0; +}; + class BlasHip { hipblasLtHandle_t ltHandle = nullptr; hipblasHandle_t handle = nullptr; @@ -97,22 +95,36 @@ class BlasHip { size_t workspaceSize = 1u << 25; // 32 MB hipStream_t stream = nullptr; - std::unordered_map, - hipblasLtMatrixLayout_t, PairHash, PairEq> - layoutStore; + enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 }; + hipblasLtMatrixLayout_t roleLayout[3] = {}; std::unordered_map descStore; - std::unordered_map - algoCache; + // algo cache entry + struct CacheEntry { + hipblasLtMatmulHeuristicResult_t h{}; + std::list::iterator lru{}; + }; + std::unordered_map algoCache; + std::list lruOrder; + // 0 = unbounded + std::size_t algoCacheLimit = 0; + + std::vector envelopes; + + LayoutStats stats; public: + const LayoutStats &layoutStats() const { return stats; } + std::size_t algoCacheSize() const { return algoCache.size(); } + BlasHip(const BlasHip &) = delete; BlasHip &operator=(const BlasHip &) = delete; BlasHip(BlasHip &&) = delete; BlasHip &operator=(BlasHip &&) = delete; - BlasHip(alpaka::QueueHipRtNonBlocking &queue) : m_queue{queue} { + BlasHip(alpaka::QueueHipRtNonBlocking &queue, std::size_t algoCacheLimit_ = 0) + : algoCacheLimit{algoCacheLimit_}, m_queue{queue} { stream = static_cast(m_queue.getNativeHandle()); CHECK_HIPBLAS(hipblasLtCreate(<Handle)); @@ -128,9 +140,9 @@ class BlasHip { } ~BlasHip() { - for (auto &[key, layout] : layoutStore) - if (layout) - hipblasLtMatrixLayoutDestroy(layout); + for (auto L : roleLayout) + if (L) + hipblasLtMatrixLayoutDestroy(L); for (auto &[key, desc] : descStore) if (desc) hipblasLtMatmulDescDestroy(desc); @@ -160,18 +172,21 @@ class BlasHip { } } - void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, - std::size_t lda, std::size_t ldb, std::size_t ldc, - char transa, char transb) { - if (transa == 'N' || transa == 'n') - checkAndAddLayout(m, k, lda); - else - checkAndAddLayout(k, m, lda); - if (transb == 'N' || transb == 'n') - checkAndAddLayout(k, n, ldb); - else - checkAndAddLayout(n, k, ldb); - checkAndAddLayout(m, n, ldc); + void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, std::size_t, + std::size_t, std::size_t, char transa, char transb) { + const auto kA = layoutKeyA(transa, m, k); + const auto kB = layoutKeyB(transb, k, n); + const std::pair kC{m, n}; + envelopes.push_back({kA.first, kA.second, kB.first, kB.second, m, n}); + + const hipblasOperation_t tA = charToHipBlasTranspose(transa); + const hipblasOperation_t tB = charToHipBlasTranspose(transb); + const hipblasLtEpilogue_t eps[] = {HIPBLASLT_EPILOGUE_DEFAULT, + HIPBLASLT_EPILOGUE_BIAS, + HIPBLASLT_EPILOGUE_RELU_BIAS}; + for (hipblasLtEpilogue_t ep : eps) { + getOrComputeAlgo(tA, tB, ep, kA, kB, kC, /*required=*/false); + } } template @@ -377,14 +392,62 @@ class BlasHip { : std::make_pair(n, k); } - void checkAndAddLayout(std::size_t rows, std::size_t cols, std::size_t ld) { - auto key = std::make_pair(rows, cols); - if (layoutStore.find(key) == layoutStore.end()) { - hipblasLtMatrixLayout_t layout = nullptr; - CHECK_HIPBLAS( - hipblasLtMatrixLayoutCreate(&layout, HIP_R_32F, rows, cols, ld)); - layoutStore.emplace(key, layout); + hipblasLtMatrixLayout_t + stampLayout(LayoutRole role, const std::pair &key) { + const uint64_t rows = key.first, cols = key.second; + const int64_t ld = static_cast(key.first); + hipblasLtMatrixLayout_t &L = roleLayout[role]; + if (!L) { + CHECK_HIPBLAS(hipblasLtMatrixLayoutCreate(&L, HIP_R_32F, rows, cols, ld)); + } else { + CHECK_HIPBLAS(hipblasLtMatrixLayoutSetAttribute( + L, HIPBLASLT_MATRIX_LAYOUT_ROWS, &rows, sizeof(rows))); + CHECK_HIPBLAS(hipblasLtMatrixLayoutSetAttribute( + L, HIPBLASLT_MATRIX_LAYOUT_COLS, &cols, sizeof(cols))); + CHECK_HIPBLAS(hipblasLtMatrixLayoutSetAttribute( + L, HIPBLASLT_MATRIX_LAYOUT_LD, &ld, sizeof(ld))); } + return L; + } + + const ShapeEnvelope * + findEnvelope(const std::pair &kA, + const std::pair &kB, + const std::pair &kC) const { + const ShapeEnvelope *best = nullptr; + std::size_t bestExcess = std::numeric_limits::max(); + for (const auto &e : envelopes) { + if (e.colsA != kA.second || e.rowsB != kB.first) + continue; + if (e.rowsA < kA.first || e.colsB < kB.second || e.rowsC < kC.first || + e.colsC < kC.second) + continue; + const std::size_t ex = (e.rowsA - kA.first) + (e.colsA - kA.second) + + (e.rowsB - kB.first) + (e.colsB - kB.second); + if (ex < bestExcess) { + bestExcess = ex; + best = &e; + } + } + return best; + } + + bool algoUsable(hipblasLtMatmulDesc_t desc, const hipblasLtMatmulAlgo_t &algo, + const std::pair &kA, + const std::pair &kB, + const std::pair &kC) { + auto lA = stampLayout(ROLE_A, kA); + auto lB = stampLayout(ROLE_B, kB); + auto lC = stampLayout(ROLE_C, kC); + // hipBLASLt has no hipblasLtMatmulAlgoCheck; the ext API rewrites the algo + // it is handed, so probe a copy. + hipblasLtMatmulAlgo_t probe = algo; + const float alpha = 1.f, beta = 0.f; + std::size_t ws = 0; + return hipblaslt_ext::matmulIsAlgoSupported(ltHandle, desc, &alpha, lA, lB, + &beta, lC, lC, probe, + ws) == HIPBLAS_STATUS_SUCCESS && + ws <= workspaceSize; } hipblasLtMatmulDesc_t &getOrCreateDesc(hipblasOperation_t transA, @@ -414,29 +477,37 @@ class BlasHip { return descStore.at(key); } - hipblasLtMatmulHeuristicResult_t & + hipblasLtMatmulHeuristicResult_t * getOrComputeAlgo(hipblasOperation_t transA, hipblasOperation_t transB, hipblasLtEpilogue_t epilogue, const std::pair &kA, const std::pair &kB, - const std::pair &kC) { + const std::pair &kC, + bool required = true) { AlgoKey key{{(int)transA, (int)transB, (int)epilogue}, kA.first, kA.second, kB.first, kB.second}; auto it = algoCache.find(key); - if (it != algoCache.end()) - return it->second; + if (it != algoCache.end()) { + if (algoCacheLimit) + lruOrder.splice(lruOrder.begin(), lruOrder, it->second.lru); + return &it->second.h; + } auto &desc = getOrCreateDesc(transA, transB, epilogue); + auto lA = stampLayout(ROLE_A, kA); + auto lB = stampLayout(ROLE_B, kB); + auto lC = stampLayout(ROLE_C, kC); hipblasLtMatmulHeuristicResult_t h{}; int returnedResults = 0; CHECK_HIPBLAS(hipblasLtMatmulAlgoGetHeuristic( - ltHandle, desc, layoutStore.at(kA), layoutStore.at(kB), - layoutStore.at(kC), layoutStore.at(kC), preference, 1, &h, - &returnedResults)); + ltHandle, desc, lA, lB, lC, lC, preference, 1, &h, &returnedResults)); + ++stats.heuristicQueries; if (returnedResults == 0) { + if (!required) + return nullptr; std::cerr << "[sofieBLAS] No suitable hipBLASLt algorithm found for " << "transA=" << transA << " transB=" << transB << " epilogue=" << epilogue << " A=[" << kA.first << "x" @@ -444,8 +515,17 @@ class BlasHip { << " B=[" << kB.first << "x" << kB.second << "]\n"; exit(EXIT_FAILURE); } - algoCache.emplace(key, h); - return algoCache.at(key); + auto ins = algoCache.emplace(key, CacheEntry{h, {}}).first; + if (algoCacheLimit) { + lruOrder.push_front(key); + ins->second.lru = lruOrder.begin(); + while (algoCache.size() > algoCacheLimit) { + algoCache.erase(lruOrder.back()); + lruOrder.pop_back(); + ++stats.evictions; + } + } + return &ins->second.h; } void executeMatmul(hipblasOperation_t transA, hipblasOperation_t transB, @@ -455,8 +535,6 @@ class BlasHip { const std::pair &kA, const std::pair &kB, const std::pair &kC) { - auto &h = getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); - auto &desc = getOrCreateDesc(transA, transB, epilogue); if (bias_ptr) { CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( @@ -464,10 +542,25 @@ class BlasHip { sizeof(bias_ptr))); } - CHECK_HIPBLAS(hipblasLtMatmul(ltHandle, desc, &alpha, A, layoutStore.at(kA), - B, layoutStore.at(kB), &beta, D_in, - layoutStore.at(kC), C_out, layoutStore.at(kC), - &h.algo, d_workspace, workspaceSize, stream)); + const ShapeEnvelope *env = findEnvelope(kA, kB, kC); + const std::pair + aA = env ? std::make_pair(env->rowsA, env->colsA) : kA, + aB = env ? std::make_pair(env->rowsB, env->colsB) : kB, + aC = env ? std::make_pair(env->rowsC, env->colsC) : kC; + hipblasLtMatmulHeuristicResult_t h = + *getOrComputeAlgo(transA, transB, epilogue, aA, aB, aC); + + if (env && !algoUsable(desc, h.algo, kA, kB, kC)) { + ++stats.envelopeRejects; + h = *getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); + } + + auto lA = stampLayout(ROLE_A, kA); + auto lB = stampLayout(ROLE_B, kB); + auto lC = stampLayout(ROLE_C, kC); + CHECK_HIPBLAS(hipblasLtMatmul(ltHandle, desc, &alpha, A, lA, B, lB, &beta, + D_in, lC, C_out, lC, &h.algo, d_workspace, + workspaceSize, stream)); } }; diff --git a/tests/test.cc b/tests/test.cc index fcb15b8..0d4002a 100644 --- a/tests/test.cc +++ b/tests/test.cc @@ -890,6 +890,116 @@ static void runHipTests() { } } +static void runHipDynamicShapeTests() { + std::cout << "\n=== HIP Dynamic-Shape Tests ===\n"; + + alpaka::PlatformHipRt platform{}; + auto dev = alpaka::getDevByIdx(platform, 0u); + alpaka::Queue queue{dev}; + sofieBLAS blas(queue); + + alpaka::PlatformCpu hostPlatform{}; + auto hostDev = alpaka::getDevByIdx(hostPlatform, 0u); + + constexpr int MCAP = 96, MENV = 64, N = 3, K = 5; + + auto hA = alpaka::allocBuf(hostDev, static_cast(MCAP * K)); + auto hB = alpaka::allocBuf(hostDev, static_cast(K * N)); + auto hC = alpaka::allocBuf(hostDev, static_cast(MCAP * N)); + float *A = alpaka::getPtrNative(hA); + float *B = alpaka::getPtrNative(hB); + float *C = alpaka::getPtrNative(hC); + fillSeq(A, MCAP * K, 0.5f, 0.25f); + fillSeq(B, K * N, 1.f, 0.5f); + + auto dA = + alpaka::allocAsyncBuf(queue, static_cast(MCAP * K)); + auto dB = alpaka::allocAsyncBuf(queue, static_cast(K * N)); + auto dC = + alpaka::allocAsyncBuf(queue, static_cast(MCAP * N)); + alpaka::memcpy(queue, dA, hA); + alpaka::memcpy(queue, dB, hB); + alpaka::wait(queue); + + blas.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), + MENV, 'N', 'N'); + + std::vector ref; + auto runAt = [&](int m, const std::string &name) { + ref.assign(static_cast(m) * N, 0.f); + refMatmul(ref.data(), A, B, m, N, K, 1.f, 0.f, false, false); + blas.matmul('N', 'N', static_cast(m), static_cast(N), + static_cast(K), 1.f, dA, dB, 0.f, dC); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + checkClose(C, ref.data(), m * N, name); + }; + + for (int m : {MENV, 37, 8, 51, 1, MENV, MCAP}) + runAt(m, "hip::dynamic m=" + std::to_string(m)); + + const int nSizes = MENV - 1; + const std::size_t cacheBefore = blas.algoCacheSize(); + const std::size_t rejBefore = blas.layoutStats().envelopeRejects; + const std::size_t searchBefore = blas.layoutStats().heuristicQueries; + for (int m = 2; m <= MENV; ++m) + blas.matmul('N', 'N', static_cast(m), static_cast(N), + static_cast(K), 1.f, dA, dB, 0.f, dC); + alpaka::wait(queue); + const std::size_t added = blas.algoCacheSize() - cacheBefore; + const std::size_t rejected = blas.layoutStats().envelopeRejects - rejBefore; + + std::cout << " " << nSizes << " sizes added " << added + << " cache entries, " << rejected << " rejected\n"; + if (added < static_cast(nSizes)) { + std::cout << " PASS hip::cache bounded\n"; + } else { + std::cerr << " FAIL [hip::cache bounded] one entry per size\n"; + ++gFailures; + } + + { + sofieBLAS capped(queue, 8); + capped.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), + MENV, 'N', 'N'); + std::vector cref; + float worst = 0.f; + for (int m = MENV + 1; m <= MCAP; ++m) { + cref.assign(static_cast(m) * N, 0.f); + refMatmul(cref.data(), A, B, m, N, K, 1.f, 0.f, false, false); + capped.matmul('N', 'N', static_cast(m), + static_cast(N), static_cast(K), 1.f, dA, + dB, 0.f, dC); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + for (std::size_t i = 0; i < cref.size(); ++i) + worst = std::max(worst, std::abs(C[i] - cref[i])); + } + std::cout << " limit=8: " << capped.algoCacheSize() << " entries, " + << capped.layoutStats().evictions << " evictions over " + << (MCAP - MENV) << " above-envelope sizes, worst err " << worst + << "\n"; + if (capped.algoCacheSize() <= 8 && worst < 1e-3f) { + std::cout << " PASS hip::cache limit honoured\n"; + } else { + std::cerr << " FAIL [hip::cache limit honoured] " + << capped.algoCacheSize() << " entries, worst err " << worst + << "\n"; + ++gFailures; + } + } + + const std::size_t searched = + blas.layoutStats().heuristicQueries - searchBefore; + if (searched <= rejected) { + std::cout << " PASS hip::no search during inference\n"; + } else { + std::cerr << " FAIL [hip::no search during inference] " << searched + << " searches, only " << rejected << " explained by rejects\n"; + ++gFailures; + } +} + #endif // ALPAKA_ACC_GPU_HIP_ENABLED // --------------------------------------------------------------------------- @@ -906,6 +1016,7 @@ int main() { #endif #ifdef ALPAKA_ACC_GPU_HIP_ENABLED runHipTests(); + runHipDynamicShapeTests(); #endif std::cout << "\n"; From 11ff88e8a73d1e46c47a1086eb2dfba1d047f656 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Thu, 20 Aug 2026 21:31:48 +0530 Subject: [PATCH 09/13] docs: describe layout roles, envelope matching and shape naming --- README.md | 6 +- .../backends/cuda/sofieBLAS_cublas.hpp | 119 +++++++++++------- .../backends/hip/sofieBLAS_hipblaslt.hpp | 118 ++++++++++------- 3 files changed, 147 insertions(+), 96 deletions(-) diff --git a/README.md b/README.md index 5ea25f3..2c96e69 100644 --- a/README.md +++ b/README.md @@ -100,9 +100,11 @@ The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgel ## Dynamic GEMM shapes and the algorithm cache -Both GPU backends behave the same way here. One instance serves GEMM calls at sizes that vary at runtime: matrix layouts are not tied to a shape, and each call stamps its dimensions into a shared per-role descriptor right before the multiply. +A GEMM call computes `C = alpha * op(A) * op(B) + beta * C`: A and B are the input matrices, C the output, and `op` an optional transpose. The CUDA backend (`BlasCuda`, over cuBLASLt) and the HIP backend (`BlasHip`, over hipBLASLt) implement the following identically. One instance serves GEMM calls at sizes that vary at runtime: for every call, each of the three matrices has its dimensions stamped into a persistent per-matrix layout descriptor right before the multiply, so no layout is ever tied to a fixed shape. -Algorithm selection is cached. `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)` declares the largest shape a call site will use (its envelope) and resolves the algorithm for it once, up front. Any later call at a size covered by an envelope reuses that entry, so sweeping sizes does not re-query the heuristic or grow the cache. A call with no covering envelope is resolved at its exact shape and cached per shape, and if an envelope's algorithm cannot run a particular size the call falls back to exact-shape resolution. +Every matmul also needs an algorithm: the concrete GEMM kernel and its configuration, chosen by querying the library's heuristic (`cublasLtMatmulAlgoGetHeuristic` / `hipblasLtMatmulAlgoGetHeuristic`) for the given transpose settings, epilogue, and matrix dimensions. That query runs on the host and is not free, so its results are cached, keyed by the same parameters. Left at that, a model with runtime-varying sizes would pay one heuristic query per distinct size it meets, and the cache would grow with every new size. + +The envelope mechanism removes that cost. An envelope is the largest shape a GEMM call site will ever run, declared through `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)`, which a generated Session constructor calls once per call site with its own construction-time dimensions. `addLayoutConfig` stores the envelope and resolves its algorithm immediately. At matmul time the call is matched to a declared envelope (the contraction dimension must match exactly, which is what ties a call to its own call site), and the algorithm is looked up at the envelope's dimensions rather than the call's, so every runtime size a call site produces shares that one cache entry: no further heuristic queries, no cache growth. A call covered by no envelope is resolved and cached at its exact shape. The library may also reject an envelope's algorithm at a specific smaller size (m=1 was found to do this in testing); the backend checks for this and falls back to exact-shape resolution for that call. The cache is unbounded by default. Pass a limit at construction to cap it with LRU eviction: diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index aa75a8d..d4b6db8 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -94,6 +94,11 @@ class BlasCuda { size_t workspaceSize = 1u << 25; // 32 MB cudaStream_t stream = nullptr; + // One persistent layout descriptor per GEMM operand: ROLE_A and ROLE_B are + // the input matrices of C = alpha * op(A) * op(B) + beta * C, ROLE_C the + // output (cublasLtMatmul takes it twice, as C and D). stampLayout rewrites + // each descriptor in place to the shape of the call at hand, which is what + // lets one instance serve any runtime size. enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 }; cublasLtMatrixLayout_t roleLayout[3] = {}; @@ -174,10 +179,11 @@ class BlasCuda { void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, std::size_t, std::size_t, std::size_t, char transa, char transb) { - const auto kA = layoutKeyA(transa, m, k); - const auto kB = layoutKeyB(transb, k, n); - const std::pair kC{m, n}; - envelopes.push_back({kA.first, kA.second, kB.first, kB.second, m, n}); + const auto shapeA = layoutKeyA(transa, m, k); + const auto shapeB = layoutKeyB(transb, k, n); + const std::pair shapeC{m, n}; + envelopes.push_back( + {shapeA.first, shapeA.second, shapeB.first, shapeB.second, m, n}); const cublasOperation_t tA = charToCuBlasTranspose(transa); const cublasOperation_t tB = charToCuBlasTranspose(transb); @@ -185,7 +191,7 @@ class BlasCuda { CUBLASLT_EPILOGUE_BIAS, CUBLASLT_EPILOGUE_RELU_BIAS}; for (cublasLtEpilogue_t ep : eps) { - getOrComputeAlgo(tA, tB, ep, kA, kB, kC, /*required=*/false); + getOrComputeAlgo(tA, tB, ep, shapeA, shapeB, shapeC, /*required=*/false); } } @@ -411,20 +417,29 @@ class BlasCuda { return L; } + // Returns the shape declared through addLayoutConfig that this call belongs + // to, or null if none covers it. Shapes are the physical (rows, cols) of + // matrices A, B and C, after any transpose is applied (transa='T' makes + // shapeA = (k, m)). The contraction dimension (colsA / rowsB) must match + // exactly: it comes from the weight tensor and never varies at runtime, so + // it identifies the call site and stops one site's declared shape from + // serving another's calls. The free dimensions only need covering; among + // candidates the least excess wins. const ShapeEnvelope * - findEnvelope(const std::pair &kA, - const std::pair &kB, - const std::pair &kC) const { + findEnvelope(const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC) const { const ShapeEnvelope *best = nullptr; std::size_t bestExcess = std::numeric_limits::max(); for (const auto &e : envelopes) { - if (e.colsA != kA.second || e.rowsB != kB.first) + if (e.colsA != shapeA.second || e.rowsB != shapeB.first) continue; - if (e.rowsA < kA.first || e.colsB < kB.second || e.rowsC < kC.first || - e.colsC < kC.second) + if (e.rowsA < shapeA.first || e.colsB < shapeB.second || + e.rowsC < shapeC.first || e.colsC < shapeC.second) continue; - const std::size_t ex = (e.rowsA - kA.first) + (e.colsA - kA.second) + - (e.rowsB - kB.first) + (e.colsB - kB.second); + const std::size_t ex = + (e.rowsA - shapeA.first) + (e.colsA - shapeA.second) + + (e.rowsB - shapeB.first) + (e.colsB - shapeB.second); if (ex < bestExcess) { bestExcess = ex; best = &e; @@ -462,12 +477,12 @@ class BlasCuda { } bool algoUsable(cublasLtMatmulDesc_t desc, const cublasLtMatmulAlgo_t &algo, - const std::pair &kA, - const std::pair &kB, - const std::pair &kC) { - auto lA = stampLayout(ROLE_A, kA); - auto lB = stampLayout(ROLE_B, kB); - auto lC = stampLayout(ROLE_C, kC); + const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC) { + auto lA = stampLayout(ROLE_A, shapeA); + auto lB = stampLayout(ROLE_B, shapeB); + auto lC = stampLayout(ROLE_C, shapeC); cublasLtMatmulHeuristicResult_t chk{}; return cublasLtMatmulAlgoCheck(ltHandle, desc, lA, lB, lC, lC, &algo, &chk) == CUBLAS_STATUS_SUCCESS && @@ -477,15 +492,15 @@ class BlasCuda { cublasLtMatmulHeuristicResult_t * getOrComputeAlgo(cublasOperation_t transA, cublasOperation_t transB, cublasLtEpilogue_t epilogue, - const std::pair &kA, - const std::pair &kB, - const std::pair &kC, + const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC, bool required = true) { AlgoKey key{{(int)transA, (int)transB, (int)epilogue}, - kA.first, - kA.second, - kB.first, - kB.second}; + shapeA.first, + shapeA.second, + shapeB.first, + shapeB.second}; auto it = algoCache.find(key); if (it != algoCache.end()) { if (algoCacheLimit) @@ -494,9 +509,9 @@ class BlasCuda { } auto &desc = getOrCreateDesc(transA, transB, epilogue); - auto lA = stampLayout(ROLE_A, kA); - auto lB = stampLayout(ROLE_B, kB); - auto lC = stampLayout(ROLE_C, kC); + auto lA = stampLayout(ROLE_A, shapeA); + auto lB = stampLayout(ROLE_B, shapeB); + auto lC = stampLayout(ROLE_C, shapeC); cublasLtMatmulHeuristicResult_t h{}; int returnedResults = 0; CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic( @@ -507,9 +522,9 @@ class BlasCuda { return nullptr; std::cerr << "[sofieBLAS] No suitable cuBLASLt algorithm found for " << "transA=" << transA << " transB=" << transB - << " epilogue=" << epilogue << " A=[" << kA.first << "x" - << kA.second << "]" - << " B=[" << kB.first << "x" << kB.second << "]\n"; + << " epilogue=" << epilogue << " A=[" << shapeA.first << "x" + << shapeA.second << "]" + << " B=[" << shapeB.first << "x" << shapeB.second << "]\n"; exit(EXIT_FAILURE); } auto ins = algoCache.emplace(key, CacheEntry{h, {}}).first; @@ -529,9 +544,9 @@ class BlasCuda { cublasLtEpilogue_t epilogue, float alpha, const float *A, const float *B, float beta, const float *D_in, float *C_out, const void *bias_ptr, - const std::pair &kA, - const std::pair &kB, - const std::pair &kC) { + const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC) { auto &desc = getOrCreateDesc(transA, transB, epilogue); if (bias_ptr) { CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( @@ -539,22 +554,32 @@ class BlasCuda { sizeof(bias_ptr))); } - const ShapeEnvelope *env = findEnvelope(kA, kB, kC); + // Resolve the algorithm at the call site's declared shape when one covers + // this call, so every runtime size the site produces shares one cache + // entry; with no covering declaration, resolve at the exact shape. + const ShapeEnvelope *env = findEnvelope(shapeA, shapeB, shapeC); const std::pair - aA = env ? std::make_pair(env->rowsA, env->colsA) : kA, - aB = env ? std::make_pair(env->rowsB, env->colsB) : kB, - aC = env ? std::make_pair(env->rowsC, env->colsC) : kC; - cublasLtMatmulHeuristicResult_t h = - *getOrComputeAlgo(transA, transB, epilogue, aA, aB, aC); - - if (env && !algoUsable(desc, h.algo, kA, kB, kC)) { + algoShapeA = env ? std::make_pair(env->rowsA, env->colsA) : shapeA, + algoShapeB = env ? std::make_pair(env->rowsB, env->colsB) : shapeB, + algoShapeC = env ? std::make_pair(env->rowsC, env->colsC) : shapeC; + cublasLtMatmulHeuristicResult_t h = *getOrComputeAlgo( + transA, transB, epilogue, algoShapeA, algoShapeB, algoShapeC); + + // Normally the resolution above is the only one. Only when cuBLASLt + // rejects the declared shape's algorithm at this call's exact size + // (returns NOT_SUPPORTED; m=1 was found to do this in testing) is the + // algorithm resolved a second time, at the exact shape, and that result + // is cached as well. + if (env && !algoUsable(desc, h.algo, shapeA, shapeB, shapeC)) { ++stats.envelopeRejects; - h = *getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); + h = *getOrComputeAlgo(transA, transB, epilogue, shapeA, shapeB, shapeC); } - auto lA = stampLayout(ROLE_A, kA); - auto lB = stampLayout(ROLE_B, kB); - auto lC = stampLayout(ROLE_C, kC); + // Stamp the exact call shape last: the algorithm resolution and the + // validity check above leave the shared descriptors at other dims. + auto lA = stampLayout(ROLE_A, shapeA); + auto lB = stampLayout(ROLE_B, shapeB); + auto lC = stampLayout(ROLE_C, shapeC); CHECK_CUBLAS(cublasLtMatmul(ltHandle, desc, &alpha, A, lA, B, lB, &beta, D_in, lC, C_out, lC, &h.algo, d_workspace, workspaceSize, stream)); diff --git a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp index 5ce5fa5..e9544ad 100644 --- a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp +++ b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp @@ -95,6 +95,11 @@ class BlasHip { size_t workspaceSize = 1u << 25; // 32 MB hipStream_t stream = nullptr; + // One persistent layout descriptor per GEMM operand: ROLE_A and ROLE_B are + // the input matrices of C = alpha * op(A) * op(B) + beta * C, ROLE_C the + // output (hipblasLtMatmul takes it twice, as C and D). stampLayout rewrites + // each descriptor in place to the shape of the call at hand, which is what + // lets one instance serve any runtime size. enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 }; hipblasLtMatrixLayout_t roleLayout[3] = {}; @@ -174,10 +179,11 @@ class BlasHip { void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, std::size_t, std::size_t, std::size_t, char transa, char transb) { - const auto kA = layoutKeyA(transa, m, k); - const auto kB = layoutKeyB(transb, k, n); - const std::pair kC{m, n}; - envelopes.push_back({kA.first, kA.second, kB.first, kB.second, m, n}); + const auto shapeA = layoutKeyA(transa, m, k); + const auto shapeB = layoutKeyB(transb, k, n); + const std::pair shapeC{m, n}; + envelopes.push_back( + {shapeA.first, shapeA.second, shapeB.first, shapeB.second, m, n}); const hipblasOperation_t tA = charToHipBlasTranspose(transa); const hipblasOperation_t tB = charToHipBlasTranspose(transb); @@ -185,7 +191,7 @@ class BlasHip { HIPBLASLT_EPILOGUE_BIAS, HIPBLASLT_EPILOGUE_RELU_BIAS}; for (hipblasLtEpilogue_t ep : eps) { - getOrComputeAlgo(tA, tB, ep, kA, kB, kC, /*required=*/false); + getOrComputeAlgo(tA, tB, ep, shapeA, shapeB, shapeC, /*required=*/false); } } @@ -410,20 +416,29 @@ class BlasHip { return L; } + // Returns the shape declared through addLayoutConfig that this call belongs + // to, or null if none covers it. Shapes are the physical (rows, cols) of + // matrices A, B and C, after any transpose is applied (transa='T' makes + // shapeA = (k, m)). The contraction dimension (colsA / rowsB) must match + // exactly: it comes from the weight tensor and never varies at runtime, so + // it identifies the call site and stops one site's declared shape from + // serving another's calls. The free dimensions only need covering; among + // candidates the least excess wins. const ShapeEnvelope * - findEnvelope(const std::pair &kA, - const std::pair &kB, - const std::pair &kC) const { + findEnvelope(const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC) const { const ShapeEnvelope *best = nullptr; std::size_t bestExcess = std::numeric_limits::max(); for (const auto &e : envelopes) { - if (e.colsA != kA.second || e.rowsB != kB.first) + if (e.colsA != shapeA.second || e.rowsB != shapeB.first) continue; - if (e.rowsA < kA.first || e.colsB < kB.second || e.rowsC < kC.first || - e.colsC < kC.second) + if (e.rowsA < shapeA.first || e.colsB < shapeB.second || + e.rowsC < shapeC.first || e.colsC < shapeC.second) continue; - const std::size_t ex = (e.rowsA - kA.first) + (e.colsA - kA.second) + - (e.rowsB - kB.first) + (e.colsB - kB.second); + const std::size_t ex = + (e.rowsA - shapeA.first) + (e.colsA - shapeA.second) + + (e.rowsB - shapeB.first) + (e.colsB - shapeB.second); if (ex < bestExcess) { bestExcess = ex; best = &e; @@ -433,12 +448,12 @@ class BlasHip { } bool algoUsable(hipblasLtMatmulDesc_t desc, const hipblasLtMatmulAlgo_t &algo, - const std::pair &kA, - const std::pair &kB, - const std::pair &kC) { - auto lA = stampLayout(ROLE_A, kA); - auto lB = stampLayout(ROLE_B, kB); - auto lC = stampLayout(ROLE_C, kC); + const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC) { + auto lA = stampLayout(ROLE_A, shapeA); + auto lB = stampLayout(ROLE_B, shapeB); + auto lC = stampLayout(ROLE_C, shapeC); // hipBLASLt has no hipblasLtMatmulAlgoCheck; the ext API rewrites the algo // it is handed, so probe a copy. hipblasLtMatmulAlgo_t probe = algo; @@ -480,15 +495,15 @@ class BlasHip { hipblasLtMatmulHeuristicResult_t * getOrComputeAlgo(hipblasOperation_t transA, hipblasOperation_t transB, hipblasLtEpilogue_t epilogue, - const std::pair &kA, - const std::pair &kB, - const std::pair &kC, + const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC, bool required = true) { AlgoKey key{{(int)transA, (int)transB, (int)epilogue}, - kA.first, - kA.second, - kB.first, - kB.second}; + shapeA.first, + shapeA.second, + shapeB.first, + shapeB.second}; auto it = algoCache.find(key); if (it != algoCache.end()) { if (algoCacheLimit) @@ -497,9 +512,9 @@ class BlasHip { } auto &desc = getOrCreateDesc(transA, transB, epilogue); - auto lA = stampLayout(ROLE_A, kA); - auto lB = stampLayout(ROLE_B, kB); - auto lC = stampLayout(ROLE_C, kC); + auto lA = stampLayout(ROLE_A, shapeA); + auto lB = stampLayout(ROLE_B, shapeB); + auto lC = stampLayout(ROLE_C, shapeC); hipblasLtMatmulHeuristicResult_t h{}; int returnedResults = 0; CHECK_HIPBLAS(hipblasLtMatmulAlgoGetHeuristic( @@ -510,9 +525,9 @@ class BlasHip { return nullptr; std::cerr << "[sofieBLAS] No suitable hipBLASLt algorithm found for " << "transA=" << transA << " transB=" << transB - << " epilogue=" << epilogue << " A=[" << kA.first << "x" - << kA.second << "]" - << " B=[" << kB.first << "x" << kB.second << "]\n"; + << " epilogue=" << epilogue << " A=[" << shapeA.first << "x" + << shapeA.second << "]" + << " B=[" << shapeB.first << "x" << shapeB.second << "]\n"; exit(EXIT_FAILURE); } auto ins = algoCache.emplace(key, CacheEntry{h, {}}).first; @@ -532,9 +547,9 @@ class BlasHip { hipblasLtEpilogue_t epilogue, float alpha, const float *A, const float *B, float beta, const float *D_in, float *C_out, const void *bias_ptr, - const std::pair &kA, - const std::pair &kB, - const std::pair &kC) { + const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC) { auto &desc = getOrCreateDesc(transA, transB, epilogue); if (bias_ptr) { CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( @@ -542,22 +557,31 @@ class BlasHip { sizeof(bias_ptr))); } - const ShapeEnvelope *env = findEnvelope(kA, kB, kC); + // Resolve the algorithm at the call site's declared shape when one covers + // this call, so every runtime size the site produces shares one cache + // entry; with no covering declaration, resolve at the exact shape. + const ShapeEnvelope *env = findEnvelope(shapeA, shapeB, shapeC); const std::pair - aA = env ? std::make_pair(env->rowsA, env->colsA) : kA, - aB = env ? std::make_pair(env->rowsB, env->colsB) : kB, - aC = env ? std::make_pair(env->rowsC, env->colsC) : kC; - hipblasLtMatmulHeuristicResult_t h = - *getOrComputeAlgo(transA, transB, epilogue, aA, aB, aC); - - if (env && !algoUsable(desc, h.algo, kA, kB, kC)) { + algoShapeA = env ? std::make_pair(env->rowsA, env->colsA) : shapeA, + algoShapeB = env ? std::make_pair(env->rowsB, env->colsB) : shapeB, + algoShapeC = env ? std::make_pair(env->rowsC, env->colsC) : shapeC; + hipblasLtMatmulHeuristicResult_t h = *getOrComputeAlgo( + transA, transB, epilogue, algoShapeA, algoShapeB, algoShapeC); + + // Normally the resolution above is the only one. Only when hipBLASLt + // rejects the declared shape's algorithm at this call's exact size is the + // algorithm resolved a second time, at the exact shape, and that result + // is cached as well. + if (env && !algoUsable(desc, h.algo, shapeA, shapeB, shapeC)) { ++stats.envelopeRejects; - h = *getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); + h = *getOrComputeAlgo(transA, transB, epilogue, shapeA, shapeB, shapeC); } - auto lA = stampLayout(ROLE_A, kA); - auto lB = stampLayout(ROLE_B, kB); - auto lC = stampLayout(ROLE_C, kC); + // Stamp the exact call shape last: the algorithm resolution and the + // validity check above leave the shared descriptors at other dims. + auto lA = stampLayout(ROLE_A, shapeA); + auto lB = stampLayout(ROLE_B, shapeB); + auto lC = stampLayout(ROLE_C, shapeC); CHECK_HIPBLAS(hipblasLtMatmul(ltHandle, desc, &alpha, A, lA, B, lB, &beta, D_in, lC, C_out, lC, &h.algo, d_workspace, workspaceSize, stream)); From 98165905b4a858bd0880a0823f256bf022f4df9e Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Thu, 20 Aug 2026 21:48:50 +0530 Subject: [PATCH 10/13] docs: explain addLayoutConfig, algo resolution and stats counters --- README.md | 2 +- .../backends/cuda/sofieBLAS_cublas.hpp | 23 ++++++++++++++++--- .../backends/hip/sofieBLAS_hipblaslt.hpp | 23 ++++++++++++++++--- 3 files changed, 41 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 2c96e69..c153824 100644 --- a/README.md +++ b/README.md @@ -104,7 +104,7 @@ A GEMM call computes `C = alpha * op(A) * op(B) + beta * C`: A and B are the inp Every matmul also needs an algorithm: the concrete GEMM kernel and its configuration, chosen by querying the library's heuristic (`cublasLtMatmulAlgoGetHeuristic` / `hipblasLtMatmulAlgoGetHeuristic`) for the given transpose settings, epilogue, and matrix dimensions. That query runs on the host and is not free, so its results are cached, keyed by the same parameters. Left at that, a model with runtime-varying sizes would pay one heuristic query per distinct size it meets, and the cache would grow with every new size. -The envelope mechanism removes that cost. An envelope is the largest shape a GEMM call site will ever run, declared through `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)`, which a generated Session constructor calls once per call site with its own construction-time dimensions. `addLayoutConfig` stores the envelope and resolves its algorithm immediately. At matmul time the call is matched to a declared envelope (the contraction dimension must match exactly, which is what ties a call to its own call site), and the algorithm is looked up at the envelope's dimensions rather than the call's, so every runtime size a call site produces shares that one cache entry: no further heuristic queries, no cache growth. A call covered by no envelope is resolved and cached at its exact shape. The library may also reject an envelope's algorithm at a specific smaller size (m=1 was found to do this in testing); the backend checks for this and falls back to exact-shape resolution for that call. +The envelope mechanism removes that cost. An envelope is the largest shape a GEMM call site will ever run, declared through `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)`, which a generated Session constructor calls once per call site with its own construction-time dimensions. `addLayoutConfig` stores the envelope and resolves its algorithm immediately. At matmul time the call is matched to a declared envelope (the contraction dimension must match exactly, which is what ties a call to its own call site), and the algorithm is looked up at the envelope's dimensions rather than the call's, so every runtime size a call site produces shares that one cache entry: no further heuristic queries, no cache growth. A call covered by no envelope is resolved and cached at its exact shape. The library may also reject an envelope's algorithm at a specific smaller size (m=1 was found to do this with cuBLASLt in testing); the backend checks for this and falls back to exact-shape resolution for that call. The cache is unbounded by default. Pass a limit at construction to cap it with LRU eviction: diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index d4b6db8..62edfb0 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -76,14 +76,15 @@ struct AlgoKeyHash { } }; +// A call site's declared maximum shape, recorded by addLayoutConfig. struct ShapeEnvelope { std::size_t rowsA, colsA, rowsB, colsB, rowsC, colsC; }; struct LayoutStats { - std::size_t heuristicQueries = 0; - std::size_t envelopeRejects = 0; - std::size_t evictions = 0; + std::size_t heuristicQueries = 0; // algorithm searches issued + std::size_t envelopeRejects = 0; // declared-shape algo unusable at a call + std::size_t evictions = 0; // entries dropped to stay under the limit }; class BlasCuda { @@ -107,6 +108,7 @@ class BlasCuda { // algo cache entry struct CacheEntry { cublasLtMatmulHeuristicResult_t h{}; + // position in lruOrder; only valid when a limit is set std::list::iterator lru{}; }; std::unordered_map algoCache; @@ -177,6 +179,12 @@ class BlasCuda { } } + // Declares a call site's largest shape (its envelope) and resolves the + // algorithm for it up front; the generated Session constructor calls this + // once per GEMM call site with its construction-time dimensions. Which + // epilogue the site will use is unknown here, so all three used by the + // generated code are resolved; unused ones cost one heuristic query each, + // off the inference path. void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, std::size_t, std::size_t, std::size_t, char transa, char transb) { const auto shapeA = layoutKeyA(transa, m, k); @@ -399,6 +407,8 @@ class BlasCuda { : std::make_pair(n, k); } + // Sets a role's layout descriptor to the given physical (rows, cols), + // creating it on first use. Matrices are dense column-major, so ld = rows. cublasLtMatrixLayout_t stampLayout(LayoutRole role, const std::pair &key) { const uint64_t rows = key.first, cols = key.second; @@ -476,6 +486,9 @@ class BlasCuda { return descStore.at(key); } + // Whether the given algorithm can run this exact shape within the + // workspace. An algorithm resolved at a declared shape is not guaranteed to + // run at every smaller size it covers. bool algoUsable(cublasLtMatmulDesc_t desc, const cublasLtMatmulAlgo_t &algo, const std::pair &shapeA, const std::pair &shapeB, @@ -489,6 +502,10 @@ class BlasCuda { chk.workspaceSize <= workspaceSize; } + // Looks up, or resolves and caches, the algorithm for the given transpose + // settings, epilogue and shapes. required=false is for constructor warmup: + // a speculatively resolved epilogue may legitimately have no algorithm, and + // returns null instead of aborting. cublasLtMatmulHeuristicResult_t * getOrComputeAlgo(cublasOperation_t transA, cublasOperation_t transB, cublasLtEpilogue_t epilogue, diff --git a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp index e9544ad..7d58c74 100644 --- a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp +++ b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp @@ -77,14 +77,15 @@ struct AlgoKeyHash { } }; +// A call site's declared maximum shape, recorded by addLayoutConfig. struct ShapeEnvelope { std::size_t rowsA, colsA, rowsB, colsB, rowsC, colsC; }; struct LayoutStats { - std::size_t heuristicQueries = 0; - std::size_t envelopeRejects = 0; - std::size_t evictions = 0; + std::size_t heuristicQueries = 0; // algorithm searches issued + std::size_t envelopeRejects = 0; // declared-shape algo unusable at a call + std::size_t evictions = 0; // entries dropped to stay under the limit }; class BlasHip { @@ -108,6 +109,7 @@ class BlasHip { // algo cache entry struct CacheEntry { hipblasLtMatmulHeuristicResult_t h{}; + // position in lruOrder; only valid when a limit is set std::list::iterator lru{}; }; std::unordered_map algoCache; @@ -177,6 +179,12 @@ class BlasHip { } } + // Declares a call site's largest shape (its envelope) and resolves the + // algorithm for it up front; the generated Session constructor calls this + // once per GEMM call site with its construction-time dimensions. Which + // epilogue the site will use is unknown here, so all three used by the + // generated code are resolved; unused ones cost one heuristic query each, + // off the inference path. void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, std::size_t, std::size_t, std::size_t, char transa, char transb) { const auto shapeA = layoutKeyA(transa, m, k); @@ -398,6 +406,8 @@ class BlasHip { : std::make_pair(n, k); } + // Sets a role's layout descriptor to the given physical (rows, cols), + // creating it on first use. Matrices are dense column-major, so ld = rows. hipblasLtMatrixLayout_t stampLayout(LayoutRole role, const std::pair &key) { const uint64_t rows = key.first, cols = key.second; @@ -447,6 +457,9 @@ class BlasHip { return best; } + // Whether the given algorithm can run this exact shape within the + // workspace. An algorithm resolved at a declared shape is not guaranteed to + // run at every smaller size it covers. bool algoUsable(hipblasLtMatmulDesc_t desc, const hipblasLtMatmulAlgo_t &algo, const std::pair &shapeA, const std::pair &shapeB, @@ -492,6 +505,10 @@ class BlasHip { return descStore.at(key); } + // Looks up, or resolves and caches, the algorithm for the given transpose + // settings, epilogue and shapes. required=false is for constructor warmup: + // a speculatively resolved epilogue may legitimately have no algorithm, and + // returns null instead of aborting. hipblasLtMatmulHeuristicResult_t * getOrComputeAlgo(hipblasOperation_t transA, hipblasOperation_t transB, hipblasLtEpilogue_t epilogue, From 4bcb1e2eab2928a3b248b6a860489c46ec5871b6 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Wed, 26 Aug 2026 05:38:45 +0530 Subject: [PATCH 11/13] fix: cache layouts and algorithms, one shared BLASLt backend --- README.md | 22 +- benchmark/bench.cc | 4 +- .../backends/cuda/sofieBLAS_cublas.hpp | 625 ++--------------- .../gpu/detail/sofieBLAS_blaslt_common.tpp | 429 ++++++++++++ .../backends/hip/sofieBLAS_hipblaslt.hpp | 628 ++---------------- tests/test.cc | 242 +------ 6 files changed, 559 insertions(+), 1391 deletions(-) create mode 100644 include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp diff --git a/README.md b/README.md index c153824..191fe2a 100644 --- a/README.md +++ b/README.md @@ -96,27 +96,7 @@ sofieBLAS blas(queue); blas.matmul('N', 'N', size, size, size, 1.0f, dA, dB, 0.0f, dC); ``` -The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addLayoutConfig` (declares a call site's shape ahead of the first call, which feeds the algorithm cache described below). - -## Dynamic GEMM shapes and the algorithm cache - -A GEMM call computes `C = alpha * op(A) * op(B) + beta * C`: A and B are the input matrices, C the output, and `op` an optional transpose. The CUDA backend (`BlasCuda`, over cuBLASLt) and the HIP backend (`BlasHip`, over hipBLASLt) implement the following identically. One instance serves GEMM calls at sizes that vary at runtime: for every call, each of the three matrices has its dimensions stamped into a persistent per-matrix layout descriptor right before the multiply, so no layout is ever tied to a fixed shape. - -Every matmul also needs an algorithm: the concrete GEMM kernel and its configuration, chosen by querying the library's heuristic (`cublasLtMatmulAlgoGetHeuristic` / `hipblasLtMatmulAlgoGetHeuristic`) for the given transpose settings, epilogue, and matrix dimensions. That query runs on the host and is not free, so its results are cached, keyed by the same parameters. Left at that, a model with runtime-varying sizes would pay one heuristic query per distinct size it meets, and the cache would grow with every new size. - -The envelope mechanism removes that cost. An envelope is the largest shape a GEMM call site will ever run, declared through `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)`, which a generated Session constructor calls once per call site with its own construction-time dimensions. `addLayoutConfig` stores the envelope and resolves its algorithm immediately. At matmul time the call is matched to a declared envelope (the contraction dimension must match exactly, which is what ties a call to its own call site), and the algorithm is looked up at the envelope's dimensions rather than the call's, so every runtime size a call site produces shares that one cache entry: no further heuristic queries, no cache growth. A call covered by no envelope is resolved and cached at its exact shape. The library may also reject an envelope's algorithm at a specific smaller size (m=1 was found to do this with cuBLASLt in testing); the backend checks for this and falls back to exact-shape resolution for that call. - -The cache is unbounded by default. Pass a limit at construction to cap it with LRU eviction: - -```cpp -sofieBLAS blas(queue); // unbounded cache (default) -sofieBLAS capped(queue, 32); // at most 32 entries, LRU eviction -sofieBLAS hipCapped(queue, 32); -``` - -`algoCacheSize()` returns the current entry count and `layoutStats()` returns counters (`heuristicQueries`, `envelopeRejects`, `evictions`) for inspecting cache behaviour. - -The exact-shape fallback uses `cublasLtMatmulAlgoCheck` on CUDA. hipBLASLt has no equivalent in its core API, so the HIP backend uses `hipblaslt_ext::matmulIsAlgoSupported` from `hipblaslt-ext.hpp` for the same check. +The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addLayoutConfig` (used to pre-register cuBLASLt/hipBLASLt matrix layouts for a given shape before the first `matmul`/`gemm` call on that shape). ## Contributing diff --git a/benchmark/bench.cc b/benchmark/bench.cc index 3263dab..8935856 100644 --- a/benchmark/bench.cc +++ b/benchmark/bench.cc @@ -139,7 +139,7 @@ static void runCudaBench(const BenchOptions &opt) { alpaka::memcpy(queue, dB, hB); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N'); + blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N', 'n'); for (int i = 0; i < opt.warmup; ++i) blas.matmul('N', 'N', M, N, K, 1.f, dA, dB, 0.f, dC); @@ -187,7 +187,7 @@ static void runHipBench(const BenchOptions &opt) { alpaka::memcpy(queue, dB, hB); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N'); + blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N', 'n'); for (int i = 0; i < opt.warmup; ++i) blas.matmul('N', 'N', M, N, K, 1.f, dA, dB, 0.f, dC); diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index 62edfb0..4717730 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -5,12 +5,11 @@ #include #include #include -#include #include #include +#include #include #include -#include #include "sofieBLAS/core.hpp" #include @@ -33,575 +32,67 @@ } \ } while (0) -struct DescKey { - int transA; // CUBLAS_OP_N / CUBLAS_OP_T encoded as int - int transB; - int epilogue; // cublasLtEpilogue_t encoded as int - bool operator==(const DescKey &o) const noexcept { - return transA == o.transA && transB == o.transB && epilogue == o.epilogue; - } -}; - -struct DescKeyHash { - std::size_t operator()(const DescKey &k) const noexcept { - std::size_t h = static_cast(k.transA) * 97u + - static_cast(k.transB) * 31u + - static_cast(k.epilogue); - return h ^ (h >> 16); - } +// The cuBLASLt spellings of everything the shared BlasLt implementation in +// backends/gpu/detail uses: types, enum values and functions. +struct CublasLtApi { + using Queue = alpaka::QueueCudaRtNonBlocking; + using Handle = cublasLtHandle_t; + using BlasHandle = cublasHandle_t; + using Preference = cublasLtMatmulPreference_t; + using Stream = cudaStream_t; + using Layout = cublasLtMatrixLayout_t; + using MatmulDesc = cublasLtMatmulDesc_t; + using HeuristicResult = cublasLtMatmulHeuristicResult_t; + using Operation = cublasOperation_t; + using Epilogue = cublasLtEpilogue_t; + + static constexpr auto OpN = CUBLAS_OP_N; + static constexpr auto OpT = CUBLAS_OP_T; + static constexpr auto OpC = CUBLAS_OP_C; + static constexpr auto EpilogueDefault = CUBLASLT_EPILOGUE_DEFAULT; + static constexpr auto EpilogueBias = CUBLASLT_EPILOGUE_BIAS; + static constexpr auto EpilogueReluBias = CUBLASLT_EPILOGUE_RELU_BIAS; + static constexpr auto EpilogueGeluBias = CUBLASLT_EPILOGUE_GELU_BIAS; + static constexpr auto ComputeF32 = CUBLAS_COMPUTE_32F; + static constexpr auto RealF32 = CUDA_R_32F; + static constexpr auto DescTransA = CUBLASLT_MATMUL_DESC_TRANSA; + static constexpr auto DescTransB = CUBLASLT_MATMUL_DESC_TRANSB; + static constexpr auto DescEpilogue = CUBLASLT_MATMUL_DESC_EPILOGUE; + static constexpr auto DescBiasPointer = CUBLASLT_MATMUL_DESC_BIAS_POINTER; + static constexpr auto PrefMaxWorkspace = + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES; + static constexpr const char *name = "cuBLASLt"; + + static constexpr auto ltCreate = cublasLtCreate; + static constexpr auto ltDestroy = cublasLtDestroy; + static constexpr auto blasCreate = cublasCreate; + static constexpr auto blasDestroy = cublasDestroy; + static constexpr auto blasSetStream = cublasSetStream; + static constexpr auto prefCreate = cublasLtMatmulPreferenceCreate; + static constexpr auto prefDestroy = cublasLtMatmulPreferenceDestroy; + static constexpr auto prefSetAttribute = cublasLtMatmulPreferenceSetAttribute; + static constexpr auto layoutCreate = cublasLtMatrixLayoutCreate; + static constexpr auto layoutDestroy = cublasLtMatrixLayoutDestroy; + static constexpr auto descCreate = cublasLtMatmulDescCreate; + static constexpr auto descDestroy = cublasLtMatmulDescDestroy; + static constexpr auto descSetAttribute = cublasLtMatmulDescSetAttribute; + static constexpr auto getHeuristic = cublasLtMatmulAlgoGetHeuristic; + static constexpr auto matmul = cublasLtMatmul; + static constexpr auto sgemmStridedBatched = cublasSgemmStridedBatched; + + // cudaMalloc has a templated C++ overload, so a pointer to it is ambiguous + static cudaError_t rtMalloc(void **ptr, std::size_t size) { + return cudaMalloc(ptr, size); + } + static cudaError_t rtFree(void *ptr) { return cudaFree(ptr); } }; -struct AlgoKey { - DescKey dk; - std::size_t rowsA, colsA; - std::size_t rowsB, colsB; - bool operator==(const AlgoKey &o) const noexcept { - return dk == o.dk && rowsA == o.rowsA && colsA == o.colsA && - rowsB == o.rowsB && colsB == o.colsB; - } -}; - -struct AlgoKeyHash { - std::size_t operator()(const AlgoKey &k) const noexcept { - std::size_t h = DescKeyHash{}(k.dk); - auto mix = [&](std::size_t v) { - h ^= std::hash{}(v) + 0x9e3779b97f4a7c15ULL + (h << 6) + - (h >> 2); - }; - mix(k.rowsA); - mix(k.colsA); - mix(k.rowsB); - mix(k.colsB); - return h; - } -}; - -// A call site's declared maximum shape, recorded by addLayoutConfig. -struct ShapeEnvelope { - std::size_t rowsA, colsA, rowsB, colsB, rowsC, colsC; -}; - -struct LayoutStats { - std::size_t heuristicQueries = 0; // algorithm searches issued - std::size_t envelopeRejects = 0; // declared-shape algo unusable at a call - std::size_t evictions = 0; // entries dropped to stay under the limit -}; - -class BlasCuda { - cublasLtHandle_t ltHandle = nullptr; - cublasHandle_t handle = nullptr; - cublasLtMatmulPreference_t preference = nullptr; - void *d_workspace = nullptr; - size_t workspaceSize = 1u << 25; // 32 MB - cudaStream_t stream = nullptr; - - // One persistent layout descriptor per GEMM operand: ROLE_A and ROLE_B are - // the input matrices of C = alpha * op(A) * op(B) + beta * C, ROLE_C the - // output (cublasLtMatmul takes it twice, as C and D). stampLayout rewrites - // each descriptor in place to the shape of the call at hand, which is what - // lets one instance serve any runtime size. - enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 }; - cublasLtMatrixLayout_t roleLayout[3] = {}; - - std::unordered_map descStore; - - // algo cache entry - struct CacheEntry { - cublasLtMatmulHeuristicResult_t h{}; - // position in lruOrder; only valid when a limit is set - std::list::iterator lru{}; - }; - std::unordered_map algoCache; - std::list lruOrder; - // 0 = unbounded - std::size_t algoCacheLimit = 0; - - std::vector envelopes; - - LayoutStats stats; - -public: - const LayoutStats &layoutStats() const { return stats; } - std::size_t algoCacheSize() const { return algoCache.size(); } +#define SOFIEBLAS_CHECK_LT(x) CHECK_CUBLAS(x) +#define SOFIEBLAS_CHECK_RT(x) CHECK_CUDA(x) - BlasCuda(const BlasCuda &) = delete; - BlasCuda &operator=(const BlasCuda &) = delete; - BlasCuda(BlasCuda &&) = delete; - BlasCuda &operator=(BlasCuda &&) = delete; - - BlasCuda(alpaka::QueueCudaRtNonBlocking &queue, - std::size_t algoCacheLimit_ = 0) - : algoCacheLimit{algoCacheLimit_}, m_queue{queue} { - stream = static_cast(m_queue.getNativeHandle()); - - CHECK_CUBLAS(cublasLtCreate(<Handle)); - - CHECK_CUBLAS(cublasCreate(&handle)); - CHECK_CUBLAS(cublasSetStream(handle, stream)); - - CHECK_CUBLAS(cublasLtMatmulPreferenceCreate(&preference)); - CHECK_CUDA(cudaMalloc(&d_workspace, workspaceSize)); - CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( - preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, - sizeof(workspaceSize))); - } - - ~BlasCuda() { - for (auto L : roleLayout) - if (L) - cublasLtMatrixLayoutDestroy(L); - for (auto &[key, desc] : descStore) - if (desc) - cublasLtMatmulDescDestroy(desc); - if (preference) - cublasLtMatmulPreferenceDestroy(preference); - if (ltHandle) - cublasLtDestroy(ltHandle); - if (handle) - cublasDestroy(handle); - if (d_workspace) - cudaFree(d_workspace); - } - - inline cublasOperation_t charToCuBlasTranspose(char trans) { - switch (trans) { - case 'N': - case 'n': - return CUBLAS_OP_N; - case 'T': - case 't': - return CUBLAS_OP_T; - case 'C': - case 'c': - return CUBLAS_OP_C; - default: - throw std::invalid_argument("Invalid transpose character for cuBLAS."); - } - } - - // Declares a call site's largest shape (its envelope) and resolves the - // algorithm for it up front; the generated Session constructor calls this - // once per GEMM call site with its construction-time dimensions. Which - // epilogue the site will use is unknown here, so all three used by the - // generated code are resolved; unused ones cost one heuristic query each, - // off the inference path. - void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, std::size_t, - std::size_t, std::size_t, char transa, char transb) { - const auto shapeA = layoutKeyA(transa, m, k); - const auto shapeB = layoutKeyB(transb, k, n); - const std::pair shapeC{m, n}; - envelopes.push_back( - {shapeA.first, shapeA.second, shapeB.first, shapeB.second, m, n}); - - const cublasOperation_t tA = charToCuBlasTranspose(transa); - const cublasOperation_t tB = charToCuBlasTranspose(transb); - const cublasLtEpilogue_t eps[] = {CUBLASLT_EPILOGUE_DEFAULT, - CUBLASLT_EPILOGUE_BIAS, - CUBLASLT_EPILOGUE_RELU_BIAS}; - for (cublasLtEpilogue_t ep : eps) { - getOrComputeAlgo(tA, tB, ep, shapeA, shapeB, shapeC, /*required=*/false); - } - } - - template - inline void - gemm(char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, alpaka::BufCudaRt, TIdx> const &A, - alpaka::BufCudaRt, TIdx> const &B, float beta, - alpaka::BufCudaRt, TIdx> &bias, - alpaka::BufCudaRt, TIdx> &C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_BIAS, alpha, alpaka::getPtrNative(A), - alpaka::getPtrNative(B), beta, alpaka::getPtrNative(bias), - alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemm( - char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, - alpaka::ViewPlainPtr, TIdx> const - &A, - alpaka::ViewPlainPtr, TIdx> const - &B, - float beta, - alpaka::ViewPlainPtr, TIdx> - &bias, - alpaka::ViewPlainPtr, TIdx> &C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_BIAS, alpha, alpaka::getPtrNative(A), - alpaka::getPtrNative(B), beta, alpaka::getPtrNative(bias), - alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemm(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, T const *A, T const *B, - float beta, T *bias, T *C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_BIAS, alpha, A, B, beta, bias, C, - static_cast(bias), layoutKeyA(transa, m, k), - layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmrelu(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, - alpaka::BufCudaRt, TIdx> const &A, - alpaka::BufCudaRt, TIdx> const &B, - float beta, - alpaka::BufCudaRt, TIdx> &bias, - alpaka::BufCudaRt, TIdx> &C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_RELU_BIAS, alpha, alpaka::getPtrNative(A), - alpaka::getPtrNative(B), beta, alpaka::getPtrNative(bias), - alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } +#include "../gpu/detail/sofieBLAS_blaslt_common.tpp" - template - inline void gemmrelu( - char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, - alpaka::ViewPlainPtr, TIdx> const - &A, - alpaka::ViewPlainPtr, TIdx> const - &B, - float beta, - alpaka::ViewPlainPtr, TIdx> - &bias, - alpaka::ViewPlainPtr, TIdx> &C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_RELU_BIAS, alpha, alpaka::getPtrNative(A), - alpaka::getPtrNative(B), beta, alpaka::getPtrNative(bias), - alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmrelu(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, T const *A, T const *B, - float beta, T *bias, T *C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_RELU_BIAS, alpha, A, B, beta, bias, C, - static_cast(bias), layoutKeyA(transa, m, k), - layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmgelu(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, - alpaka::BufCudaRt, TIdx> const &A, - alpaka::BufCudaRt, TIdx> const &B, - float beta, - alpaka::BufCudaRt, TIdx> &bias, - alpaka::BufCudaRt, TIdx> &C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_GELU_BIAS, alpha, alpaka::getPtrNative(A), - alpaka::getPtrNative(B), beta, alpaka::getPtrNative(bias), - alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmgelu( - char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, - alpaka::ViewPlainPtr, TIdx> const - &A, - alpaka::ViewPlainPtr, TIdx> const - &B, - float beta, - alpaka::ViewPlainPtr, TIdx> - &bias, - alpaka::ViewPlainPtr, TIdx> &C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_GELU_BIAS, alpha, alpaka::getPtrNative(A), - alpaka::getPtrNative(B), beta, alpaka::getPtrNative(bias), - alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmgelu(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, T const *A, T const *B, - float beta, T *bias, T *C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_GELU_BIAS, alpha, A, B, beta, bias, C, - static_cast(bias), layoutKeyA(transa, m, k), - layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void matmul(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, - alpaka::BufCudaRt, TIdx> const &A, - alpaka::BufCudaRt, TIdx> const &B, - float beta, - alpaka::BufCudaRt, TIdx> &C) { - float *c = alpaka::getPtrNative(C); - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_DEFAULT, alpha, alpaka::getPtrNative(A), - alpaka::getPtrNative(B), beta, c, c, nullptr, - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void matmul( - char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, - alpaka::ViewPlainPtr, TIdx> const - &A, - alpaka::ViewPlainPtr, TIdx> const - &B, - float beta, - alpaka::ViewPlainPtr, TIdx> &C) { - T *c = alpaka::getPtrNative(C); - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_DEFAULT, alpha, alpaka::getPtrNative(A), - alpaka::getPtrNative(B), beta, c, c, nullptr, - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - // Raw-pointer overload - template - inline void matmul(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, T const *A, T const *B, - float beta, T *C) { - executeMatmul(charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), - CUBLASLT_EPILOGUE_DEFAULT, alpha, A, B, beta, C, C, nullptr, - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - inline void gemmStridedBatched(char transa, char transb, int m, int n, int k, - float alpha, const float *A, int lda, - long long strideA, const float *B, int ldb, - long long strideB, float beta, float *C, - int ldc, long long strideC, int batchCount) { - CHECK_CUBLAS(cublasSgemmStridedBatched( - handle, charToCuBlasTranspose(transa), charToCuBlasTranspose(transb), m, - n, k, &alpha, A, lda, strideA, B, ldb, strideB, &beta, C, ldc, strideC, - batchCount)); - } - -private: - alpaka::QueueCudaRtNonBlocking m_queue; - - static std::pair - layoutKeyA(char trans, std::size_t m, std::size_t k) { - return (trans == 'N' || trans == 'n') ? std::make_pair(m, k) - : std::make_pair(k, m); - } - - static std::pair - layoutKeyB(char trans, std::size_t k, std::size_t n) { - return (trans == 'N' || trans == 'n') ? std::make_pair(k, n) - : std::make_pair(n, k); - } - - // Sets a role's layout descriptor to the given physical (rows, cols), - // creating it on first use. Matrices are dense column-major, so ld = rows. - cublasLtMatrixLayout_t - stampLayout(LayoutRole role, const std::pair &key) { - const uint64_t rows = key.first, cols = key.second; - const int64_t ld = static_cast(key.first); - cublasLtMatrixLayout_t &L = roleLayout[role]; - if (!L) { - CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&L, CUDA_R_32F, rows, cols, ld)); - } else { - CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( - L, CUBLASLT_MATRIX_LAYOUT_ROWS, &rows, sizeof(rows))); - CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( - L, CUBLASLT_MATRIX_LAYOUT_COLS, &cols, sizeof(cols))); - CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( - L, CUBLASLT_MATRIX_LAYOUT_LD, &ld, sizeof(ld))); - } - return L; - } - - // Returns the shape declared through addLayoutConfig that this call belongs - // to, or null if none covers it. Shapes are the physical (rows, cols) of - // matrices A, B and C, after any transpose is applied (transa='T' makes - // shapeA = (k, m)). The contraction dimension (colsA / rowsB) must match - // exactly: it comes from the weight tensor and never varies at runtime, so - // it identifies the call site and stops one site's declared shape from - // serving another's calls. The free dimensions only need covering; among - // candidates the least excess wins. - const ShapeEnvelope * - findEnvelope(const std::pair &shapeA, - const std::pair &shapeB, - const std::pair &shapeC) const { - const ShapeEnvelope *best = nullptr; - std::size_t bestExcess = std::numeric_limits::max(); - for (const auto &e : envelopes) { - if (e.colsA != shapeA.second || e.rowsB != shapeB.first) - continue; - if (e.rowsA < shapeA.first || e.colsB < shapeB.second || - e.rowsC < shapeC.first || e.colsC < shapeC.second) - continue; - const std::size_t ex = - (e.rowsA - shapeA.first) + (e.colsA - shapeA.second) + - (e.rowsB - shapeB.first) + (e.colsB - shapeB.second); - if (ex < bestExcess) { - bestExcess = ex; - best = &e; - } - } - return best; - } - - cublasLtMatmulDesc_t &getOrCreateDesc(cublasOperation_t transA, - cublasOperation_t transB, - cublasLtEpilogue_t epilogue) { - DescKey key{(int)transA, (int)transB, (int)epilogue}; - auto it = descStore.find(key); - if (it != descStore.end()) - return it->second; - - cublasLtMatmulDesc_t desc = nullptr; - CHECK_CUBLAS( - cublasLtMatmulDescCreate(&desc, CUBLAS_COMPUTE_32F, CUDA_R_32F)); - CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( - desc, CUBLASLT_MATMUL_DESC_TRANSA, &transA, sizeof(transA))); - CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( - desc, CUBLASLT_MATMUL_DESC_TRANSB, &transB, sizeof(transB))); - CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( - desc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue))); - // For bias epilogues: set a non-null dummy pointer so the descriptor is - // valid for cublasLtMatmulAlgoGetHeuristic. - if (epilogue != CUBLASLT_EPILOGUE_DEFAULT) { - const void *dummy = d_workspace; - CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( - desc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &dummy, sizeof(dummy))); - } - descStore.emplace(key, desc); - return descStore.at(key); - } - - // Whether the given algorithm can run this exact shape within the - // workspace. An algorithm resolved at a declared shape is not guaranteed to - // run at every smaller size it covers. - bool algoUsable(cublasLtMatmulDesc_t desc, const cublasLtMatmulAlgo_t &algo, - const std::pair &shapeA, - const std::pair &shapeB, - const std::pair &shapeC) { - auto lA = stampLayout(ROLE_A, shapeA); - auto lB = stampLayout(ROLE_B, shapeB); - auto lC = stampLayout(ROLE_C, shapeC); - cublasLtMatmulHeuristicResult_t chk{}; - return cublasLtMatmulAlgoCheck(ltHandle, desc, lA, lB, lC, lC, &algo, - &chk) == CUBLAS_STATUS_SUCCESS && - chk.workspaceSize <= workspaceSize; - } - - // Looks up, or resolves and caches, the algorithm for the given transpose - // settings, epilogue and shapes. required=false is for constructor warmup: - // a speculatively resolved epilogue may legitimately have no algorithm, and - // returns null instead of aborting. - cublasLtMatmulHeuristicResult_t * - getOrComputeAlgo(cublasOperation_t transA, cublasOperation_t transB, - cublasLtEpilogue_t epilogue, - const std::pair &shapeA, - const std::pair &shapeB, - const std::pair &shapeC, - bool required = true) { - AlgoKey key{{(int)transA, (int)transB, (int)epilogue}, - shapeA.first, - shapeA.second, - shapeB.first, - shapeB.second}; - auto it = algoCache.find(key); - if (it != algoCache.end()) { - if (algoCacheLimit) - lruOrder.splice(lruOrder.begin(), lruOrder, it->second.lru); - return &it->second.h; - } - - auto &desc = getOrCreateDesc(transA, transB, epilogue); - auto lA = stampLayout(ROLE_A, shapeA); - auto lB = stampLayout(ROLE_B, shapeB); - auto lC = stampLayout(ROLE_C, shapeC); - cublasLtMatmulHeuristicResult_t h{}; - int returnedResults = 0; - CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic( - ltHandle, desc, lA, lB, lC, lC, preference, 1, &h, &returnedResults)); - ++stats.heuristicQueries; - if (returnedResults == 0) { - if (!required) - return nullptr; - std::cerr << "[sofieBLAS] No suitable cuBLASLt algorithm found for " - << "transA=" << transA << " transB=" << transB - << " epilogue=" << epilogue << " A=[" << shapeA.first << "x" - << shapeA.second << "]" - << " B=[" << shapeB.first << "x" << shapeB.second << "]\n"; - exit(EXIT_FAILURE); - } - auto ins = algoCache.emplace(key, CacheEntry{h, {}}).first; - if (algoCacheLimit) { - lruOrder.push_front(key); - ins->second.lru = lruOrder.begin(); - while (algoCache.size() > algoCacheLimit) { - algoCache.erase(lruOrder.back()); - lruOrder.pop_back(); - ++stats.evictions; - } - } - return &ins->second.h; - } - - void executeMatmul(cublasOperation_t transA, cublasOperation_t transB, - cublasLtEpilogue_t epilogue, float alpha, const float *A, - const float *B, float beta, const float *D_in, - float *C_out, const void *bias_ptr, - const std::pair &shapeA, - const std::pair &shapeB, - const std::pair &shapeC) { - auto &desc = getOrCreateDesc(transA, transB, epilogue); - if (bias_ptr) { - CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( - desc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, - sizeof(bias_ptr))); - } - - // Resolve the algorithm at the call site's declared shape when one covers - // this call, so every runtime size the site produces shares one cache - // entry; with no covering declaration, resolve at the exact shape. - const ShapeEnvelope *env = findEnvelope(shapeA, shapeB, shapeC); - const std::pair - algoShapeA = env ? std::make_pair(env->rowsA, env->colsA) : shapeA, - algoShapeB = env ? std::make_pair(env->rowsB, env->colsB) : shapeB, - algoShapeC = env ? std::make_pair(env->rowsC, env->colsC) : shapeC; - cublasLtMatmulHeuristicResult_t h = *getOrComputeAlgo( - transA, transB, epilogue, algoShapeA, algoShapeB, algoShapeC); - - // Normally the resolution above is the only one. Only when cuBLASLt - // rejects the declared shape's algorithm at this call's exact size - // (returns NOT_SUPPORTED; m=1 was found to do this in testing) is the - // algorithm resolved a second time, at the exact shape, and that result - // is cached as well. - if (env && !algoUsable(desc, h.algo, shapeA, shapeB, shapeC)) { - ++stats.envelopeRejects; - h = *getOrComputeAlgo(transA, transB, epilogue, shapeA, shapeB, shapeC); - } - - // Stamp the exact call shape last: the algorithm resolution and the - // validity check above leave the shared descriptors at other dims. - auto lA = stampLayout(ROLE_A, shapeA); - auto lB = stampLayout(ROLE_B, shapeB); - auto lC = stampLayout(ROLE_C, shapeC); - CHECK_CUBLAS(cublasLtMatmul(ltHandle, desc, &alpha, A, lA, B, lB, &beta, - D_in, lC, C_out, lC, &h.algo, d_workspace, - workspaceSize, stream)); - } -}; +using BlasCuda = BlasLt; namespace traits { diff --git a/include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp b/include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp new file mode 100644 index 0000000..1e83b9c --- /dev/null +++ b/include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp @@ -0,0 +1,429 @@ +// Shared implementation of the cuBLASLt and hipBLASLt backends. The two +// vendor APIs have the same shape under different names, so the backend is +// written once against an Api table. A vendor header defines that table +// (the types, constants and functions of its library), defines the check +// macros SOFIEBLAS_CHECK_LT and SOFIEBLAS_CHECK_RT, includes the vendor and +// standard headers (, , , , +// , , , , alpaka), and then +// includes this file. + +struct PairHash { + std::size_t + operator()(const std::pair &p) const noexcept { + std::size_t h1 = std::hash{}(p.first); + std::size_t h2 = std::hash{}(p.second); + return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2)); + } +}; + +struct PairEq { + bool operator()(const std::pair &a, + const std::pair &b) const noexcept { + return a.first == b.first && a.second == b.second; + } +}; + +struct DescKey { + int transA; // backend transpose enum encoded as int + int transB; + int epilogue; // backend epilogue enum encoded as int + bool operator==(const DescKey &o) const noexcept { + return transA == o.transA && transB == o.transB && epilogue == o.epilogue; + } +}; + +struct DescKeyHash { + std::size_t operator()(const DescKey &k) const noexcept { + std::size_t h = static_cast(k.transA) * 97u + + static_cast(k.transB) * 31u + + static_cast(k.epilogue); + return h ^ (h >> 16); + } +}; + +struct AlgoKey { + DescKey dk; + std::size_t rowsA, colsA; // physical dimensions of A in layoutStore + std::size_t rowsB, colsB; // physical dimensions of B in layoutStore + bool operator==(const AlgoKey &o) const noexcept { + return dk == o.dk && rowsA == o.rowsA && colsA == o.colsA && + rowsB == o.rowsB && colsB == o.colsB; + } +}; + +struct AlgoKeyHash { + std::size_t operator()(const AlgoKey &k) const noexcept { + std::size_t h = DescKeyHash{}(k.dk); + auto mix = [&](std::size_t v) { + h ^= std::hash{}(v) + 0x9e3779b97f4a7c15ULL + (h << 6) + + (h >> 2); + }; + mix(k.rowsA); + mix(k.colsA); + mix(k.rowsB); + mix(k.colsB); + return h; + } +}; + +template class BlasLt { + typename Api::Handle ltHandle = nullptr; + typename Api::BlasHandle handle = nullptr; + typename Api::Preference preference = nullptr; + void *d_workspace = nullptr; + size_t workspaceSize = 1u << 25; // 32 MB + typename Api::Stream stream = nullptr; + + std::unordered_map, typename Api::Layout, + PairHash, PairEq> + layoutStore; + + std::unordered_map descStore; + + // One cache entry per exact GEMM configuration: the heuristic result to + // reuse, plus this entry's position in the recency list so a hit can mark + // itself most-recently-used in O(1). The position is only maintained when a + // cache limit is set; with no limit the list stays empty. + struct CacheEntry { + typename Api::HeuristicResult h{}; + std::list::iterator lru{}; + }; + std::unordered_map algoCache; + // entries ordered most- to least-recently used; drives eviction + std::list lruOrder; + // 0 = unbounded + std::size_t algoCacheLimit = 0; + +public: + std::size_t algoCacheSize() const { return algoCache.size(); } + + BlasLt(const BlasLt &) = delete; + BlasLt &operator=(const BlasLt &) = delete; + BlasLt(BlasLt &&) = delete; + BlasLt &operator=(BlasLt &&) = delete; + + BlasLt(typename Api::Queue &queue, std::size_t cacheLimit = 0) + : algoCacheLimit{cacheLimit}, m_queue{queue} { + stream = static_cast(m_queue.getNativeHandle()); + + SOFIEBLAS_CHECK_LT(Api::ltCreate(<Handle)); + + SOFIEBLAS_CHECK_LT(Api::blasCreate(&handle)); + SOFIEBLAS_CHECK_LT(Api::blasSetStream(handle, stream)); + + SOFIEBLAS_CHECK_LT(Api::prefCreate(&preference)); + SOFIEBLAS_CHECK_RT(Api::rtMalloc(&d_workspace, workspaceSize)); + SOFIEBLAS_CHECK_LT(Api::prefSetAttribute(preference, Api::PrefMaxWorkspace, + &workspaceSize, + sizeof(workspaceSize))); + } + + ~BlasLt() { + for (auto &[key, layout] : layoutStore) + if (layout) + Api::layoutDestroy(layout); + for (auto &[key, desc] : descStore) + if (desc) + Api::descDestroy(desc); + if (preference) + Api::prefDestroy(preference); + if (ltHandle) + Api::ltDestroy(ltHandle); + if (handle) + Api::blasDestroy(handle); + if (d_workspace) + Api::rtFree(d_workspace); + } + + inline typename Api::Operation charToTranspose(char trans) { + switch (trans) { + case 'N': + case 'n': + return Api::OpN; + case 'T': + case 't': + return Api::OpT; + case 'C': + case 'c': + return Api::OpC; + default: + throw std::invalid_argument( + std::string("Invalid transpose character for ") + Api::name + "."); + } + } + + // An epilogue is the extra step the library fuses into the multiply kernel + // after the product: nothing, adding the bias vector, or adding it and + // applying the activation. Which one a call site uses is decided by the + // function it calls (matmul, gemm, gemmrelu, gemmgelu); addLayoutConfig + // receives the same choice as a character so it can resolve the site's + // algorithm up front for the right configuration. + inline typename Api::Epilogue charToEpilogue(char epilogue) { + switch (epilogue) { + case 'N': + case 'n': + return Api::EpilogueDefault; + case 'B': + case 'b': + return Api::EpilogueBias; + case 'R': + case 'r': + return Api::EpilogueReluBias; + case 'G': + case 'g': + return Api::EpilogueGeluBias; + default: + throw std::invalid_argument( + std::string("Invalid epilogue character for ") + Api::name + "."); + } + } + + // Registers a call site's construction-time shape: creates the three matrix + // layouts and resolves the multiply algorithm for them up front, so the + // first call at this shape finds everything cached. + void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, + std::size_t lda, std::size_t ldb, std::size_t ldc, + char transa, char transb, char epilogue) { + const auto shapeA = layoutKeyA(transa, m, k); + const auto shapeB = layoutKeyB(transb, k, n); + const std::pair shapeC{m, n}; + getOrCreateLayout(shapeA, lda); + getOrCreateLayout(shapeB, ldb); + getOrCreateLayout(shapeC, ldc); + getOrComputeAlgo(charToTranspose(transa), charToTranspose(transb), + charToEpilogue(epilogue), shapeA, shapeB, shapeC); + } + + // Each multiply variant comes as one generic overload, where A, B, bias and + // C are any alpaka buffers or views (anything alpaka::getPtrNative + // accepts), and one raw device-pointer overload, which generated code + // calls. + template + inline void gemm(char transa, char transb, unsigned int m, unsigned int n, + unsigned int k, float alpha, TA const &A, TB const &B, + float beta, TBias &bias, TC &C) { + executeMatmul(charToTranspose(transa), charToTranspose(transb), + Api::EpilogueBias, alpha, alpaka::getPtrNative(A), + alpaka::getPtrNative(B), beta, alpaka::getPtrNative(bias), + alpaka::getPtrNative(C), + static_cast(alpaka::getPtrNative(bias)), + layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); + } + + template + inline void gemm(char transa, char transb, unsigned int m, unsigned int n, + unsigned int k, float alpha, T const *A, T const *B, + float beta, T *bias, T *C) { + executeMatmul(charToTranspose(transa), charToTranspose(transb), + Api::EpilogueBias, alpha, A, B, beta, bias, C, + static_cast(bias), layoutKeyA(transa, m, k), + layoutKeyB(transb, k, n), {m, n}); + } + + template + inline void gemmrelu(char transa, char transb, unsigned int m, unsigned int n, + unsigned int k, float alpha, TA const &A, TB const &B, + float beta, TBias &bias, TC &C) { + executeMatmul(charToTranspose(transa), charToTranspose(transb), + Api::EpilogueReluBias, alpha, alpaka::getPtrNative(A), + alpaka::getPtrNative(B), beta, alpaka::getPtrNative(bias), + alpaka::getPtrNative(C), + static_cast(alpaka::getPtrNative(bias)), + layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); + } + + template + inline void gemmrelu(char transa, char transb, unsigned int m, unsigned int n, + unsigned int k, float alpha, T const *A, T const *B, + float beta, T *bias, T *C) { + executeMatmul(charToTranspose(transa), charToTranspose(transb), + Api::EpilogueReluBias, alpha, A, B, beta, bias, C, + static_cast(bias), layoutKeyA(transa, m, k), + layoutKeyB(transb, k, n), {m, n}); + } + + template + inline void gemmgelu(char transa, char transb, unsigned int m, unsigned int n, + unsigned int k, float alpha, TA const &A, TB const &B, + float beta, TBias &bias, TC &C) { + executeMatmul(charToTranspose(transa), charToTranspose(transb), + Api::EpilogueGeluBias, alpha, alpaka::getPtrNative(A), + alpaka::getPtrNative(B), beta, alpaka::getPtrNative(bias), + alpaka::getPtrNative(C), + static_cast(alpaka::getPtrNative(bias)), + layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); + } + + template + inline void gemmgelu(char transa, char transb, unsigned int m, unsigned int n, + unsigned int k, float alpha, T const *A, T const *B, + float beta, T *bias, T *C) { + executeMatmul(charToTranspose(transa), charToTranspose(transb), + Api::EpilogueGeluBias, alpha, A, B, beta, bias, C, + static_cast(bias), layoutKeyA(transa, m, k), + layoutKeyB(transb, k, n), {m, n}); + } + + template + inline void matmul(char transa, char transb, unsigned int m, unsigned int n, + unsigned int k, float alpha, TA const &A, TB const &B, + float beta, TC &C) { + auto *c = alpaka::getPtrNative(C); + executeMatmul(charToTranspose(transa), charToTranspose(transb), + Api::EpilogueDefault, alpha, alpaka::getPtrNative(A), + alpaka::getPtrNative(B), beta, c, c, nullptr, + layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); + } + + template + inline void matmul(char transa, char transb, unsigned int m, unsigned int n, + unsigned int k, float alpha, T const *A, T const *B, + float beta, T *C) { + executeMatmul(charToTranspose(transa), charToTranspose(transb), + Api::EpilogueDefault, alpha, A, B, beta, C, C, nullptr, + layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); + } + + inline void gemmStridedBatched(char transa, char transb, int m, int n, int k, + float alpha, const float *A, int lda, + long long strideA, const float *B, int ldb, + long long strideB, float beta, float *C, + int ldc, long long strideC, int batchCount) { + SOFIEBLAS_CHECK_LT(Api::sgemmStridedBatched( + handle, charToTranspose(transa), charToTranspose(transb), m, n, k, + &alpha, A, lda, strideA, B, ldb, strideB, &beta, C, ldc, strideC, + batchCount)); + } + +private: + typename Api::Queue m_queue; + + static std::pair + layoutKeyA(char trans, std::size_t m, std::size_t k) { + return (trans == 'N' || trans == 'n') ? std::make_pair(m, k) + : std::make_pair(k, m); + } + + static std::pair + layoutKeyB(char trans, std::size_t k, std::size_t n) { + return (trans == 'N' || trans == 'n') ? std::make_pair(k, n) + : std::make_pair(n, k); + } + + // Returns the layout describing a (rows, cols) matrix, creating and caching + // it on first use. Every caller passes ld = rows (dense column-major). + typename Api::Layout + getOrCreateLayout(const std::pair &shape, + std::size_t ld) { + auto it = layoutStore.find(shape); + if (it != layoutStore.end()) + return it->second; + typename Api::Layout layout = nullptr; + SOFIEBLAS_CHECK_LT(Api::layoutCreate(&layout, Api::RealF32, shape.first, + shape.second, ld)); + layoutStore.emplace(shape, layout); + return layout; + } + + typename Api::MatmulDesc &getOrCreateDesc(typename Api::Operation transA, + typename Api::Operation transB, + typename Api::Epilogue epilogue) { + DescKey key{(int)transA, (int)transB, (int)epilogue}; + auto it = descStore.find(key); + if (it != descStore.end()) + return it->second; + + typename Api::MatmulDesc desc = nullptr; + SOFIEBLAS_CHECK_LT(Api::descCreate(&desc, Api::ComputeF32, Api::RealF32)); + SOFIEBLAS_CHECK_LT( + Api::descSetAttribute(desc, Api::DescTransA, &transA, sizeof(transA))); + SOFIEBLAS_CHECK_LT( + Api::descSetAttribute(desc, Api::DescTransB, &transB, sizeof(transB))); + SOFIEBLAS_CHECK_LT(Api::descSetAttribute(desc, Api::DescEpilogue, &epilogue, + sizeof(epilogue))); + // For bias epilogues: set a non-null dummy pointer so the descriptor is + // valid for the heuristic query. + if (epilogue != Api::EpilogueDefault) { + const void *dummy = d_workspace; + SOFIEBLAS_CHECK_LT(Api::descSetAttribute(desc, Api::DescBiasPointer, + &dummy, sizeof(dummy))); + } + descStore.emplace(key, desc); + return descStore.at(key); + } + + typename Api::HeuristicResult & + getOrComputeAlgo(typename Api::Operation transA, + typename Api::Operation transB, + typename Api::Epilogue epilogue, + const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC) { + AlgoKey key{{(int)transA, (int)transB, (int)epilogue}, + shapeA.first, + shapeA.second, + shapeB.first, + shapeB.second}; + auto it = algoCache.find(key); + if (it != algoCache.end()) { + if (algoCacheLimit) + lruOrder.splice(lruOrder.begin(), lruOrder, it->second.lru); + return it->second.h; + } + + auto &desc = getOrCreateDesc(transA, transB, epilogue); + auto lA = getOrCreateLayout(shapeA, shapeA.first); + auto lB = getOrCreateLayout(shapeB, shapeB.first); + auto lC = getOrCreateLayout(shapeC, shapeC.first); + typename Api::HeuristicResult h{}; + int returnedResults = 0; + SOFIEBLAS_CHECK_LT(Api::getHeuristic(ltHandle, desc, lA, lB, lC, lC, + preference, 1, &h, &returnedResults)); + if (returnedResults == 0) { + std::cerr << "[sofieBLAS] No suitable " << Api::name + << " algorithm found for " + << "transA=" << transA << " transB=" << transB + << " epilogue=" << epilogue << " A=[" << shapeA.first << "x" + << shapeA.second << "]" + << " B=[" << shapeB.first << "x" << shapeB.second << "]\n"; + exit(EXIT_FAILURE); + } + auto ins = algoCache.emplace(key, CacheEntry{h, {}}).first; + if (algoCacheLimit) { + lruOrder.push_front(key); + ins->second.lru = lruOrder.begin(); + while (algoCache.size() > algoCacheLimit) { + algoCache.erase(lruOrder.back()); + lruOrder.pop_back(); + } + } + return ins->second.h; + } + + void executeMatmul(typename Api::Operation transA, + typename Api::Operation transB, + typename Api::Epilogue epilogue, float alpha, + const float *A, const float *B, float beta, + const float *D_in, float *C_out, const void *bias_ptr, + const std::pair &shapeA, + const std::pair &shapeB, + const std::pair &shapeC) { + // Retrieve (or lazily compute) the cached algorithm for this shape + auto &h = + getOrComputeAlgo(transA, transB, epilogue, shapeA, shapeB, shapeC); + + // Retrieve the cached descriptor and patch the real bias pointer in-place + auto &desc = getOrCreateDesc(transA, transB, epilogue); + if (bias_ptr) { + SOFIEBLAS_CHECK_LT(Api::descSetAttribute(desc, Api::DescBiasPointer, + &bias_ptr, sizeof(bias_ptr))); + } + + auto lA = getOrCreateLayout(shapeA, shapeA.first); + auto lB = getOrCreateLayout(shapeB, shapeB.first); + auto lC = getOrCreateLayout(shapeC, shapeC.first); + SOFIEBLAS_CHECK_LT(Api::matmul(ltHandle, desc, &alpha, A, lA, B, lB, &beta, + D_in, lC, C_out, lC, &h.algo, d_workspace, + workspaceSize, stream)); + } +}; diff --git a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp index 7d58c74..d876eea 100644 --- a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp +++ b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp @@ -5,17 +5,15 @@ #include #include #include -#include #include #include +#include #include #include -#include #include "sofieBLAS/core.hpp" #include #include -#include #include #define CHECK_HIP(err) \ @@ -34,576 +32,68 @@ } \ } while (0) -struct DescKey { - int transA; // HIPBLAS_OP_N / HIPBLAS_OP_T encoded as int - int transB; - int epilogue; // hipblasLtEpilogue_t encoded as int - bool operator==(const DescKey &o) const noexcept { - return transA == o.transA && transB == o.transB && epilogue == o.epilogue; - } -}; - -struct DescKeyHash { - std::size_t operator()(const DescKey &k) const noexcept { - std::size_t h = static_cast(k.transA) * 97u + - static_cast(k.transB) * 31u + - static_cast(k.epilogue); - return h ^ (h >> 16); - } +// The hipBLASLt spellings of everything the shared BlasLt implementation in +// backends/gpu/detail uses: types, enum values and functions. +struct HipblasLtApi { + using Queue = alpaka::QueueHipRtNonBlocking; + using Handle = hipblasLtHandle_t; + using BlasHandle = hipblasHandle_t; + using Preference = hipblasLtMatmulPreference_t; + using Stream = hipStream_t; + using Layout = hipblasLtMatrixLayout_t; + using MatmulDesc = hipblasLtMatmulDesc_t; + using HeuristicResult = hipblasLtMatmulHeuristicResult_t; + using Operation = hipblasOperation_t; + using Epilogue = hipblasLtEpilogue_t; + + static constexpr auto OpN = HIPBLAS_OP_N; + static constexpr auto OpT = HIPBLAS_OP_T; + static constexpr auto OpC = HIPBLAS_OP_C; + static constexpr auto EpilogueDefault = HIPBLASLT_EPILOGUE_DEFAULT; + static constexpr auto EpilogueBias = HIPBLASLT_EPILOGUE_BIAS; + static constexpr auto EpilogueReluBias = HIPBLASLT_EPILOGUE_RELU_BIAS; + static constexpr auto EpilogueGeluBias = HIPBLASLT_EPILOGUE_GELU_BIAS; + static constexpr auto ComputeF32 = HIPBLAS_COMPUTE_32F; + static constexpr auto RealF32 = HIP_R_32F; + static constexpr auto DescTransA = HIPBLASLT_MATMUL_DESC_TRANSA; + static constexpr auto DescTransB = HIPBLASLT_MATMUL_DESC_TRANSB; + static constexpr auto DescEpilogue = HIPBLASLT_MATMUL_DESC_EPILOGUE; + static constexpr auto DescBiasPointer = HIPBLASLT_MATMUL_DESC_BIAS_POINTER; + static constexpr auto PrefMaxWorkspace = + HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES; + static constexpr const char *name = "hipBLASLt"; + + static constexpr auto ltCreate = hipblasLtCreate; + static constexpr auto ltDestroy = hipblasLtDestroy; + static constexpr auto blasCreate = hipblasCreate; + static constexpr auto blasDestroy = hipblasDestroy; + static constexpr auto blasSetStream = hipblasSetStream; + static constexpr auto prefCreate = hipblasLtMatmulPreferenceCreate; + static constexpr auto prefDestroy = hipblasLtMatmulPreferenceDestroy; + static constexpr auto prefSetAttribute = + hipblasLtMatmulPreferenceSetAttribute; + static constexpr auto layoutCreate = hipblasLtMatrixLayoutCreate; + static constexpr auto layoutDestroy = hipblasLtMatrixLayoutDestroy; + static constexpr auto descCreate = hipblasLtMatmulDescCreate; + static constexpr auto descDestroy = hipblasLtMatmulDescDestroy; + static constexpr auto descSetAttribute = hipblasLtMatmulDescSetAttribute; + static constexpr auto getHeuristic = hipblasLtMatmulAlgoGetHeuristic; + static constexpr auto matmul = hipblasLtMatmul; + static constexpr auto sgemmStridedBatched = hipblasSgemmStridedBatched; + + // hipMalloc has a templated C++ overload, so a pointer to it is ambiguous + static hipError_t rtMalloc(void **ptr, std::size_t size) { + return hipMalloc(ptr, size); + } + static hipError_t rtFree(void *ptr) { return hipFree(ptr); } }; -struct AlgoKey { - DescKey dk; - std::size_t rowsA, colsA; - std::size_t rowsB, colsB; - bool operator==(const AlgoKey &o) const noexcept { - return dk == o.dk && rowsA == o.rowsA && colsA == o.colsA && - rowsB == o.rowsB && colsB == o.colsB; - } -}; - -struct AlgoKeyHash { - std::size_t operator()(const AlgoKey &k) const noexcept { - std::size_t h = DescKeyHash{}(k.dk); - auto mix = [&](std::size_t v) { - h ^= std::hash{}(v) + 0x9e3779b97f4a7c15ULL + (h << 6) + - (h >> 2); - }; - mix(k.rowsA); - mix(k.colsA); - mix(k.rowsB); - mix(k.colsB); - return h; - } -}; - -// A call site's declared maximum shape, recorded by addLayoutConfig. -struct ShapeEnvelope { - std::size_t rowsA, colsA, rowsB, colsB, rowsC, colsC; -}; - -struct LayoutStats { - std::size_t heuristicQueries = 0; // algorithm searches issued - std::size_t envelopeRejects = 0; // declared-shape algo unusable at a call - std::size_t evictions = 0; // entries dropped to stay under the limit -}; - -class BlasHip { - hipblasLtHandle_t ltHandle = nullptr; - hipblasHandle_t handle = nullptr; - hipblasLtMatmulPreference_t preference = nullptr; - void *d_workspace = nullptr; - size_t workspaceSize = 1u << 25; // 32 MB - hipStream_t stream = nullptr; - - // One persistent layout descriptor per GEMM operand: ROLE_A and ROLE_B are - // the input matrices of C = alpha * op(A) * op(B) + beta * C, ROLE_C the - // output (hipblasLtMatmul takes it twice, as C and D). stampLayout rewrites - // each descriptor in place to the shape of the call at hand, which is what - // lets one instance serve any runtime size. - enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 }; - hipblasLtMatrixLayout_t roleLayout[3] = {}; - - std::unordered_map descStore; - - // algo cache entry - struct CacheEntry { - hipblasLtMatmulHeuristicResult_t h{}; - // position in lruOrder; only valid when a limit is set - std::list::iterator lru{}; - }; - std::unordered_map algoCache; - std::list lruOrder; - // 0 = unbounded - std::size_t algoCacheLimit = 0; - - std::vector envelopes; - - LayoutStats stats; - -public: - const LayoutStats &layoutStats() const { return stats; } - std::size_t algoCacheSize() const { return algoCache.size(); } - - BlasHip(const BlasHip &) = delete; - BlasHip &operator=(const BlasHip &) = delete; - BlasHip(BlasHip &&) = delete; - BlasHip &operator=(BlasHip &&) = delete; - - BlasHip(alpaka::QueueHipRtNonBlocking &queue, std::size_t algoCacheLimit_ = 0) - : algoCacheLimit{algoCacheLimit_}, m_queue{queue} { - stream = static_cast(m_queue.getNativeHandle()); - - CHECK_HIPBLAS(hipblasLtCreate(<Handle)); - - CHECK_HIPBLAS(hipblasCreate(&handle)); - CHECK_HIPBLAS(hipblasSetStream(handle, stream)); - - CHECK_HIPBLAS(hipblasLtMatmulPreferenceCreate(&preference)); - CHECK_HIP(hipMalloc(&d_workspace, workspaceSize)); - CHECK_HIPBLAS(hipblasLtMatmulPreferenceSetAttribute( - preference, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, - sizeof(workspaceSize))); - } - - ~BlasHip() { - for (auto L : roleLayout) - if (L) - hipblasLtMatrixLayoutDestroy(L); - for (auto &[key, desc] : descStore) - if (desc) - hipblasLtMatmulDescDestroy(desc); - if (preference) - hipblasLtMatmulPreferenceDestroy(preference); - if (ltHandle) - hipblasLtDestroy(ltHandle); - if (handle) - hipblasDestroy(handle); - if (d_workspace) - hipFree(d_workspace); - } - - inline hipblasOperation_t charToHipBlasTranspose(char trans) { - switch (trans) { - case 'N': - case 'n': - return HIPBLAS_OP_N; - case 'T': - case 't': - return HIPBLAS_OP_T; - case 'C': - case 'c': - return HIPBLAS_OP_C; - default: - throw std::invalid_argument("Invalid transpose character for hipBLAS."); - } - } - - // Declares a call site's largest shape (its envelope) and resolves the - // algorithm for it up front; the generated Session constructor calls this - // once per GEMM call site with its construction-time dimensions. Which - // epilogue the site will use is unknown here, so all three used by the - // generated code are resolved; unused ones cost one heuristic query each, - // off the inference path. - void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, std::size_t, - std::size_t, std::size_t, char transa, char transb) { - const auto shapeA = layoutKeyA(transa, m, k); - const auto shapeB = layoutKeyB(transb, k, n); - const std::pair shapeC{m, n}; - envelopes.push_back( - {shapeA.first, shapeA.second, shapeB.first, shapeB.second, m, n}); - - const hipblasOperation_t tA = charToHipBlasTranspose(transa); - const hipblasOperation_t tB = charToHipBlasTranspose(transb); - const hipblasLtEpilogue_t eps[] = {HIPBLASLT_EPILOGUE_DEFAULT, - HIPBLASLT_EPILOGUE_BIAS, - HIPBLASLT_EPILOGUE_RELU_BIAS}; - for (hipblasLtEpilogue_t ep : eps) { - getOrComputeAlgo(tA, tB, ep, shapeA, shapeB, shapeC, /*required=*/false); - } - } - - template - inline void - gemm(char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, alpaka::BufHipRt, TIdx> const &A, - alpaka::BufHipRt, TIdx> const &B, float beta, - alpaka::BufHipRt, TIdx> &bias, - alpaka::BufHipRt, TIdx> &C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_BIAS, - alpha, alpaka::getPtrNative(A), alpaka::getPtrNative(B), beta, - alpaka::getPtrNative(bias), alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemm( - char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, - alpaka::ViewPlainPtr, TIdx> const - &A, - alpaka::ViewPlainPtr, TIdx> const - &B, - float beta, - alpaka::ViewPlainPtr, TIdx> &bias, - alpaka::ViewPlainPtr, TIdx> &C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_BIAS, - alpha, alpaka::getPtrNative(A), alpaka::getPtrNative(B), beta, - alpaka::getPtrNative(bias), alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemm(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, T const *A, T const *B, - float beta, T *bias, T *C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_BIAS, - alpha, A, B, beta, bias, C, static_cast(bias), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmrelu(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, - alpaka::BufHipRt, TIdx> const &A, - alpaka::BufHipRt, TIdx> const &B, - float beta, - alpaka::BufHipRt, TIdx> &bias, - alpaka::BufHipRt, TIdx> &C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_RELU_BIAS, - alpha, alpaka::getPtrNative(A), alpaka::getPtrNative(B), beta, - alpaka::getPtrNative(bias), alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmrelu( - char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, - alpaka::ViewPlainPtr, TIdx> const - &A, - alpaka::ViewPlainPtr, TIdx> const - &B, - float beta, - alpaka::ViewPlainPtr, TIdx> &bias, - alpaka::ViewPlainPtr, TIdx> &C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_RELU_BIAS, - alpha, alpaka::getPtrNative(A), alpaka::getPtrNative(B), beta, - alpaka::getPtrNative(bias), alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmrelu(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, T const *A, T const *B, - float beta, T *bias, T *C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_RELU_BIAS, - alpha, A, B, beta, bias, C, static_cast(bias), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmgelu(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, - alpaka::BufHipRt, TIdx> const &A, - alpaka::BufHipRt, TIdx> const &B, - float beta, - alpaka::BufHipRt, TIdx> &bias, - alpaka::BufHipRt, TIdx> &C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_GELU_BIAS, - alpha, alpaka::getPtrNative(A), alpaka::getPtrNative(B), beta, - alpaka::getPtrNative(bias), alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmgelu( - char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, - alpaka::ViewPlainPtr, TIdx> const - &A, - alpaka::ViewPlainPtr, TIdx> const - &B, - float beta, - alpaka::ViewPlainPtr, TIdx> &bias, - alpaka::ViewPlainPtr, TIdx> &C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_GELU_BIAS, - alpha, alpaka::getPtrNative(A), alpaka::getPtrNative(B), beta, - alpaka::getPtrNative(bias), alpaka::getPtrNative(C), - static_cast(alpaka::getPtrNative(bias)), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void gemmgelu(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, T const *A, T const *B, - float beta, T *bias, T *C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_GELU_BIAS, - alpha, A, B, beta, bias, C, static_cast(bias), - layoutKeyA(transa, m, k), layoutKeyB(transb, k, n), {m, n}); - } +#define SOFIEBLAS_CHECK_LT(x) CHECK_HIPBLAS(x) +#define SOFIEBLAS_CHECK_RT(x) CHECK_HIP(x) - template - inline void matmul(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, - alpaka::BufHipRt, TIdx> const &A, - alpaka::BufHipRt, TIdx> const &B, - float beta, - alpaka::BufHipRt, TIdx> &C) { - float *c = alpaka::getPtrNative(C); - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_DEFAULT, - alpha, alpaka::getPtrNative(A), alpaka::getPtrNative(B), beta, - c, c, nullptr, layoutKeyA(transa, m, k), - layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void matmul( - char transa, char transb, unsigned int m, unsigned int n, unsigned int k, - float alpha, - alpaka::ViewPlainPtr, TIdx> const - &A, - alpaka::ViewPlainPtr, TIdx> const - &B, - float beta, - alpaka::ViewPlainPtr, TIdx> &C) { - T *c = alpaka::getPtrNative(C); - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_DEFAULT, - alpha, alpaka::getPtrNative(A), alpaka::getPtrNative(B), beta, - c, c, nullptr, layoutKeyA(transa, m, k), - layoutKeyB(transb, k, n), {m, n}); - } - - template - inline void matmul(char transa, char transb, unsigned int m, unsigned int n, - unsigned int k, float alpha, T const *A, T const *B, - float beta, T *C) { - executeMatmul(charToHipBlasTranspose(transa), - charToHipBlasTranspose(transb), HIPBLASLT_EPILOGUE_DEFAULT, - alpha, A, B, beta, C, C, nullptr, layoutKeyA(transa, m, k), - layoutKeyB(transb, k, n), {m, n}); - } - - inline void gemmStridedBatched(char transa, char transb, int m, int n, int k, - float alpha, const float *A, int lda, - long long strideA, const float *B, int ldb, - long long strideB, float beta, float *C, - int ldc, long long strideC, int batchCount) { - CHECK_HIPBLAS(hipblasSgemmStridedBatched( - handle, charToHipBlasTranspose(transa), charToHipBlasTranspose(transb), - m, n, k, &alpha, A, lda, strideA, B, ldb, strideB, &beta, C, ldc, - strideC, batchCount)); - } - -private: - alpaka::QueueHipRtNonBlocking m_queue; - - static std::pair - layoutKeyA(char trans, std::size_t m, std::size_t k) { - return (trans == 'N' || trans == 'n') ? std::make_pair(m, k) - : std::make_pair(k, m); - } +#include "../gpu/detail/sofieBLAS_blaslt_common.tpp" - static std::pair - layoutKeyB(char trans, std::size_t k, std::size_t n) { - return (trans == 'N' || trans == 'n') ? std::make_pair(k, n) - : std::make_pair(n, k); - } - - // Sets a role's layout descriptor to the given physical (rows, cols), - // creating it on first use. Matrices are dense column-major, so ld = rows. - hipblasLtMatrixLayout_t - stampLayout(LayoutRole role, const std::pair &key) { - const uint64_t rows = key.first, cols = key.second; - const int64_t ld = static_cast(key.first); - hipblasLtMatrixLayout_t &L = roleLayout[role]; - if (!L) { - CHECK_HIPBLAS(hipblasLtMatrixLayoutCreate(&L, HIP_R_32F, rows, cols, ld)); - } else { - CHECK_HIPBLAS(hipblasLtMatrixLayoutSetAttribute( - L, HIPBLASLT_MATRIX_LAYOUT_ROWS, &rows, sizeof(rows))); - CHECK_HIPBLAS(hipblasLtMatrixLayoutSetAttribute( - L, HIPBLASLT_MATRIX_LAYOUT_COLS, &cols, sizeof(cols))); - CHECK_HIPBLAS(hipblasLtMatrixLayoutSetAttribute( - L, HIPBLASLT_MATRIX_LAYOUT_LD, &ld, sizeof(ld))); - } - return L; - } - - // Returns the shape declared through addLayoutConfig that this call belongs - // to, or null if none covers it. Shapes are the physical (rows, cols) of - // matrices A, B and C, after any transpose is applied (transa='T' makes - // shapeA = (k, m)). The contraction dimension (colsA / rowsB) must match - // exactly: it comes from the weight tensor and never varies at runtime, so - // it identifies the call site and stops one site's declared shape from - // serving another's calls. The free dimensions only need covering; among - // candidates the least excess wins. - const ShapeEnvelope * - findEnvelope(const std::pair &shapeA, - const std::pair &shapeB, - const std::pair &shapeC) const { - const ShapeEnvelope *best = nullptr; - std::size_t bestExcess = std::numeric_limits::max(); - for (const auto &e : envelopes) { - if (e.colsA != shapeA.second || e.rowsB != shapeB.first) - continue; - if (e.rowsA < shapeA.first || e.colsB < shapeB.second || - e.rowsC < shapeC.first || e.colsC < shapeC.second) - continue; - const std::size_t ex = - (e.rowsA - shapeA.first) + (e.colsA - shapeA.second) + - (e.rowsB - shapeB.first) + (e.colsB - shapeB.second); - if (ex < bestExcess) { - bestExcess = ex; - best = &e; - } - } - return best; - } - - // Whether the given algorithm can run this exact shape within the - // workspace. An algorithm resolved at a declared shape is not guaranteed to - // run at every smaller size it covers. - bool algoUsable(hipblasLtMatmulDesc_t desc, const hipblasLtMatmulAlgo_t &algo, - const std::pair &shapeA, - const std::pair &shapeB, - const std::pair &shapeC) { - auto lA = stampLayout(ROLE_A, shapeA); - auto lB = stampLayout(ROLE_B, shapeB); - auto lC = stampLayout(ROLE_C, shapeC); - // hipBLASLt has no hipblasLtMatmulAlgoCheck; the ext API rewrites the algo - // it is handed, so probe a copy. - hipblasLtMatmulAlgo_t probe = algo; - const float alpha = 1.f, beta = 0.f; - std::size_t ws = 0; - return hipblaslt_ext::matmulIsAlgoSupported(ltHandle, desc, &alpha, lA, lB, - &beta, lC, lC, probe, - ws) == HIPBLAS_STATUS_SUCCESS && - ws <= workspaceSize; - } - - hipblasLtMatmulDesc_t &getOrCreateDesc(hipblasOperation_t transA, - hipblasOperation_t transB, - hipblasLtEpilogue_t epilogue) { - DescKey key{(int)transA, (int)transB, (int)epilogue}; - auto it = descStore.find(key); - if (it != descStore.end()) - return it->second; - - hipblasLtMatmulDesc_t desc = nullptr; - CHECK_HIPBLAS( - hipblasLtMatmulDescCreate(&desc, HIPBLAS_COMPUTE_32F, HIP_R_32F)); - CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( - desc, HIPBLASLT_MATMUL_DESC_TRANSA, &transA, sizeof(transA))); - CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( - desc, HIPBLASLT_MATMUL_DESC_TRANSB, &transB, sizeof(transB))); - CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( - desc, HIPBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue))); - - if (epilogue != HIPBLASLT_EPILOGUE_DEFAULT) { - const void *dummy = d_workspace; - CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( - desc, HIPBLASLT_MATMUL_DESC_BIAS_POINTER, &dummy, sizeof(dummy))); - } - descStore.emplace(key, desc); - return descStore.at(key); - } - - // Looks up, or resolves and caches, the algorithm for the given transpose - // settings, epilogue and shapes. required=false is for constructor warmup: - // a speculatively resolved epilogue may legitimately have no algorithm, and - // returns null instead of aborting. - hipblasLtMatmulHeuristicResult_t * - getOrComputeAlgo(hipblasOperation_t transA, hipblasOperation_t transB, - hipblasLtEpilogue_t epilogue, - const std::pair &shapeA, - const std::pair &shapeB, - const std::pair &shapeC, - bool required = true) { - AlgoKey key{{(int)transA, (int)transB, (int)epilogue}, - shapeA.first, - shapeA.second, - shapeB.first, - shapeB.second}; - auto it = algoCache.find(key); - if (it != algoCache.end()) { - if (algoCacheLimit) - lruOrder.splice(lruOrder.begin(), lruOrder, it->second.lru); - return &it->second.h; - } - - auto &desc = getOrCreateDesc(transA, transB, epilogue); - auto lA = stampLayout(ROLE_A, shapeA); - auto lB = stampLayout(ROLE_B, shapeB); - auto lC = stampLayout(ROLE_C, shapeC); - hipblasLtMatmulHeuristicResult_t h{}; - int returnedResults = 0; - CHECK_HIPBLAS(hipblasLtMatmulAlgoGetHeuristic( - ltHandle, desc, lA, lB, lC, lC, preference, 1, &h, &returnedResults)); - ++stats.heuristicQueries; - if (returnedResults == 0) { - if (!required) - return nullptr; - std::cerr << "[sofieBLAS] No suitable hipBLASLt algorithm found for " - << "transA=" << transA << " transB=" << transB - << " epilogue=" << epilogue << " A=[" << shapeA.first << "x" - << shapeA.second << "]" - << " B=[" << shapeB.first << "x" << shapeB.second << "]\n"; - exit(EXIT_FAILURE); - } - auto ins = algoCache.emplace(key, CacheEntry{h, {}}).first; - if (algoCacheLimit) { - lruOrder.push_front(key); - ins->second.lru = lruOrder.begin(); - while (algoCache.size() > algoCacheLimit) { - algoCache.erase(lruOrder.back()); - lruOrder.pop_back(); - ++stats.evictions; - } - } - return &ins->second.h; - } - - void executeMatmul(hipblasOperation_t transA, hipblasOperation_t transB, - hipblasLtEpilogue_t epilogue, float alpha, const float *A, - const float *B, float beta, const float *D_in, - float *C_out, const void *bias_ptr, - const std::pair &shapeA, - const std::pair &shapeB, - const std::pair &shapeC) { - auto &desc = getOrCreateDesc(transA, transB, epilogue); - if (bias_ptr) { - CHECK_HIPBLAS(hipblasLtMatmulDescSetAttribute( - desc, HIPBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, - sizeof(bias_ptr))); - } - - // Resolve the algorithm at the call site's declared shape when one covers - // this call, so every runtime size the site produces shares one cache - // entry; with no covering declaration, resolve at the exact shape. - const ShapeEnvelope *env = findEnvelope(shapeA, shapeB, shapeC); - const std::pair - algoShapeA = env ? std::make_pair(env->rowsA, env->colsA) : shapeA, - algoShapeB = env ? std::make_pair(env->rowsB, env->colsB) : shapeB, - algoShapeC = env ? std::make_pair(env->rowsC, env->colsC) : shapeC; - hipblasLtMatmulHeuristicResult_t h = *getOrComputeAlgo( - transA, transB, epilogue, algoShapeA, algoShapeB, algoShapeC); - - // Normally the resolution above is the only one. Only when hipBLASLt - // rejects the declared shape's algorithm at this call's exact size is the - // algorithm resolved a second time, at the exact shape, and that result - // is cached as well. - if (env && !algoUsable(desc, h.algo, shapeA, shapeB, shapeC)) { - ++stats.envelopeRejects; - h = *getOrComputeAlgo(transA, transB, epilogue, shapeA, shapeB, shapeC); - } - - // Stamp the exact call shape last: the algorithm resolution and the - // validity check above leave the shared descriptors at other dims. - auto lA = stampLayout(ROLE_A, shapeA); - auto lB = stampLayout(ROLE_B, shapeB); - auto lC = stampLayout(ROLE_C, shapeC); - CHECK_HIPBLAS(hipblasLtMatmul(ltHandle, desc, &alpha, A, lA, B, lB, &beta, - D_in, lC, C_out, lC, &h.algo, d_workspace, - workspaceSize, stream)); - } -}; +using BlasHip = BlasLt; namespace traits { diff --git a/tests/test.cc b/tests/test.cc index 0d4002a..06b27d1 100644 --- a/tests/test.cc +++ b/tests/test.cc @@ -394,7 +394,7 @@ static void runCudaTests() { // ---- matmul NN ---- blas.addLayoutConfig(M, N, K, ldaFor('N', M, K), ldbFor('N', K, N), M, 'N', - 'N'); + 'N', 'n'); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), A, B, M, N, K, 1.f, 0.f, false, false); blas.matmul('N', 'N', M, N, K, 1.f, dA, dB, 0.f, dC); @@ -410,7 +410,7 @@ static void runCudaTests() { alpaka::memcpy(queue, dAt, hAt); alpaka::wait(queue); blas.addLayoutConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, 'T', - 'N'); + 'N', 'n'); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), At, B, M, N, K, 1.f, 0.f, true, false); blas.matmul('T', 'N', M, N, K, 1.f, dAt, dB, 0.f, dC); @@ -427,7 +427,7 @@ static void runCudaTests() { alpaka::memcpy(queue, dBt, hBt); alpaka::wait(queue); blas.addLayoutConfig(M, N, K, ldaFor('N', M, K), ldbFor('T', K, N), M, 'N', - 'T'); + 'T', 'n'); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), A, Bt, M, N, K, 1.f, 0.f, false, true); blas.matmul('N', 'T', M, N, K, 1.f, dA, dBt, 0.f, dC); @@ -466,7 +466,7 @@ static void runCudaTests() { alpaka::memcpy(queue, dAt, hAt); alpaka::wait(queue); blas.addLayoutConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, 'T', - 'N'); + 'N', 'b'); std::fill(ref.begin(), ref.end(), 0.f); refGemm(ref.data(), At, B, bias, M, N, K, 1.f, 0.f, true, false); blas.gemm('T', 'N', M, N, K, 1.f, dAt, dB, 0.f, dBias, dC); @@ -494,7 +494,7 @@ static void runCudaTests() { alpaka::memcpy(queue, dBp, hBp); alpaka::memcpy(queue, dBiasz, hBiasz); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N'); + blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N', 'r'); std::fill(ref.begin(), ref.end(), 0.f); refGemmRelu(ref.data(), Ap, Bp, alpaka::getPtrNative(hBiasz), M, N, K, 1.f, 0.f, false, false); @@ -559,116 +559,6 @@ static void runCudaTests() { } } -static void runDynamicShapeTests() { - std::cout << "\n=== CUDA Dynamic-Shape Tests ===\n"; - - alpaka::PlatformCudaRt platform{}; - auto dev = alpaka::getDevByIdx(platform, 0u); - alpaka::Queue queue{dev}; - sofieBLAS blas(queue); - - alpaka::PlatformCpu hostPlatform{}; - auto hostDev = alpaka::getDevByIdx(hostPlatform, 0u); - - constexpr int MCAP = 96, MENV = 64, N = 3, K = 5; - - auto hA = alpaka::allocBuf(hostDev, static_cast(MCAP * K)); - auto hB = alpaka::allocBuf(hostDev, static_cast(K * N)); - auto hC = alpaka::allocBuf(hostDev, static_cast(MCAP * N)); - float *A = alpaka::getPtrNative(hA); - float *B = alpaka::getPtrNative(hB); - float *C = alpaka::getPtrNative(hC); - fillSeq(A, MCAP * K, 0.5f, 0.25f); - fillSeq(B, K * N, 1.f, 0.5f); - - auto dA = - alpaka::allocAsyncBuf(queue, static_cast(MCAP * K)); - auto dB = alpaka::allocAsyncBuf(queue, static_cast(K * N)); - auto dC = - alpaka::allocAsyncBuf(queue, static_cast(MCAP * N)); - alpaka::memcpy(queue, dA, hA); - alpaka::memcpy(queue, dB, hB); - alpaka::wait(queue); - - blas.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), - MENV, 'N', 'N'); - - std::vector ref; - auto runAt = [&](int m, const std::string &name) { - ref.assign(static_cast(m) * N, 0.f); - refMatmul(ref.data(), A, B, m, N, K, 1.f, 0.f, false, false); - blas.matmul('N', 'N', static_cast(m), static_cast(N), - static_cast(K), 1.f, dA, dB, 0.f, dC); - alpaka::memcpy(queue, hC, dC); - alpaka::wait(queue); - checkClose(C, ref.data(), m * N, name); - }; - - for (int m : {MENV, 37, 8, 51, 1, MENV, MCAP}) - runAt(m, "cuda::dynamic m=" + std::to_string(m)); - - const int nSizes = MENV - 1; - const std::size_t cacheBefore = blas.algoCacheSize(); - const std::size_t rejBefore = blas.layoutStats().envelopeRejects; - const std::size_t searchBefore = blas.layoutStats().heuristicQueries; - for (int m = 2; m <= MENV; ++m) - blas.matmul('N', 'N', static_cast(m), static_cast(N), - static_cast(K), 1.f, dA, dB, 0.f, dC); - alpaka::wait(queue); - const std::size_t added = blas.algoCacheSize() - cacheBefore; - const std::size_t rejected = blas.layoutStats().envelopeRejects - rejBefore; - - std::cout << " " << nSizes << " sizes added " << added - << " cache entries, " << rejected << " rejected\n"; - if (added < static_cast(nSizes)) { - std::cout << " PASS cuda::cache bounded\n"; - } else { - std::cerr << " FAIL [cuda::cache bounded] one entry per size\n"; - ++gFailures; - } - - { - sofieBLAS capped(queue, 8); - capped.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), - MENV, 'N', 'N'); - std::vector cref; - float worst = 0.f; - for (int m = MENV + 1; m <= MCAP; ++m) { - cref.assign(static_cast(m) * N, 0.f); - refMatmul(cref.data(), A, B, m, N, K, 1.f, 0.f, false, false); - capped.matmul('N', 'N', static_cast(m), - static_cast(N), static_cast(K), 1.f, dA, - dB, 0.f, dC); - alpaka::memcpy(queue, hC, dC); - alpaka::wait(queue); - for (std::size_t i = 0; i < cref.size(); ++i) - worst = std::max(worst, std::abs(C[i] - cref[i])); - } - std::cout << " limit=8: " << capped.algoCacheSize() << " entries, " - << capped.layoutStats().evictions << " evictions over " - << (MCAP - MENV) << " above-envelope sizes, worst err " << worst - << "\n"; - if (capped.algoCacheSize() <= 8 && worst < 1e-3f) { - std::cout << " PASS cuda::cache limit honoured\n"; - } else { - std::cerr << " FAIL [cuda::cache limit honoured] " - << capped.algoCacheSize() << " entries, worst err " << worst - << "\n"; - ++gFailures; - } - } - - const std::size_t searched = - blas.layoutStats().heuristicQueries - searchBefore; - if (searched <= rejected) { - std::cout << " PASS cuda::no search during inference\n"; - } else { - std::cerr << " FAIL [cuda::no search during inference] " << searched - << " searches, only " << rejected << " explained by rejects\n"; - ++gFailures; - } -} - #endif // ALPAKA_ACC_GPU_CUDA_ENABLED // --------------------------------------------------------------------------- @@ -725,7 +615,7 @@ static void runHipTests() { // ---- matmul NN ---- blas.addLayoutConfig(M, N, K, ldaFor('N', M, K), ldbFor('N', K, N), M, 'N', - 'N'); + 'N', 'n'); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), A, B, M, N, K, 1.f, 0.f, false, false); blas.matmul('N', 'N', M, N, K, 1.f, dA, dB, 0.f, dC); @@ -741,7 +631,7 @@ static void runHipTests() { alpaka::memcpy(queue, dAt, hAt); alpaka::wait(queue); blas.addLayoutConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, 'T', - 'N'); + 'N', 'n'); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), At, B, M, N, K, 1.f, 0.f, true, false); blas.matmul('T', 'N', M, N, K, 1.f, dAt, dB, 0.f, dC); @@ -758,7 +648,7 @@ static void runHipTests() { alpaka::memcpy(queue, dBt, hBt); alpaka::wait(queue); blas.addLayoutConfig(M, N, K, ldaFor('N', M, K), ldbFor('T', K, N), M, 'N', - 'T'); + 'T', 'n'); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), A, Bt, M, N, K, 1.f, 0.f, false, true); blas.matmul('N', 'T', M, N, K, 1.f, dA, dBt, 0.f, dC); @@ -797,7 +687,7 @@ static void runHipTests() { alpaka::memcpy(queue, dAt, hAt); alpaka::wait(queue); blas.addLayoutConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, 'T', - 'N'); + 'N', 'b'); std::fill(ref.begin(), ref.end(), 0.f); refGemm(ref.data(), At, B, bias, M, N, K, 1.f, 0.f, true, false); blas.gemm('T', 'N', M, N, K, 1.f, dAt, dB, 0.f, dBias, dC); @@ -825,7 +715,7 @@ static void runHipTests() { alpaka::memcpy(queue, dBp, hBp); alpaka::memcpy(queue, dBiasz, hBiasz); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N'); + blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N', 'r'); std::fill(ref.begin(), ref.end(), 0.f); refGemmRelu(ref.data(), Ap, Bp, alpaka::getPtrNative(hBiasz), M, N, K, 1.f, 0.f, false, false); @@ -890,116 +780,6 @@ static void runHipTests() { } } -static void runHipDynamicShapeTests() { - std::cout << "\n=== HIP Dynamic-Shape Tests ===\n"; - - alpaka::PlatformHipRt platform{}; - auto dev = alpaka::getDevByIdx(platform, 0u); - alpaka::Queue queue{dev}; - sofieBLAS blas(queue); - - alpaka::PlatformCpu hostPlatform{}; - auto hostDev = alpaka::getDevByIdx(hostPlatform, 0u); - - constexpr int MCAP = 96, MENV = 64, N = 3, K = 5; - - auto hA = alpaka::allocBuf(hostDev, static_cast(MCAP * K)); - auto hB = alpaka::allocBuf(hostDev, static_cast(K * N)); - auto hC = alpaka::allocBuf(hostDev, static_cast(MCAP * N)); - float *A = alpaka::getPtrNative(hA); - float *B = alpaka::getPtrNative(hB); - float *C = alpaka::getPtrNative(hC); - fillSeq(A, MCAP * K, 0.5f, 0.25f); - fillSeq(B, K * N, 1.f, 0.5f); - - auto dA = - alpaka::allocAsyncBuf(queue, static_cast(MCAP * K)); - auto dB = alpaka::allocAsyncBuf(queue, static_cast(K * N)); - auto dC = - alpaka::allocAsyncBuf(queue, static_cast(MCAP * N)); - alpaka::memcpy(queue, dA, hA); - alpaka::memcpy(queue, dB, hB); - alpaka::wait(queue); - - blas.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), - MENV, 'N', 'N'); - - std::vector ref; - auto runAt = [&](int m, const std::string &name) { - ref.assign(static_cast(m) * N, 0.f); - refMatmul(ref.data(), A, B, m, N, K, 1.f, 0.f, false, false); - blas.matmul('N', 'N', static_cast(m), static_cast(N), - static_cast(K), 1.f, dA, dB, 0.f, dC); - alpaka::memcpy(queue, hC, dC); - alpaka::wait(queue); - checkClose(C, ref.data(), m * N, name); - }; - - for (int m : {MENV, 37, 8, 51, 1, MENV, MCAP}) - runAt(m, "hip::dynamic m=" + std::to_string(m)); - - const int nSizes = MENV - 1; - const std::size_t cacheBefore = blas.algoCacheSize(); - const std::size_t rejBefore = blas.layoutStats().envelopeRejects; - const std::size_t searchBefore = blas.layoutStats().heuristicQueries; - for (int m = 2; m <= MENV; ++m) - blas.matmul('N', 'N', static_cast(m), static_cast(N), - static_cast(K), 1.f, dA, dB, 0.f, dC); - alpaka::wait(queue); - const std::size_t added = blas.algoCacheSize() - cacheBefore; - const std::size_t rejected = blas.layoutStats().envelopeRejects - rejBefore; - - std::cout << " " << nSizes << " sizes added " << added - << " cache entries, " << rejected << " rejected\n"; - if (added < static_cast(nSizes)) { - std::cout << " PASS hip::cache bounded\n"; - } else { - std::cerr << " FAIL [hip::cache bounded] one entry per size\n"; - ++gFailures; - } - - { - sofieBLAS capped(queue, 8); - capped.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), - MENV, 'N', 'N'); - std::vector cref; - float worst = 0.f; - for (int m = MENV + 1; m <= MCAP; ++m) { - cref.assign(static_cast(m) * N, 0.f); - refMatmul(cref.data(), A, B, m, N, K, 1.f, 0.f, false, false); - capped.matmul('N', 'N', static_cast(m), - static_cast(N), static_cast(K), 1.f, dA, - dB, 0.f, dC); - alpaka::memcpy(queue, hC, dC); - alpaka::wait(queue); - for (std::size_t i = 0; i < cref.size(); ++i) - worst = std::max(worst, std::abs(C[i] - cref[i])); - } - std::cout << " limit=8: " << capped.algoCacheSize() << " entries, " - << capped.layoutStats().evictions << " evictions over " - << (MCAP - MENV) << " above-envelope sizes, worst err " << worst - << "\n"; - if (capped.algoCacheSize() <= 8 && worst < 1e-3f) { - std::cout << " PASS hip::cache limit honoured\n"; - } else { - std::cerr << " FAIL [hip::cache limit honoured] " - << capped.algoCacheSize() << " entries, worst err " << worst - << "\n"; - ++gFailures; - } - } - - const std::size_t searched = - blas.layoutStats().heuristicQueries - searchBefore; - if (searched <= rejected) { - std::cout << " PASS hip::no search during inference\n"; - } else { - std::cerr << " FAIL [hip::no search during inference] " << searched - << " searches, only " << rejected << " explained by rejects\n"; - ++gFailures; - } -} - #endif // ALPAKA_ACC_GPU_HIP_ENABLED // --------------------------------------------------------------------------- @@ -1012,11 +792,9 @@ int main() { #endif #ifdef ALPAKA_ACC_GPU_CUDA_ENABLED runCudaTests(); - runDynamicShapeTests(); #endif #ifdef ALPAKA_ACC_GPU_HIP_ENABLED runHipTests(); - runHipDynamicShapeTests(); #endif std::cout << "\n"; From 8b6c3e57864543e7b91cebdf2aff57901e4f9839 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Fri, 4 Sep 2026 15:58:49 +0530 Subject: [PATCH 12/13] feat: addOperationConfig with a epilogue, tests for runtime sizes --- benchmark/bench.cc | 4 +- .../backends/cuda/sofieBLAS_cublas.hpp | 4 +- .../gpu/detail/sofieBLAS_blaslt_common.tpp | 49 ++-- .../backends/hip/sofieBLAS_hipblaslt.hpp | 4 +- include/sofieBLAS/core.hpp | 2 + tests/test.cc | 220 ++++++++++++++++-- 6 files changed, 229 insertions(+), 54 deletions(-) diff --git a/benchmark/bench.cc b/benchmark/bench.cc index 8935856..687c312 100644 --- a/benchmark/bench.cc +++ b/benchmark/bench.cc @@ -139,7 +139,7 @@ static void runCudaBench(const BenchOptions &opt) { alpaka::memcpy(queue, dB, hB); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N', 'n'); + blas.addOperationConfig(M, N, K, M, K, M, 'N', 'N', Epilogue::Default); for (int i = 0; i < opt.warmup; ++i) blas.matmul('N', 'N', M, N, K, 1.f, dA, dB, 0.f, dC); @@ -187,7 +187,7 @@ static void runHipBench(const BenchOptions &opt) { alpaka::memcpy(queue, dB, hB); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N', 'n'); + blas.addOperationConfig(M, N, K, M, K, M, 'N', 'N', Epilogue::Default); for (int i = 0; i < opt.warmup; ++i) blas.matmul('N', 'N', M, N, K, 1.f, dA, dB, 0.f, dC); diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index 4717730..b7c0e06 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -87,8 +87,8 @@ struct CublasLtApi { static cudaError_t rtFree(void *ptr) { return cudaFree(ptr); } }; -#define SOFIEBLAS_CHECK_LT(x) CHECK_CUBLAS(x) -#define SOFIEBLAS_CHECK_RT(x) CHECK_CUDA(x) +#define SOFIEBLAS_CHECK_LT(status) CHECK_CUBLAS(status) +#define SOFIEBLAS_CHECK_RT(err) CHECK_CUDA(err) #include "../gpu/detail/sofieBLAS_blaslt_common.tpp" diff --git a/include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp b/include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp index 1e83b9c..b688b76 100644 --- a/include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp +++ b/include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp @@ -152,46 +152,35 @@ public: } } - // An epilogue is the extra step the library fuses into the multiply kernel - // after the product: nothing, adding the bias vector, or adding it and - // applying the activation. Which one a call site uses is decided by the - // function it calls (matmul, gemm, gemmrelu, gemmgelu); addLayoutConfig - // receives the same choice as a character so it can resolve the site's - // algorithm up front for the right configuration. - inline typename Api::Epilogue charToEpilogue(char epilogue) { - switch (epilogue) { - case 'N': - case 'n': - return Api::EpilogueDefault; - case 'B': - case 'b': - return Api::EpilogueBias; - case 'R': - case 'r': - return Api::EpilogueReluBias; - case 'G': - case 'g': - return Api::EpilogueGeluBias; - default: - throw std::invalid_argument( - std::string("Invalid epilogue character for ") + Api::name + "."); - } - } - // Registers a call site's construction-time shape: creates the three matrix // layouts and resolves the multiply algorithm for them up front, so the // first call at this shape finds everything cached. - void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, - std::size_t lda, std::size_t ldb, std::size_t ldc, - char transa, char transb, char epilogue) { + void addOperationConfig(std::size_t m, std::size_t n, std::size_t k, + std::size_t lda, std::size_t ldb, std::size_t ldc, + char transa, char transb, Epilogue epilogue) { const auto shapeA = layoutKeyA(transa, m, k); const auto shapeB = layoutKeyB(transb, k, n); const std::pair shapeC{m, n}; getOrCreateLayout(shapeA, lda); getOrCreateLayout(shapeB, ldb); getOrCreateLayout(shapeC, ldc); + + typename Api::Epilogue apiEpilogue = Api::EpilogueDefault; + switch (epilogue) { + case Epilogue::Bias: + apiEpilogue = Api::EpilogueBias; + break; + case Epilogue::ReluBias: + apiEpilogue = Api::EpilogueReluBias; + break; + case Epilogue::GeluBias: + apiEpilogue = Api::EpilogueGeluBias; + break; + case Epilogue::Default: + break; + } getOrComputeAlgo(charToTranspose(transa), charToTranspose(transb), - charToEpilogue(epilogue), shapeA, shapeB, shapeC); + apiEpilogue, shapeA, shapeB, shapeC); } // Each multiply variant comes as one generic overload, where A, B, bias and diff --git a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp index d876eea..7c379ce 100644 --- a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp +++ b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp @@ -88,8 +88,8 @@ struct HipblasLtApi { static hipError_t rtFree(void *ptr) { return hipFree(ptr); } }; -#define SOFIEBLAS_CHECK_LT(x) CHECK_HIPBLAS(x) -#define SOFIEBLAS_CHECK_RT(x) CHECK_HIP(x) +#define SOFIEBLAS_CHECK_LT(status) CHECK_HIPBLAS(status) +#define SOFIEBLAS_CHECK_RT(err) CHECK_HIP(err) #include "../gpu/detail/sofieBLAS_blaslt_common.tpp" diff --git a/include/sofieBLAS/core.hpp b/include/sofieBLAS/core.hpp index f2f3da9..1fb8890 100644 --- a/include/sofieBLAS/core.hpp +++ b/include/sofieBLAS/core.hpp @@ -6,3 +6,5 @@ template class sofieBLAS; template using sofieBLAS = typename traits::sofieBLAS::Impl; + +enum class Epilogue { Default, Bias, ReluBias, GeluBias }; diff --git a/tests/test.cc b/tests/test.cc index 06b27d1..512f59b 100644 --- a/tests/test.cc +++ b/tests/test.cc @@ -393,8 +393,8 @@ static void runCudaTests() { }; // ---- matmul NN ---- - blas.addLayoutConfig(M, N, K, ldaFor('N', M, K), ldbFor('N', K, N), M, 'N', - 'N', 'n'); + blas.addOperationConfig(M, N, K, ldaFor('N', M, K), ldbFor('N', K, N), M, 'N', + 'N', Epilogue::Default); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), A, B, M, N, K, 1.f, 0.f, false, false); blas.matmul('N', 'N', M, N, K, 1.f, dA, dB, 0.f, dC); @@ -409,8 +409,8 @@ static void runCudaTests() { alpaka::allocAsyncBuf(queue, static_cast(K * M)); alpaka::memcpy(queue, dAt, hAt); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, 'T', - 'N', 'n'); + blas.addOperationConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, + 'T', 'N', Epilogue::Default); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), At, B, M, N, K, 1.f, 0.f, true, false); blas.matmul('T', 'N', M, N, K, 1.f, dAt, dB, 0.f, dC); @@ -426,8 +426,8 @@ static void runCudaTests() { alpaka::allocAsyncBuf(queue, static_cast(N * K)); alpaka::memcpy(queue, dBt, hBt); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, ldaFor('N', M, K), ldbFor('T', K, N), M, 'N', - 'T', 'n'); + blas.addOperationConfig(M, N, K, ldaFor('N', M, K), ldbFor('T', K, N), M, + 'N', 'T', Epilogue::Default); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), A, Bt, M, N, K, 1.f, 0.f, false, true); blas.matmul('N', 'T', M, N, K, 1.f, dA, dBt, 0.f, dC); @@ -465,8 +465,8 @@ static void runCudaTests() { alpaka::allocAsyncBuf(queue, static_cast(K * M)); alpaka::memcpy(queue, dAt, hAt); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, 'T', - 'N', 'b'); + blas.addOperationConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, + 'T', 'N', Epilogue::Bias); std::fill(ref.begin(), ref.end(), 0.f); refGemm(ref.data(), At, B, bias, M, N, K, 1.f, 0.f, true, false); blas.gemm('T', 'N', M, N, K, 1.f, dAt, dB, 0.f, dBias, dC); @@ -494,7 +494,7 @@ static void runCudaTests() { alpaka::memcpy(queue, dBp, hBp); alpaka::memcpy(queue, dBiasz, hBiasz); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N', 'r'); + blas.addOperationConfig(M, N, K, M, K, M, 'N', 'N', Epilogue::ReluBias); std::fill(ref.begin(), ref.end(), 0.f); refGemmRelu(ref.data(), Ap, Bp, alpaka::getPtrNative(hBiasz), M, N, K, 1.f, 0.f, false, false); @@ -559,6 +559,97 @@ static void runCudaTests() { } } +static void runDynamicShapeTests() { + std::cout << "\n=== CUDA Dynamic-Shape Tests ===\n"; + + alpaka::PlatformCudaRt platform{}; + auto dev = alpaka::getDevByIdx(platform, 0u); + alpaka::Queue queue{dev}; + + alpaka::PlatformCpu hostPlatform{}; + auto hostDev = alpaka::getDevByIdx(hostPlatform, 0u); + + // M0 is the construction-time size given to addOperationConfig; the buffers + // hold MCAP rows so sizes above M0 are exercised too. + constexpr int MCAP = 96, M0 = 64, N = 3, K = 5; + + auto hA = alpaka::allocBuf(hostDev, static_cast(MCAP * K)); + auto hB = alpaka::allocBuf(hostDev, static_cast(K * N)); + auto hC = alpaka::allocBuf(hostDev, static_cast(MCAP * N)); + float *A = alpaka::getPtrNative(hA); + float *B = alpaka::getPtrNative(hB); + float *C = alpaka::getPtrNative(hC); + fillSeq(A, MCAP * K, 0.5f, 0.25f); + fillSeq(B, K * N, 1.f, 0.5f); + + auto dA = + alpaka::allocAsyncBuf(queue, static_cast(MCAP * K)); + auto dB = alpaka::allocAsyncBuf(queue, static_cast(K * N)); + auto dC = + alpaka::allocAsyncBuf(queue, static_cast(MCAP * N)); + alpaka::memcpy(queue, dA, hA); + alpaka::memcpy(queue, dB, hB); + alpaka::wait(queue); + + // One instance serving sizes never passed to addOperationConfig (issue #10), + // including m=1 and a size above the construction-time one. + sofieBLAS blas(queue); + blas.addOperationConfig(M0, N, K, ldaFor('N', M0, K), ldbFor('N', K, N), M0, + 'N', 'N', Epilogue::Default); + + std::vector ref; + auto runAt = [&](int m, const std::string &name) { + ref.assign(static_cast(m) * N, 0.f); + refMatmul(ref.data(), A, B, m, N, K, 1.f, 0.f, false, false); + blas.matmul('N', 'N', static_cast(m), static_cast(N), + static_cast(K), 1.f, dA, dB, 0.f, dC); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + checkClose(C, ref.data(), m * N, name); + }; + + for (int m : {M0, 37, 8, 51, 1, M0, MCAP}) + runAt(m, "cuda::dynamic m=" + std::to_string(m)); + + // Generated code calls the raw-pointer overloads; one call keeps them + // compiled and resolving to the right overload. + ref.assign(static_cast(45) * N, 0.f); + refMatmul(ref.data(), A, B, 45, N, K, 1.f, 0.f, false, false); + blas.matmul('N', 'N', 45u, static_cast(N), static_cast(K), + 1.f, alpaka::getPtrNative(dA), alpaka::getPtrNative(dB), 0.f, + alpaka::getPtrNative(dC)); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + checkClose(C, ref.data(), 45 * N, "cuda::dynamic raw pointers m=45"); + + // 32 distinct sizes through a cache limited to 8 entries. + { + sofieBLAS capped(queue, 8); + capped.addOperationConfig(M0, N, K, ldaFor('N', M0, K), ldbFor('N', K, N), + M0, 'N', 'N', Epilogue::Default); + float worst = 0.f; + for (int m = M0 + 1; m <= MCAP; ++m) { + ref.assign(static_cast(m) * N, 0.f); + refMatmul(ref.data(), A, B, m, N, K, 1.f, 0.f, false, false); + capped.matmul('N', 'N', static_cast(m), + static_cast(N), static_cast(K), 1.f, dA, + dB, 0.f, dC); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + for (std::size_t i = 0; i < ref.size(); ++i) + worst = std::max(worst, std::abs(C[i] - ref[i])); + } + if (capped.algoCacheSize() <= 8 && worst < 1e-3f) { + std::cout << " PASS cuda::cache limit honoured\n"; + } else { + std::cerr << " FAIL [cuda::cache limit honoured] " + << capped.algoCacheSize() << " entries, worst err " << worst + << "\n"; + ++gFailures; + } + } +} + #endif // ALPAKA_ACC_GPU_CUDA_ENABLED // --------------------------------------------------------------------------- @@ -614,8 +705,8 @@ static void runHipTests() { }; // ---- matmul NN ---- - blas.addLayoutConfig(M, N, K, ldaFor('N', M, K), ldbFor('N', K, N), M, 'N', - 'N', 'n'); + blas.addOperationConfig(M, N, K, ldaFor('N', M, K), ldbFor('N', K, N), M, 'N', + 'N', Epilogue::Default); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), A, B, M, N, K, 1.f, 0.f, false, false); blas.matmul('N', 'N', M, N, K, 1.f, dA, dB, 0.f, dC); @@ -630,8 +721,8 @@ static void runHipTests() { alpaka::allocAsyncBuf(queue, static_cast(K * M)); alpaka::memcpy(queue, dAt, hAt); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, 'T', - 'N', 'n'); + blas.addOperationConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, + 'T', 'N', Epilogue::Default); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), At, B, M, N, K, 1.f, 0.f, true, false); blas.matmul('T', 'N', M, N, K, 1.f, dAt, dB, 0.f, dC); @@ -647,8 +738,8 @@ static void runHipTests() { alpaka::allocAsyncBuf(queue, static_cast(N * K)); alpaka::memcpy(queue, dBt, hBt); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, ldaFor('N', M, K), ldbFor('T', K, N), M, 'N', - 'T', 'n'); + blas.addOperationConfig(M, N, K, ldaFor('N', M, K), ldbFor('T', K, N), M, + 'N', 'T', Epilogue::Default); std::fill(ref.begin(), ref.end(), 0.f); refMatmul(ref.data(), A, Bt, M, N, K, 1.f, 0.f, false, true); blas.matmul('N', 'T', M, N, K, 1.f, dA, dBt, 0.f, dC); @@ -686,8 +777,8 @@ static void runHipTests() { alpaka::allocAsyncBuf(queue, static_cast(K * M)); alpaka::memcpy(queue, dAt, hAt); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, 'T', - 'N', 'b'); + blas.addOperationConfig(M, N, K, ldaFor('T', M, K), ldbFor('N', K, N), M, + 'T', 'N', Epilogue::Bias); std::fill(ref.begin(), ref.end(), 0.f); refGemm(ref.data(), At, B, bias, M, N, K, 1.f, 0.f, true, false); blas.gemm('T', 'N', M, N, K, 1.f, dAt, dB, 0.f, dBias, dC); @@ -715,7 +806,7 @@ static void runHipTests() { alpaka::memcpy(queue, dBp, hBp); alpaka::memcpy(queue, dBiasz, hBiasz); alpaka::wait(queue); - blas.addLayoutConfig(M, N, K, M, K, M, 'N', 'N', 'r'); + blas.addOperationConfig(M, N, K, M, K, M, 'N', 'N', Epilogue::ReluBias); std::fill(ref.begin(), ref.end(), 0.f); refGemmRelu(ref.data(), Ap, Bp, alpaka::getPtrNative(hBiasz), M, N, K, 1.f, 0.f, false, false); @@ -780,6 +871,97 @@ static void runHipTests() { } } +static void runHipDynamicShapeTests() { + std::cout << "\n=== HIP Dynamic-Shape Tests ===\n"; + + alpaka::PlatformHipRt platform{}; + auto dev = alpaka::getDevByIdx(platform, 0u); + alpaka::Queue queue{dev}; + + alpaka::PlatformCpu hostPlatform{}; + auto hostDev = alpaka::getDevByIdx(hostPlatform, 0u); + + // M0 is the construction-time size given to addOperationConfig; the buffers + // hold MCAP rows so sizes above M0 are exercised too. + constexpr int MCAP = 96, M0 = 64, N = 3, K = 5; + + auto hA = alpaka::allocBuf(hostDev, static_cast(MCAP * K)); + auto hB = alpaka::allocBuf(hostDev, static_cast(K * N)); + auto hC = alpaka::allocBuf(hostDev, static_cast(MCAP * N)); + float *A = alpaka::getPtrNative(hA); + float *B = alpaka::getPtrNative(hB); + float *C = alpaka::getPtrNative(hC); + fillSeq(A, MCAP * K, 0.5f, 0.25f); + fillSeq(B, K * N, 1.f, 0.5f); + + auto dA = + alpaka::allocAsyncBuf(queue, static_cast(MCAP * K)); + auto dB = alpaka::allocAsyncBuf(queue, static_cast(K * N)); + auto dC = + alpaka::allocAsyncBuf(queue, static_cast(MCAP * N)); + alpaka::memcpy(queue, dA, hA); + alpaka::memcpy(queue, dB, hB); + alpaka::wait(queue); + + // One instance serving sizes never passed to addOperationConfig (issue #10), + // including m=1 and a size above the construction-time one. + sofieBLAS blas(queue); + blas.addOperationConfig(M0, N, K, ldaFor('N', M0, K), ldbFor('N', K, N), M0, + 'N', 'N', Epilogue::Default); + + std::vector ref; + auto runAt = [&](int m, const std::string &name) { + ref.assign(static_cast(m) * N, 0.f); + refMatmul(ref.data(), A, B, m, N, K, 1.f, 0.f, false, false); + blas.matmul('N', 'N', static_cast(m), static_cast(N), + static_cast(K), 1.f, dA, dB, 0.f, dC); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + checkClose(C, ref.data(), m * N, name); + }; + + for (int m : {M0, 37, 8, 51, 1, M0, MCAP}) + runAt(m, "hip::dynamic m=" + std::to_string(m)); + + // Generated code calls the raw-pointer overloads; one call keeps them + // compiled and resolving to the right overload. + ref.assign(static_cast(45) * N, 0.f); + refMatmul(ref.data(), A, B, 45, N, K, 1.f, 0.f, false, false); + blas.matmul('N', 'N', 45u, static_cast(N), static_cast(K), + 1.f, alpaka::getPtrNative(dA), alpaka::getPtrNative(dB), 0.f, + alpaka::getPtrNative(dC)); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + checkClose(C, ref.data(), 45 * N, "hip::dynamic raw pointers m=45"); + + // 32 distinct sizes through a cache limited to 8 entries. + { + sofieBLAS capped(queue, 8); + capped.addOperationConfig(M0, N, K, ldaFor('N', M0, K), ldbFor('N', K, N), + M0, 'N', 'N', Epilogue::Default); + float worst = 0.f; + for (int m = M0 + 1; m <= MCAP; ++m) { + ref.assign(static_cast(m) * N, 0.f); + refMatmul(ref.data(), A, B, m, N, K, 1.f, 0.f, false, false); + capped.matmul('N', 'N', static_cast(m), + static_cast(N), static_cast(K), 1.f, dA, + dB, 0.f, dC); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + for (std::size_t i = 0; i < ref.size(); ++i) + worst = std::max(worst, std::abs(C[i] - ref[i])); + } + if (capped.algoCacheSize() <= 8 && worst < 1e-3f) { + std::cout << " PASS hip::cache limit honoured\n"; + } else { + std::cerr << " FAIL [hip::cache limit honoured] " + << capped.algoCacheSize() << " entries, worst err " << worst + << "\n"; + ++gFailures; + } + } +} + #endif // ALPAKA_ACC_GPU_HIP_ENABLED // --------------------------------------------------------------------------- @@ -792,9 +974,11 @@ int main() { #endif #ifdef ALPAKA_ACC_GPU_CUDA_ENABLED runCudaTests(); + runDynamicShapeTests(); #endif #ifdef ALPAKA_ACC_GPU_HIP_ENABLED runHipTests(); + runHipDynamicShapeTests(); #endif std::cout << "\n"; From d6395751ba061e288e865c0370854dd65e744a01 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Fri, 4 Sep 2026 16:20:16 +0530 Subject: [PATCH 13/13] docs: describe addOperationConfig, the epilogue enum and the cache limit --- README.md | 42 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 191fe2a..1952f90 100644 --- a/README.md +++ b/README.md @@ -96,7 +96,47 @@ sofieBLAS blas(queue); blas.matmul('N', 'N', size, size, size, 1.0f, dA, dB, 0.0f, dC); ``` -The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addLayoutConfig` (used to pre-register cuBLASLt/hipBLASLt matrix layouts for a given shape before the first `matmul`/`gemm` call on that shape). +The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addOperationConfig` (creates the matrix layouts and resolves the multiply algorithm for a call site's shape ahead of its first call, see below). + +## Dynamic GEMM shapes and the algorithm cache + +A GEMM call computes `C = alpha * op(A) * op(B) + beta * C`, where A and B are the input matrices, C the output, and `op` an optional transpose. To run one, cuBLASLt and hipBLASLt need three kinds of objects besides the data: + +- a **matrix layout** per matrix: a descriptor holding its rows, columns and leading dimension; +- a **matmul descriptor**: the operation settings (the transposes and the epilogue); +- an **algorithm**: the concrete GEMM kernel the library selects for the given settings and dimensions, obtained by querying its heuristic (`cublasLtMatmulAlgoGetHeuristic` / `hipblasLtMatmulAlgoGetHeuristic`). The query runs on the host and is not free. + +The CUDA backend (`BlasCuda`, over cuBLASLt) and the HIP backend (`BlasHip`, over hipBLASLt) behave identically: all three objects are created the first time a combination appears and cached, keyed by the exact dimensions plus, for descriptors and algorithms, the transposes and the epilogue. One instance therefore serves GEMM calls at sizes that vary at runtime: a size seen for the first time creates and caches its objects, and a repeated size reuses them without another heuristic query. + +### Warming the cache with addOperationConfig + +`addOperationConfig(m, n, k, lda, ldb, ldc, transa, transb, epilogue)` fills the cache for one call site ahead of its first call: it creates the three layouts and resolves the algorithm for the given dimensions. The `epilogue` argument is the `Epilogue` enum from `sofieBLAS/core.hpp` and names which call the site will make, because the fused epilogue is part of the selected kernel: + +| `Epilogue` value | call it configures | +| --- | --- | +| `Epilogue::Default` | `matmul` (no bias) | +| `Epilogue::Bias` | `gemm` (adds the bias vector) | +| `Epilogue::ReluBias` | `gemmrelu` (bias, then ReLU) | +| `Epilogue::GeluBias` | `gemmgelu` (bias, then GELU) | + +A generated Session constructor calls it once per GEMM call site with the construction-time dimensions, so the first inference pays no heuristic queries at those sizes. + +### Initializing the cache limit + +The algorithm cache is unbounded by default. Passing a limit as the second constructor argument caps the number of cached algorithms; when an insertion would exceed the limit, the least recently used entries are evicted. Choose a limit at least as large as the number of distinct shapes the workload uses regularly, or leave it unbounded. `algoCacheSize()` returns the current number of entries. + +```cpp +sofieBLAS blas(queue); // unbounded algorithm cache (default) +sofieBLAS capped(queue, 32); // at most 32 entries, LRU eviction + +blas.addOperationConfig(64, 3, 5, 64, 5, 64, 'N', 'N', Epilogue::Default); +blas.matmul('N', 'N', 64, 3, 5, 1.0f, dA, dB, 0.0f, dC); // warmed: cache hit +blas.matmul('N', 'N', 37, 3, 5, 1.0f, dA, dB, 0.0f, dC); // new size: created on first use +``` + +### One implementation for both GPU backends + +The CUDA and HIP backends share one implementation: `include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp` contains the class template `BlasLt`, and each backend header defines a table with its library's types, constants and functions (`CublasLtApi` in `sofieBLAS_cublas.hpp`, `HipblasLtApi` in `sofieBLAS_hipblaslt.hpp`) and instantiates the template with it. This is the same arrangement the CPU backends use with `backends/cpu/detail/sofieBLAS_cblas_common.hpp`. ## Contributing