feat(hipblaslt): re-select the top solution when the adaptive hint changes

projects/hipblaslt/CHANGELOG.md

@@ -26,6 +26,7 @@ Full documentation for hipBLASLt is available at [rocm.docs.amd.com/projects/hip
 * Opt-in runtime tuning cache, off by default. `HIPBLASLT_TUNING_MODE=off|cache|tune` and `HIPBLASLT_TUNING_CACHE_PATH=<file>` let the library benchmark candidate kernels for an unseen problem at its first `hipblasLtMatmul`, record the winner, and replay it afterwards. `tune` measures candidates the way `hipblaslt-bench` does (warm-up and timed launches over rotating buffers, with an instruction-cache flush), serialises tuning within a process, and bounds each shape's search to five minutes by default (`HIPBLASLT_TUNING_BUDGET_MS_PER_SHAPE`). A search the budget stops records its best candidate as incomplete, which is never slower than the untuned kernel because that kernel is measured first, and tune mode revisits it when a later run can finish or widen the search. `cache` only replays, validating each entry as `HIPBLASLT_TUNING_OVERRIDE_FILE` does, and also serves `hipblasLtMatmul` calls that pass no `algo`. An explicit `algo` always launches as given. Both modes report what they loaded, each tuned shape's start and result, and a closing summary without needing a log level. With a tuning mode set, `HIPBLASLT_TUNING_OVERRIDE_FILE` is ignored, and a process in a secure execution context ignores every tuning variable.
 * `TENSILELITE_EMIT_OCCUPANCY_PROBE` CMake option and `invoke build --emit-occupancy-probe` (off by default) to emit a CU-occupancy probe into gfx950 Stream-K Hybrid kernels, needed by `HIPBLASLT_ADAPTIVE_SM_COUNT`.
 * Adaptive `sm_count_target` mode, off by default. With `HIPBLASLT_ADAPTIVE_SM_COUNT=1`, `hipblasLtMatmul` periodically measures the CUs available to a stream with the occupancy probe and uses the count as an implicit `sm_count_target`; `HIPBLASLT_ADAPTIVE_SM_COUNT_*` variables tune or force it. In this mode the top heuristic result reports `workspaceSize = min(MAX_WORKSPACE_BYTES, max(required, 128 MiB))`, and `hipblaslt-bench`, `hipblaslt-sequence`, and `hipblaslt-api-overhead` refuse to time unless `HIPBLASLT_ADAPTIVE_SM_COUNT_FORCE` is set.
+* In adaptive `sm_count_target` mode, `hipblasLtMatmul` re-selects the top heuristic result at launch when the adaptive hint differs from the one it was ranked under and the new solution fits the workspace. A result from `HIPBLASLT_TUNING_OVERRIDE_FILE` or the tuning cache, or a launch on a capturing stream, is never re-selected.
 
 ### Changed
 

projects/hipblaslt/clients/tests/src/adaptive_sm_count_gtest.cpp

@@ -1,11 +1,12 @@
 // Copyright Advanced Micro Devices, Inc., or its affiliates.
 // SPDX-License-Identifier: MIT
 
-// Adaptive sm_count_target through the C API: the heuristic-query workspace,
-// launches under a forced hint, probe readings, graph capture and the stream
-// lock. The mode is read from the environment once per process, so
-// the cases for a mode the environment does not select skip; run with
-// HIPBLASLT_ADAPTIVE_SM_COUNT=1 or HIPBLASLT_ADAPTIVE_SM_COUNT_FORCE=<cus> (and
+// Adaptive sm_count_target through the C API: heuristic-query tagging and
+// workspace, launches under a forced hint, probe readings, graph capture, the
+// stream lock and launch-time re-selection. The mode is read from the
+// environment once per process, so the cases for a mode the environment does
+// not select skip; run with HIPBLASLT_ADAPTIVE_SM_COUNT=1 or
+// HIPBLASLT_ADAPTIVE_SM_COUNT_FORCE=<cus> (and
 // HIPBLASLT_ADAPTIVE_SM_COUNT_FORCE_PROBE=1) to cover them. The cotenant cases
 // are Linux only; the exact one is manual (HIPBLASLT_TEST_COTENANT_PROBE=1,
 // idle GPU). Host-only cases are in
@@ -37,6 +38,8 @@ extern "C" bool hipblaslt_adaptive_sm_state_for_test(void*       handle,
                                                      size_t      size);
 extern "C" bool hipblaslt_adaptive_sm_launch_lock_for_test(
     void* handle, hipStream_t stream, bool tiny, bool* probe, bool* heldBetween, bool* heldAfter);
+extern "C" size_t hipblaslt_adaptive_reselections_for_test(void* matmulDesc);
+extern "C" int    hipblaslt_tuning_last_launch_for_test();
 
 namespace rocblaslt::adaptive_sm
 {
@@ -165,7 +168,13 @@ namespace
             return st;
         }
 
-        // Under FORCE every launch takes the forced hint. Probes run only with
+        size_t reselections()
+        {
+            return hipblaslt_adaptive_reselections_for_test(r.desc);
+        }
+
+        // Under FORCE the query is ranked and tagged for the forced hint and
+        // every launch takes it, so nothing re-selects. Probes run only with
         // FORCE_PROBE, and their readings never move the hint.
         void checkForcedHint(int32_t sk)
         {
@@ -176,22 +185,43 @@ namespace
             if(results.empty())
                 GTEST_SKIP() << "No solution for this problem in the loaded library";
             const auto& result = results[0];
+            uint32_t    tag    = 0;
+            ASSERT_TRUE(algoTag(result.algo.data, &tag));
+            EXPECT_EQ(tag, cfg.force);
 
             for(uint32_t i = 0; i < cfg.probeEvery + 1; ++i)
                 ASSERT_EQ(run(result.algo, result.workspaceSize), HIPBLAS_STATUS_SUCCESS);
             ASSERT_EQ(hipStreamSynchronize(stream), hipSuccess);
             ASSERT_EQ(run(result.algo, result.workspaceSize), HIPBLAS_STATUS_SUCCESS);
+            EXPECT_EQ(reselections(), 0u);
             const StreamStateForTest st = state();
             EXPECT_EQ(st.published, 0u) << "nothing is published under FORCE";
             if(cfg.forceProbe)
                 EXPECT_GT(st.epoch, 0u);
             else
                 EXPECT_EQ(st.epoch, 0u);
         }
+
+        // The top result retagged as if ranked for `hint`; skips without
+        // FORCE or a solution.
+        void retaggedResult(uint32_t hint, hipblasLtMatmulHeuristicResult_t* out)
+        {
+            if(!cfg.force)
+                GTEST_SKIP() << "set HIPBLASLT_ADAPTIVE_SM_COUNT_FORCE";
+            const auto results = query();
+            if(results.empty())
+                GTEST_SKIP() << "No solution for this problem in the loaded library";
+            *out         = results[0];
+            uint32_t tag = 0;
+            ASSERT_TRUE(algoTag(out->algo.data, &tag));
+            ASSERT_NE(hint, cfg.force);
+            tagAlgo(out->algo.data, hint);
+        }
     };
 
     // Adaptive mode reports min(limit, max(plain, c_adaptiveWorkspace)) for the
-    // top result; with it off the result reports the plain size.
+    // tagged top result; with it off the result is untagged and reports the
+    // plain size.
     TEST_F(AdaptiveSmCount_pre_checkin, QueryReportsTheAdaptiveWorkspace)
     {
         for(uint64_t limit : {kWsBudgetBytes / 4, kWsBudgetBytes * 2})
@@ -200,8 +230,10 @@ namespace
             const auto results = query(1, limit);
             if(results.empty())
                 GTEST_SKIP() << "No solution for this problem in the loaded library";
-            const auto&  result = results[0];
-            const size_t plain  = required(result.algo);
+            const auto& result = results[0];
+            uint32_t    tag    = 0;
+            EXPECT_EQ(algoTag(result.algo.data, &tag), cfg.enabled);
+            const size_t plain = required(result.algo);
             ASSERT_LE(plain, limit);
             const size_t expected
                 = cfg.enabled ? std::min<size_t>(limit, std::max(plain, c_adaptiveWorkspace))
@@ -210,6 +242,19 @@ namespace
         }
     }
 
+    // StreamK OFF never takes the adaptive hint.
+    TEST_F(AdaptiveSmCount_pre_checkin, StreamKOffIsUntagged)
+    {
+        if(!cfg.enabled)
+            GTEST_SKIP() << "adaptive mode is off";
+        setStreamK(HIPBLASLT_STREAMK_TILE_SCHEDULING_OFF);
+        const auto results = query();
+        if(results.empty())
+            GTEST_SKIP() << "No solution for this problem in the loaded library";
+        uint32_t tag = 0;
+        EXPECT_FALSE(algoTag(results[0].algo.data, &tag));
+    }
+
     TEST_F(AdaptiveSmCount_pre_checkin, ForcedHintForQueryAndLaunchAuto)
     {
         checkForcedHint(HIPBLASLT_STREAMK_TILE_SCHEDULING_AUTO);
@@ -482,4 +527,136 @@ namespace
             EXPECT_FALSE(heldAfter);
         }
     }
