Add binary cache

docs/dev/triage-migraphx.rst

@@ -151,4 +151,41 @@ Debugging and Tracing
 
 - ``MIGRAPHX_TRACE_BENCHMARKING=3``     # Kernel benchmarking process
 
-This systematic approach helps maintainers quickly understand and fix root causes.
+This systematic approach helps maintainers quickly understand and fix root causes.
+
+Binary Cache
+============
+
+Compiled kernels are shared within a single compile, so a kernel that appears many times in a
+model is compiled only once. Setting a directory makes that reuse survive across runs:
+
+.. code-block:: bash
+
+   export MIGRAPHX_BINARY_CACHE=$HOME/.cache/migraphx
+
+Entries are grouped by a directory naming the entry format, the HIP compiler, a digest of the
+embedded kernel headers, and the rocMLIR build, so entries a build cannot use are never
+consulted. The compiler is identified by compiling a small probe that records
+``__clang_version__`` into the object and reading it back, because the device compiler is loaded
+at runtime and need not be the one MIGraphX was built with. Reclaim space by deleting the
+directories for builds you no longer use:
+
+.. code-block:: bash
+
+   ls $HOME/.cache/migraphx        # each directory has a cache.info describing the build
+   rm -r $HOME/.cache/migraphx/v1-hip22.0.*
+
+The same settings are available as backend options, which take precedence over the environment
+and are how tests configure the cache:
+
+.. code-block:: python
+
+   model.compile(migraphx.get_target("gpu"),
+                 advance_backend_options={"binary_cache": "/tmp/cache",
+                                          "binary_cache_verify": True})
+
+Entries are keyed on everything handed to the backend compiler. If a kernel is suspected of
+being reused when it should not be, set ``binary_cache_verify``. That compiles even when a
+result could be reused and fails loudly if the two disagree, which is the only way an
+incomplete key shows up other than as wrong results. It is slower than compiling without a
+cache at all, so use it to diagnose rather than routinely.

src/include/migraphx/program.hpp

@@ -179,6 +179,11 @@ struct MIGRAPHX_EXPORT program
     // bumped if any changes occur to the format of the MXR file.
     static constexpr int program_file_version = 8;
 };
+
+/// Lets a program be a member of a reflected type, using the same conversion as saving one.
+inline void migraphx_to_value(value& v, const program& p) { v = p.to_value(); }
+inline void migraphx_from_value(const value& v, program& p) { p.from_value(v); }
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx
 

src/include/migraphx/tmp_dir.hpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -33,6 +33,9 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
+/// A name unlikely to collide with one made by another process or thread.
+MIGRAPHX_EXPORT std::string unique_string(const std::string& prefix);
+
 struct MIGRAPHX_EXPORT tmp_dir
 {
     fs::path path;
@@ -45,7 +48,7 @@ struct MIGRAPHX_EXPORT tmp_dir
         execute(std::string_view{cmd.string()}, args);
     }
 
-    tmp_dir(tmp_dir const&) = delete;
+    tmp_dir(tmp_dir const&)            = delete;
     tmp_dir& operator=(tmp_dir const&) = delete;
 
     ~tmp_dir();

src/targets/gpu/CMakeLists.txt

@@ -24,6 +24,25 @@
 
 find_package(hip REQUIRED)
 
+# Produce a short identity for a toolchain file, used to separate binary cache entries that
+# different toolchains cannot share. The content hash is exact but slow on the larger libraries,
+# so fall back to size and timestamp, which still changes whenever the file is replaced.
+function(migraphx_hash_file FILE_PATH OUT_VAR)
+    if(NOT FILE_PATH OR NOT EXISTS "${FILE_PATH}")
+        set(${OUT_VAR} "unknown" PARENT_SCOPE)
+        return()
+    endif()
+    file(SIZE "${FILE_PATH}" HASH_FILE_SIZE)
+    if(HASH_FILE_SIZE LESS 268435456)
+        file(MD5 "${FILE_PATH}" HASH_RESULT)
+    else()
+        file(TIMESTAMP "${FILE_PATH}" HASH_FILE_TIME "%s" UTC)
+        string(MD5 HASH_RESULT "${FILE_PATH}-${HASH_FILE_SIZE}-${HASH_FILE_TIME}")
+    endif()
+    string(SUBSTRING "${HASH_RESULT}" 0 12 HASH_RESULT)
+    set(${OUT_VAR} "${HASH_RESULT}" PARENT_SCOPE)
+endfunction()
+
 # MIGRAPHX_USE_AMDMLSS is resolved to a plain boolean at the top-level CMakeLists.
 if(MIGRAPHX_USE_AMDMLSS)
     message(STATUS "MIGraphX is using AMDMLSS")
