Use object-based DLPack interchange and test numeric dtypes in CI

onnxruntime/core/dlpack/dlpack_converter.cc

@@ -90,7 +90,7 @@ const char* GetOrtDeviceName(const OrtDevice& device) {
     case OrtDevice::CPU:
       return CPU;
     case OrtDevice::GPU:
-      return CUDA;
+      return device.Vendor() == OrtDevice::VendorIds::AMD ? HIP : CUDA;
     default:
       ORT_THROW("Unknown device type: ", device.Type());
   }
@@ -136,7 +136,15 @@ DLDevice GetDlpackDevice(const OrtValue& ort_value, const int64_t& device_id) {
       device.device_type = DLDeviceType::kDLCPU;
       break;
     case OrtDevice::GPU:
-      device.device_type = DLDeviceType::kDLCUDA;
+      if (location.device.Vendor() == OrtDevice::VendorIds::NVIDIA) {
+        device.device_type = DLDeviceType::kDLCUDA;
+      } else if (location.device.Vendor() == OrtDevice::VendorIds::AMD) {
+        device.device_type = DLDeviceType::kDLROCM;
+      } else {
+        ORT_THROW(
+            "DLPack GPU export supports only NVIDIA CUDA and AMD ROCm tensors. "
+            "Copy this tensor to CPU before exporting it.");
+      }
       break;
     default:
       ORT_THROW("Cannot pack tensors on this device.");
@@ -176,12 +184,11 @@ DLManagedTensor* OrtValueToDlpack(OrtValue& ort_value) {
   ort_dlmanaged_tensor->tensor.dl_tensor.device = GetDlpackDevice(ort_value, tensor.Location().device.Id());
   ort_dlmanaged_tensor->tensor.dl_tensor.ndim = static_cast<int>(tensor.Shape().NumDimensions());
   ort_dlmanaged_tensor->tensor.dl_tensor.dtype = GetDlpackDataType(ort_value);
-  ort_dlmanaged_tensor->tensor.dl_tensor.shape =
-      tensor.Shape().NumDimensions() > 0 ? &const_cast<TensorShape&>(tensor.Shape())[0] : nullptr;
+  const auto dims = tensor.Shape().GetDims();
+  ort_dlmanaged_tensor->tensor.dl_tensor.shape = dims.empty() ? nullptr : const_cast<int64_t*>(dims.data());
   auto& strides = ort_dlmanaged_tensor->strides;
-  strides.reserve(tensor.Shape().NumDimensions());
+  strides.reserve(dims.size());
   // Empty tensors have no addressable elements, so zero strides avoid overflowing dimension products.
-  const auto dims = tensor.Shape().GetDims();
   const bool is_empty = std::find(dims.begin(), dims.end(), 0) != dims.end();
   SafeInt<int64_t> stride = 1;
   for (size_t i = tensor.Shape().NumDimensions(); i > 0; --i) {

onnxruntime/core/dlpack/dlpack_converter.h

@@ -11,10 +11,8 @@
 namespace onnxruntime {
 namespace dlpack {
 
-// This convertor will take an OrtValue and wrap it as a DLPack tensor
-// This may create a new ownership to the underlying tensor in OrtValue,
-// so we do pass-by-value here. We don't use pass-by-reference because
-// it implies no new ownership.
+// Wrap an OrtValue as a DLPack tensor, retaining a copy of the OrtValue
+// until the consumer calls the DLPack deleter.
 DLManagedTensor* OrtValueToDlpack(OrtValue& ort_value);
 
 // Imports a contiguous tensor, preserving its DLPack dtype (including kDLBool).

onnxruntime/core/providers/cuda/cuda_provider_factory.cc

@@ -228,6 +228,20 @@ struct ProviderInfo_CUDA_Impl final : ProviderInfo_CUDA {
     return num_devices;
   }
 
+  Status SynchronizeDevice(int device_id) override {
+    int original_device;
+    ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetDevice(&original_device)));
+    if (original_device != device_id) {
+      ORT_RETURN_IF_ERROR(CUDA_CALL(cudaSetDevice(device_id)));
+    }
+    const auto sync_error = cudaDeviceSynchronize();
+    // Restore the caller's device even when synchronization reports an error.
+    const auto restore_error = original_device != device_id ? cudaSetDevice(original_device) : cudaSuccess;
+    ORT_RETURN_IF_ERROR(CUDA_CALL(sync_error));
+    ORT_RETURN_IF_ERROR(CUDA_CALL(restore_error));
+    return Status::OK();
+  }
+
   void CUDAExecutionProviderInfo__FromProviderOptions(const ProviderOptions& options, CUDAExecutionProviderInfo& info) override {
     info = CUDAExecutionProviderInfo::FromProviderOptions(options);
   }

onnxruntime/core/providers/cuda/cuda_provider_factory.h

@@ -61,6 +61,8 @@ struct ProviderInfo_CUDA {
     ORT_NOT_IMPLEMENTED(__FUNCTION__, " is only implements in test code path.");
   }
 
+  virtual Status SynchronizeDevice(int device_id) = 0;
+
  protected:
   ~ProviderInfo_CUDA() = default;  // Can only be destroyed through a subclass instance
 };

onnxruntime/core/providers/migraphx/migraphx_provider_factory.cc

@@ -69,6 +69,29 @@ struct ProviderInfo_MIGraphX_Impl final : ProviderInfo_MIGraphX {
     HIP_CALL_THROW(hipMemcpy(dst, src, count, hipMemcpyDeviceToHost));
   }
 
+  Status SynchronizeDevice(int device_id) override {
+    int original_device;
+    HIP_RETURN_IF_ERROR(hipGetDevice(&original_device));
+    if (original_device != device_id) {
+      HIP_RETURN_IF_ERROR(hipSetDevice(device_id));
+    }
+    const auto sync_error = hipDeviceSynchronize();
+    const auto restore_error = original_device != device_id ? hipSetDevice(original_device) : hipSuccess;
+    HIP_RETURN_IF_ERROR(sync_error);
+    HIP_RETURN_IF_ERROR(restore_error);
+    return Status::OK();
+  }
+
+  int GetCurrentGpuDeviceId() override {
+    int device_id;
+    HIP_CALL_THROW(hipGetDevice(&device_id));
+    return device_id;
+  }
+
+  void SetCurrentGpuDeviceId(int device_id) override {
+    HIP_CALL_THROW(hipSetDevice(device_id));
+  }
+
   std::shared_ptr<IAllocator> CreateMIGraphXAllocator(OrtDevice::DeviceId device_id, size_t mem_limit, onnxruntime::ArenaExtendStrategy arena_extend_strategy,
                                                       void* alloc_fn, void* free_fn, void* empty_cache_fn, const OrtArenaCfg* default_memory_arena_cfg) override {
     if (alloc_fn != nullptr && free_fn != nullptr) {

onnxruntime/core/providers/migraphx/migraphx_provider_factory.h

@@ -5,6 +5,7 @@
 
 #include <memory>
 
+#include "core/common/status.h"
 #include "core/framework/arena_extend_strategy.h"
 #include "core/framework/ortdevice.h"
 
@@ -18,6 +19,9 @@ struct ProviderInfo_MIGraphX {
   virtual void MIGraphXMemcpy_DeviceToHost(void* dst, const void* src, size_t count) = 0;
   virtual std::shared_ptr<IAllocator> CreateMIGraphXAllocator(OrtDevice::DeviceId device_id, size_t mem_limit,
                                                               ArenaExtendStrategy arena_extend_strategy, void* alloc_fn, void* free_fn, void* empty_cache_fn, const OrtArenaCfg* default_memory_arena_cfg) = 0;
+  virtual Status SynchronizeDevice(int device_id) = 0;
+  virtual int GetCurrentGpuDeviceId() = 0;
+  virtual void SetCurrentGpuDeviceId(int device_id) = 0;
 
  protected:
   ~ProviderInfo_MIGraphX() = default;  // Can only be destroyed through a subclass instance

onnxruntime/python/onnxruntime_inference_collection.py

@@ -1308,8 +1308,10 @@ def ortvalue_from_numpy(
         cls, numpy_obj: np.ndarray, /, device_type="cpu", device_id=0, vendor_id: int | OrtDeviceVendorId = -1
     ) -> OrtValue:
         """
-        Factory method to construct an OrtValue (which holds a Tensor) from a given Numpy object
-        A copy of the data in the Numpy object is held by the OrtValue only if the device is NOT cpu
+        Construct an OrtValue (which holds a Tensor) from a NumPy array.
+        Contiguous CPU arrays are shared without copying; noncontiguous arrays
+        are copied to contiguous storage. Other devices copy the data to the
+        target device. The native OrtValue retains any NumPy storage it uses.
 
         :param numpy_obj: The Numpy object to construct the OrtValue from
         :param device_type: e.g. cpu, cuda, cann, cpu by default
@@ -1331,10 +1333,10 @@ def ortvalue_from_numpy(
     def ortvalue_from_numpy_with_onnx_type(cls, data: np.ndarray, /, onnx_element_type: int) -> OrtValue:
         """
         This method creates an instance of OrtValue on top of the numpy array.
-        No data copy is made and the lifespan of the resulting OrtValue should never
-        exceed the lifespan of bytes object. The API attempts to reinterpret
-        the data type which is expected to be the same size. This is useful
+        No data copy is made; the native OrtValue retains the NumPy array.
+        The API reinterprets the data type, which must have the same element size. This is useful
         when we want to use an ONNX data type that is not supported by numpy.
+        The array must be C-contiguous; use ``numpy.ascontiguousarray`` otherwise.
 
         :param data: numpy.ndarray.
         :param onnx_element_type: a valid onnx TensorProto::DataType enum value
@@ -1529,14 +1531,20 @@ def __dlpack__(self, *, stream=None):
         ``from_dlpack(ortvalue)`` (e.g. ``torch.from_dlpack``,
         ``jax.dlpack.from_dlpack``, ``numpy.from_dlpack``).
 
-        The OrtValue must hold a contiguous tensor. No data is copied;
-        the consumer shares memory with this OrtValue, which must remain
-        alive while the capsule is in use.
+        The OrtValue must hold a contiguous tensor. No data is copied.
+        The capsule and its consumer retain the tensor allocation, so this
+        OrtValue can be released. NumPy-backed OrtValues also retain the NumPy
+        allocation through the consumer's lifetime.
         WebGPU OrtValues cannot be exported through DLPack because their data
         is stored in opaque WebGPU buffers rather than CUDA-addressable memory.
 
-        :param stream: Optional stream on which the tensor data is accessible.
-            Currently unused; included for protocol compliance.
+        :param stream: GPU consumer stream: for CUDA, ``None`` or ``1`` for the legacy
+            default, ``2`` for the per-thread default, or a positive stream handle.
+            ROCm also accepts ``0`` for its default stream.
+            Since OrtValues do not track their producer stream, GPU exports
+            synchronize the entire tensor device before returning. ``-1`` skips
+            synchronization; the caller must ensure the data is ready. CPU exports
+            require ``None``.
         :return: A PyCapsule holding a DLManagedTensor.
         """
         return self._ortvalue.__dlpack__(stream=stream)
@@ -1555,7 +1563,7 @@ def __dlpack_device__(self) -> tuple[int, int]:
         return self._ortvalue.__dlpack_device__()
 
     @classmethod
-    def from_dlpack(cls, data, /) -> OrtValue:
+    def from_dlpack(cls, data, /, *, synchronize: bool = True) -> OrtValue:
         """
         Construct an OrtValue from an object that implements the DLPack protocol.
 
@@ -1571,20 +1579,31 @@ def from_dlpack(cls, data, /) -> OrtValue:
         including float8, packed float4, int4/uint4, and int2/uint2. String,
         complex, and float6 tensors are not supported.
 
+        Interchange also depends on the other framework's representation.
+        PyTorch 2.10 exchanges all byte-aligned ORT numeric and boolean types,
+        including bfloat16 and the five float8 types. Its float4 uses two lanes (unsupported here),
+        and its int2/int4 shell dtypes export as int8/uint8, not packed types.
+        PyTorch cannot consume ORT's packed float4, int4/uint4, or int2/uint2.
+
         No data is copied; the new OrtValue shares memory with the source.
+        GPU imports synchronize the tensor's device before returning, so ORT
+        can safely use the storage on its own streams. CUDA and ROCm require
+        their corresponding CUDA or MIGraphX provider for synchronization.
+        When that provider is available, GPU protocol exports run with the producer's
+        device selected and restore the caller's current device, including on failure.
+        The capsule is consumed before synchronization; if that wait fails, the
+        imported storage is released and the capsule cannot be retried.
 
         :param data: A tensor object supporting the DLPack protocol, or a raw
             DLPack PyCapsule.
+        :param synchronize: Set False only when the caller arranges producer/consumer
+            stream ordering before using the imported tensor, for example during
+            graph capture or when preparing a batch followed by one stream wait.
+            GPU protocol producers are asked to export with ``stream=-1`` in this mode.
         :return: An OrtValue wrapping the tensor data.
         :raises TypeError: If the object does not provide a valid DLPack capsule.
         """
-        # If the input supports the __dlpack__ protocol, call it to get the capsule.
-        if hasattr(data, "__dlpack__"):
-            capsule = data.__dlpack__()
-        else:
-            capsule = data
-
-        return cls(C.OrtValue.from_dlpack(capsule))
+        return cls(C.OrtValue.from_dlpack(data, synchronize=synchronize))
 
     def update_inplace(self, data) -> None:
         """