Display error message when too many arguments are passed (#526)

Author: oulgen

helion/runtime/kernel.py

@@ -150,6 +150,10 @@ def bind(self, args: tuple[object, ...]) -> BoundKernel[_R]:
         if not isinstance(args, tuple):
             assert isinstance(args, list), "args must be a tuple or list"
             args = tuple(args)
+        if len(args) > len(self.signature.parameters):
+            raise TypeError(
+                f"Too many arguments passed to the kernel, expected: {len(self.signature.parameters)} got: {len(args)}."
+            )
         signature = self.specialization_key(args)
         cache_key = self._get_bound_kernel_cache_key(args, signature)
         bound_kernel = (

test/test_errors.py

@@ -214,6 +214,20 @@ def fn(a: torch.Tensor) -> torch.Tensor:
         ):
             code_and_output(fn, (torch.randn(4, device=DEVICE),))
 
+    def test_too_many_args(self):
+        @helion.kernel()
+        def kernel(x: torch.Tensor) -> torch.Tensor:
+            result = torch.zeros_like(x)
+            for i in hl.tile(x.size()):
+                result[i] = x[i]
+            return result
+
+        with self.assertRaisesRegex(
+            TypeError, r"Too many arguments passed to the kernel, expected: 1 got: 2."
+        ):
+            a = torch.randn(8, device=DEVICE)
+            code_and_output(kernel, (a, a))
+
 
 if __name__ == "__main__":
     unittest.main()