Fix data transfer routing to the selected backend

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) {

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
+}  // namespace gpu_ep

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

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 {

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<DataTransfer>(*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<DataTransfer>(*this).release();
     return STATUS_OK;
 }