Merge pull request #129000 from kokol16:fix/space-to-depth

PiperOrigin-RevId: 1005574314

tensorflow/core/kernels/BUILD

@@ -5028,7 +5028,22 @@ tf_kernel_library(
     deps = [
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
-        "//tensorflow/core/util:overflow",
+        "//tensorflow/core/platform:errors",
+        "@com_google_absl//absl/status",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+tf_cuda_only_cc_test(
+    name = "depth_space_ops_gpu_test",
+    size = "small",
+    srcs = ["depth_space_ops_gpu_test.cu.cc"],
+    deps = [
+        ":depth_space_ops",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:status_matchers",
         "@com_google_absl//absl/status",
         "@eigen_archive//:eigen3",
     ],

tensorflow/core/kernels/depth_space_ops_gpu_test.cu.cc

@@ -0,0 +1,97 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#define EIGEN_USE_GPU
+
+#include "absl/status/status.h"
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/core/framework/tensor_types.h"
+#include "tensorflow/core/kernels/depthtospace_op.h"
+#include "tensorflow/core/kernels/spacetodepth_op.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/status_matchers.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace tensorflow {
+namespace functor {
+namespace {
+
+using GPUDevice = Eigen::GpuDevice;
+using Dims = Eigen::DSizes<Eigen::DenseIndex, 4>;
+using ::testing::HasSubstr;
+
+class DepthSpaceOpGpuTest : public ::testing::Test {
+ protected:
+  template <typename Functor>
+  absl::Status RunFunctor(const Dims& input_dims, const Dims& output_dims) {
+    TTypes<float, 4>::ConstTensor input(&dummy_, input_dims);
+    TTypes<float, 4>::Tensor output(&dummy_, output_dims);
+    return Functor()(device_, input, /*block_size=*/2, output);
+  }
+
+  Eigen::GpuStreamDevice stream_;
+  GPUDevice device_{&stream_};
+  float dummy_ = 0.0f;
+};
+
+TEST_F(DepthSpaceOpGpuTest, SpaceToDepthRejectsInputOverInt32Max) {
+  EXPECT_THAT(
+      (RunFunctor<SpaceToDepthOpFunctor<GPUDevice, float, FORMAT_NHWC>>(
+          Dims(1, 2, 2, 600000000), Dims(1, 1, 1, 2400000000))),
+      absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             HasSubstr("input size exceeds int32 bounds")));
+  EXPECT_THAT(
+      (RunFunctor<SpaceToDepthOpFunctor<GPUDevice, float, FORMAT_NCHW>>(
+          Dims(1, 600000000, 2, 2), Dims(1, 2400000000, 1, 1))),
+      absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             HasSubstr("input size exceeds int32 bounds")));
+}
+
+TEST_F(DepthSpaceOpGpuTest, DepthToSpaceRejectsInputOverInt32Max) {
+  EXPECT_THAT(
+      (RunFunctor<DepthToSpaceOpFunctor<GPUDevice, float, FORMAT_NHWC>>(
+          Dims(1, 1, 1, 2400000000), Dims(1, 2, 2, 600000000))),
+      absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             HasSubstr("input size exceeds int32 bounds")));
+  EXPECT_THAT(
+      (RunFunctor<DepthToSpaceOpFunctor<GPUDevice, float, FORMAT_NCHW>>(
+          Dims(1, 2400000000, 1, 1), Dims(1, 600000000, 2, 2))),
+      absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             HasSubstr("input size exceeds int32 bounds")));
+}
+
+TEST_F(DepthSpaceOpGpuTest, EmptyInputWithDimensionOverInt32MaxSucceeds) {
+  TF_EXPECT_OK(
+      (RunFunctor<SpaceToDepthOpFunctor<GPUDevice, float, FORMAT_NHWC>>(
+          Dims(3000000000, 0, 2, 2), Dims(3000000000, 0, 1, 8))));
+  TF_EXPECT_OK(
+      (RunFunctor<SpaceToDepthOpFunctor<GPUDevice, float, FORMAT_NCHW>>(
+          Dims(3000000000, 0, 2, 2), Dims(3000000000, 0, 1, 1))));
+  TF_EXPECT_OK(
+      (RunFunctor<DepthToSpaceOpFunctor<GPUDevice, float, FORMAT_NHWC>>(
+          Dims(3000000000, 0, 1, 4), Dims(3000000000, 0, 2, 1))));
+  TF_EXPECT_OK(
+      (RunFunctor<DepthToSpaceOpFunctor<GPUDevice, float, FORMAT_NCHW>>(
+          Dims(3000000000, 0, 1, 1), Dims(3000000000, 0, 2, 2))));
+}
+
+}  // namespace
+}  // namespace functor
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

