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 ded6ba8..4717730 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 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; // 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(x) CHECK_CUBLAS(x) +#define SOFIEBLAS_CHECK_RT(x) CHECK_CUDA(x) - 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..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 0e41e70..d876eea 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 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; // 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(x) CHECK_HIPBLAS(x) +#define SOFIEBLAS_CHECK_RT(x) CHECK_HIP(x) -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/tests/test.cc b/tests/test.cc index 11d6e4c..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); @@ -615,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); @@ -631,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); @@ -648,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); @@ -687,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); @@ -715,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);