@@ -202,7 +221,9 @@ endif()
 add_library(migraphx_gpu
     analyze_streams.cpp
     allocation_model.cpp
+    binary_cache.cpp
     code_object_op.cpp
+    compiled_code.cpp
     compile_ops.cpp
     compile_gen.cpp
     compile_hip.cpp
@@ -313,6 +334,21 @@ if(MIGRAPHX_ENABLE_MLIR)
     find_package(rocMLIR 1.0.0 CONFIG REQUIRED)
     message(STATUS "Build with rocMLIR::rockCompiler ${rocMLIR_VERSION}")
     target_compile_definitions(migraphx_gpu PRIVATE "-DMIGRAPHX_MLIR")
+
+    # Identify the rocMLIR build for the binary cache by hashing the library that gets linked.
+    # That covers a locally built or substituted rocMLIR, which a recorded version or commit
+    # would not.
+    get_target_property(MIGRAPHX_ROCKCOMPILER_LIB rocMLIR::rockCompiler IMPORTED_LOCATION)
+    if(NOT MIGRAPHX_ROCKCOMPILER_LIB)
+        # rockCompiler is an interface target, so the archive is named as a link dependency.
+        get_target_property(MIGRAPHX_ROCKCOMPILER_LIB rocMLIR::rockCompiler INTERFACE_LINK_LIBRARIES)
+        list(FILTER MIGRAPHX_ROCKCOMPILER_LIB INCLUDE REGEX "rockCompiler")
+        list(GET MIGRAPHX_ROCKCOMPILER_LIB 0 MIGRAPHX_ROCKCOMPILER_LIB)
+    endif()
+    migraphx_hash_file("${MIGRAPHX_ROCKCOMPILER_LIB}" MIGRAPHX_ROCMLIR_ID)
+    message(STATUS "rocMLIR binary cache id: ${MIGRAPHX_ROCMLIR_ID}")
+    target_compile_definitions(migraphx_gpu
+                               PRIVATE "-DMIGRAPHX_ROCMLIR_ID=\"${MIGRAPHX_ROCMLIR_ID}\"")
     # Make this private to avoid multiple inclusions of LLVM symbols.
     target_link_libraries(migraphx_gpu PRIVATE rocMLIR::rockCompiler)
     # Hide LLVM internals that come from rocMLIR.

src/targets/gpu/binary_cache.cpp

@@ -0,0 +1,208 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2026 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/gpu/binary_cache.hpp>
+#include <migraphx/gpu/context.hpp>
+#include <migraphx/gpu/compile_hip.hpp>
+#include <migraphx/file_buffer.hpp>
+#include <migraphx/filesystem.hpp>
+#include <migraphx/logger.hpp>
+#include <migraphx/md5.hpp>
+#include <migraphx/msgpack.hpp>
+#include <migraphx/serialize.hpp>
+#include <migraphx/stringutils.hpp>
+#include <migraphx/tmp_dir.hpp>
+#include <migraphx_kernels.hpp>
+#include <sstream>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+// Bump when the shape of a stored fragment changes, which happens when a compiler changes the
+// instructions it replaces with or when a serialized operator gains or loses a field. Such a
+// change is invisible to the key, since the source handed to the backend is unaffected.
+static constexpr const char* binary_cache_format = "v1";
+
+#ifdef MIGRAPHX_ROCMLIR_ID
+static constexpr const char* rocmlir_id = MIGRAPHX_ROCMLIR_ID;
+#else
+static constexpr const char* rocmlir_id = "nomlir";
+#endif
+
+std::shared_ptr<binary_cache> make_binary_cache() { return std::make_shared<binary_cache>(); }
+
+bool binary_cache::verify() const { return settings.verify; }
+
+static std::string short_digest(const std::string& s) { return md5(s).substr(0, 12); }
+
+/// A digest of the kernel headers compiled into this build. Taken from the embedded sources
+/// rather than the files on disk, so it tracks what is actually compiled even when the build
+/// system has not reconfigured.
+static const std::string& kernels_digest()
+{
+    static const std::string digest = [] {
+        std::stringstream ss;
+        for(const auto& [path, content] : ::migraphx_kernels())
+        {
+            ss << path << "\n" << content << "\n";
+        }
+        return short_digest(ss.str());
+    }();
+    return digest;
+}
+
+const std::string& binary_cache::version_dir()
+{
+    static const std::string dir = [] {
+        const auto& compiler = hip_compiler_version();
+        if(compiler.empty())
+            return std::string{};
+        // The version numbers make the directory readable; the hash of the full version string
+        // separates builds that share them, since it also covers the source revision.
+        return std::string{binary_cache_format} + "-hip" + compiler.major + "." + compiler.minor +
+               "." + short_digest(compiler.version) + "-kernels" + kernels_digest() + "-rocmlir" +
+               rocmlir_id;
+    }();
+    return dir;
+}
+
+/// Entries are grouped by the device they were compiled for. This keeps the directory
+/// self-describing; the arch, core count and wavefront size already reach the key through the
+/// arch line, the launch bounds and the -D defines.
+static std::string device_dir(const context& ctx)
+{
+    const auto& device = ctx.get_current_device();
+    return to_c_id(device.get_device_name()) + "_cu" + std::to_string(device.get_cu_count()) +
+           "_wf" + std::to_string(device.get_wavefront_size());
+}
+
+/// Where an entry lives, or an empty path when the toolchain cannot be identified and entries
+/// from different toolchains would be indistinguishable.
+static fs::path entry_path(const fs::path& root, const context& ctx, const std::string& key)
+{
+    const auto& version = binary_cache::version_dir();
+    if(version.empty())
+        return {};
+    return root / version / device_dir(ctx) / (md5(key) + ".mxr");
+}
+
+/// Record what this build is, so a directory full of hashes can be identified later.
+static void write_stamp(const fs::path& dir)
+{
+    auto stamp = dir / "cache.info";
+    if(fs::exists(stamp))
+        return;
+    std::stringstream ss;
+    ss << "format: " << binary_cache_format << "\n";
+    ss << "hip: " << hip_compiler_version().version << "\n";
+    ss << "kernels: " << kernels_digest() << "\n";
+    ss << "rocmlir: " << rocmlir_id << "\n";
+    // Publish by rename like the entries, so concurrent writers cannot tear the file.
+    auto tmp = stamp;
+    tmp += "." + unique_string("tmp");
+    write_string(tmp, ss.str());
+    fs::rename(tmp, stamp);
+}
+
+optional<compiled_code> binary_cache::get(const context& ctx, const std::string& key)
+{
+    if(key.empty())
+        return nullopt;
+    auto it = memo.find(key);
+    if(it != memo.end())
+    {
+        counters.reused++;
+        return it->second;
+    }
+    const auto& root = settings.path;
+    if(root.empty())
+    {
+        counters.misses++;
+        return nullopt;
+    }
+    auto path = entry_path(root, ctx, key);
+    if(path.empty() or not fs::exists(path))
+    {
+        counters.misses++;
+        return nullopt;
+    }
+    entry e;
+    try
+    {
+        migraphx::from_value(from_msgpack(read_buffer(path)), e);
+    }
+    catch(const std::exception& ex)
+    {
+        // A damaged or stale entry is only worth a recompile, so treat every failure as a miss
+        // and let the result be written over the top.
+        log::warn() << "Ignoring unreadable binary cache entry " << path << ": " << ex.what();
+        counters.misses++;
+        return nullopt;
+    }
+    if(e.key != key)
+    {
+        log::warn() << "Ignoring binary cache entry with mismatched key: " << path;
+        counters.misses++;
+        return nullopt;
+    }
+
+    counters.hits++;
+    return memo[key] = std::move(e.code);
+}
+
+void binary_cache::insert(const context& ctx, entry e)
+{
+    if(e.key.empty())
+        return;
+    counters.compiled++;
+    const auto& root = settings.path;
+    auto path        = root.empty() ? fs::path{} : entry_path(root, ctx, e.key);
+    if(not path.empty())
+    {
+        // Publish by rename so a reader never sees a half-written entry. The content is decided
+        // entirely by the key, so a writer that loses the race replaces the file with the same
+        // bytes and no locking is needed.
+        auto tmp = path;
+        tmp += "." + unique_string("tmp");
+        try
+        {
+            fs::create_directories(path.parent_path());
+            write_stamp(fs::path(root) / version_dir());
+            write_buffer(tmp, to_msgpack(migraphx::to_value(e)));
+            fs::rename(tmp, path);
+        }
+        catch(const std::exception& ex)
+        {
+            // Leaving the temporary behind would accumulate in the cache directory.
+            std::error_code ec;
+            fs::remove(tmp, ec);
+            log::warn() << "Failed to write binary cache entry " << path << ": " << ex.what();
+        }
+    }
+    memo[std::move(e.key)] = std::move(e.code);
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx