Fix issue w/ DML batched readback size

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

@@ -30,6 +30,7 @@
 #define ERROR 0
 
 #include "core/session/onnxruntime_c_api.h"
+#include <iomanip>
 #include <wil/wrl.h>
 #ifndef _GAMING_XBOX
 #include <dxgi1_6.h>
@@ -43,6 +44,25 @@
 {
     using namespace onnxruntime::common;
 
+#ifndef ORT_NO_EXCEPTIONS
+    static Status HResultToStatus(HRESULT hr, const char* operation, const char* details)
+    {
+        const StatusCode status_code = hr == E_INVALIDARG ? INVALID_ARGUMENT : FAIL;
+        return Status(
+            ONNXRUNTIME,
+            status_code,
+            onnxruntime::MakeString(
+                operation,
+                " failed with HRESULT 0x",
+                std::setfill('0'),
+                std::uppercase,
+                std::hex,
+                std::setw(8),
+                static_cast<uint32_t>(hr), ": ", details
+            ));
+    }
+#endif
+
     ExecutionProvider::~ExecutionProvider()
     {
         if (m_impl)
@@ -543,12 +563,12 @@
         // Source and destination for batched GPU -> CPU copies
         std::vector<ID3D12Resource*> srcDatas;
         std::vector<void*> dstDatas;
-        std::vector<uint32_t> dataSizesInBytes;
+        std::vector<size_t> dataSizesInBytes;
 
         assert(!m_closed);
         auto provider = const_cast<ExecutionProviderImpl*>(this);
 
-        for (uint32_t i = 0; i < dst.size(); ++i)
+        for (size_t i = 0; i < dst.size(); ++i)
         {
             // This batching implementation only handles GPU -> CPU copies.  Other copies do not require synchronization
             // and are batched across multiple calls to CopyTensor.
@@ -559,15 +579,15 @@
             }
 
             const size_t dataSizeInBytes = ComputeByteSizeFromTensor(*dst[i]);
-            ORT_THROW_HR_IF(E_INVALIDARG, dataSizeInBytes != ComputeByteSizeFromTensor(*src[i])); // Tensors must be the same size
+            const size_t srcSizeInBytes = ComputeByteSizeFromTensor(*src[i]);
+            ORT_THROW_HR_IF(E_INVALIDARG, dataSizeInBytes != srcSizeInBytes); // Tensors must be the same size
 
             if (dataSizeInBytes == 0)
             {
                 continue;
             }
 
-            dataSizesInBytes.push_back(static_cast<uint32_t>(ComputeByteSizeFromTensor(*dst[i])));
-            ORT_THROW_HR_IF(E_INVALIDARG, dataSizesInBytes.back() != ComputeByteSizeFromTensor(*src[i])); // Tensors must be the same size
+            dataSizesInBytes.push_back(dataSizeInBytes);
 
             dstDatas.push_back(dst[i]->GetData());
             const AllocationInfo* srcAllocInfo = m_allocator->DecodeDataHandle(MLOperatorTensor(src[i]).GetDataInterface().Get());
@@ -959,12 +979,12 @@
         // Source and destination for batched GPU -> CPU copies
         std::vector<ID3D12Resource*> srcDatas;
         std::vector<void*> dstDatas;
-        std::vector<uint32_t> dataSizesInBytes;
+        std::vector<size_t> dataSizesInBytes;
 
         assert(!m_closed);
         auto provider = const_cast<ExecutionProviderImpl*>(this);
 
-        for (uint32_t i = 0; i < src_dst_pairs.size(); ++i)
+        for (size_t i = 0; i < src_dst_pairs.size(); ++i)
         {
             // This batching implementation only handles GPU -> CPU copies.  Other copies do not require synchronization
             // and are batched across multiple calls to CopyTensor.
@@ -987,15 +1007,19 @@
                 true);
 
             const size_t dataSizeInBytes = ComputeByteSizeFromTensor(dstWrapper);
-            ORT_THROW_HR_IF(E_INVALIDARG, dataSizeInBytes != ComputeByteSizeFromTensor(srcWrapper)); // Tensors must be the same size
+            const size_t srcSizeInBytes = ComputeByteSizeFromTensor(srcWrapper);
+            if (dataSizeInBytes != srcSizeInBytes)
+            {
+                return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "DML tensor size mismatch. src size: ",
+                                       srcSizeInBytes, " dst size: ", dataSizeInBytes);
+            }
 
             if (dataSizeInBytes == 0)
             {
-                return onnxruntime::common::Status::OK();
+                continue;
             }
 
-            dataSizesInBytes.push_back(static_cast<uint32_t>(ComputeByteSizeFromTensor(dstWrapper)));
-            ORT_THROW_HR_IF(E_INVALIDARG, dataSizesInBytes[i] != ComputeByteSizeFromTensor(srcWrapper)); // Tensors must be the same size
+            dataSizesInBytes.push_back(dataSizeInBytes);
 
             dstDatas.push_back(dstWrapper.GetData());
             const AllocationInfo* srcAllocInfo = m_allocator->DecodeDataHandle(MLOperatorTensor(&srcWrapper).GetDataInterface().Get());
@@ -1006,7 +1030,18 @@
         const auto srcState = D3D12_RESOURCE_STATE_UNORDERED_ACCESS; // GPU resources are always kept in UAV state
 
         // Performs a blocking call to synchronize and read back data from the GPU into the destination buffer
-        m_readbackHeap->ReadbackFromGpu(dstDatas, dataSizesInBytes, srcDatas, srcState);
+#ifndef ORT_NO_EXCEPTIONS
+        ORT_TRY
+        {
+#endif
+            m_readbackHeap->ReadbackFromGpu(dstDatas, dataSizesInBytes, srcDatas, srcState);
+#ifndef ORT_NO_EXCEPTIONS
+        }
+        ORT_CATCH(const wil::ResultException& ex)
+        {
+            return HResultToStatus(ex.GetErrorCode(), "DML GPU readback", ex.what());
+        }
+#endif
 
         return onnxruntime::common::Status::OK();
     }

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

@@ -7,6 +7,22 @@
 
 namespace Dml
 {
+    namespace detail
+    {
+        size_t ComputeTotalReadbackSize(gsl::span<const size_t> sizes)
+        {
+            // Batched readback heap sizing must not wrap at the former uint32_t boundary.
+            size_t totalSize = 0;
+            for (auto size : sizes)
+            {
+                ORT_THROW_HR_IF(E_INVALIDARG, size > std::numeric_limits<size_t>::max() - totalSize);
+                totalSize += size;
+            }
+
+            return totalSize;
+        }
+    }
+
     static ComPtr<ID3D12Resource> CreateReadbackHeap(ID3D12Device* device, size_t size)
     {
         ComPtr<ID3D12Resource> readbackHeap;
@@ -103,30 +119,28 @@
 
     void ReadbackHeap::ReadbackFromGpu(
         gsl::span<void*> dst,
-        gsl::span<const uint32_t > dstSizes,
+        gsl::span<const size_t> dstSizes,
         gsl::span<ID3D12Resource*> src,
         D3D12_RESOURCE_STATES srcState)
     {
         assert(dst.size() == src.size());
         assert(dstSizes.size() == src.size());
+        ORT_THROW_HR_IF(E_INVALIDARG, dst.size() != src.size());
+        ORT_THROW_HR_IF(E_INVALIDARG, dstSizes.size() != src.size());
 
         if (dst.empty())
         {
             return;
         }
 
-        uint32_t totalSize = 0;
-        for (auto size : dstSizes)
-        {
-            totalSize += size;
-        }
-
+        const size_t totalSize = detail::ComputeTotalReadbackSize(dstSizes);
         EnsureReadbackHeap(totalSize);
 
         // Copy from the source resource into the readback heap
-        uint32_t offset = 0;
-        for (uint32_t i = 0; i < dst.size(); ++i)
+        size_t offset = 0;
+        for (size_t i = 0; i < dst.size(); ++i)
         {
+            ORT_THROW_HR_IF(E_INVALIDARG, src[i] == nullptr);
             m_executionContext->CopyBufferRegion(
                 m_readbackHeap.Get(),
                 offset,
@@ -147,15 +161,15 @@
         // Map the readback heap and copy it into the destination
         void* readbackHeapData = nullptr;
         ORT_THROW_IF_FAILED(m_readbackHeap->Map(0, nullptr, &readbackHeapData));
+        [[maybe_unused]] auto unmapReadbackHeap = gsl::finally([this]() { m_readbackHeap->Unmap(0, nullptr); });
 
-        // Copy from the source resource into the readback heap
+        // Copy from the readback heap into the destination buffers.
         offset = 0;
-        for (uint32_t i = 0; i < dst.size(); ++i)
+        for (size_t i = 0; i < dst.size(); ++i)
         {
+            ORT_THROW_HR_IF(E_INVALIDARG, dst[i] == nullptr);
             memcpy(dst[i], static_cast<uint8_t*>(readbackHeapData) + offset, dstSizes[i]);
             offset += dstSizes[i];
         }
-
-        m_readbackHeap->Unmap(0, nullptr);
     }
 } // namespace Dml

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

@@ -7,6 +7,11 @@ namespace Dml
 {
     class ExecutionContext;
 
+    namespace detail
+    {
+        size_t ComputeTotalReadbackSize(gsl::span<const size_t> sizes);
+    }
+
     // Because we never perform more than one readback at a time, we don't need anything fancy for managing the
     // readback heap - just maintain a single resource and reallocate it if it's not big enough.
     class ReadbackHeap
@@ -25,7 +30,7 @@ namespace Dml
         // Overload supporting batching
         void ReadbackFromGpu(
             gsl::span<void*> dst,
-            gsl::span<const uint32_t > dstSizes,
+            gsl::span<const size_t> dstSizes,
             gsl::span<ID3D12Resource*> src,
             D3D12_RESOURCE_STATES srcState);
 

onnxruntime/test/providers/dml_readback_heap_test.cc

@@ -0,0 +1,65 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#ifdef USE_DML
+
+#include "gtest/gtest.h"
+
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+
+#ifdef _GAMING_XBOX_SCARLETT
+#include <d3d12_xs.h>
+#elif defined(_GAMING_XBOX_XBOXONE)
+#include <d3d12_x.h>
+#else
+#include "directx/d3d12.h"
+#endif
+#include <gsl/gsl>
+#include <wrl/client.h>
+
+using Microsoft::WRL::ComPtr;
+
+#include "core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.h"
+
+namespace onnxruntime {
+namespace test {
+
+TEST(DmlReadbackHeapTest, ComputeTotalReadbackSizeDoesNotWrapAtUint32Max) {
+  if (sizeof(size_t) <= sizeof(uint32_t)) {
+    GTEST_SKIP() << "This regression needs a size_t wider than uint32_t.";
+  }
+
+  constexpr size_t kUint32Max = std::numeric_limits<uint32_t>::max();
+  const std::array<size_t, 3> sizes = {kUint32Max, 1, 7};
+
+  EXPECT_EQ(Dml::detail::ComputeTotalReadbackSize(gsl::make_span(sizes.data(), sizes.size())), kUint32Max + 8);
+}
+
+TEST(DmlReadbackHeapTest, ComputeTotalReadbackSizeIncludesSizesAfterZero) {
+  const std::array<size_t, 3> sizes = {5, 0, 7};
+
+  EXPECT_EQ(Dml::detail::ComputeTotalReadbackSize(gsl::make_span(sizes.data(), sizes.size())), 12u);
+}
+
+#ifndef ORT_NO_EXCEPTIONS
+TEST(DmlReadbackHeapTest, ComputeTotalReadbackSizeRejectsSizeTOverflow) {
+  const std::array<size_t, 2> sizes = {std::numeric_limits<size_t>::max(), 1};
+
+  EXPECT_ANY_THROW((void)Dml::detail::ComputeTotalReadbackSize(gsl::make_span(sizes.data(), sizes.size())));
+}
+
+TEST(DmlReadbackHeapTest, ComputeTotalReadbackSizeRejectsMidBatchOverflow) {
+  const size_t half_max = std::numeric_limits<size_t>::max() / 2;
+  const std::array<size_t, 2> sizes = {half_max + 1, half_max + 1};
+
+  EXPECT_ANY_THROW((void)Dml::detail::ComputeTotalReadbackSize(gsl::make_span(sizes.data(), sizes.size())));
+}
+#endif
+
+}  // namespace test
+}  // namespace onnxruntime
+
+#endif  // USE_DML