diff --git a/README.md b/README.md index 191fe2a..d7df695 100644 --- a/README.md +++ b/README.md @@ -96,7 +96,44 @@ 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` for batched gemm operations through strides +- `addOperationConfig` that 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. + +### addOperationConfig + +`addOperationConfig(m, n, k, lda, ldb, ldc, transa, transb, epilogue)` creates all three objects for one operation (the matrix layouts, the matmul descriptor and the algorithm) for the given dimensions, transposes and epilogue, before the corresponding call is made. It is optional: a combination that was never configured is created and cached on its first call. 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) | + +### 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); // created by addOperationConfig: cache hit +blas.matmul('N', 'N', 37, 3, 5, 1.0f, dA, dB, 0.0f, dC); // new size: created on first use +``` ## Contributing diff --git a/benchmark/bench.cc b/benchmark/bench.cc index 3263dab..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'); + 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'); + 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 ded6ba8..5b3b92a 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -5,7 +5,9 @@ #include #include #include +#include #include +#include #include #include @@ -30,453 +32,67 @@ } \ } 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; - 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 forwarding of the shared BlasLt implementation in +// backends/gpu/detail +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; // 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; - } -}; - -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; - - std::unordered_map, - cublasLtMatrixLayout_t, PairHash, PairEq> - layoutStore; - - std::unordered_map descStore; - - std::unordered_map - algoCache; - -public: - BlasCuda(const BlasCuda &) = delete; - BlasCuda &operator=(const BlasCuda &) = delete; - BlasCuda(BlasCuda &&) = delete; - BlasCuda &operator=(BlasCuda &&) = delete; - - BlasCuda(alpaka::QueueCudaRtNonBlocking &queue) : 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 &[key, layout] : layoutStore) - if (layout) - cublasLtMatrixLayoutDestroy(layout); - 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."); - } - } - - 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); - } - - 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}); - } - - 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}); - } +#define SOFIEBLAS_CHECK_LT(status) CHECK_CUBLAS(status) +#define SOFIEBLAS_CHECK_RT(err) CHECK_CUDA(err) - 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}); - } +#include "../gpu/detail/sofieBLAS_blaslt_common.tpp" - 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); - } - - 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); - } - } - - 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); - } - - cublasLtMatmulHeuristicResult_t & - getOrComputeAlgo(cublasOperation_t transA, cublasOperation_t transB, - cublasLtEpilogue_t epilogue, - const std::pair &kA, - const std::pair &kB, - const std::pair &kC) { - 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; - - auto &desc = getOrCreateDesc(transA, transB, epilogue); - cublasLtMatmulHeuristicResult_t h{}; - int returnedResults = 0; - CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic( - ltHandle, desc, layoutStore.at(kA), layoutStore.at(kB), - layoutStore.at(kC), layoutStore.at(kC), preference, 1, &h, - &returnedResults)); - if (returnedResults == 0) { - 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"; - exit(EXIT_FAILURE); - } - algoCache.emplace(key, h); - return algoCache.at(key); - } - - 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 &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( - desc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, - sizeof(bias_ptr))); - } - - 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), - &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..b688b76 --- /dev/null +++ b/include/sofieBLAS/backends/gpu/detail/sofieBLAS_blaslt_common.tpp @@ -0,0 +1,418 @@ +// 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 + "."); + } + } + + // 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 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), + apiEpilogue, 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 0e41e70..4e178df 100644 --- a/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp +++ b/include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp @@ -5,7 +5,9 @@ #include #include #include +#include #include +#include #include #include @@ -30,446 +32,68 @@ } \ } 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; - 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 forwarding of the shared BlasLt implementation in +// backends/gpu/detail +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; // 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; - } -}; - -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; - - std::unordered_map, - hipblasLtMatrixLayout_t, PairHash, PairEq> - layoutStore; - - std::unordered_map descStore; - - std::unordered_map - algoCache; +#define SOFIEBLAS_CHECK_LT(status) CHECK_HIPBLAS(status) +#define SOFIEBLAS_CHECK_RT(err) CHECK_HIP(err) -public: - BlasHip(const BlasHip &) = delete; - BlasHip &operator=(const BlasHip &) = delete; - BlasHip(BlasHip &&) = delete; - BlasHip &operator=(BlasHip &&) = delete; - - BlasHip(alpaka::QueueHipRtNonBlocking &queue) : 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 &[key, layout] : layoutStore) - if (layout) - hipblasLtMatrixLayoutDestroy(layout); - 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."); - } - } +#include "../gpu/detail/sofieBLAS_blaslt_common.tpp" - 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); - } - - 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}); - } - - 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); - } - - 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); - } - - 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); - } - } - - 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); - } - - hipblasLtMatmulHeuristicResult_t & - getOrComputeAlgo(hipblasOperation_t transA, hipblasOperation_t transB, - hipblasLtEpilogue_t epilogue, - const std::pair &kA, - const std::pair &kB, - const std::pair &kC) { - 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; - - auto &desc = getOrCreateDesc(transA, transB, epilogue); - 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)); - if (returnedResults == 0) { - 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"; - exit(EXIT_FAILURE); - } - algoCache.emplace(key, h); - return algoCache.at(key); - } - - 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 &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( - desc, HIPBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, - 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)); - } -}; +using BlasHip = BlasLt; namespace traits { 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 11d6e4c..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'); + 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'); + 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'); + 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'); + 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'); + 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'); + 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'); + 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'); + 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'); + 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'); + 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";