Print error message for torch.chunk / torch.unbind to redirect users to hl.split (#921)

Author: yf225

helion/_compiler/type_propagation.py

@@ -773,6 +773,10 @@ def propagate_call(  # pyright: ignore[reportIncompatibleMethodOverride]
     ) -> TypeInfo | None:
         if self.value in (torch.nonzero, torch.Tensor.nonzero) and origin.is_device():
             raise exc.DataDependentOutputShapeNotSupported(op_desc="torch.nonzero")
+        if self.value in (torch.chunk, torch.Tensor.chunk) and origin.is_device():
+            raise exc.UnsupportedSplitOperation(op="torch.chunk")
+        if self.value in (torch.unbind, torch.Tensor.unbind) and origin.is_device():
+            raise exc.UnsupportedSplitOperation(op="torch.unbind")
         if is_api_func(fn := self.value):
             if fn._is_device_only and origin.is_host():
                 raise exc.DeviceAPIOnHost(fn.__qualname__)

helion/exc.py

@@ -114,6 +114,14 @@ class DataDependentOutputShapeNotSupported(BaseError):
     )
 
 
+class UnsupportedSplitOperation(BaseError):
+    message = (
+        "{op} is not supported in Helion device loops. "
+        "For splitting the last dimension with size 2, use hl.split(). "
+        "For other splitting operations, consider reshaping or using other hl.* operations."
+    )
+
+
 class RequiresTensorInAssignment(BaseError):
     message = "Expected tensor in right-hand side of assignment, got {0!s}."
 

test/test_errors.py

@@ -293,6 +293,74 @@ def torch_nonzero_in_device_code(x: torch.Tensor) -> torch.Tensor:
                 torch_nonzero_in_device_code, (torch.randn(2, 2, device=DEVICE),)
             )
 
+    def test_torch_chunk_device_error(self):
+        """Test that torch.chunk raises error in device loops and suggests hl.split()."""
+
+        @helion.kernel(use_default_config=True, static_shapes=True)
+        def kernel_with_chunk(q: torch.Tensor) -> torch.Tensor:
+            _, _, M, D = q.shape
+            D = hl.specialize(D)
+            M = hl.specialize(M)
+            q = q.reshape(-1, D)
+            total_rows = q.shape[0]
+            block_m = hl.register_block_size(M)
+            result = hl.zeros([total_rows, D])
+            for tile_m in hl.tile(total_rows, block_size=block_m):
+                acc = hl.zeros([tile_m, D])
+
+                for _tile_n in hl.tile(M, block_size=block_m):
+                    acc = torch.stack(torch.chunk(acc, 2, dim=-1), dim=-2).reshape(
+                        acc.shape
+                    )
+                    acc = acc + 0
+
+                result[tile_m, :] = acc
+
+            return result
+
+        with self.assertRaisesRegex(
+            helion.exc.UnsupportedSplitOperation,
+            r"torch\.chunk is not supported in Helion device loops.*hl\.split\(\)",
+        ):
+            code_and_output(
+                kernel_with_chunk,
+                (torch.randn(1, 1, 128, 128, device=DEVICE, dtype=torch.bfloat16),),
+            )
+
+    def test_torch_unbind_device_error(self):
+        """Test that torch.unbind raises error in device loops and suggests hl.split()."""
+
+        @helion.kernel(use_default_config=True, static_shapes=True)
+        def kernel_with_unbind(q: torch.Tensor) -> torch.Tensor:
+            _, _, M, D = q.shape
+            D = hl.specialize(D)
+            M = hl.specialize(M)
+            q = q.reshape(-1, D)
+            total_rows = q.shape[0]
+            block_m = hl.register_block_size(M)
+            result = hl.zeros([total_rows, D])
+            for tile_m in hl.tile(total_rows, block_size=block_m):
+                acc = hl.zeros([tile_m, D])
+
+                for _tile_n in hl.tile(M, block_size=block_m):
+                    reshaped = acc.reshape(tile_m, 2, D // 2)
+                    acc0, acc1 = torch.unbind(reshaped, dim=1)
+                    acc = torch.stack((acc0, acc1), dim=1).reshape(tile_m, D)
+                    acc = acc + 0
+
+                result[tile_m, :] = acc
+
+            return result
+
+        with self.assertRaisesRegex(
+            helion.exc.UnsupportedSplitOperation,
+            r"torch\.unbind is not supported in Helion device loops.*hl\.split\(\)",
+        ):
+            code_and_output(
+                kernel_with_unbind,
+                (torch.randn(1, 1, 128, 128, device=DEVICE, dtype=torch.bfloat16),),
+            )
+
     def test_closure_fn(self):
         @helion.kernel()
         def bad_fn(x: torch.Tensor) -> torch.Tensor: