Cache HIP code object compilation - #5270
dhernandez0 wants to merge 1 commit into
Conversation
| try | ||
| { | ||
| auto compiled = compile(); | ||
| auto bytes = std::accumulate(compiled.begin(), |
There was a problem hiding this comment.
[format.py] reported by reviewdog 🐶
| auto bytes = std::accumulate(compiled.begin(), | |
| auto bytes = std::accumulate(compiled.begin(), |
There was a problem hiding this comment.
🟡 Changes recommended
Critical ABI and cache-key correctness issues, along with additional cache, eviction, diagnostics, testing, and changelog fixes, remain unresolved.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
This PR adds a per-device, thread-safe HIP compilation cache to reduce repeated kernel compilation.
Changes:
- Adds shared-future deduplication, LRU limits, and retry-on-failure.
- Integrates caching into HIP code-object compilation.
- Adds tests and documents cache configuration variables.
File summaries
| File | Summary and final findings |
|---|---|
test/gpu/jit.cpp |
Tests concurrency, retries, and eviction. |
src/targets/gpu/include/migraphx/gpu/context.hpp |
Stores the device compilation cache. critical (1 vote), line 268: Preserve hip_device copyability by avoiding a by-value mutex-containing cache. |
src/targets/gpu/include/migraphx/gpu/compile_hip.hpp |
Declares cache APIs. critical (3 votes), line 108: Preserve the existing five-argument exported overload for ABI compatibility. |
src/targets/gpu/compile_hip.cpp |
Implements caching and key generation. moderate (1 vote), line 82: Trim completed LRU entries using lru.size(). critical (3 votes), line 499: Avoid relying solely on an MD5 digest as the cache key. nit (1 vote), line 105: Add zero-limit cache tests. nit (3 votes), line 93: Replace the disallowed std::for_each side-effect wrapper. moderate (1 vote), line 494: Account for diagnostic controls in caching behavior. nit (1 vote), line 64: Add the cache environment variables to the changelog. moderate (1 vote), line 146: Prevent unresolved entries from premature eviction and clean up failed publication. |
src/targets/gpu/compile_hip_code_object.cpp |
Supplies the device cache during compilation. |
docs/reference/MIGraphX-dev-env-vars.rst |
Documents cache limits. nit (2 votes), line 739: Add the optimization to the appropriate CHANGELOG.md section. |
Review details
Suppressed comments (5)
src/targets/gpu/compile_hip.cpp:82
- [agent]:
trim()counts all map entries, includingready == falsecompilations that are intentionally absent fromlru. With more thanmax_entriesconcurrent compiles, each early completion can be inserted and immediately evicted, so a later request recompiles it even thoughmax_entriescompleted binaries could have been retained; count completed LRU entries (lru.size()) instead.
auto count = entries.size();
src/targets/gpu/compile_hip.cpp:106
- [agent]: The new cache tests cover positive limits, eviction, and retry, but not the documented zero-limit behavior. A regression in this early-return branch would go unnoticed; add cases for
max_bytes == 0andmax_entries == 0that invoke the same key twice and verify the callback runs twice.
if(max_bytes == 0 or max_entries == 0)
return compile();
src/targets/gpu/compile_hip.cpp:497
- The cache key omits the diagnostic controls that are read inside
compile_hip_src_impl(MIGRAPHX_GPU_DUMP_SRC,MIGRAPHX_GPU_DUMP_ASM, and, for HIPRTC,MIGRAPHX_TRACE_HIPRTC). With any of these enabled, the first compile prints the diagnostic, but later identical requests hit the cache and silently skip it; bypass caching while diagnostics are enabled or move those side effects outside the cached compilation.
key["debug"] = enabled(MIGRAPHX_GPU_DEBUG{});
key["debug_symbols"] = enabled(MIGRAPHX_GPU_DEBUG_SYM{});
key["optimize"] = string_value_of(MIGRAPHX_GPU_OPTIMIZE{}, "3");
key["extra_flags"] = string_value_of(MIGRAPHX_GPU_HIP_FLAGS{}, "");
src/targets/gpu/compile_hip.cpp:65
- The PR description selects the Optimized changelog category and explicitly requires a
CHANGELOG.mdentry, but these new user-visible cache environment variables are not recorded there. Add an entry underDevelop/Optimizedwith the PR number before merging.
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_GPU_HIP_CACHE_MAX_BYTES);
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_GPU_HIP_CACHE_MAX_ENTRIES);
src/targets/gpu/compile_hip.cpp:150
- The entry is marked ready and exposed to LRU eviction before
producer->set_valuepublishes its future. With a small limit, another compilation can trim this key in that gap, so a concurrent caller inserts a second producer and recompiles the same source; if publication throws, the catch path can also leave a stale LRU node. Publish the future before making the entry evictable, or otherwise keep unresolved entries out oftrimand clean up their LRU state on failure.
it->second.ready = true;
current_bytes += bytes;
trim();
}
producer->set_value(compiled);
- Files reviewed: 6/6 changed files
- Comments generated: 5
- Review effort level: Lite
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| return cache->get_or_compile(md5(std::string_view{packed_key.data(), packed_key.size()}), | ||
| compile); |
| bool disable_processes = false, | ||
| bool quiet = false, | ||
| hip_compile_cache* cache = nullptr); |
| device_description desc = {}; | ||
| hip_compile_cache compile_cache{}; |
| * - | ``MIGRAPHX_GPU_HIP_CACHE_MAX_BYTES`` | ||
| | Sets the maximum total size, in bytes, of compiled binaries retained in each GPU context's HIP compilation cache. | ||
| | Least-recently-used binaries are evicted when the cache exceeds this limit. | ||
|
|
||
| - | Takes a non-negative integer. Setting the value to ``0`` disables the cache. |
| std::for_each(first_to_evict, lru.end(), [&](const std::string& cache_key) { | ||
| auto it = entries.find(cache_key); | ||
| assert(it != entries.end()); | ||
| current_bytes -= it->second.bytes; | ||
| entries.erase(it); |
There was a problem hiding this comment.
🟡 Changes recommended
The exported ABI change and cache correctness issues must be addressed before approval.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (2)
src/targets/gpu/include/migraphx/gpu/compile_hip.hpp:110
- Adding the cache parameter changes the exported C++ symbol from a five-argument function to a six-argument function; default arguments preserve source compatibility but not ABI compatibility, so existing clients will fail to link. Keep the five-argument overload and add a separate cache-aware overload that it forwards to.
bool disable_processes = false,
bool quiet = false,
hip_compile_cache* cache = nullptr);
src/targets/gpu/compile_hip.cpp:500
- The MD5 digest is being used as the complete cache identity without checking the serialized request on a hit. Two distinct source/option requests with the same digest will therefore return the first request's code object; retain and compare the full serialized key (using the digest only as an index) before reusing a binary.
return cache->get_or_compile(md5(std::string_view{packed_key.data(), packed_key.size()}),
compile);
- Files reviewed: 6/6 changed files
- Comments generated: 3
- Review effort level: Balanced
| producer->set_exception(std::current_exception()); | ||
| std::lock_guard<std::mutex> lock{mutex}; | ||
| entries.erase(key); |
| MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_GPU_HIP_CACHE_MAX_BYTES); | ||
| MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_GPU_HIP_CACHE_MAX_ENTRIES); |
| value key; | ||
| key["srcs"] = to_value(std::vector<hiprtc_src_file>{srcs.begin(), srcs.end()}); | ||
| key["params"] = to_value(params); |
There was a problem hiding this comment.
hiprtc already provides a cache when you enable it with AMD_COMGR_CACHE and AMD_COMGR_CACHE_DIR. No need to reinvent the wheel here.
Also #5100 adds a persistent binary cache as well which works for MLIR and hip kernels and doesnt require locking.
|
Actually from here it says the cache is enabled by default, so I am not sure why its not working for you. |
I didn't know about this cache. I've done some experiments and measured COMGR warm hit at ~326 ms. So, looks like |
sounds good, should I close this PR? Is that one ready for reviews? Edit: note that the motivation of this PR is to avoid recompiling the pointwise and fill kernels after split-k perf_config candidates in the tuning list. Because all of them are followed by the same pointwise (and the same fill kernel). Is your branch able to cache those as well? Edit2: Doesn’t this code create one key for the entire MLIR candidate, including its pointwise and fill? Since mlir_compile_key includes the tuning solution, each candidate appears to have a different key, so identical pointwise and fill kernels across candidates would still be recompiled. Is there another cache lookup for those sub-kernels that I’m missing? I'm not familiar with migraphx, so maybe I misunderstood the code. |
Motivation
Split-K tuning may split each candidate into an MLIR kernel plus identical pointwise-tail and fill kernels. These HIP kernels were recompiled for every candidate, significantly increasing compilation time.
With this branch combined with fix/compile-benchmark-device-ops-parallel and justinr-additional-timing-info, using rocmlirTriton, APNB v2 FP32 512×512 compilation improves:
Parallel: 24.712s → 9.124s
Serial: 369.789s → 117.274s
Technical Details
Add a per-device, thread-safe HIP compilation cache keyed by source, compiler options, architecture, and compilation mode. Concurrent requests share one compilation through a shared future.
Completed binaries use LRU eviction, defaulting to 256 entries and 256 MiB. Both limits are configurable through environment variables, and failed compilations are removed for retry.
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot ApplicableFollow the LLVM AI Tool Use Policy for contributions using AI.