[WebGPU] Handle aligned copies within the same buffer (#25193)

WebGPU doesn't allow copying within the same buffer, even when the
ranges don't overlap. We can use a temporary buffer for these aligned
copies and clean it up through the existing path.

---------

Signed-off-by: anutosh491 <andersonbhat491@gmail.com>

runtime/src/iree/hal/cts/command_buffer/copy_buffer_test.cc

diff --git a/runtime/src/iree/hal/cts/command_buffer/copy_buffer_test.cc b/runtime/src/iree/hal/cts/command_buffer/copy_buffer_test.cc
index 96f3208..f90665a 100644
--- a/runtime/src/iree/hal/cts/command_buffer/copy_buffer_test.cc
+++ b/runtime/src/iree/hal/cts/command_buffer/copy_buffer_test.cc
@@ -264,6 +264,40 @@
   RunCopyBufferTest(/*source_offset=*/0, /*target_offset=*/0, 1024 * 1024);
 }
 
+TEST_P(CommandBufferCopyBufferTest, CopyAlignedDisjointRangesWithinSameBuffer) {
+  for (iree_device_size_t length : {4, 16, 256, 1024}) {
+    for (iree_device_size_t source_offset : {iree_device_size_t{0}, length}) {
+      iree_device_size_t target_offset = length - source_offset;
+      SCOPED_TRACE(::testing::Message()
+                   << "length=" << length << " source_offset=" << source_offset
+                   << " target_offset=" << target_offset);
+      std::vector<uint8_t> expected = MakeDeterministicBytes(2 * length + 16);
+      Ref<iree_hal_buffer_t> buffer;
+      IREE_ASSERT_OK(CreateDeviceBufferWithData(expected.data(),
+                                                expected.size(), buffer.out()));
+
+      Ref<iree_hal_command_buffer_t> command_buffer;
+      IREE_ASSERT_OK(iree_hal_command_buffer_create(
+          device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
+          IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
+          binding_capacity(), command_buffer.out()));
+      IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
+      IREE_ASSERT_OK(iree_hal_command_buffer_copy_buffer(
+          command_buffer,
+          MakeRef(buffer, kSourceBufferSlot, source_offset, length),
+          MakeRef(buffer, kTargetBufferSlot, target_offset, length),
+          IREE_HAL_COPY_FLAG_NONE));
+      IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
+      IREE_ASSERT_OK(SubmitWithBindings(command_buffer, buffer, buffer));
+
+      std::memcpy(expected.data() + target_offset,
+                  expected.data() + source_offset, length);
+      EXPECT_THAT(ReadBufferBytes(buffer, 0, expected.size()),
+                  ContainerEq(expected));
+    }
+  }
+}
+
 CTS_REGISTER_COMMAND_BUFFER_TEST_SUITE(CommandBufferCopyBufferTest);
 
 }  // namespace iree::hal::cts

runtime/src/iree/hal/drivers/webgpu/webgpu_imports.mjs

diff --git a/runtime/src/iree/hal/drivers/webgpu/webgpu_imports.mjs b/runtime/src/iree/hal/drivers/webgpu/webgpu_imports.mjs
index f8e4ba8..0e94286 100644
--- a/runtime/src/iree/hal/drivers/webgpu/webgpu_imports.mjs
+++ b/runtime/src/iree/hal/drivers/webgpu/webgpu_imports.mjs
@@ -270,10 +270,21 @@
 
           if ((effectiveSrc & 3) === 0 && (effectiveDst & 3) === 0 &&
               (length & 3) === 0) {
-            // Aligned: use native copyBufferToBuffer.
-            encoder.copyBufferToBuffer(
-                buffers[srcSlot], effectiveSrc, buffers[dstSlot], effectiveDst,
-                length);
+            if (buffers[srcSlot] === buffers[dstSlot]) {
+              const scratch = device.createBuffer({
+                usage: 0x0004 | 0x0008,  // COPY_SRC | COPY_DST
+                size: length,
+              });
+              temporaryBuffers.push(scratch);
+              encoder.copyBufferToBuffer(
+                  buffers[srcSlot], effectiveSrc, scratch, 0, length);
+              encoder.copyBufferToBuffer(
+                  scratch, 0, buffers[dstSlot], effectiveDst, length);
+            } else {
+              encoder.copyBufferToBuffer(
+                  buffers[srcSlot], effectiveSrc, buffers[dstSlot],
+                  effectiveDst, length);
+            }
           } else {
             // Unaligned: dispatch copy builtin shader.
             const params = device.createBuffer({