tensorflow/core/kernels/depthtospace_op_gpu.cu.cc

@@ -17,13 +17,15 @@ limitations under the License.
 
 #define EIGEN_USE_GPU
 
+#include <cstdint>
 #include <limits>
 
+#include "absl/status/status.h"
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/kernels/depthtospace_op.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/gpu_kernel_helper.h"
-#include "tensorflow/core/util/overflow.h"
 
 namespace tensorflow {
 namespace {
@@ -152,6 +154,13 @@ struct DepthToSpaceOpFunctor<GPUDevice, T, FORMAT_NHWC> {
                           typename TTypes<T, 4>::ConstTensor input,
                           int block_size,
                           typename TTypes<T, 4>::Tensor output) {
+    if (input.size() == 0) {
+      return absl::OkStatus();
+    }
+    if (input.size() > std::numeric_limits<int32_t>::max()) {
+      return absl::InvalidArgumentError(
+          "DepthToSpaceOpFunctor NHWC: input size exceeds int32 bounds");
+    }
     const int batch_size = output.dimension(0);
     const int input_height = input.dimension(1);
     const int input_width = input.dimension(2);
@@ -160,22 +169,10 @@ struct DepthToSpaceOpFunctor<GPUDevice, T, FORMAT_NHWC> {
     const int output_width = output.dimension(2);
     const int output_depth = output.dimension(3);
 
-    int64_t total_count_64 = MultiplyWithoutOverflow(
-        MultiplyWithoutOverflow(
-            MultiplyWithoutOverflow(batch_size, output_height), output_width),
-        output_depth);
-    if (total_count_64 < 0 ||
-        total_count_64 > std::numeric_limits<int32_t>::max()) {
-      return absl::InternalError(
-          "DepthToSpaceOpFunctor NHWC: total_count exceeds int32 bounds");
-    }
-    const int total_count = total_count_64;
+    const int total_count = static_cast<int>(input.size());
 
-    if (total_count == 0) {
-      return absl::OkStatus();
-    }
     GpuLaunchConfig config = GetGpuLaunchConfig(total_count, d);
-    TF_CHECK_OK(GpuLaunchKernel(
+    TF_RETURN_IF_ERROR(GpuLaunchKernel(
         D2S_NHWC<T>, config.block_count, config.thread_per_block, 0, d.stream(),
         config.virtual_thread_count, input.data(), block_size, batch_size,
         input_height, input_width, input_depth, output_height, output_width,
@@ -198,60 +195,33 @@ struct DepthToSpaceOpFunctor<GPUDevice, T, FORMAT_NCHW> {
                           typename TTypes<T, 4>::ConstTensor input,
                           int block_size,
                           typename TTypes<T, 4>::Tensor output) {
+    if (input.size() == 0) {
+      return absl::OkStatus();
+    }
+    if (input.size() > std::numeric_limits<int32_t>::max()) {
+      return absl::InvalidArgumentError(
+          "DepthToSpaceOpFunctor NCHW: input size exceeds int32 bounds");
+    }
     const int batch_size = input.dimension(0);
     const int input_depth = input.dimension(1);
     const int input_height = input.dimension(2);
     const int input_width = input.dimension(3);
     const int output_depth = output.dimension(1);
 
-    int64_t input_area_64 = MultiplyWithoutOverflow(input_width, input_height);
-    if (input_area_64 < 0 ||
-        input_area_64 > std::numeric_limits<int32_t>::max()) {
-      return absl::InternalError(
-          "DepthToSpaceOpFunctor NCHW: input_area exceeds int32 bounds");
-    }
-    const int input_area = input_area_64;
-
-    int64_t input_depth_by_input_area_64 =
-        MultiplyWithoutOverflow(input_depth, input_area);
-    if (input_depth_by_input_area_64 < 0 ||
-        input_depth_by_input_area_64 > std::numeric_limits<int32_t>::max()) {
-      return absl::InternalError(
-          "DepthToSpaceOpFunctor NCHW: input_depth_by_input_area exceeds int32 "
-          "bounds");
-    }
-    const int input_depth_by_input_area = input_depth_by_input_area_64;
+    const int input_area = input_width * input_height;
+    const int input_depth_by_input_area = input_depth * input_area;
 
     // We improve performance by generating instantiations of the loop kernel
     // for the most common block sizes.
     if (block_size <= 4) {
       const int output_width = output.dimension(3);
-      int64_t output_depth_by_input_area_64 =
-          MultiplyWithoutOverflow(output_depth, input_area);
-      if (output_depth_by_input_area_64 < 0 ||
-          output_depth_by_input_area_64 > std::numeric_limits<int32_t>::max()) {
-        return absl::InternalError(
-            "DepthToSpaceOpFunctor NCHW: output_depth_by_input_area exceeds "
-            "int32 bounds");
-      }
-      const int output_depth_by_input_area = output_depth_by_input_area_64;
+      const int output_depth_by_input_area = output_depth * input_area;
+      const int total_count = batch_size * output_depth_by_input_area;
 
-      int64_t total_count_64 =
-          MultiplyWithoutOverflow(batch_size, output_depth_by_input_area);
-      if (total_count_64 < 0 ||
-          total_count_64 > std::numeric_limits<int32_t>::max()) {
-        return absl::InternalError(
-            "DepthToSpaceOpFunctor NCHW: total_count exceeds int32 bounds");
-      }
-      const int total_count = total_count_64;
-
-      if (total_count == 0) {
-        return absl::OkStatus();
-      }
       GpuLaunchConfig config = GetGpuLaunchConfig(total_count, d);
       switch (block_size) {
         case 2:
-          TF_CHECK_OK(GpuLaunchKernel(
+          TF_RETURN_IF_ERROR(GpuLaunchKernel(
               D2S_NCHW_LOOP<T, 2>, config.block_count, config.thread_per_block,
               0, d.stream(), total_count, input.data(), input_width,
               output_width, output_depth_by_input_area,
@@ -259,7 +229,7 @@ struct DepthToSpaceOpFunctor<GPUDevice, T, FORMAT_NCHW> {
           return d.ok() ? absl::OkStatus()
                         : absl::InternalError("GPU execution failed");
         case 3:
-          TF_CHECK_OK(GpuLaunchKernel(
+          TF_RETURN_IF_ERROR(GpuLaunchKernel(
               D2S_NCHW_LOOP<T, 3>, config.block_count, config.thread_per_block,
               0, d.stream(), total_count, input.data(), input_width,
               output_width, output_depth_by_input_area,
@@ -267,7 +237,7 @@ struct DepthToSpaceOpFunctor<GPUDevice, T, FORMAT_NCHW> {
           return d.ok() ? absl::OkStatus()
                         : absl::InternalError("GPU execution failed");
         case 4:
-          TF_CHECK_OK(GpuLaunchKernel(
+          TF_RETURN_IF_ERROR(GpuLaunchKernel(
               D2S_NCHW_LOOP<T, 4>, config.block_count, config.thread_per_block,
               0, d.stream(), total_count, input.data(), input_width,
               output_width, output_depth_by_input_area,
@@ -278,20 +248,10 @@ struct DepthToSpaceOpFunctor<GPUDevice, T, FORMAT_NCHW> {
     }
 
     // Other block sizes are processed by the generic kernel.
-    int64_t total_count_64 =
-        MultiplyWithoutOverflow(batch_size, input_depth_by_input_area);
-    if (total_count_64 < 0 ||
-        total_count_64 > std::numeric_limits<int32_t>::max()) {
-      return absl::InternalError(
-          "DepthToSpaceOpFunctor NCHW: total_count exceeds int32 bounds");
-    }
-    const int total_count = total_count_64;
+    const int total_count = batch_size * input_depth_by_input_area;
 
-    if (total_count == 0) {
-      return absl::OkStatus();
-    }
     auto config = GetGpuLaunchConfig(total_count, d);
-    TF_CHECK_OK(GpuLaunchKernel(
+    TF_RETURN_IF_ERROR(GpuLaunchKernel(
         D2S_NCHW<T>, config.block_count, config.thread_per_block, 0, d.stream(),
         config.virtual_thread_count, input.data(), block_size, input_width,
         output_depth * input_height, output.data()));

tensorflow/core/kernels/spacetodepth_op_gpu.cu.cc

@@ -17,13 +17,15 @@ limitations under the License.
 
 #define EIGEN_USE_GPU
 
+#include <cstdint>
 #include <limits>
 
+#include "absl/status/status.h"
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/kernels/spacetodepth_op.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/util/gpu_kernel_helper.h"
-#include "tensorflow/core/util/overflow.h"
 
 namespace tensorflow {
 
@@ -149,6 +151,13 @@ struct SpaceToDepthOpFunctor<GPUDevice, T, FORMAT_NHWC> {
                           typename TTypes<T, 4>::ConstTensor input,
                           int block_size,
                           typename TTypes<T, 4>::Tensor output) {
+    if (input.size() == 0) {
+      return absl::OkStatus();
+    }
+    if (input.size() > std::numeric_limits<int32_t>::max()) {
+      return absl::InvalidArgumentError(
+          "SpaceToDepthOpFunctor NHWC: input size exceeds int32 bounds");
+    }
     const int batch_size = output.dimension(0);
     const int input_height = input.dimension(1);
     const int input_width = input.dimension(2);
@@ -157,22 +166,10 @@ struct SpaceToDepthOpFunctor<GPUDevice, T, FORMAT_NHWC> {
     const int output_width = output.dimension(2);
     const int output_depth = output.dimension(3);
 
-    int64_t total_count_64 = MultiplyWithoutOverflow(
-        MultiplyWithoutOverflow(
-            MultiplyWithoutOverflow(batch_size, input_height), input_width),
-        input_depth);
-    if (total_count_64 < 0 ||
-        total_count_64 > std::numeric_limits<int32_t>::max()) {
-      return absl::InternalError(
-          "SpaceToDepthOpFunctor NHWC: total_count exceeds int32 bounds");
-    }
-    const int total_count = total_count_64;
+    const int total_count = static_cast<int>(input.size());
 
-    if (total_count == 0) {
-      return absl::OkStatus();
-    }
     GpuLaunchConfig config = GetGpuLaunchConfig(total_count, d);
-    TF_CHECK_OK(GpuLaunchKernel(
+    TF_RETURN_IF_ERROR(GpuLaunchKernel(
         S2D_NHWC<T>, config.block_count, config.thread_per_block, 0, d.stream(),
         config.virtual_thread_count, input.data(), block_size, batch_size,
         input_height, input_width, input_depth, output_height, output_width,
@@ -195,61 +192,33 @@ struct SpaceToDepthOpFunctor<GPUDevice, T, FORMAT_NCHW> {
                           typename TTypes<T, 4>::ConstTensor input,
                           int block_size,
                           typename TTypes<T, 4>::Tensor output) {
+    if (input.size() == 0) {
+      return absl::OkStatus();
+    }
+    if (input.size() > std::numeric_limits<int32_t>::max()) {
+      return absl::InvalidArgumentError(
+          "SpaceToDepthOpFunctor NCHW: input size exceeds int32 bounds");
+    }
     const int batch_size = output.dimension(0);
     const int input_depth = input.dimension(1);
     const int output_depth = output.dimension(1);
     const int output_height = output.dimension(2);
     const int output_width = output.dimension(3);
 
-    int64_t output_area_64 =
-        MultiplyWithoutOverflow(output_width, output_height);
-    if (output_area_64 < 0 ||
-        output_area_64 > std::numeric_limits<int32_t>::max()) {
-      return absl::InternalError(
-          "SpaceToDepthOpFunctor NCHW: output_area exceeds int32 bounds");
-    }
-    const int output_area = output_area_64;
-
-    int64_t output_depth_by_output_area_64 =
-        MultiplyWithoutOverflow(output_depth, output_area);
-    if (output_depth_by_output_area_64 < 0 ||
-        output_depth_by_output_area_64 > std::numeric_limits<int32_t>::max()) {
-      return absl::InternalError(
-          "SpaceToDepthOpFunctor NCHW: output_depth_by_output_area exceeds "
-          "int32 bounds");
-    }
-    const int output_depth_by_output_area = output_depth_by_output_area_64;
+    const int output_area = output_width * output_height;
+    const int output_depth_by_output_area = output_depth * output_area;
 
     // We improve performance by generating instantiations of the loop kernel
     // for the most common block sizes.
     if (block_size <= 4) {
       const int input_width = input.dimension(3);
-      int64_t input_depth_by_output_area_64 =
-          MultiplyWithoutOverflow(input_depth, output_area);
-      if (input_depth_by_output_area_64 < 0 ||
-          input_depth_by_output_area_64 > std::numeric_limits<int32_t>::max()) {
-        return absl::InternalError(
-            "SpaceToDepthOpFunctor NCHW: input_depth_by_output_area exceeds "
-            "int32 bounds");
-      }
-      const int input_depth_by_output_area = input_depth_by_output_area_64;
+      const int input_depth_by_output_area = input_depth * output_area;
+      const int total_count = batch_size * input_depth_by_output_area;
 
-      int64_t total_count_64 =
-          MultiplyWithoutOverflow(batch_size, input_depth_by_output_area);
-      if (total_count_64 < 0 ||
-          total_count_64 > std::numeric_limits<int32_t>::max()) {
-        return absl::InternalError(
-            "SpaceToDepthOpFunctor NCHW: total_count exceeds int32 bounds");
-      }
-      const int total_count = total_count_64;
-
-      if (total_count == 0) {
-        return absl::OkStatus();
-      }
       GpuLaunchConfig config = GetGpuLaunchConfig(total_count, d);
       switch (block_size) {
         case 2:
-          TF_CHECK_OK(GpuLaunchKernel(
+          TF_RETURN_IF_ERROR(GpuLaunchKernel(
               S2D_NCHW_LOOP<T, 2>, config.block_count, config.thread_per_block,
               0, d.stream(), total_count, input.data(), output_width,
               input_width, input_depth_by_output_area,
@@ -257,7 +226,7 @@ struct SpaceToDepthOpFunctor<GPUDevice, T, FORMAT_NCHW> {
           return d.ok() ? absl::OkStatus()
                         : absl::InternalError("GPU execution failed");
         case 3:
-          TF_CHECK_OK(GpuLaunchKernel(
+          TF_RETURN_IF_ERROR(GpuLaunchKernel(
               S2D_NCHW_LOOP<T, 3>, config.block_count, config.thread_per_block,
               0, d.stream(), total_count, input.data(), output_width,
               input_width, input_depth_by_output_area,
@@ -265,7 +234,7 @@ struct SpaceToDepthOpFunctor<GPUDevice, T, FORMAT_NCHW> {
           return d.ok() ? absl::OkStatus()
                         : absl::InternalError("GPU execution failed");
         case 4:
-          TF_CHECK_OK(GpuLaunchKernel(
+          TF_RETURN_IF_ERROR(GpuLaunchKernel(
               S2D_NCHW_LOOP<T, 4>, config.block_count, config.thread_per_block,
               0, d.stream(), total_count, input.data(), output_width,
               input_width, input_depth_by_output_area,
@@ -276,20 +245,10 @@ struct SpaceToDepthOpFunctor<GPUDevice, T, FORMAT_NCHW> {
     }
 
     // Other block sizes are processed by the generic kernel.
-    int64_t total_count_64 =
-        MultiplyWithoutOverflow(batch_size, output_depth_by_output_area);
-    if (total_count_64 < 0 ||
-        total_count_64 > std::numeric_limits<int32_t>::max()) {
-      return absl::InternalError(
-          "SpaceToDepthOpFunctor NCHW: total_count exceeds int32 bounds");
-    }
-    const int total_count = total_count_64;
+    const int total_count = batch_size * output_depth_by_output_area;
 
-    if (total_count == 0) {
-      return absl::OkStatus();
-    }
     GpuLaunchConfig config = GetGpuLaunchConfig(total_count, d);
-    TF_CHECK_OK(GpuLaunchKernel(
+    TF_RETURN_IF_ERROR(GpuLaunchKernel(
         S2D_NCHW<T>, config.block_count, config.thread_per_block, 0, d.stream(),
         config.virtual_thread_count, input.data(), block_size, output_width,
         input_depth * output_height, output.data()));