From 99946464eecd300977c3a864af4f04d32c44f414 Mon Sep 17 00:00:00 2001 From: Antti Hyttinen Date: Fri, 2 Oct 2026 06:04:13 -0500 Subject: [PATCH] [MIOpen] Reduce convolution heuristic overhead Rebased onto develop (squashed). Adds the candidate-selection result cache (B-RAM) and the per-row kernel-config embedding cache (C1) to cut two-tower heuristic overhead on repeated calls. Conflict resolution vs develop: - ai_candidate_selection.cpp: kept develop's LgbmOnly() bypass and ConfigEncoder-ranked logging; added the RAM result-cache short-circuit and GetOrComputeKernelEmbeddings() embedding cache. - ai_heuristics.cpp: took develop's unified TunaNet->LGBM->WTI predictor; dropped the PR's timing log on the removed ND-first block. - conv_ai_candidate_selection_model.cpp: unioned develop's decode tests with the PR's 3 cache tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../heuristics/ai_candidate_selection.cpp | 181 +++++++++++++++++- .../src/conv/heuristics/ai_heuristics.cpp | 1 + .../heuristics/ai_candidate_selection.hpp | 4 + .../ai_conv_nd_kernel_tuning_utils.hpp | 14 +- .../conv_ai_candidate_selection_model.cpp | 162 ++++++++++++++++ 5 files changed, 358 insertions(+), 4 deletions(-) diff --git a/projects/miopen/src/conv/heuristics/ai_candidate_selection.cpp b/projects/miopen/src/conv/heuristics/ai_candidate_selection.cpp index 0c615c3ce56c..e7f4dde3e043 100644 --- a/projects/miopen/src/conv/heuristics/ai_candidate_selection.cpp +++ b/projects/miopen/src/conv/heuristics/ai_candidate_selection.cpp @@ -34,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -46,6 +47,8 @@ #include #include #include +#include +#include #include #if MIOPEN_ENABLE_AI_KERNEL_TUNING @@ -1480,10 +1483,173 @@ EncodeKernelParams(const std::vector>& valid_kernel_par return encoded_candidates; } +// ---- Per-process RAM cache for CandidateSelectionResult -------------- +// +// Mirrors the TunaNet RAM cache (GetCachedPrediction / StorePredictionCache in +// ai_heuristics.cpp). A cache hit at this level bypasses all fdeep inference. +namespace { + +struct CandidateSelectionCache +{ + std::unordered_map map; + std::mutex mtx; + + static CandidateSelectionCache& Get(const std::string& arch) + { + static std::mutex instances_mtx; + static std::unordered_map> instances; + std::lock_guard lk(instances_mtx); + auto& ptr = instances[arch]; + if(!ptr) + ptr = std::make_unique(); + return *ptr; + } +}; + +std::string MakeCandidateSelectionKey(const std::string& solver, + const conv::ProblemDescription& problem, + bool use_split_k) +{ + std::ostringstream ss; + ss << solver << "|split_k=" << use_split_k << "|"; + problem.Serialize(ss); + return ss.str(); +} + +std::optional +GetCachedCandidateSelection(const std::string& arch, + const std::string& solver, + const conv::ProblemDescription& problem, + bool use_split_k) +{ + auto& cache = CandidateSelectionCache::Get(arch); + std::lock_guard lk(cache.mtx); + const auto it = cache.map.find(MakeCandidateSelectionKey(solver, problem, use_split_k)); + if(it == cache.map.end()) + return std::nullopt; + MIOPEN_LOG_I2("CandidateSelection RAM cache hit for solver: " << solver); + return it->second; +} + +void StoreCandidateSelectionCache(const std::string& arch, + const std::string& solver, + const conv::ProblemDescription& problem, + bool use_split_k, + const CandidateSelectionResult& result) +{ + if(result.IsEmpty()) + return; + auto& cache = CandidateSelectionCache::Get(arch); + std::lock_guard lk(cache.mtx); + cache.map.emplace(MakeCandidateSelectionKey(solver, problem, use_split_k), result); +} + +// ---- Kernel-config embedding cache ------------------------------------------ +// +// The config encoder processes each kernel independently, so its output for +// kernel K is the same regardless of what other kernels are in the batch. +// Cache at individual kernel-row granularity so embeddings are reused across +// different shapes that share a subset of valid kernels. +struct KernelEmbeddingCache +{ + std::unordered_map> map; // one entry per kernel row + std::mutex mtx; + + static KernelEmbeddingCache& Get() + { + static KernelEmbeddingCache instance; + return instance; + } +}; + +std::string +KernelRowKey(const std::string& arch, const std::string& solver, const std::vector& row) +{ + // Use the exact bit pattern of each float so the key is collision-free. + // Encoded candidate rows contain small integer values from EncodeKernelParams; + // a hash-only approach risks structured collisions for such inputs. + std::ostringstream ss; + ss << arch << '|' << solver << '|' << row.size(); + for(float v : row) + { + uint32_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + ss << '|' << std::hex << bits; + } + return ss.str(); +} + +std::vector> +GetOrComputeKernelEmbeddings(const std::string& arch, + const std::string& solver, + const std::vector>& encoded_candidates, + const CandidateSelectionModel& model) +{ + auto& cache = KernelEmbeddingCache::Get(); + + std::vector> result(encoded_candidates.size()); + std::vector miss_indices; + std::vector miss_keys; + + // Phase 1: look up each kernel row individually. + { + std::lock_guard lk(cache.mtx); + for(size_t i = 0; i < encoded_candidates.size(); ++i) + { + auto key = KernelRowKey(arch, solver, encoded_candidates[i]); + auto it = cache.map.find(key); + if(it != cache.map.end()) + { + result[i] = it->second; + } + else + { + miss_indices.push_back(i); + miss_keys.push_back(std::move(key)); + } + } + } + + if(miss_indices.empty()) + { + MIOPEN_LOG_I2("Kernel-config embedding cache: all " + << encoded_candidates.size() << " kernels hit for solver: " << solver); + return result; + } + + MIOPEN_LOG_I2("Kernel-config embedding cache: " + << (encoded_candidates.size() - miss_indices.size()) << "/" + << encoded_candidates.size() << " kernels hit, encoding " << miss_indices.size() + << " new kernels for solver: " << solver); + + // Phase 2: batch-encode only the cache misses. + std::vector> miss_candidates; + miss_candidates.reserve(miss_indices.size()); + for(size_t idx : miss_indices) + miss_candidates.push_back(encoded_candidates[idx]); + + auto miss_embeddings = model.EncodeKernelConfigs(miss_candidates); + + // Phase 3: store new embeddings and assemble the full result. + { + std::lock_guard lk(cache.mtx); + for(size_t i = 0; i < miss_indices.size(); ++i) + { + result[miss_indices[i]] = miss_embeddings[i]; + cache.map.emplace(miss_keys[i], miss_embeddings[i]); + } + } + + return result; +} + +} // anonymous namespace + MIOPEN_INTERNALS_EXPORT CandidateSelectionResult ModelSelectBestCandidate(const std::string& arch, const std::string& solver, + const conv::ProblemDescription& problem, const std::map& features, const std::vector>& valid_kernel_params, const bool use_split_k, @@ -1495,10 +1661,17 @@ ModelSelectBestCandidate(const std::string& arch, << solver << " on " << arch); return CandidateSelectionResult{{}, {}}; } + + // RAM cache: bypass all model inference on repeated calls for the same problem + { + auto cached = GetCachedCandidateSelection(arch, solver, problem, use_split_k); + if(cached) + return *cached; + } + try { const auto& model = GetCandidateSelectionModel(arch, solver); - // debug: show that we successfully retrieved the model MIOPEN_LOG_I2("Retrieved CandidateSelectionModel for arch: " << arch << ", solver: " << solver); std::vector> expanded_params = valid_kernel_params; @@ -1558,8 +1731,9 @@ ModelSelectBestCandidate(const std::string& arch, for(const auto& candidate : expanded_params) candidate_kernel_names.push_back(candidate.empty() ? std::string{} : candidate.front()); - const auto& encoded_configs = - model.EncodeKernelConfigs(encoded_candidates, &candidate_kernel_names); + // Reuse kernel-config embeddings if the kernel set hasn't changed + const auto encoded_configs = + GetOrComputeKernelEmbeddings(arch, solver, encoded_candidates, model); { std::ostringstream encoded_configs_log; encoded_configs_log << "Encoded configs: ["; @@ -1605,6 +1779,7 @@ ModelSelectBestCandidate(const std::string& arch, } } + StoreCandidateSelectionCache(arch, solver, problem, use_split_k, result); return result; } catch(const miopen::Exception& ex) diff --git a/projects/miopen/src/conv/heuristics/ai_heuristics.cpp b/projects/miopen/src/conv/heuristics/ai_heuristics.cpp index a03549743d2d..8e066af80823 100644 --- a/projects/miopen/src/conv/heuristics/ai_heuristics.cpp +++ b/projects/miopen/src/conv/heuristics/ai_heuristics.cpp @@ -40,6 +40,7 @@ #include #include #include +#include #include #include diff --git a/projects/miopen/src/include/miopen/conv/heuristics/ai_candidate_selection.hpp b/projects/miopen/src/include/miopen/conv/heuristics/ai_candidate_selection.hpp index f6fe308508ac..8e6e2535f62f 100644 --- a/projects/miopen/src/include/miopen/conv/heuristics/ai_candidate_selection.hpp +++ b/projects/miopen/src/include/miopen/conv/heuristics/ai_candidate_selection.hpp @@ -34,6 +34,9 @@ #include namespace miopen { +namespace conv { +struct ProblemDescription; +} // namespace conv namespace ai { namespace tuning { namespace candidate_selection { @@ -224,6 +227,7 @@ MIOPEN_INTERNALS_EXPORT CandidateSelectionResult ModelSelectBestCandidate(const std::string& arch, const std::string& solver, + const conv::ProblemDescription& problem, const std::map& features, const std::vector>& valid_kernel_params, bool use_split_k, diff --git a/projects/miopen/src/include/miopen/conv/heuristics/ai_conv_nd_kernel_tuning_utils.hpp b/projects/miopen/src/include/miopen/conv/heuristics/ai_conv_nd_kernel_tuning_utils.hpp index fca525c6fe14..c04bd93af115 100644 --- a/projects/miopen/src/include/miopen/conv/heuristics/ai_conv_nd_kernel_tuning_utils.hpp +++ b/projects/miopen/src/include/miopen/conv/heuristics/ai_conv_nd_kernel_tuning_utils.hpp @@ -259,6 +259,7 @@ RunParameterPredictionModel( auto result = ai::tuning::candidate_selection::ModelSelectBestCandidate( arch, solver_name, + problem, features, heuristic_kernels, use_split_k, @@ -411,6 +412,10 @@ bool RunAIHeuristics(const SolverHeuristicConfig& solver_cfg, MIOPEN_LOG_I2(solver_cfg.solver_name << ": RunAIHeuristics on " << arch << " (use_tf32=" << std::boolalpha << use_tf32 << ")"); + const bool logging = miopen::IsLogging(LoggingLevel::Info2); + const auto start = logging ? std::chrono::high_resolution_clock::now() + : std::chrono::high_resolution_clock::time_point{}; + // Only applicable for supported architectures if(arch != "gfx90a" && arch != "gfx942" && arch != "gfx950") { @@ -484,7 +489,14 @@ bool RunAIHeuristics(const SolverHeuristicConfig& solver_cfg, }; auto result = ai::tuning::candidate_selection::ModelSelectBestCandidate( - arch, solver_name, features, heuristic_kernels, use_split_k, is_valid); + arch, solver_name, problem, features, heuristic_kernels, use_split_k, is_valid); + + if(logging) + { + const auto end = std::chrono::high_resolution_clock::now(); + MIOPEN_LOG_I2("Heuristic candidate selection inference time: " + << (end - start).count() * .000001f << " ms"); + } if(!result.IsEmpty()) { diff --git a/projects/miopen/test/gtest/conv_ai_candidate_selection_model.cpp b/projects/miopen/test/gtest/conv_ai_candidate_selection_model.cpp index 0135f5684e23..93d98975d446 100644 --- a/projects/miopen/test/gtest/conv_ai_candidate_selection_model.cpp +++ b/projects/miopen/test/gtest/conv_ai_candidate_selection_model.cpp @@ -30,12 +30,16 @@ #include #include #include +#include +#include +#include #include #include #include #include #include #include +#include #include #if MIOPEN_ENABLE_AI_KERNEL_TUNING @@ -87,6 +91,19 @@ std::string FormatMissingFiles(const std::vector& missing_file return os.str(); } +// Build a minimal 2D conv ProblemDescription for use as a cache key in tests. +miopen::conv::ProblemDescription +MakeTestProblem(int n = 1, int c = 4, int h = 8, int w = 8, int k = 8, int y = 3, int x = 3) +{ + miopen::TensorDescriptor in_desc(miopenFloat, {n, c, h, w}); + miopen::TensorDescriptor wei_desc(miopenFloat, {k, c, y, x}); + miopen::TensorDescriptor out_desc(miopenFloat, {n, k, h - y + 1, w - x + 1}); + miopen::ConvolutionDescriptor conv_desc( + 2, miopenConvolution, miopenPaddingDefault, {0, 0}, {1, 1}, {1, 1}); + return miopen::conv::ProblemDescription( + in_desc, wei_desc, out_desc, conv_desc, miopen::conv::Direction::BackwardWeights); +} + std::vector> GenerateValidKernelParams( const CandidateSelectionMetadata& meta, const std::string& kernel_name, int num_candidates = 3) { @@ -330,8 +347,10 @@ TEST_P(GPU_CandidateSelection_FP32, ModelSelectBestCandidate_Test) for(const auto& name : meta.input_params()) features[name] = 1.0f; auto valid_kernel_params = GenerateValidKernelParams(meta, params.kernel_name, 3); + auto problem = MakeTestProblem(); auto result = ModelSelectBestCandidate(params.arch, params.solver, + problem, features, valid_kernel_params, /*use_split_k=*/false, @@ -1025,6 +1044,149 @@ TEST_P(GPU_CandidateSelection_FP32, StaticKernelDecodes_Test) meta.output_params().size() - meta.GetConstantOutputIndices().size()); } +// === CACHE TESTS === + +// Calling ModelSelectBestCandidate twice with the same arguments must return +// identical results. The second call hits the in-process RAM cache. +TEST_P(GPU_CandidateSelection_FP32, CandidateSelectionRamCache_Test) +{ + const auto& params = GetParam(); + CandidateSelectionMetadata meta(params.arch, params.solver); + std::map features; + for(const auto& name : meta.input_params()) + features[name] = 1.0f; + auto valid_kernel_params = GenerateValidKernelParams(meta, params.kernel_name, 3); + auto problem = MakeTestProblem(); + + const auto t0 = std::chrono::high_resolution_clock::now(); + auto result1 = ModelSelectBestCandidate(params.arch, + params.solver, + problem, + features, + valid_kernel_params, + /*use_split_k=*/false, + accept_all_combinations); + const auto t1 = std::chrono::high_resolution_clock::now(); + auto result2 = ModelSelectBestCandidate(params.arch, + params.solver, + problem, + features, + valid_kernel_params, + /*use_split_k=*/false, + accept_all_combinations); + const auto t2 = std::chrono::high_resolution_clock::now(); + + ASSERT_FALSE(result1.IsEmpty()) << "First call returned empty result"; + ASSERT_FALSE(result2.IsEmpty()) << "Second call (cache hit) returned empty result"; + ASSERT_EQ(result1.kernel_indices, result2.kernel_indices) << "Cache returned different indices"; + ASSERT_EQ(result1.split_k_values, result2.split_k_values) + << "Cache returned different split_k values"; + + const auto first_ms = std::chrono::duration(t1 - t0).count(); + const auto second_ms = std::chrono::duration(t2 - t1).count(); + std::cout << "[CandidateSelectionRamCache] first=" << first_ms << " ms, second=" << second_ms + << " ms\n"; +} + +// Kernel-config embeddings are cached per row via GetOrComputeKernelEmbeddings, +// which is exercised through ModelSelectBestCandidate. Two calls with identical +// features and kernel params but different problems bypass the B-RAM cache yet +// share the per-row embedding cache. Because both inputs to the scorer are then +// identical (same Tower 1 output, same cached Tower 2 embeddings), the rankings +// must be deterministically the same. +TEST_P(GPU_CandidateSelection_FP32, KernelEmbeddingCache_Test) +{ + const auto& params = GetParam(); + CandidateSelectionMetadata meta(params.arch, params.solver); + std::map features; + for(const auto& name : meta.input_params()) + features[name] = 1.0f; + auto valid_kernel_params = GenerateValidKernelParams(meta, params.kernel_name, 3); + + // Call 1: populates the B-RAM cache (keyed on problem_a) and the per-row + // embedding cache for all kernel rows. + auto problem_a = MakeTestProblem(1, 4, 8, 8, 8, 3, 3); + auto result_a = ModelSelectBestCandidate(params.arch, + params.solver, + problem_a, + features, + valid_kernel_params, + /*use_split_k=*/false, + accept_all_combinations); + + // Call 2: different N → different B-RAM key → B-RAM cache miss. + // Same features and same kernel params → embedding cache hit for all rows. + // Identical scorer inputs must produce identical rankings. + auto problem_b = MakeTestProblem(2, 4, 8, 8, 8, 3, 3); + auto result_b = ModelSelectBestCandidate(params.arch, + params.solver, + problem_b, + features, + valid_kernel_params, + /*use_split_k=*/false, + accept_all_combinations); + + ASSERT_FALSE(result_a.IsEmpty()) << "First call returned empty result"; + ASSERT_FALSE(result_b.IsEmpty()) + << "Second call returned empty result (embedding cache may be corrupt)"; + ASSERT_EQ(result_a.kernel_indices, result_b.kernel_indices) + << "Rankings differ despite identical features and kernel set"; + ASSERT_EQ(result_a.split_k_values, result_b.split_k_values); +} + +// use_split_k=false and use_split_k=true must produce independent cache entries. +TEST_P(GPU_CandidateSelection_FP32, CandidateSelectionRamCacheIsolation_Test) +{ + const auto& params = GetParam(); + CandidateSelectionMetadata meta(params.arch, params.solver); + std::map features; + for(const auto& name : meta.input_params()) + features[name] = 1.0f; + auto valid_kernel_params = GenerateValidKernelParams(meta, params.kernel_name, 3); + auto problem = MakeTestProblem(); + + auto result_no_splitk = ModelSelectBestCandidate(params.arch, + params.solver, + problem, + features, + valid_kernel_params, + /*use_split_k=*/false, + accept_all_combinations); + auto result_splitk = ModelSelectBestCandidate(params.arch, + params.solver, + problem, + features, + valid_kernel_params, + /*use_split_k=*/true, + accept_all_combinations); + + ASSERT_FALSE(result_no_splitk.IsEmpty()) << "use_split_k=false returned empty"; + ASSERT_FALSE(result_splitk.IsEmpty()) << "use_split_k=true returned empty"; + + auto result_no_splitk_cached = ModelSelectBestCandidate(params.arch, + params.solver, + problem, + features, + valid_kernel_params, + /*use_split_k=*/false, + accept_all_combinations); + auto result_splitk_cached = ModelSelectBestCandidate(params.arch, + params.solver, + problem, + features, + valid_kernel_params, + /*use_split_k=*/true, + accept_all_combinations); + + ASSERT_EQ(result_no_splitk.kernel_indices, result_no_splitk_cached.kernel_indices) + << "Cache isolation failed: use_split_k=false entry was overwritten"; + ASSERT_EQ(result_splitk.kernel_indices, result_splitk_cached.kernel_indices) + << "Cache isolation failed: use_split_k=true entry was overwritten"; + + EXPECT_NE(result_no_splitk.split_k_values, result_splitk.split_k_values) + << "Expected use_split_k=false and use_split_k=true to produce different split_k_values"; +} + // === INSTANTIATION === // Helper function to generate test parameters for both 2D and 3D solvers