[DML EP] Add BFC allocator

include/onnxruntime/core/framework/execution_provider.h

@@ -324,6 +324,13 @@
     return default_device_;
   };
 
+  /**
+   * Return the appropriate OrtDevice object given OrtMemType that can be used directly by external callers.
+   */
+  virtual OrtDevice GetExternalOrtDeviceByMemType(OrtMemType mem_type) const {
+    return GetOrtDeviceByMemType(mem_type);
+  };
+
   /**
    * Create Preferred allocators for the current Execution Provider
    * This function is a stateless function which creates new instances of Allocator, without storing them in EP.

include/onnxruntime/core/framework/ortdevice.h

@@ -24,6 +24,7 @@ struct OrtDevice {
     static const MemoryType CUDA_PINNED = 1;
     static const MemoryType HIP_PINNED = 2;
     static const MemoryType CANN_PINNED = 3;
+    static const MemoryType DML_EXTERNAL = 4;
   };
 
   constexpr OrtDevice(DeviceType device_type_, MemoryType memory_type_, DeviceId device_id_)

onnxruntime/core/framework/allocator.cc

@@ -139,12 +139,17 @@ ORT_API_STATUS_IMPL(OrtApis::CreateMemoryInfo, _In_ const char* name1, enum OrtA
     *out = new OrtMemoryInfo(onnxruntime::CPU, type, OrtDevice(), id1, mem_type1);
   } else if (strcmp(name1, onnxruntime::CUDA) == 0 ||
              strcmp(name1, onnxruntime::OpenVINO_GPU) == 0 ||
-             strcmp(name1, onnxruntime::DML) == 0 ||
              strcmp(name1, onnxruntime::HIP) == 0 ||
              strcmp(name1, onnxruntime::WEBGPU_BUFFER) == 0) {
     *out = new OrtMemoryInfo(
         name1, type, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast<OrtDevice::DeviceId>(id1)), id1,
         mem_type1);
+  } else if (strcmp(name1, onnxruntime::DML) == 0) {
+    // Since EPs cannot have 2 allocators with the same OrtMemType and Memory ID,
+    // we use -1 as the memory ID to represent external allocations that don't have any allocator.
+    *out = new OrtMemoryInfo(
+        name1, type, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DML_EXTERNAL, static_cast<OrtDevice::DeviceId>(id1)),
+        -1, mem_type1);
   } else if (strcmp(name1, onnxruntime::OpenVINO_RT_NPU) == 0) {
     *out = new OrtMemoryInfo(
         name1, type, OrtDevice(OrtDevice::NPU, OrtDevice::MemType::DEFAULT, static_cast<OrtDevice::DeviceId>(id1)), id1,

onnxruntime/core/framework/bfc_arena.cc

@@ -41,23 +41,8 @@
   memory_limit_ = total_memory;
   stats_.bytes_limit = static_cast<int64_t>(total_memory);
 
-  arena_extend_strategy_ = arena_extend_strategy;
+  SetArenaExtendStrategy(arena_extend_strategy);
 
-  // We never want to shrink the initial allocation if the arena extend strategy is kNextPowerOfTwo.
-  // This could seem confusingly arbitrary but the rationale is as follows:
-  // The user selected initial allocation chunk is only valid for the arena extend strategy kNextPowerOfTwo
-  // and the user has likely chosen this initial value so that any ad-hoc arena extensions/shrinkages could potentially
-  // be avoided. So we do not consider the initial allocation for shrinkage whatever its usage status.
-  // On the other hand, if the arena extension strategy is kSameAsRequested, any initial chunk set by the user or otherwise,
-  // is moot and the arena will only extend based on the request size. In these cases, we consider any allocation for shrinkage
-  // if it is left unused (even if it is the first allocation).
-  if (arena_extend_strategy_ == ArenaExtendStrategy::kSameAsRequested) {
-    // Consider all allocation regions (including first allocation region) for shrinkage
-    consider_first_allocation_region_for_shrinkage_ = true;
-  } else {  // arena_extend_strategy_ == kNextPowerOfTwo
-    // Do not consider the first allocation region for shrinkage
-    consider_first_allocation_region_for_shrinkage_ = false;
-  }
   // Create a bunch of bins of various good sizes.
 
   // We create bins to fit all possible ranges that cover the
@@ -91,6 +76,29 @@
   }
 }
 
+void BFCArena::UpdateFirstAllocationShrinkageLogic() {
+  // We never want to shrink the initial allocation if the arena extend strategy is kNextPowerOfTwo.
+  // This could seem confusingly arbitrary but the rationale is as follows:
+  // The user selected initial allocation chunk is only valid for the arena extend strategy kNextPowerOfTwo
+  // and the user has likely chosen this initial value so that any ad-hoc arena extensions/shrinkages could potentially
+  // be avoided. So we do not consider the initial allocation for shrinkage whatever its usage status.
+  // On the other hand, if the arena extension strategy is kSameAsRequested, any initial chunk set by the user or otherwise,
+  // is moot and the arena will only extend based on the request size. In these cases, we consider any allocation for shrinkage
+  // if it is left unused (even if it is the first allocation).
+  if (arena_extend_strategy_ == ArenaExtendStrategy::kSameAsRequested) {
+    // Consider all allocation regions (including first allocation region) for shrinkage
+    consider_first_allocation_region_for_shrinkage_ = true;
+  } else {  // arena_extend_strategy_ == kNextPowerOfTwo
+    // Do not consider the first allocation region for shrinkage
+    consider_first_allocation_region_for_shrinkage_ = false;
+  }
+}
+
+void BFCArena::SetArenaExtendStrategy(ArenaExtendStrategy arena_extend_strategy) {
+  arena_extend_strategy_ = arena_extend_strategy;
+  UpdateFirstAllocationShrinkageLogic();
+}
+
 BFCArena::Chunk* BFCArena::ChunkFromHandle(ChunkHandle h) {
   ORT_ENFORCE(h < chunks_.size());
   return &(chunks_[h]);

onnxruntime/core/framework/bfc_arena.h

@@ -77,6 +77,11 @@ class BFCArena : public IAllocator {
 
   ~BFCArena() override;
 
+  // Allows the caller to change the arena extend strategy after the allocator is done initializing.
+  // For example, kSameAsRequested may be desirable in certain situations and kNextPowerOfTwo may be
+  // desirable in others.
+  void SetArenaExtendStrategy(ArenaExtendStrategy arena_extend_strategy);
+
   // If size is 0, then this function returns either NULL,
   // or a unique pointer value that can later be successfully
   // passed to free(). Whatever, do not dereference that pointer
@@ -123,6 +128,9 @@ class BFCArena : public IAllocator {
  private:
   void DeallocateRawInternal(void* ptr);
 
+  // Updates whether the first allocation should be considered for shrinkage depending on the strategy type.
+  void UpdateFirstAllocationShrinkageLogic();
+
   // A ChunkHandle is an index into the chunks_ vector in BFCAllocator
   // kInvalidChunkHandle means an invalid chunk
   using ChunkHandle = size_t;

onnxruntime/core/framework/utils.cc

@@ -163,6 +163,19 @@
     return Status::OK();
   }
 
+#ifdef USE_DML
+  const bool bothValuesOnGPU = copy_info.source_device.Type() == OrtDevice::GPU && copy_info.target_device.Type() == OrtDevice::GPU;
+  const bool sourceIsDmlAlloc = copy_info.source_device.MemType() == OrtDevice::MemType::DEFAULT || copy_info.source_device.MemType() == OrtDevice::MemType::DML_EXTERNAL;
+  const bool targetIsInternalAlloc = copy_info.target_device.MemType() == OrtDevice::MemType::DEFAULT;
+  const bool bothValuesOnSameDevice = copy_info.source_device.Id() == copy_info.target_device.Id();
+
+  // The DML EP supports binding external allocations directly, even if the memory types don't match, as long as they are on the same D3D12 device
+  if (bothValuesOnGPU && sourceIsDmlAlloc && targetIsInternalAlloc && bothValuesOnSameDevice) {
+    target_mlvalue = source_mlvalue;
+    return Status::OK();
+  }
+#endif
+
   auto allocator = session_state.GetAllocator(copy_info.target_device);
   if (!target_mlvalue.IsAllocated()) {
     ORT_ENFORCE(allocator != nullptr, "Failed to find allocator for device ", copy_info.target_device.ToString());

onnxruntime/core/providers/dml/DmlExecutionProvider/inc/DmlExecutionProvider.h

@@ -10,6 +10,7 @@ interface ID3D12Resource;
 #include "core/common/status.h"
 #include "core/framework/data_transfer.h"
 #include "IWinmlExecutionProvider.h"
+#include "core/providers/dml/DmlExecutionProvider/src/DmlBufferRegion.h"
 #include "core/providers/dml/DmlExecutionProvider/src/ExecutionContext.h"
 
 namespace onnxruntime
@@ -21,12 +22,6 @@ namespace onnxruntime
     class KernelRegistry;
 }
 
-enum class AllocatorRoundingMode
-{
-    Disabled = 0,
-    Enabled = 1,
-};
-
 namespace Dml
 {
     std::unique_ptr<onnxruntime::IExecutionProvider> CreateExecutionProvider(
@@ -35,9 +30,9 @@ namespace Dml
         bool enableMetacommands,
         bool enableGraphCapture,
         bool enableCpuSyncSpinning,
-        bool disableMemoryArena);
+        bool disableMemoryArena,
+        bool enableBfcAllocator);
 
-    ID3D12Resource* GetD3D12ResourceFromAllocation(onnxruntime::IAllocator* allocator, void* ptr);
     void FlushContext(onnxruntime::IExecutionProvider* provider);
     void ReleaseCompletedReferences(onnxruntime::IExecutionProvider* provider);
 

onnxruntime/core/providers/dml/DmlExecutionProvider/inc/IWinmlExecutionProvider.h

@@ -10,6 +10,7 @@
 #include <wrl/client.h>
 
 #include "core/framework/op_kernel.h"
+#include "core/providers/dml/DmlExecutionProvider/src/DmlBufferRegion.h"
 #include "core/providers/dml/DmlExecutionProvider/src/DmlEdgeShapes.h"
 
 struct AbstractOperatorDesc;
@@ -25,6 +26,11 @@ namespace onnxruntime
     class Node;
 }
 
+namespace Dml
+{
+    struct TaggedPointer;
+}
+
 namespace Windows::AI::MachineLearning::Adapter
 {
     interface __declspec(uuid("5b19a18a-5ed5-4df2-a363-21b89380a698"))
@@ -37,19 +43,9 @@ namespace Windows::AI::MachineLearning::Adapter
         // the provider's underlying queues.
         virtual void QueueReference(IUnknown *object) = 0;
 
-        virtual void GetShadowCopyIfRequired(
-            bool isInternalOperator,
-            IUnknown* data,
-            IUnknown** dataCopy) const = 0;
-
-        virtual void GetABIDataInterface(
-            bool isInternalOperator,
-            IUnknown* data,
-            IUnknown** abiData) const = 0;
+        virtual Dml::D3D12BufferRegion GetBufferRegion(void* opaquePointer, uint64_t size) const = 0;
 
-        virtual uint64_t TryGetPooledAllocationId(
-            IUnknown* data,
-            bool isInternalOperator) = 0;
+        virtual uint64_t GetUniqueId(void* opaquePointer) = 0;
 
         virtual void GetABIExecutionInterfaceAndInvalidateState(
             bool isInternalOperator,

onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.cpp

@@ -564,11 +564,17 @@ HRESULT STDMETHODCALLTYPE AbiCustomRegistry::RegisterOperatorKernel(
         //
         // For backward compatibility, this does not propagate errors for external operators
         static_cast<void>(m_kernelRegistry->RegisterCustomKernel(create_info));  // ignore result
+        m_hasExternalOperators = true;
     }
 
     return S_OK;
     }
     ORT_CATCH_RETURN
 }
 
+bool STDMETHODCALLTYPE AbiCustomRegistry::HasExternalOperators() const noexcept
+{
+    return m_hasExternalOperators;
+}
+
 }

onnxruntime/core/providers/dml/DmlExecutionProvider/src/AbiCustomRegistry.h

@@ -50,6 +50,8 @@
         IMLOperatorKernelFactory* operatorKernelFactory,
         _In_opt_ IMLOperatorShapeInferrer* shapeInferrer) const noexcept override;
 
+    bool STDMETHODCALLTYPE HasExternalOperators() const noexcept override;
+
     std::list<std::shared_ptr<onnxruntime::CustomRegistry>> GetRegistries()
     {
         std::list<std::shared_ptr<onnxruntime::CustomRegistry>> registries;
@@ -108,6 +110,8 @@
     // Map between Lotus KernelDefs and extended data used during partitioning
     mutable std::shared_ptr<InternalRegistrationInfoMap> m_internalRegInfoMap;
 
+    mutable bool m_hasExternalOperators = false;
+
 };
 
 }    // namespace Windows::AI::MachineLearning::Adapter