make device copy operator dynamic shape support

Gasoonjia · web-flow · commit e0dfec58bfdd · 2026-06-08T23:49:38.000Z
Differential Revision: D107901331 Pull Request resolved: pytorch#20116
diff --git a/kernels/portable/cpu/op__device_copy.cpp b/kernels/portable/cpu/op__device_copy.cpp
@@ -56,15 +56,15 @@ _h2d_copy_out(KernelRuntimeContext& ctx, const Tensor& self, Tensor& out) {
       out,
       "_h2d_copy: destination tensor must be on a non-CPU device");
 
-  auto nbytes = self.nbytes();
   ET_KERNEL_CHECK_MSG(
       ctx,
-      nbytes == out.nbytes(),
+      resize_tensor(out, self.sizes()) == Error::Ok,
       InvalidArgument,
       out,
-      "_h2d_copy: size mismatch: self.nbytes()=%zu, out.nbytes()=%zu",
-      nbytes,
+      "_h2d_copy: cannot resize out to self sizes (self.nbytes()=%zu exceeds out planned capacity %zu?)",
+      self.nbytes(),
       out.nbytes());
+  auto nbytes = self.nbytes();
 
   DeviceAllocator* allocator =
       executorch::runtime::get_device_allocator(device_type);
@@ -117,15 +117,15 @@ _d2h_copy_out(KernelRuntimeContext& ctx, const Tensor& self, Tensor& out) {
       "_d2h_copy: destination tensor must be on CPU, got device_type=%d",
       static_cast<int>(out.unsafeGetTensorImpl()->device_type()));
 
-  auto nbytes = self.nbytes();
   ET_KERNEL_CHECK_MSG(
       ctx,
-      nbytes == out.nbytes(),
+      resize_tensor(out, self.sizes()) == Error::Ok,
       InvalidArgument,
       out,
-      "_d2h_copy: size mismatch: self.nbytes()=%zu, out.nbytes()=%zu",
-      nbytes,
+      "_d2h_copy: cannot resize out to self sizes (self.nbytes()=%zu exceeds out planned capacity %zu?)",
+      self.nbytes(),
       out.nbytes());
+  auto nbytes = self.nbytes();
 
   DeviceAllocator* allocator =
       executorch::runtime::get_device_allocator(device_type);
diff --git a/kernels/test/op__device_copy_test.cpp b/kernels/test/op__device_copy_test.cpp
@@ -246,3 +246,246 @@ TEST_F(OpDeviceCopyTest, H2dCopyMultidimensionalTensor) {
     EXPECT_EQ(dst_data[i], src_data[i]);
   }
 }
