diff --git a/src/amdgpu/gpu_allocator.cc b/src/amdgpu/gpu_allocator.cc index 8ab11d7..7175de6 100644 --- a/src/amdgpu/gpu_allocator.cc +++ b/src/amdgpu/gpu_allocator.cc @@ -30,6 +30,9 @@ Allocator::~Allocator() { } } +// Always lazy. Sessions no longer reach this wrapper — they resolve allocators per-session +// through OrtEp::CreateAllocator (gpu_ep.cc). What is left is the factory-level fallback, +// notably ORT's environment shared allocator, which is shared and has no session to pin to. OrtAllocator* Allocator::GetBackendAllocator() const noexcept { const auto backend_factory{factory_.GetBackendFactory()}; if (backend_factory == nullptr) { diff --git a/src/amdgpu/gpu_data_transfer.cc b/src/amdgpu/gpu_data_transfer.cc index f1df7d3..a9b6794 100644 --- a/src/amdgpu/gpu_data_transfer.cc +++ b/src/amdgpu/gpu_data_transfer.cc @@ -6,8 +6,8 @@ namespace gpu_ep { -DataTransfer::DataTransfer(const ProviderFactory& factory, OrtEpFactory* backend_factory) - : OrtDataTransferImpl{ORT_API_VERSION}, factory_{factory}, backend_factory_{backend_factory} +DataTransfer::DataTransfer(const ProviderFactory& factory) + : OrtDataTransferImpl{ORT_API_VERSION}, factory_{factory} { OrtDataTransferImpl::Release = [](OrtDataTransferImpl* this_) noexcept { // ORT creates one DataTransfer per session and owns it, delete on Release. @@ -28,36 +28,65 @@ DataTransfer::DataTransfer(const ProviderFactory& factory, OrtEpFactory* backend API_CALL_S(DataTransfer, this_, CopyTensors, src_tensors, dst_tensors, streams, num_tensors); }; - // Snapshot the backend transfer once from the backend this session selected. - // A null backend_factory_ (library-registration-time creation) leaves this - // instance inert: CanCopy returns false, CopyTensors returns an error. - if (backend_factory_ != nullptr) { - if (backend_factory_->CreateDataTransfer(backend_factory_, &backend_data_transfer_) != nullptr) { + // A backend here means we are the per-session instance (see header): freeze on it, so + // a later session overwriting the process-global slot cannot redirect this session's + // copies. No backend means the registration-time instance: stay lazy. + // + // frozen_ is still false, so this call takes the lazy path and populates + // backend_data_transfer_ before we freeze on the result. Two statements to keep that + // ordering explicit. A failed CreateDataTransfer leaves frozen_ false, so the instance + // retries later instead of latching the failure. + const bool backend_already_selected{GetBackendDataTransfer() != nullptr}; + frozen_ = backend_already_selected; +} + +OrtDataTransferImpl* DataTransfer::GetBackendDataTransfer() const noexcept { + if (frozen_) { + // Pinned to this session's own backend; the global slot may since have changed. + return backend_data_transfer_; + } + const auto backend_factory{factory_.GetBackendFactory()}; + if (backend_factory == nullptr) { + // No backend selected yet (e.g. called before any CreateEp). Not an error — + // this instance resolves once a backend exists. + return nullptr; + } + if (backend_factory != backend_factory_) { + // First use, or the backend changed (e.g. profile switch) — (re-)query the + // transfer from the currently selected backend factory. The backend owns it + // (Release is a no-op there), so there is nothing to release for the old one. + backend_data_transfer_ = nullptr; + backend_factory_ = nullptr; + if (backend_factory->CreateDataTransfer(backend_factory, &backend_data_transfer_) != nullptr) { backend_data_transfer_ = nullptr; - backend_factory_ = nullptr; + return nullptr; } + backend_factory_ = backend_factory; } + return backend_data_transfer_; } bool DataTransfer::CanCopy(const OrtMemoryDevice* src_memory_device, const OrtMemoryDevice* dst_memory_device) const noexcept { - if (backend_data_transfer_ == nullptr) { + const auto backend_data_transfer{GetBackendDataTransfer()}; + if (backend_data_transfer == nullptr) { return false; } - return backend_data_transfer_->CanCopy(backend_data_transfer_, + return backend_data_transfer->CanCopy(backend_data_transfer, src_memory_device, dst_memory_device); } Ort::Status DataTransfer::CopyTensors(const OrtValue** src_tensors, OrtValue** dst_tensors, OrtSyncStream** streams, size_t num_tensors) const noexcept { - if (backend_data_transfer_ == nullptr) { + const auto backend_data_transfer{GetBackendDataTransfer()}; + if (backend_data_transfer == nullptr) { return MAKE_STATUS(ORT_EP_FAIL, "invalid backend factory"); } - RETURN_IF_ERROR(backend_data_transfer_->CopyTensors(backend_data_transfer_, + RETURN_IF_ERROR(backend_data_transfer->CopyTensors(backend_data_transfer, src_tensors, dst_tensors, streams, num_tensors)); return STATUS_OK; } -} // namespace gpu_ep \ No newline at end of file +} // namespace gpu_ep diff --git a/src/amdgpu/gpu_data_transfer.h b/src/amdgpu/gpu_data_transfer.h index 2f486aa..2b7040d 100644 --- a/src/amdgpu/gpu_data_transfer.h +++ b/src/amdgpu/gpu_data_transfer.h @@ -11,13 +11,17 @@ struct ProviderFactory; struct DataTransfer : OrtDataTransferImpl { DataTransfer() = delete; - // Per-session: the backend is snapshotted at construction from the factory that - // the session's CreateEp just selected, and never re-queried. ORT creates one - // DataTransfer per session (CreateDataTransfer -> GetDataTransfer) and owns it, - // Releasing (deleting) it at session teardown. backend_factory may be null when - // ORT creates a transfer at library-registration time (before any backend is - // selected); such an instance is inert. - DataTransfer(const ProviderFactory& factory, OrtEpFactory* backend_factory); + // ORT creates one instance per session (after that session's CreateEp) and one per + // factory at library registration (before any CreateEp). These need opposite + // behaviour, so the constructor branches on whether a backend exists yet: + // + // - backend selected -> per-session: snapshot and freeze. Re-resolving would follow + // ProviderFactory's process-global slot into a *later* session's backend. + // - no backend -> registration-time instance, the only one the env-level + // OrtApi::CopyTensors can reach. Stay lazy; resolving here breaks it permanently. + // + // ORT owns each instance and Releases (deletes) it at teardown. + explicit DataTransfer(const ProviderFactory& factory); private: bool CanCopy(const OrtMemoryDevice* src_memory_device, @@ -26,9 +30,20 @@ struct DataTransfer : OrtDataTransferImpl { [[nodiscard]] Ort::Status CopyTensors(const OrtValue** src_tensors, OrtValue** dst_tensors, OrtSyncStream** streams, size_t num_tensors) const noexcept; + // Returns the backend's data transfer. When frozen_, the snapshot taken in the + // constructor. Otherwise re-queried if the selected backend changed (e.g. profile + // switch), and null until a backend has been selected. + OrtDataTransferImpl* GetBackendDataTransfer() const noexcept; + + // True when a backend already existed at construction — i.e. the per-session + // instance. Not mutable: written only in the constructor, by design. The decision + // belongs there, not in the const GetBackendDataTransfer(). + bool frozen_{}; + + mutable OrtEpFactory* backend_factory_{}; + mutable OrtDataTransferImpl* backend_data_transfer_{}; + const ProviderFactory& factory_; - OrtEpFactory* backend_factory_{}; - OrtDataTransferImpl* backend_data_transfer_{}; }; } // namespace gpu_ep diff --git a/src/amdgpu/gpu_ep.cc b/src/amdgpu/gpu_ep.cc index e765fe7..94ae557 100644 --- a/src/amdgpu/gpu_ep.cc +++ b/src/amdgpu/gpu_ep.cc @@ -102,6 +102,12 @@ ExecutionProvider::ExecutionProvider(ProviderFactory& factory, std::string_view OrtEp::OnRunEnd = [](OrtEp* this_, const OrtRunOptions* run_options, bool sync_stream) noexcept { API_CALL_S(ExecutionProvider, this_, OnRunEnd, run_options, sync_stream); }; + // Wired for every profile so allocators resolve through this EP rather than factory_'s + // process-global backend slot, which the next session's CreateEp overwrites. + OrtEp::CreateAllocator = [](OrtEp* this_, const OrtMemoryInfo* memory_info, + OrtAllocator** allocator) noexcept { + API_CALL_S(ExecutionProvider, this_, CreateAllocator, memory_info, allocator); + }; OrtEp::CreateSyncStreamForDevice = [](OrtEp* this_, const OrtMemoryDevice* memory_device, OrtSyncStreamImpl** stream) noexcept { API_CALL_S(ExecutionProvider, this_, CreateSyncStreamForDevice, memory_device, stream); @@ -189,21 +195,10 @@ ExecutionProvider::ExecutionProvider(ProviderFactory& factory, std::string_view #ifdef USE_DML const auto create_directx_backend = [&] { THROW_IF_ERROR(factory.CreateDirectXBackend(local_session_options, logger, backend_ep_)); - // DirectML manages its own per-session GPU allocator (DmlBucketizedBufferAllocator) - // via EP-level CreateAllocator. Wire it now that we know the backend is DirectML. - // MIGraphX allocators are handled at factory level — leave OrtEp::CreateAllocator null - // so ORT falls back to ep_factory_.CreateAllocator (the Allocator wrapper). - OrtEp::CreateAllocator = [](OrtEp* this_, const OrtMemoryInfo* memory_info, - OrtAllocator** allocator) noexcept { - API_CALL_S(ExecutionProvider, this_, CreateAllocator, memory_info, allocator); - }; }; #endif const auto create_hip_backend = [&] { - // hip backend manages allocator/data-transfer at the backend factory level, - // reached through the amdgpu Allocator/DataTransfer wrappers — leave - // OrtEp::CreateAllocator null so ORT falls back to ep_factory_.CreateAllocator. THROW_IF_ERROR(factory.CreateHipBackend(local_session_options, logger, backend_ep_)); }; @@ -407,7 +402,18 @@ Ort::Status ExecutionProvider::OnRunStart(const OrtRunOptions* run_options) cons Ort::Status ExecutionProvider::CreateAllocator(const OrtMemoryInfo* memory_info, OrtAllocator** allocator) const noexcept { - EP_CALL_S(backend_ep_, CreateAllocator, memory_info, allocator); + // DirectML implements this on its OrtEp; migraphx and hip implement it on their factory. + // Prefer the former, as ORT itself does. The explicit null check is needed because + // EP_CALL_S reports a missing function pointer as success, leaving *allocator unset. + if (backend_ep_ != nullptr && backend_ep_->CreateAllocator != nullptr) { + EP_CALL_S(backend_ep_, CreateAllocator, memory_info, allocator); + } + if (backend_ep_factory_ == nullptr) { + return MAKE_STATUS(ORT_EP_FAIL, "CreateAllocator: invalid backend factory"); + } + RETURN_IF_ERROR(backend_ep_factory_->CreateAllocator(backend_ep_factory_, memory_info, + nullptr /*allocator_options*/, allocator)); + return STATUS_OK; } Ort::Status ExecutionProvider::OnRunEnd(const OrtRunOptions* run_options, bool sync_stream) const noexcept { diff --git a/src/amdgpu/gpu_factory.cc b/src/amdgpu/gpu_factory.cc index e259fd6..939f138 100644 --- a/src/amdgpu/gpu_factory.cc +++ b/src/amdgpu/gpu_factory.cc @@ -326,11 +326,11 @@ void ProviderFactory::ReleaseAllocator(OrtAllocator*) const { } Ort::Status ProviderFactory::CreateDataTransfer(OrtDataTransferImpl** data_transfer) { - // Per-session: hand ORT a fresh DataTransfer bound to the backend this session - // selected (GetBackendFactory() reflects the CreateEp that just ran; it is null at - // library-registration time, yielding an inert instance). ORT owns it and calls - // Release (which deletes it) at session teardown. - *data_transfer = std::make_unique(*this, GetBackendFactory()).release(); + // Called once per session (after that session's CreateEp) and once per factory at + // library-registration time, before any backend exists. DataTransfer's constructor + // tells the two apart and freezes or stays lazy accordingly — see its header. + // ORT owns each instance and calls Release (which deletes it) at teardown. + *data_transfer = std::make_unique(*this).release(); return STATUS_OK; }