[Metal] Fix copy_buffer_1byte dispatch grid (#24640)
copy_buffer_1byte silently dropped the trailing ~3/4 of every copy that
fell back to the compute-kernel path. The kernel copies one byte per
thread, but its dispatch grid was sized as ceil(length / (workgroup_size
* 4)) -- a quarter of the required ceil(length / workgroup_size)
workgroups -- so only the first min(ceil(length/128)*32, length) bytes
were written and the tail was left untouched. The bounds guard `if (id
>= spec.length) return` prevents over-run but cannot launch the missing
threads.
On macOS the compute path is taken whenever the source offset, target
offset, or length is not a multiple of 4 (the Metal blit fast-path
requires all three 4-byte aligned). It is also the path behind the
staging->device upload
(iree_hal_metal_command_buffer_prepare_update_buffer), so non-4-aligned
source uploads were truncated too -- a device copy then read a zeroed
source tail.
The *4 multiplier is correct for the fill kernels (fill_buffer_16byte /
fill_buffer_4byte / fill_buffer_1byte process 16/4/4 bytes per thread
respectively), which is where it was copied from; copy_buffer_1byte is a
flat 1-byte-per-thread loop with no such expansion.
Fix (dispatch only -- the copy_buffer_1byte.metal kernel is unchanged):
size the grid as ceil(length / workgroup_size) so the launched thread
count (ceil(length/32) * 32) covers every byte. The existing `if (id >=
spec.length) return` guard masks the surplus threads in the final
workgroup.
Verified:
CTS/CommandBufferCopyBufferTest.CopySizeAndAlignmentClasses/metal now
passes all alignment/size sub-cases that previously truncated
(offset/length not a multiple of 4; lengths 31/32/33/.../65536 across
all alignment classes). The aligned16_mib sub-case still fails,
independently, from a staging-buffer capacity issue (128 KiB staging
buffer vs. a 1 MiB single-command-buffer upload), not from this dispatch
sizing. Signed-off-by: Alex Vasile
<48962821+Alex-Vasile@users.noreply.github.com>
Signed-off-by: Alex Vasile <48962821+Alex-Vasile@users.noreply.github.com>
diff --git a/runtime/src/iree/hal/drivers/metal/builtin_executables.m b/runtime/src/iree/hal/drivers/metal/builtin_executables.m
index a14ece3..8ed1cc4 100644
--- a/runtime/src/iree/hal/drivers/metal/builtin_executables.m
+++ b/runtime/src/iree/hal/drivers/metal/builtin_executables.m
@@ -303,7 +303,9 @@
// Encode the dispatch.
const iree_device_size_t workgroup_size = 32;
- iree_device_size_t workgroup_count = iree_device_size_ceil_div(length, workgroup_size * 4);
+ // copy_buffer_1byte copies 1 byte/thread (no 4x expansion like the fill
+ // kernels), so the grid must span ceil(length/workgroup_size) workgroups.
+ iree_device_size_t workgroup_count = iree_device_size_ceil_div(length, workgroup_size);
[encoder dispatchThreadgroups:MTLSizeMake(workgroup_count, 1, 1)
threadsPerThreadgroup:MTLSizeMake(workgroup_size, 1, 1)];