Add opt-in co-resident program cache to avoid reloads on input-shape switches

src/migraphx/mgx_ep.cc

@@ -456,6 +456,8 @@ ExecutionProvider::ExecutionProvider(const ProviderFactory& factory, std::string
     PARSE_ENV_VAR(env_var::kDumpSubgraphs, dump_subgraphs_);
     PARSE_ENV_VAR(env_var::kDumpEpContextModel, context_enable_);
     PARSE_ENV_VAR(env_var::kExhaustiveTune, exhaustive_tune_);
+    PARSE_ENV_VAR(env_var::kCoresidentPrograms, coresident_programs_);
+    PARSE_ENV_VAR(env_var::kMaxResidentPrograms, max_resident_programs_);
     PARSE_ENV_VAR(env_var::kMlssUseSpecificOps, mlss_use_specific_ops_);
 
     platform::SetEnvironmentVar("MIGRAPHX_MLSS_USE_SPECIFIC_OPS", mlss_use_specific_ops_);
@@ -798,6 +800,11 @@ Ort::Status ExecutionProvider::CreateNodeComputeInfoFromGraph(const Ort::ConstGr
             t_,
             onnx_options,
             program,
+            coresident_programs_,
+            max_resident_programs_,
+            {},  // resident_programs
+            {},  // lru_order
+            {},  // active_shape_key
             enable_fp16_,
             enable_bf16_,
             enable_fp8_,
@@ -1017,6 +1024,50 @@ Ort::Status NodeComputeInfo::Compute(ComputeState& compute_state, const Ort::Ker
         }
     }
 
+    // Registers a program in the resident cache as the most-recently-used entry
+    // and evicts least-recently-used entries until the bound is met. The key just
+    // registered is never the eviction victim.
+    const auto register_resident = [&compute_state](const std::string& key, const migraphx::program& prog) {
+        compute_state.active_shape_key = key;
+        compute_state.resident_programs[key] = prog;
+        compute_state.lru_order.remove(key);
+        compute_state.lru_order.push_back(key);
+        while (compute_state.resident_programs.size() > compute_state.max_resident_programs &&
+               !compute_state.lru_order.empty()) {
+            const auto victim{compute_state.lru_order.front()};
+            compute_state.lru_order.pop_front();
+            if (victim != key) {
+                compute_state.resident_programs.erase(victim);
+            }
+        }
+    };
+
+    // On a shape switch, stash the active program in the resident cache under its
+    // (previous) shape key, then check whether the incoming shape is already
+    // resident. On a hit, restore it and skip the reload entirely (a lookup that
+    // replaces a multi-second reload and finalize). Default off: falls through to
+    // the legacy single-slot overwrite path below.
+    if (!input_shapes_match && compute_state.coresident_programs) {
+        const std::string shape_key{hash::ToHex(input_shapes_hash)};
+        // Save the currently-active program under the previously-active key
+        // (if we have one and it isn't already stored).
+        if (!compute_state.active_shape_key.empty() &&
+            compute_state.resident_programs.find(compute_state.active_shape_key) ==
+                compute_state.resident_programs.end()) {
+            compute_state.resident_programs.emplace(compute_state.active_shape_key, program);
+            compute_state.lru_order.push_back(compute_state.active_shape_key);
+        }
+        if (const auto it{compute_state.resident_programs.find(shape_key)};
+            it != compute_state.resident_programs.end()) {
+            program = it->second;  // shared_ptr handle copy, no reload/finalize
+            param_shapes = program.get_parameter_shapes();
+            input_shapes_match = true;
+            // Mark the restored program most-recently-used and enforce the bound
+            // (the stash above may have pushed the cache one past max).
+            register_resident(shape_key, program);
+        }
+    }
+
     // If the input shapes are different (e.g., LLMs), the EP needs to reparse and recompile the program
     if (!input_shapes_match) {
         migraphx::program_parameters compile_params{};
@@ -1061,6 +1112,12 @@ Ort::Status NodeComputeInfo::Compute(ComputeState& compute_state, const Ort::Ker
             }
         }
         param_shapes = program.get_parameter_shapes();
+
+        // The freshly loaded/compiled program is now the active one; register it
+        // in the resident cache and enforce the LRU bound.
+        if (compute_state.coresident_programs) {
+            register_resident(hash::ToHex(input_shapes_hash), program);
+        }
     }
 
     migraphx::program_parameters compute_params;

src/migraphx/mgx_ep.h

@@ -4,8 +4,10 @@
 #pragma once
 
 #include <ciso646>
+#include <list>
 #include <set>
 #include <mutex>
+#include <unordered_map>
 
 #include <hip/hip_runtime_api.h>
 #include <migraphx/migraphx.hpp>
@@ -34,14 +36,31 @@ constexpr auto kComputeMode = "ORT_MIGRAPHX_COMPUTE_MODE"sv;
 constexpr auto kINT8UseNativeCalibrationTable = "ORT_MIGRAPHX_INT8_USE_NATIVE_CALIBRATION_TABLE"sv;
 constexpr auto kExhaustiveTune = "ORT_MIGRAPHX_EXHAUSTIVE_TUNE"sv;
 constexpr auto kMlssUseSpecificOps = "ORT_MIGRAPHX_MLSS_USE_SPECIFIC_OPS"sv;
+// Keep multiple shape-specialized programs (for example an LLM's prefill and
+// decode programs) resident at once so a shape switch is a map lookup instead of
+// a full reload and finalize. Default off (single-slot, legacy overwrite
+// behavior). When on, the cache is bounded by kMaxResidentPrograms.
+constexpr auto kCoresidentPrograms = "ORT_MIGRAPHX_CORESIDENT_PROGRAMS"sv;
+constexpr auto kMaxResidentPrograms = "ORT_MIGRAPHX_MAX_RESIDENT_PROGRAMS"sv;
 }  // namespace env_vars
 
 struct ComputeState {
     std::mutex& mutex;
     int device_id;
     const migraphx::target& t;
     migraphx::onnx_options onnx_options;
-    migraphx::program program;
+    migraphx::program program;  // currently-active program (last shape run)
+    // shape-hash -> resident program cache. When coresident_programs is on,
+    // programs for distinct input shapes are kept here so switching is a lookup,
+    // not a reload and finalize. lru_order is most-recent-last; on overflow past
+    // max_resident_programs the front (least-recently-used) entry is evicted. The
+    // program run every step stays at the back and so is never the victim.
+    // active_shape_key is the key of the program currently in the `program` slot.
+    bool coresident_programs{};
+    size_t max_resident_programs{4};
+    std::unordered_map<std::string, migraphx::program> resident_programs;
+    std::list<std::string> lru_order;
+    std::string active_shape_key;
     bool enable_fp16{};
     bool enable_bf16{};
     bool enable_fp8{};
@@ -146,6 +165,8 @@ struct ExecutionProvider : OrtEp, ApiPtrs {
     bool enable_fp8_{};
     bool enable_int8_{};
     bool exhaustive_tune_{};
+    bool coresident_programs_{};
+    size_t max_resident_programs_{4};
     std::string mlss_use_specific_ops_{};
     bool int8_calibration_cache_available_{};
     bool int8_use_native_calibration_table_{};