[ExecuTorch][WebGPU] split_with_sizes_copy op tests

backends/webgpu/test/op_tests/cases.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -758,6 +758,50 @@
         atol=1e-3,
         rtol=1e-3,
     )
+from executorch.backends.webgpu.test.ops.test_split_with_sizes_copy import (
+    _det_input as _split_det_input,
+    SplitWithSizesModule,
+)
+
+
+@register_op_test("split_with_sizes_copy")
+def _split_with_sizes_copy_suite() -> WebGPUTestSuite:
+    # YOLO Detect-head split of concatenated predictions. Multi-output: the
+    # framework compares chunk 0 (out_index 0) while each case runs all N
+    # per-chunk dispatches. Covers a 3-way channel split, a dim-0 split, and
+    # a last-dim split. copy is bit-exact -> float32 golden.
+    return WebGPUTestSuite(
+        module_factory=lambda sizes, dim, out_order=None: SplitWithSizesModule(
+            sizes, dim, out_order
+        ),
+        cases=[
+            Case(
+                name="three_dim1",
+                construct={"sizes": [2, 3, 3], "dim": 1},
+                inputs=(InputSpec(shape=(1, 8, 4, 4), gen=_split_det_input),),
+            ),
+            Case(
+                name="two_dim0",
+                construct={"sizes": [3, 2], "dim": 0},
+                inputs=(InputSpec(shape=(5, 4), gen=_split_det_input),),
+            ),
+            Case(
+                name="dim_last",
+                construct={"sizes": [4, 4], "dim": -1},
+                inputs=(InputSpec(shape=(2, 8), gen=_split_det_input),),
+            ),
+            # Reorder so chunk 1 (running offset > 0) is output 0 -> verifies the
+            # per-chunk start accumulation (out_index 0 is all the framework checks).
+            Case(
+                name="offset_chunk1_first",
+                construct={"sizes": [2, 3, 3], "dim": 1, "out_order": [1, 0, 2]},
+                inputs=(InputSpec(shape=(1, 8, 4, 4), gen=_split_det_input),),
+            ),
+        ],
+        golden_dtype="float32",
+        atol=1e-4,
+        rtol=1e-3,
+    )
 from executorch.backends.webgpu.test.ops.test_max_pool2d import (
     _det_input as _maxpool_det_input,
     MaxPool2dModule,

backends/webgpu/test/ops/test_split_with_sizes_copy.py

@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.split_with_sizes_copy.default` module for the WebGPU op-test framework.
+
+split_with_sizes is on YOLO's Detect head (splitting the concatenated
+box/objectness/class predictions). `torch.split` by a size list lowers to
+`split_with_sizes_copy` in the edge dialect.
+"""
+
+import torch
+
+
+class SplitWithSizesModule(torch.nn.Module):
+    def __init__(self, sizes, dim: int, out_order=None) -> None:
+        super().__init__()
+        self.sizes = sizes
+        self.dim = dim
+        # Reorder the returned chunks so a non-first chunk can be output 0 (the
+        # op-test framework compares out_index 0) -- exercises the running offset.
+        self.out_order = out_order
+
+    def forward(self, x: torch.Tensor):
+        chunks = torch.split(x, self.sizes, self.dim)
+        if self.out_order is not None:
+            return tuple(chunks[i] for i in self.out_order)
+        return chunks
+
+
+def _det_input(shape) -> torch.Tensor:
+    # ((i % 23) - 11) / 16: exact in fp32, spans negatives through positives.
+    n = 1
+    for s in shape:
+        n *= s
+    idx = torch.arange(n, dtype=torch.int64)
+    return (((idx % 23) - 11).to(torch.float32) / 16.0).reshape(shape)