From 36318392448a8dca2e000131904b3aaf4ab25d64 Mon Sep 17 00:00:00 2001 From: Tony Davis Date: Fri, 9 Oct 2026 15:18:07 -0400 Subject: [PATCH] fix(hipblaslt): clamp workspace limits above 4 GiB instead of rejecting #12856 made any workspace limit or size above UINT32_MAX return HIPBLAS_STATUS_INVALID_VALUE. Callers pass SIZE_MAX to mean "no limit", so driftcheck now fails at hipblasLtMatmulPreferenceSetAttribute before any GEMM runs. A limit is a ceiling, so clamp it to UINT32_MAX at the same four entry points: the MAX_WORKSPACE_BYTES preference, hipblasLtMatmul, and the extension API's algoGetHeuristic and initialize. Kernels still never see more than a 32-bit range, and the int truncation fix from #12856 stays. The clamp is logged at the hints level. The reject test becomes two clamp tests. One checks the preference and needs no solution library. The other compares heuristic results at 4 GiB and SIZE_MAX against UINT32_MAX, and runs hipblasLtMatmul and extension initialize with an oversized workspace size. JIRA ID : ROCM-32535 Co-authored-by: Cursor --- projects/hipblaslt/CHANGELOG.md | 1 + .../tests/src/heuristic_workspace_gtest.cpp | 61 +++++++++++++++---- .../include/hipblaslt/hipblaslt-ext.hpp | 8 +-- .../library/include/hipblaslt/hipblaslt.h | 8 +-- .../src/include/rocblaslt_mat_utils.hpp | 16 ++--- .../rocblaslt/src/rocblaslt_auxiliary.cpp | 16 ++--- .../rocblaslt/src/rocblaslt_mat.cpp | 9 +-- 7 files changed, 73 insertions(+), 46 deletions(-) diff --git a/projects/hipblaslt/CHANGELOG.md b/projects/hipblaslt/CHANGELOG.md index bf85ce68a856..d59440bf5ffc 100644 --- a/projects/hipblaslt/CHANGELOG.md +++ b/projects/hipblaslt/CHANGELOG.md @@ -35,6 +35,7 @@ Full documentation for hipBLASLt is available at [rocm.docs.amd.com/projects/hip * Solution cache key now includes `HIPBLASLT_MATMUL_DESC_SM_COUNT_TARGET` and the StreamK tile scheduling mode, so the same problem can select a different kernel than before. * A tuning file is now trusted one entry at a time instead of all or nothing. `HIPBLASLT_TUNING_OVERRIDE_FILE` records solution indices, which are positions in one build's kernel library. Previously the C API ignored the whole file when its build-version line did not match the running build, while the C++ API applied it regardless and could run kernels it was never tuned on. `hipblaslt-bench` now records a `kernel_name` beside each `solution_index` in `HIPBLASLT_TUNING_FILE`, and on both APIs a row that records a name is checked at replay and dropped only if its index no longer names that kernel. A row without a name is used only when the file's `Git Version` line matches the running build; a build made outside a git checkout has no version, so it uses no such rows. Problems whose rows are dropped fall back to normal kernel selection. * Extension API `initialize` compares against the size set with `setMaxWorkspaceBytes` on the `Gemm` or `GroupedGemm` instance, which defaults to 0. A caller that passes a large enough buffer without setting it is now rejected with `HIPBLAS_STATUS_INVALID_VALUE` for solutions with a fixed workspace requirement, such as GSU. Stream-K and zero-workspace solutions are unaffected. +* Workspace limits and sizes above `UINT32_MAX` (4 GiB - 1), such as `SIZE_MAX` for no limit, are clamped to `UINT32_MAX`, because kernels address the workspace with 32-bit offsets. This applies to `HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES` (which then reads back as `UINT32_MAX`), the `hipblasLtMatmul` workspace size, and the extension API `setMaxWorkspaceBytes`. ### Removed diff --git a/projects/hipblaslt/clients/tests/src/heuristic_workspace_gtest.cpp b/projects/hipblaslt/clients/tests/src/heuristic_workspace_gtest.cpp index e6c34bb978e5..27cc2204b477 100644 --- a/projects/hipblaslt/clients/tests/src/heuristic_workspace_gtest.cpp +++ b/projects/hipblaslt/clients/tests/src/heuristic_workspace_gtest.cpp @@ -2,7 +2,8 @@ // SPDX-License-Identifier: MIT // The C++ heuristic query must apply a workspace limit from 2 GiB to UINT32_MAX as the C query -// does. Kernels address the workspace with 32-bit offsets, so both APIs reject anything larger. +// does. Kernels address the workspace with 32-bit offsets, so both APIs clamp anything larger to +// UINT32_MAX. Callers pass SIZE_MAX to mean "no limit", so a larger value must not be an error. #include #include @@ -40,7 +41,7 @@ namespace return out; } - // The buffers are never read: only the heuristic queries run. + // The buffers hold no data: the tests check statuses and heuristic results only. class HeuristicWorkspaceLimit : public ::testing::Test { protected: @@ -173,19 +174,52 @@ namespace } } - // Every entry point that takes a workspace size or limit rejects it before selecting or - // launching a kernel, so no solution library is needed. - TEST_F(HeuristicWorkspaceLimit, smoke_RejectsWorkspaceAboveUint32Max) + // The preference stores a limit above UINT32_MAX as UINT32_MAX, so no solution library is + // needed. + TEST_F(HeuristicWorkspaceLimit, smoke_PreferenceClampsLimitAboveUint32Max) { for(uint64_t limit : {4 * kGiB, std::numeric_limits::max()}) { SCOPED_TRACE(limit); - EXPECT_EQ(setLimit(limit), HIPBLAS_STATUS_INVALID_VALUE); + hipblasLtMatmulPreference_t pref = nullptr; + ASSERT_EQ(hipblasLtMatmulPreferenceCreate(&pref), HIPBLAS_STATUS_SUCCESS); + EXPECT_EQ(hipblasLtMatmulPreferenceSetAttribute( + pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &limit, sizeof(limit)), + HIPBLAS_STATUS_SUCCESS); + uint64_t stored = 0; + size_t written = 0; + EXPECT_EQ( + hipblasLtMatmulPreferenceGetAttribute(pref, + HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &stored, + sizeof(stored), + &written), + HIPBLAS_STATUS_SUCCESS); + EXPECT_EQ(stored, kUint32Max); + hipblasLtMatmulPreferenceDestroy(pref); + } + } + + // Every entry point that takes a workspace size or limit treats a value above UINT32_MAX + // as UINT32_MAX. + TEST_F(HeuristicWorkspaceLimit, smoke_ClampsWorkspaceAboveUint32Max) + { + // Results at the cap are the reference. Solutions that fit the real buffer are run. + std::vector capped, fitting; + if(cHeuristic(kUint32Max, capped) != HIPBLAS_STATUS_SUCCESS || capped.empty() + || cHeuristic(kWorkspace, fitting) != HIPBLAS_STATUS_SUCCESS || fitting.empty()) + GTEST_SKIP() << "No solution for this problem in the loaded library"; - std::vector results; - EXPECT_EQ(cppHeuristic(limit, results), HIPBLAS_STATUS_INVALID_VALUE); - EXPECT_TRUE(results.empty()); + for(uint64_t limit : {4 * kGiB, std::numeric_limits::max()}) + { + SCOPED_TRACE(limit); + std::vector c, cpp; + ASSERT_EQ(cHeuristic(limit, c), HIPBLAS_STATUS_SUCCESS); + EXPECT_EQ(summary(c), summary(capped)); + ASSERT_EQ(cppHeuristic(limit, cpp), HIPBLAS_STATUS_SUCCESS); + EXPECT_EQ(summary(cpp), summary(capped)); + // The algo's own workspace limit keeps the kernel inside the real buffer. EXPECT_EQ(hipblasLtMatmul(handle, desc, &alpha, @@ -198,11 +232,12 @@ namespace layoutD, d, layoutD, - nullptr, + &fitting.front().algo, ws, limit, nullptr), - HIPBLAS_STATUS_INVALID_VALUE); + HIPBLAS_STATUS_SUCCESS); + EXPECT_EQ(hipStreamSynchronize(nullptr), hipSuccess); hipblaslt_ext::Gemm gemm(handle, HIPBLAS_OP_N, @@ -214,8 +249,8 @@ namespace HIPBLAS_COMPUTE_32F); ASSERT_EQ(setProblem(gemm), HIPBLAS_STATUS_SUCCESS); gemm.setMaxWorkspaceBytes(limit); - hipblasLtMatmulAlgo_t algo{}; - EXPECT_EQ(gemm.initialize(algo, ws), HIPBLAS_STATUS_INVALID_VALUE); + hipblasLtMatmulAlgo_t algo = fitting.front().algo; + EXPECT_EQ(gemm.initialize(algo, ws), HIPBLAS_STATUS_SUCCESS); } } } // namespace diff --git a/projects/hipblaslt/library/include/hipblaslt/hipblaslt-ext.hpp b/projects/hipblaslt/library/include/hipblaslt/hipblaslt-ext.hpp index 51a128ab5af4..6ca7ebc68783 100644 --- a/projects/hipblaslt/library/include/hipblaslt/hipblaslt-ext.hpp +++ b/projects/hipblaslt/library/include/hipblaslt/hipblaslt-ext.hpp @@ -74,8 +74,8 @@ namespace hipblaslt_ext * \brief This function sets the maximum workspace size. * * @param[in] - * workspaceBytes Set the maximum workspace size in bytes. Above ``UINT32_MAX`` - * (4 GiB - 1), algoGetHeuristic returns HIPBLAS_STATUS_INVALID_VALUE. + * workspaceBytes Set the maximum workspace size in bytes. algoGetHeuristic clamps + * values above ``UINT32_MAX`` (4 GiB - 1) to ``UINT32_MAX``. */ HIPBLASLT_EXPORT void setMaxWorkspaceBytes(size_t workspaceBytes); @@ -490,8 +490,8 @@ namespace hipblaslt_ext * \brief This function sets the maximum workspace size. * * @param[in] - * workspaceBytes Sets the maximum workspace size in bytes. Above ``UINT32_MAX`` - * (4 GiB - 1), initialize returns HIPBLAS_STATUS_INVALID_VALUE. + * workspaceBytes Sets the maximum workspace size in bytes. initialize clamps + * values above ``UINT32_MAX`` (4 GiB - 1) to ``UINT32_MAX``. */ HIPBLASLT_EXPORT void setMaxWorkspaceBytes(size_t workspaceBytes); diff --git a/projects/hipblaslt/library/include/hipblaslt/hipblaslt.h b/projects/hipblaslt/library/include/hipblaslt/hipblaslt.h index bb77848eaea1..d939a3656c01 100644 --- a/projects/hipblaslt/library/include/hipblaslt/hipblaslt.h +++ b/projects/hipblaslt/library/include/hipblaslt/hipblaslt.h @@ -384,7 +384,7 @@ typedef enum { */ typedef enum { HIPBLASLT_MATMUL_PREF_SEARCH_MODE = 0, /** /******************************************************************************* - * Validate Workspace Size + * Clamp Workspace Size * Kernels and helper kernels address the workspace with 32-bit byte offsets - * and buffer sizes, so a larger workspace cannot be used. + * and buffer sizes, so a larger workspace cannot be used. A workspace size or + * limit is a ceiling, so larger values (e.g. SIZE_MAX for "no limit") are + * clamped. ******************************************************************************/ constexpr size_t rocblaslt_max_workspace_bytes = std::numeric_limits::max(); -inline rocblaslt_status validateWorkspaceSize(const char* func, size_t workspaceBytes) +inline size_t clampWorkspaceSize(const char* func, size_t workspaceBytes) { if(workspaceBytes <= rocblaslt_max_workspace_bytes) - return rocblaslt_status_success; - log_error(func, + return workspaceBytes; + log_hints(func, "workspace size", workspaceBytes, - "exceeds the maximum of", + "clamped to the maximum of", rocblaslt_max_workspace_bytes, "bytes"); - return rocblaslt_status_invalid_value; + return rocblaslt_max_workspace_bytes; } inline bool isValidOrderForDatatype(hipDataType datatype, hipblasLtOrder_t order) diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp index dfd0cb283f18..94cd6b5fa0a7 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp @@ -2124,10 +2124,7 @@ rocblaslt_status pref->search_mode); break; case ROCBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES: - if(auto status = validateWorkspaceSize(__func__, *(uint64_t*)data); - status != rocblaslt_status_success) - return status; - pref->max_workspace_bytes = *(uint64_t*)data; + pref->max_workspace_bytes = clampWorkspaceSize(__func__, *(uint64_t*)data); log_api(__func__, "matmulPref", pref, @@ -2655,9 +2652,7 @@ rocblaslt_status log_error(__func__, "invalid requested count", requestedAlgoCount); return rocblaslt_status_invalid_value; } - if(auto status = validateWorkspaceSize(__func__, maxWorkspaceBytes); - status != rocblaslt_status_success) - return status; + const size_t workspaceLimit = clampWorkspaceSize(__func__, maxWorkspaceBytes); if(gemmType == rocblaslt::RocGemmType::ROCBLASLT_GROUPED_GEMM) { log_api( @@ -2685,7 +2680,7 @@ rocblaslt_status if(override.env_mode) { override_success = problem_override_from_file_cpp( - handle, gemmType, gemmData, override_result, override.file_path, maxWorkspaceBytes); + handle, gemmType, gemmData, override_result, override.file_path, workspaceLimit); log_api(__func__, "OverrideAlgoCount", override_success ? 1 : 0); } @@ -2695,7 +2690,7 @@ rocblaslt_status = getBestSolutions(handle, gemmType, gemmData, - maxWorkspaceBytes, + workspaceLimit, override_success ? requestedAlgoCount - 1 : requestedAlgoCount, results); @@ -2723,8 +2718,7 @@ rocblaslt_status std::vector allSolutionsResults; size_t workspaceSizeInBytes = 0; if(rocblaslt_status_success - == getAllSolutions( - gemmData, handle, gemmType, allSolutionsResults, maxWorkspaceBytes)) + == getAllSolutions(gemmData, handle, gemmType, allSolutionsResults, workspaceLimit)) { status = rocblaslt_status_success; int oriReturnAlgoCount = results.size(); diff --git a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp index 749ca8d1a25f..ce859c8c3541 100644 --- a/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp +++ b/projects/hipblaslt/library/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp @@ -829,9 +829,7 @@ rocblaslt_status rocblaslt_matmul(rocblaslt_handle handle, log_error(__func__, "invalid workspace pointer"); return rocblaslt_status_invalid_pointer; } - if(auto status = validateWorkspaceSize(__func__, workspaceSizeInBytes); - status != rocblaslt_status_success) - return status; + workspaceSizeInBytes = clampWorkspaceSize(__func__, workspaceSizeInBytes); if(matC->type != matD->type) { @@ -1617,15 +1615,12 @@ rocblaslt_status rocblaslt_makeArgument_cpp(rocblaslt_handle handle hipStream_t stream, std::shared_ptr gemmData) { - if(auto status = validateWorkspaceSize(__func__, workspaceSizeInBytes); - status != rocblaslt_status_success) - return status; return makeArgument(handle, gemmType, algo, tuning, workspace, - workspaceSizeInBytes, + clampWorkspaceSize(__func__, workspaceSizeInBytes), useUserArgs, stream, gemmData);