diff --git a/onnxruntime/core/providers/webgpu/tensor/pad.cc b/onnxruntime/core/providers/webgpu/tensor/pad.cc index 7a576c4b53ecf..f7be1e2587b48 100644 --- a/onnxruntime/core/providers/webgpu/tensor/pad.cc +++ b/onnxruntime/core/providers/webgpu/tensor/pad.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include @@ -66,8 +67,12 @@ Status Pad::ComputeInternal(ComputeContext& context) const { for (size_t i = 0; i < dimension_count; i++) { int64_t lower_pad = (*p_pads)[i] + (*p_slices)[i]; int64_t upper_pad = (*p_pads)[i + dimension_count] + (*p_slices)[i + dimension_count]; + ORT_RETURN_IF_NOT(lower_pad >= std::numeric_limits::min() && + lower_pad <= std::numeric_limits::max(), + "WebGPU Pad only supports lower pads in the int32 range. Got ", lower_pad, + " for axis ", i); lower_pads[i] = static_cast(lower_pad); - output_dims[i] += lower_pad + upper_pad; + output_dims[i] += SafeInt(lower_pad) + upper_pad; } TensorShape output_shape(output_dims); diff --git a/onnxruntime/core/providers/webgpu/tensor/pad.wgsl.template b/onnxruntime/core/providers/webgpu/tensor/pad.wgsl.template index 6a7480c2197ee..a2356fc9d2684 100644 --- a/onnxruntime/core/providers/webgpu/tensor/pad.wgsl.template +++ b/onnxruntime/core/providers/webgpu/tensor/pad.wgsl.template @@ -62,6 +62,10 @@ $MAIN { in_coord = output_index - lower_pads; } +#if pad_mode == PAD_MODE_WRAP + in_coord = ((in_coord % data_shape) + data_shape) % data_shape; +#endif + input_index += select(u32(in_coord) #if output.rank > 1 * getElementAt(uniforms.data_stride, dim, output.rank - 1) diff --git a/onnxruntime/test/providers/cpu/tensor/pad_test.cc b/onnxruntime/test/providers/cpu/tensor/pad_test.cc index 990e4354c3626..34c725132a2ba 100644 --- a/onnxruntime/test/providers/cpu/tensor/pad_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/pad_test.cc @@ -219,6 +219,41 @@ TYPED_TEST(PadOpTest, Pad_Wrap_1D) { "wrap"); } +#ifdef USE_WEBGPU +TEST(PadOpTest, Pad_Wrap_WebGpu_LowerPadExceedsInt32Fails) { + if (DefaultWebGpuExecutionProvider().get() == nullptr) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + + OpTester test("Pad", 19); + test.AddAttribute("mode", "wrap"); + test.AddInput("data", {4}, {1.0f, 2.0f, 3.0f, 4.0f}); + test.AddInput("pads", {2}, {2147483648LL, -2147483647LL}, true); + test.AddOutput("output", {5}, {0.0f, 0.0f, 0.0f, 0.0f, 0.0f}); + + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, + "WebGPU Pad only supports lower pads in the int32 range", {}, nullptr, &eps); +} + +TEST(PadOpTest, Pad_Wrap_WebGpu_PadGreaterThanInputDimension) { + if (DefaultWebGpuExecutionProvider().get() == nullptr) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + + OpTester test("Pad", 19); + test.AddAttribute("mode", "wrap"); + test.AddInput("data", {3}, {1.0f, 2.0f, 3.0f}); + test.AddInput("pads", {2}, {5, 0}, true); + test.AddOutput("output", {8}, {2.0f, 3.0f, 1.0f, 2.0f, 3.0f, 1.0f, 2.0f, 3.0f}); + + std::vector> eps; + eps.push_back(DefaultWebGpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps); +} +#endif + TYPED_TEST(PadOpTest, Pad_Edge_1D) { using T = TypeParam; RunAllOpsetAllDomainPadTests({3, 2},