diff --git a/src/migraphx/mgx_ep.cc b/src/migraphx/mgx_ep.cc index 791fba5..219370b 100644 --- a/src/migraphx/mgx_ep.cc +++ b/src/migraphx/mgx_ep.cc @@ -1102,17 +1102,15 @@ try { return Ort::Status{e.what(), ORT_EP_FAIL}; } -Ort::Status ExecutionProvider::OnRunEnd(const OrtRunOptions* /* run_options */, bool /* sync_stream */) noexcept -try { - HIP_RETURN_IF_ERROR(hipSetDevice(device_id_)); - if (const auto status{hipStreamQuery(stream_)}; status != hipSuccess) { - HIP_RETURN_IF_ERROR(hipStreamSynchronize(stream_)); - } +Ort::Status ExecutionProvider::OnRunEnd(const OrtRunOptions* /* run_options */, bool /* sync_stream */) noexcept { + // Nothing to synchronize here. Compute runs on ORT's per-run device stream + // (created in CreateSyncStreamForDevice and handed to the kernel via + // GetGPUComputeStream()), and ORT flushes that stream at run end through + // DeviceStreamCollection::CleanUp -> SyncStream::Flush, honoring the run's + // sync_stream setting. The previous implementation synchronized the EP's + // default stream, which never carried any compute work, so it was both + // incorrect and a redundant CPU/GPU serialization point. return STATUS_OK; -} catch (const Ort::Exception& e) { - return Ort::Status{e}; -} catch (const std::exception& e) { - return Ort::Status{e.what(), ORT_EP_FAIL}; } Ort::Status ExecutionProvider::CreateSyncStreamForDevice(const OrtMemoryDevice* memory_device, OrtSyncStreamImpl** stream) @@ -1232,7 +1230,24 @@ Ort::Status NodeComputeInfo::Compute(ComputeState& compute_state, const Ort::Ker return shape; }}; - if (!compute_state.has_input_shapes) { + // Item 4: gather the actual input shapes once. If they are identical to the + // previous call we can skip both the shape-compare loop and the rehash loop: + // identical actual shapes imply an identical effective shape, dyn bucket, and + // hash, and (since the program was already compiled for them last time) a match. + std::vector current_input_shapes; + current_input_shapes.reserve(input_name_indices.size() * 4); + for (const auto& [name, index] : input_name_indices) { + const auto shape{kernel_context.GetInput(index).GetTensorTypeAndShapeInfo().GetShape()}; + current_input_shapes.insert(current_input_shapes.end(), shape.begin(), shape.end()); + } + const bool shapes_unchanged{compute_state.has_last_input_shapes && + current_input_shapes == compute_state.last_input_shapes}; + + if (shapes_unchanged) { + input_shapes_hash = compute_state.last_input_shapes_hash; + input_shapes_match = true; + param_shapes = program.get_parameter_shapes(); + } else if (!compute_state.has_input_shapes) { for (auto& [name, index] : input_name_indices) { auto value{kernel_context.GetInput(index)}; auto shape{effective_shape(value.GetTensorTypeAndShapeInfo().GetShape(), name)}; @@ -1347,25 +1362,50 @@ Ort::Status NodeComputeInfo::Compute(ComputeState& compute_state, const Ort::Ker if (dyn.active) { compute_state.cached_programs.emplace(current_hash, program); } + // The program for this hash was (re)built, so any binding cached + // against the previous program for the same hash is stale. + compute_state.staging_bind_cache.erase(current_hash); } param_shapes = program.get_parameter_shapes(); } - // Staging path: required for hipGraph capture (pointer stability), for dynamic - // batching (batch-axis pad/slice), and for static seq-padding (token-axis - // pad/slice). Stage I/O into EP-owned buffers, bind scratch, then replay/capture - // a graph or run eagerly. + // Item 4: remember these shapes so the next identical call takes the fast path. + // Set only after the (possible) recompile above, so the recorded hash always + // corresponds to a program that is compiled and ready for these shapes. + if (!shapes_unchanged) { + compute_state.last_input_shapes = std::move(current_input_shapes); + compute_state.last_input_shapes_hash = input_shapes_hash; + compute_state.has_last_input_shapes = true; + } + + // Staging path: required for hipGraph capture (pointer stability) and for + // dynamic batching (input padding / output slicing). Stage I/O into EP-owned + // buffers, bind scratch, then replay/capture a graph or run eagerly if ((compute_state.hip_graph_enable || dyn.active || seq.active) && param_shapes.size() > 0) { const auto shape_hash{hash::ToHex(input_shapes_hash)}; const auto hip_stream{static_cast(kernel_context.GetGPUComputeStream())}; HIP_RETURN_IF_ERROR(hipSetDevice(compute_state.device_id)); AllocateStaging(compute_state, param_shapes, hip_stream, dyn); - CopyInputsToStaging(compute_state, param_shapes, kernel_context, hip_stream, dyn, seq); - auto bind{BindStagingParams(compute_state, param_shapes, shape_hash, hip_stream)}; + CopyInputsToStaging(compute_state, param_shapes, kernel_context, hip_stream, dyn); + // Item 3: reuse a cached binding for this shape hash. Staging buffers and + // scratch are pointer-stable until FreeStaging, so binding once and replaying + // avoids re-doing N program_parameters.add() calls, string work, and a + // scratch lookup on every inference. + auto bind_it{compute_state.staging_bind_cache.find(shape_hash)}; + if (bind_it == compute_state.staging_bind_cache.end()) { + bind_it = compute_state.staging_bind_cache.emplace( + shape_hash, + BindStagingParams(compute_state, param_shapes, shape_hash, hip_stream)).first; + } + auto& bind{bind_it->second}; RunProgramOrHipGraph(compute_state, hip_stream, kernel_context, program, bind.params, bind.prog_output_indices, shape_hash, dyn); - CopyStagingOutputsToOrt(compute_state, bind, kernel_context, hip_stream, dyn, seq); - HIP_RETURN_IF_ERROR(hipStreamSynchronize(hip_stream)); + CopyStagingOutputsToOrt(compute_state, bind, kernel_context, hip_stream, dyn); + // No per-Compute sync: all work above is enqueued on ORT's compute stream, + // which ORT flushes at run end (DeviceStreamCollection::CleanUp -> + // SyncStream::Flush) honoring the run's sync_stream setting. Cross-stream + // consumers are ordered via ORT notifications. Syncing here would only + // serialize CPU/GPU and defeat that overlap. return STATUS_OK; } @@ -1414,7 +1454,7 @@ Ort::Status NodeComputeInfo::Compute(ComputeState& compute_state, const Ort::Ker HIP_RETURN_IF_ERROR(hipSetDevice(compute_state.device_id)); auto hip_stream{static_cast(kernel_context.GetGPUComputeStream())}; auto prog_outputs{program.run_async(compute_params, hip_stream)}; - HIP_RETURN_IF_ERROR(hipStreamSynchronize(hip_stream)); + // Enqueue only; ORT flushes the compute stream at run end (see OnRunEnd). if (auto output_size{prog_outputs.size()}; output_indices.size() < output_size) { for (size_t i{}; i < output_size; ++i) { @@ -1430,7 +1470,6 @@ Ort::Status NodeComputeInfo::Compute(ComputeState& compute_state, const Ort::Ker HIP_CALL_THROW(hipMemcpyWithStream(output_data, gpu_resource.data(), resource_shape.bytes(), hipMemcpyDeviceToDevice, hip_stream)); } - HIP_RETURN_IF_ERROR(hipStreamSynchronize(hip_stream)); } } return STATUS_OK; @@ -1503,7 +1542,7 @@ try { HIP_RETURN_IF_ERROR(hipSetDevice(compute_state.device_id)); auto hip_stream{static_cast(kernel_context.GetGPUComputeStream())}; auto prog_outputs{program.run_async(compute_params, hip_stream)}; - HIP_RETURN_IF_ERROR(hipStreamSynchronize(hip_stream)); + // Enqueue only; ORT flushes the compute stream at run end (see OnRunEnd). if (auto output_size{prog_outputs.size()}; output_indices.size() < output_size) { for (size_t i{}; i < output_size; ++i) { @@ -1519,7 +1558,6 @@ try { HIP_CALL_THROW(hipMemcpyWithStream(output_data, gpu_resource.data(), resource_shape.bytes(), hipMemcpyDeviceToDevice, hip_stream)); } - HIP_RETURN_IF_ERROR(hipStreamSynchronize(hip_stream)); } } return STATUS_OK; diff --git a/src/migraphx/mgx_ep.h b/src/migraphx/mgx_ep.h index fe7a96b..78b7e1e 100644 --- a/src/migraphx/mgx_ep.h +++ b/src/migraphx/mgx_ep.h @@ -18,6 +18,7 @@ #include "common/path_string.h" #include "common/plugin_ep_utils.h" +#include "common/murmurhash3.h" #include "mgx_factory.h" #include "mgx_info.h" @@ -99,6 +100,17 @@ struct CapturedHipGraph { std::vector> captured_output_zeroes{}; }; +// Result of binding staging buffers (and the EP-owned scratch) as program +// parameters for a given compiled shape. Cached per shape hash in ComputeState: +// staging buffers and scratch are pointer-stable until FreeStaging, so a binding +// built once can be replayed unchanged instead of rebuilt every Compute call. +struct StagingBindResult { + migraphx::program_parameters params{}; + std::vector prog_output_indices{}; // ORT output index per bound output + std::vector bound_output_names{}; // staging key per bound output + std::vector bound_output_shapes{}; // current bucket shape per bound output +}; + struct ComputeState { std::mutex& mutex; int device_id; @@ -170,6 +182,17 @@ struct ComputeState { Map scratch_bufs{}; // Captured graphs keyed by shape hash. Map hip_graph_cache{}; + + // ── Binding / shape-hash fast-path caches ──────────────────────────────── + // Staging parameter bindings keyed by shape hash (multi-entry, so alternating + // dynamic-batch buckets each keep their binding). Invalidated by FreeStaging + // (staging pointers change) and per-hash on recompile. + Map staging_bind_cache{}; + // Last call's actual input shapes and their hash, for skipping the shape-compare + // and rehash loops when the shapes are unchanged from the previous Compute call. + std::vector last_input_shapes{}; + hash::Value last_input_shapes_hash{}; + bool has_last_input_shapes{}; }; struct EpContextComputeState { @@ -245,7 +268,6 @@ struct ExecutionProvider : OrtEp, ApiPtrs { Map ep_context_compute_states_; Map compute_states_; - hipStream_t stream_{}; hipDeviceProp_t device_prop_{}; int device_id_{}; diff --git a/src/migraphx/mgx_hip_graph.cc b/src/migraphx/mgx_hip_graph.cc index ed46c6f..5bd3f02 100644 --- a/src/migraphx/mgx_hip_graph.cc +++ b/src/migraphx/mgx_hip_graph.cc @@ -665,6 +665,9 @@ void FreeStaging(ComputeState& cs) { cs.staging_inputs.clear(); cs.staging_outputs.clear(); cs.staging_allocated = false; + // Cached bindings reference the staging buffers just freed; drop them so they + // are rebuilt against the next allocation. + cs.staging_bind_cache.clear(); } void DestroyHipGraphs(ComputeState& cs) { diff --git a/src/migraphx/mgx_hip_graph.h b/src/migraphx/mgx_hip_graph.h index 90ba8e1..093e9a7 100644 --- a/src/migraphx/mgx_hip_graph.h +++ b/src/migraphx/mgx_hip_graph.h @@ -91,13 +91,7 @@ void CopyInputsToStaging(ComputeState& cs, const DynamicBatchContext& dyn, const StaticSeqContext& seq); -// Result of binding staging buffers as program parameters. -struct StagingBindResult { - migraphx::program_parameters params{}; - std::vector prog_output_indices{}; // ORT output index per bound output - std::vector bound_output_names{}; // staging key per bound output - std::vector bound_output_shapes{}; // current bucket shape per bound output -}; +// StagingBindResult is defined in mgx_ep.h (cached per shape hash on ComputeState). // Bind staging input/output buffers and the EP-owned scratch buffer as program // parameters for the given compiled shape.