+
+    // Under FORCE a result tagged with another hint is re-selected at launch,
+    // once: the repeat launch reuses the memoised choice.
+    TEST_F(AdaptiveSmCount_pre_checkin, LaunchReselectsOnceForAnotherHint)
+    {
+        hipblasLtMatmulHeuristicResult_t result{};
+        retaggedResult(cfg.force ^ 1, &result);
+        if(IsSkipped() || HasFatalFailure())
+            return;
+        EXPECT_EQ(reselections(), 0u);
+        ASSERT_EQ(run(result.algo, result.workspaceSize), HIPBLAS_STATUS_SUCCESS);
+        EXPECT_EQ(reselections(), 1u);
+        ASSERT_EQ(run(result.algo, result.workspaceSize), HIPBLAS_STATUS_SUCCESS);
+        EXPECT_EQ(reselections(), 1u);
+    }
+
+    // A result ranked at the full device and tagged as an adaptive query with
+    // no hint is swapped for the solution the heuristic picks at the forced
+    // hint, and that solution is the one that launches. Skips when the pick
+    // does not change at this hint.
+    TEST_F(AdaptiveSmCount_pre_checkin, LaunchSwapsToTheSolutionPickedAtTheLaunchHint)
+    {
+        if(!cfg.force)
+            GTEST_SKIP() << "set HIPBLASLT_ADAPTIVE_SM_COUNT_FORCE";
+        auto setTarget = [&](int32_t target) {
+            ASSERT_EQ(hipblasLtMatmulDescSetAttribute(
+                          r.desc, HIPBLASLT_MATMUL_DESC_SM_COUNT_TARGET, &target, sizeof(target)),
+                      HIPBLAS_STATUS_SUCCESS);
+        };
+        // An explicit descriptor hint disables adaptive mode, so this query is
+        // untagged and ranked at the full device.
+        setTarget(1024);
+        if(HasFatalFailure())
+            return;
+        const auto atFull = query();
+        setTarget(0);
+        if(HasFatalFailure())
+            return;
+        const auto atLaunch = query();
+        if(atFull.empty() || atLaunch.empty())
+            GTEST_SKIP() << "No solution for this problem in the loaded library";
+        hipblasLtMatmulHeuristicResult_t result   = atFull[0];
+        hipblasLtMatmulHeuristicResult_t pick     = atLaunch[0];
+        const int                        tagged   = hipblaslt_ext::getIndexFromAlgo(result.algo);
+        const int                        expected = hipblaslt_ext::getIndexFromAlgo(pick.algo);
+        if(tagged == expected)
+            GTEST_SKIP() << "the heuristic picks solution " << tagged << " at both hints";
+
+        uint32_t tag = 0;
+        ASSERT_FALSE(algoTag(result.algo.data, &tag));
+        tagAlgo(result.algo.data, 0);
+        static_cast<void>(hipblaslt_tuning_last_launch_for_test());
+        ASSERT_EQ(run(result.algo, kWsBudgetBytes), HIPBLAS_STATUS_SUCCESS);
+        EXPECT_EQ(reselections(), 1u);
+        EXPECT_EQ(hipblaslt_tuning_last_launch_for_test(), expected);
+        EXPECT_NE(expected, tagged);
+    }
+
+    // A launch on a capturing stream runs a result tagged with another hint
+    // as given; the next live launch re-selects it.
+    TEST_F(AdaptiveSmCount_pre_checkin, CapturedLaunchDoesNotReselect)
+    {
+        hipblasLtMatmulHeuristicResult_t result{};
+        retaggedResult(cfg.force ^ 1, &result);
+        if(IsSkipped() || HasFatalFailure())
+            return;
+        ASSERT_EQ(hipStreamBeginCapture(stream, hipStreamCaptureModeGlobal), hipSuccess);
+        const hipblasStatus_t captured = run(result.algo, result.workspaceSize);
+        hipGraph_t            graph    = nullptr;
+        ASSERT_EQ(hipStreamEndCapture(stream, &graph), hipSuccess);
+        hipGraphExec_t   exec = nullptr;
+        const hipError_t inst = hipGraphInstantiate(&exec, graph, nullptr, nullptr, 0);
+        if(inst == hipSuccess)
+        {
+            EXPECT_EQ(hipGraphLaunch(exec, stream), hipSuccess);
+            EXPECT_EQ(hipStreamSynchronize(stream), hipSuccess);
+            static_cast<void>(hipGraphExecDestroy(exec));
+        }
+        static_cast<void>(hipGraphDestroy(graph));
+        ASSERT_EQ(captured, HIPBLAS_STATUS_SUCCESS);
+        ASSERT_EQ(inst, hipSuccess);
+        EXPECT_EQ(reselections(), 0u);
+
+        ASSERT_EQ(run(result.algo, result.workspaceSize), HIPBLAS_STATUS_SUCCESS);
+        EXPECT_EQ(reselections(), 1u);
+    }
+
+    // Under uniform summation order a result tagged with another hint is
+    // re-selected too, and the launch passes the order's launch check.
+    TEST_F(AdaptiveSmCount_pre_checkin, UniformSummationOrderReselects)
+    {
+        const int32_t uso = 1;
+        ASSERT_EQ(hipblasLtMatmulDescSetAttribute(
+                      r.desc, HIPBLASLT_MATMUL_DESC_UNIFORM_SUMMATION_ORDER_EXT, &uso, sizeof(uso)),
+                  HIPBLAS_STATUS_SUCCESS);
+        hipblasLtMatmulHeuristicResult_t result{};
+        retaggedResult(cfg.force ^ 1, &result);
+        if(IsSkipped() || HasFatalFailure())
+            return;
+        ASSERT_EQ(run(result.algo, result.workspaceSize), HIPBLAS_STATUS_SUCCESS);
+        EXPECT_EQ(reselections(), 1u);
+    }
+
+    // A long K, so that the top result uses a workspace when it has one.
+    class AdaptiveSmCountLongK_pre_checkin : public AdaptiveSmCount_pre_checkin
+    {
+    protected:
+        AdaptiveSmCountLongK_pre_checkin()
+        {
+            m = 256;
+            n = 256;
+            k = 16384;
+        }
+    };
+
+    // The pick is checked against the caller's workspace, not the query's: a
+    // StreamK solution needs less workspace when given less, so even without
+    // one the launch re-selects to it and runs.
+    TEST_F(AdaptiveSmCountLongK_pre_checkin, ReselectsWithoutWorkspace)
+    {
+        hipblasLtMatmulHeuristicResult_t result{};
+        retaggedResult(cfg.force ^ 1, &result);
+        if(IsSkipped() || HasFatalFailure())
+            return;
+        static_cast<void>(hipblaslt_tuning_last_launch_for_test());
+        ASSERT_EQ(run(result.algo, 0), HIPBLAS_STATUS_SUCCESS);
+        EXPECT_EQ(reselections(), 1u);
+        // Under FORCE the launch ranks as the query did, so its pick is the
+        // query's top result.
+        EXPECT_EQ(hipblaslt_tuning_last_launch_for_test(),
+                  hipblaslt_ext::getIndexFromAlgo(result.algo));
+    }
 } // namespace

projects/hipblaslt/clients/tests/src/tuning_cache_test.cpp

@@ -18,6 +18,8 @@
 #include <hipblaslt/hipblaslt-ext.hpp>
 #include <hipblaslt/hipblaslt.h>
 
+#include "../../../library/src/amd_detail/rocblaslt/src/include/rocblaslt_adaptive_sm_count.hpp"
+
 #include <chrono>
 #include <cstdint>
 #include <cstdio>
@@ -92,7 +94,9 @@ namespace
         hipblasLtMatmulDesc_t       desc    = nullptr;
         hipblasLtMatmulPreference_t pref    = nullptr;
 
-        bool create()
+        // `streamK`: a HIPBLASLT_MATMUL_DESC_STREAMK_TILE_SCHEDULING_EXT value,
+        // or none to keep the desc's default.
+        bool create(std::optional<int32_t> streamK = std::nullopt)
         {
             if(hipblasLtMatrixLayoutCreate(&layoutA, HIP_R_16F, kM, kK, kM)
                    != HIPBLAS_STATUS_SUCCESS
@@ -105,6 +109,14 @@ namespace
                || hipblasLtMatmulPreferenceCreate(&pref) != HIPBLAS_STATUS_SUCCESS)
                 return false;
 
+            if(streamK
+               && hipblasLtMatmulDescSetAttribute(desc,
+                                                  HIPBLASLT_MATMUL_DESC_STREAMK_TILE_SCHEDULING_EXT,
+                                                  &*streamK,
+                                                  sizeof(*streamK))
+                      != HIPBLAS_STATUS_SUCCESS)
+                return false;
+
             const uint64_t maxWs = kWorkspaceBytes;
             return hipblasLtMatmulPreferenceSetAttribute(
                        pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &maxWs, sizeof(maxWs))
@@ -1508,4 +1520,57 @@ namespace
                 unsetenv(name.c_str());
         }
     }
+
+    // The top heuristic result through the C API with StreamK AUTO.
+    bool autoModeHeuristic(hipblasLtMatmulHeuristicResult_t* result)
+    {
+        hipblasLtHandle_t handle = nullptr;
+        if(hipblasLtCreate(&handle) != HIPBLAS_STATUS_SUCCESS)
+            return false;
+        bool ok = false;
+        {
+            Problem problem;
+            int     returned = 0;
+            ok               = problem.create(HIPBLASLT_STREAMK_TILE_SCHEDULING_AUTO)
+                 && hipblasLtMatmulAlgoGetHeuristic(handle,
+                                                    problem.desc,
+                                                    problem.layoutA,
+                                                    problem.layoutB,
+                                                    problem.layoutC,
+                                                    problem.layoutC,
+                                                    problem.pref,
+                                                    1,
+                                                    result,
+                                                    &returned)
+                        == HIPBLAS_STATUS_SUCCESS
+                 && returned > 0;
+        }
+        hipblasLtDestroy(handle);
+        return ok;
+    }
+
+    // In adaptive mode a heuristic result is tagged for launch-time
+    // re-selection, but not one the tuning file supplied: the file's choice
+    // is launched as recorded.
+    TEST_F(TuningCache_pre_checkin, FileEntryIsNotAdaptiveTagged)
+    {
+        if(!rocblaslt::adaptive_sm::Config::fromEnv().enabled)
+            GTEST_SKIP() << "adaptive mode is off";
+        if(!haveSolutions(2))
+            GTEST_SKIP() << "the heuristic offers one solution for this problem";
+
+        uint32_t                         hint = 0;
+        hipblasLtMatmulHeuristicResult_t r{};
+        ASSERT_TRUE(autoModeHeuristic(&r));
+        ASSERT_TRUE(rocblaslt::adaptive_sm::algoTag(r.algo.data, &hint));
+
+        const auto& recorded = m_identities[1];
+        ASSERT_TRUE(writeTuningFile(m_path, m_stamp, {{recorded.index, recorded.kernelName}}));
+        useTuningFile();
+
+        r = {};
+        ASSERT_TRUE(autoModeHeuristic(&r));
+        EXPECT_EQ(hipblaslt_ext::getIndexFromAlgo(r.algo), recorded.index);
+        EXPECT_FALSE(rocblaslt::adaptive_sm::algoTag(r.algo.data, &hint));
+    }
 } // namespace

projects/hipblaslt/clients/tests/unit/adaptive_sm_count_gtest.cpp

@@ -14,6 +14,7 @@
 #include <algorithm>
 #include <cstddef>
 #include <cstdlib>
+#include <cstring>
 #include <functional>
 #include <optional>
 #include <string>
@@ -543,6 +544,30 @@ namespace
         EXPECT_FALSE(tinyProblem(1, 1, 1, 1, cfg)) << "threshold 0";
     }
 
+    TEST(HostUnitAdaptiveSmCountPure, smoke_AlgoTagRoundTrip)
+    {
+        const int32_t index   = 0x1234;
+        uint8_t       data[8] = {};
+        std::memcpy(data, &index, sizeof(index));
+        uint32_t hint = 1;
+        EXPECT_FALSE(algoTag(data, &hint));
+        tagAlgo(data, 192);
+        ASSERT_TRUE(algoTag(data, &hint));
+        EXPECT_EQ(hint, 192u);
+        int32_t kept = 0;
+        std::memcpy(&kept, data, sizeof(kept));
+        EXPECT_EQ(kept, index);
+    }
+
+    TEST(HostUnitAdaptiveSmCountPure, smoke_AlgoTagVersionMismatchIsUntagged)
+    {
+        uint8_t data[8] = {};
+        tagAlgo(data, 192);
+        data[5]       = c_algoTagVersion + 1;
+        uint32_t hint = 0;
+        EXPECT_FALSE(algoTag(data, &hint));
+    }
+
     using Env = std::vector<std::pair<const char*, const char*>>;
 
     // The environment parseEnv() reads from.