+
+// H2D: out has a LARGER upper-bound capacity + dynamic shape, self is SMALLER.
+// After the op, out is resized down to self's shape and holds self's values.
+TEST_F(OpDeviceCopyTest, H2dCopyDynamicShapeResizesOutDownToInput) {
+  // CPU source: actual (smaller) shape [4].
+  float src_data[] = {1.0f, 2.0f, 3.0f, 4.0f};
+  int32_t src_sizes[] = {4};
+  uint8_t src_dim_order[] = {0};
+  int32_t src_strides[] = {1};
+  TensorImpl src_impl(
+      ScalarType::Float,
+      1,
+      src_sizes,
+      src_data,
+      src_dim_order,
+      src_strides,
+      TensorShapeDynamism::STATIC,
+      DeviceType::CPU,
+      0);
+  Tensor src(&src_impl);
+
+  // CUDA destination: planned at upper bound [8] (capacity = 8 elems), dynamic.
+  float dst_data[] = {0, 0, 0, 0, 0, 0, 0, 0};
+  int32_t dst_sizes[] = {8};
+  uint8_t dst_dim_order[] = {0};
+  int32_t dst_strides[] = {1};
+  TensorImpl dst_impl(
+      ScalarType::Float,
+      1,
+      dst_sizes,
+      dst_data,
+      dst_dim_order,
+      dst_strides,
+      TensorShapeDynamism::DYNAMIC_BOUND,
+      DeviceType::CUDA,
+      0);
+  Tensor dst(&dst_impl);
+
+  Tensor& result = op_h2d_copy_out(src, dst);
+
+  // out was resized down to match self.
+  EXPECT_EQ(dst.dim(), 1);
+  EXPECT_EQ(dst.size(0), 4);
+  EXPECT_EQ(dst.numel(), 4);
+
+  // Only self.nbytes() worth of data was copied.
+  EXPECT_EQ(g_mock_cuda.h2d_count_, 1);
+  EXPECT_EQ(g_mock_cuda.last_h2d_size_, 4 * sizeof(float));
+
+  // out values equal self values.
+  EXPECT_EQ(dst_data[0], 1.0f);
+  EXPECT_EQ(dst_data[1], 2.0f);
+  EXPECT_EQ(dst_data[2], 3.0f);
+  EXPECT_EQ(dst_data[3], 4.0f);
+
+  EXPECT_EQ(&result, &dst);
+}
+
+// D2H: mirror of the above, device -> host with a larger planned out buffer.
+TEST_F(OpDeviceCopyTest, D2hCopyDynamicShapeResizesOutDownToInput) {
+  // CUDA source: actual (smaller) shape [4].
+  float src_data[] = {5.0f, 6.0f, 7.0f, 8.0f};
+  int32_t src_sizes[] = {4};
+  uint8_t src_dim_order[] = {0};
+  int32_t src_strides[] = {1};
+  TensorImpl src_impl(
+      ScalarType::Float,
+      1,
+      src_sizes,
+      src_data,
+      src_dim_order,
+      src_strides,
+      TensorShapeDynamism::STATIC,
+      DeviceType::CUDA,
+      0);
+  Tensor src(&src_impl);
+
+  // CPU destination: planned at upper bound [8] (capacity = 8 elems), dynamic.
+  float dst_data[] = {0, 0, 0, 0, 0, 0, 0, 0};
+  int32_t dst_sizes[] = {8};
+  uint8_t dst_dim_order[] = {0};
+  int32_t dst_strides[] = {1};
+  TensorImpl dst_impl(
+      ScalarType::Float,
+      1,
+      dst_sizes,
+      dst_data,
+      dst_dim_order,
+      dst_strides,
+      TensorShapeDynamism::DYNAMIC_BOUND,
+      DeviceType::CPU,
+      0);
+  Tensor dst(&dst_impl);
+
+  Tensor& result = op_d2h_copy_out(src, dst);
+
+  EXPECT_EQ(dst.dim(), 1);
+  EXPECT_EQ(dst.size(0), 4);
+  EXPECT_EQ(dst.numel(), 4);
+
+  EXPECT_EQ(g_mock_cuda.d2h_count_, 1);
+  EXPECT_EQ(g_mock_cuda.last_d2h_size_, 4 * sizeof(float));
+
+  EXPECT_EQ(dst_data[0], 5.0f);
+  EXPECT_EQ(dst_data[1], 6.0f);
+  EXPECT_EQ(dst_data[2], 7.0f);
+  EXPECT_EQ(dst_data[3], 8.0f);
+
+  EXPECT_EQ(&result, &dst);
+}
+
+// H2D: self LARGER than out's planned capacity -> resize fails -> op errors
+// with InvalidArgument and does NOT copy.
+TEST_F(OpDeviceCopyTest, H2dCopyFailsWhenInputExceedsOutCapacity) {
+  // CPU source: shape [4].
+  float src_data[] = {1.0f, 2.0f, 3.0f, 4.0f};
+  int32_t src_sizes[] = {4};
+  uint8_t src_dim_order[] = {0};
+  int32_t src_strides[] = {1};
+  TensorImpl src_impl(
+      ScalarType::Float,
+      1,
+      src_sizes,
+      src_data,
+      src_dim_order,
+      src_strides,
+      TensorShapeDynamism::STATIC,
+      DeviceType::CPU,
+      0);
+  Tensor src(&src_impl);
+
+  // CUDA destination: planned capacity only [2], smaller than self.
+  float dst_data[] = {0, 0};
+  int32_t dst_sizes[] = {2};
+  uint8_t dst_dim_order[] = {0};
+  int32_t dst_strides[] = {1};
+  TensorImpl dst_impl(
+      ScalarType::Float,
+      1,
+      dst_sizes,
+      dst_data,
+      dst_dim_order,
+      dst_strides,
+      TensorShapeDynamism::DYNAMIC_BOUND,
+      DeviceType::CUDA,
+      0);
+  Tensor dst(&dst_impl);
+
+  ET_EXPECT_KERNEL_FAILURE(context_, op_h2d_copy_out(src, dst));
+
+#ifndef USE_ATEN_LIB
+  EXPECT_EQ(context_.failure_state(), Error::InvalidArgument);
+#endif
+  // The kernel bailed before copying.
+  EXPECT_EQ(g_mock_cuda.h2d_count_, 0);
+}
+
+// D2H: self LARGER than out's planned capacity -> resize fails -> op errors
+// with InvalidArgument and does NOT copy.
+TEST_F(OpDeviceCopyTest, D2hCopyFailsWhenInputExceedsOutCapacity) {
+  // CUDA source: shape [4].
+  float src_data[] = {5.0f, 6.0f, 7.0f, 8.0f};
+  int32_t src_sizes[] = {4};
+  uint8_t src_dim_order[] = {0};
+  int32_t src_strides[] = {1};
+  TensorImpl src_impl(
+      ScalarType::Float,
+      1,
+      src_sizes,
+      src_data,
+      src_dim_order,
+      src_strides,
+      TensorShapeDynamism::STATIC,
+      DeviceType::CUDA,
+      0);
+  Tensor src(&src_impl);
+
+  // CPU destination: planned capacity only [2], smaller than self.
+  float dst_data[] = {0, 0};
+  int32_t dst_sizes[] = {2};
+  uint8_t dst_dim_order[] = {0};
+  int32_t dst_strides[] = {1};
+  TensorImpl dst_impl(
+      ScalarType::Float,
+      1,
+      dst_sizes,
+      dst_data,
+      dst_dim_order,
+      dst_strides,
+      TensorShapeDynamism::DYNAMIC_BOUND,
+      DeviceType::CPU,
+      0);
+  Tensor dst(&dst_impl);
+
+  ET_EXPECT_KERNEL_FAILURE(context_, op_d2h_copy_out(src, dst));
+
+#ifndef USE_ATEN_LIB
+  EXPECT_EQ(context_.failure_state(), Error::InvalidArgument);
+#endif
+  EXPECT_EQ(g_mock_cuda.d2h_count_, 0);
+}
+
+// Equal-size case under the dynamic-bound path: capacity == input size still
+// copies correctly (confirms existing behavior is preserved by the resize).
+TEST_F(OpDeviceCopyTest, H2dCopyDynamicBoundEqualSizeStillCopies) {
+  float src_data[] = {1.0f, 2.0f, 3.0f, 4.0f};
+  int32_t sizes[] = {4};
+  uint8_t dim_order[] = {0};
+  int32_t strides[] = {1};
+  TensorImpl src_impl(
+      ScalarType::Float,
+      1,
+      sizes,
+      src_data,
+      dim_order,
+      strides,
+      TensorShapeDynamism::STATIC,
+      DeviceType::CPU,
+      0);
+  Tensor src(&src_impl);
+
+  float dst_data[] = {0, 0, 0, 0};
+  TensorImpl dst_impl(
+      ScalarType::Float,
+      1,
+      sizes,
+      dst_data,
+      dim_order,
+      strides,
+      TensorShapeDynamism::DYNAMIC_BOUND,
+      DeviceType::CUDA,
+      0);
+  Tensor dst(&dst_impl);
+
+  op_h2d_copy_out(src, dst);
+
+  EXPECT_EQ(dst.size(0), 4);
+  EXPECT_EQ(g_mock_cuda.h2d_count_, 1);
+  EXPECT_EQ(g_mock_cuda.last_h2d_size_, 4 * sizeof(float));
+  for (int i = 0; i < 4; ++i) {
+    EXPECT_EQ(dst_data[i], src_data[i]);
+  }
+}
