[Error] Add error message when the number of elements in kernel arguments exceed

python/taichi/lang/kernel_impl.py

@@ -12,7 +12,7 @@
 from taichi.lang.ast import (ASTTransformerContext, KernelSimplicityASTChecker,
                              transform_tree)
 from taichi.lang.enums import Layout
-from taichi.lang.exception import (TaichiCompilationError,
+from taichi.lang.exception import (TaichiCompilationError, TaichiRuntimeError,
                                    TaichiRuntimeTypeError, TaichiSyntaxError)
 from taichi.lang.expr import Expr
 from taichi.lang.matrix import Matrix, MatrixType
@@ -608,6 +608,20 @@ def func__(*args):
             if not self.is_grad and self.runtime.target_tape and not self.runtime.grad_replaced:
                 self.runtime.target_tape.insert(self, args)
 
+            if actual_argument_slot > 8 and (
+                    _ti_core.arch_name(impl.current_cfg().arch) == 'opengl'
+                    or _ti_core.arch_name(impl.current_cfg().arch) == 'cc'):
+                raise TaichiRuntimeError(
+                    f"The number of elements in kernel arguments is too big! Do not exceed 8 on {_ti_core.arch_name(impl.current_cfg().arch)} backend."
+                )
+
+            if actual_argument_slot > 64 and (
+                (_ti_core.arch_name(impl.current_cfg().arch) != 'opengl'
+                 and _ti_core.arch_name(impl.current_cfg().arch) != 'cc')):
+                raise TaichiRuntimeError(
+                    f"The number of elements in kernel arguments is too big! Do not exceed 64 on {_ti_core.arch_name(impl.current_cfg().arch)} backend."
+                )
+
             t_kernel(launch_ctx)
 
             ret = None

tests/python/test_argument.py

@@ -0,0 +1,183 @@
+import pytest
+
+import taichi as ti
+from tests import test_utils
+
+
+@test_utils.test(arch=ti.opengl)
+def test_exceed_max_eight():
+    @ti.kernel
+    def foo1(a: ti.i32, b: ti.i32, c: ti.i32, d: ti.i32, e: ti.i32, f: ti.i32,
+             g: ti.i32, h: ti.i32) -> ti.i32:
+        return a + b + c + d + e + f + g + h
+
+    assert foo1(1, 2, 3, 4, 5, 6, 7, 8) == 36
+
+    @ti.kernel
+    def foo2(a: ti.i32, b: ti.i32, c: ti.i32, d: ti.i32, e: ti.i32, f: ti.i32,
+             g: ti.i32, h: ti.i32, i: ti.i32) -> ti.i32:
+        return a + b + c + d + e + f + g + h + i
+
+    with pytest.raises(
+            ti.TaichiRuntimeError,
+            match=
+            "The number of elements in kernel arguments is too big! Do not exceed 8 on opengl backend."
+    ):
+        foo2(1, 2, 3, 4, 5, 6, 7, 8, 9)
+
+
+@test_utils.test(arch=ti.cc)
+def test_exceed_max_eight():
+    @ti.kernel
+    def foo1(a: ti.i32, b: ti.i32, c: ti.i32, d: ti.i32, e: ti.i32, f: ti.i32,
+             g: ti.i32, h: ti.i32) -> ti.i32:
+        return a + b + c + d + e + f + g + h
+
+    assert foo1(1, 2, 3, 4, 5, 6, 7, 8) == 36
+
+    @ti.kernel
+    def foo2(a: ti.i32, b: ti.i32, c: ti.i32, d: ti.i32, e: ti.i32, f: ti.i32,
+             g: ti.i32, h: ti.i32, i: ti.i32) -> ti.i32:
+        return a + b + c + d + e + f + g + h + i
+
+    with pytest.raises(
+            ti.TaichiRuntimeError,
+            match=
+            "The number of elements in kernel arguments is too big! Do not exceed 8 on cc backend."
+    ):
+        foo2(1, 2, 3, 4, 5, 6, 7, 8, 9)
+
+
+@test_utils.test(arch=ti.cuda)
+def test_exceed_max_64():
+    N = 64
+
+    @ti.kernel
+    def foo1(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+    assert foo1(A) == 64
+
+    N = 65
+
+    @ti.kernel
+    def foo2(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+
+    with pytest.raises(
+            ti.TaichiRuntimeError,
+            match=
+            "The number of elements in kernel arguments is too big! Do not exceed 64 on cuda backend."
+    ):
+        foo2(A)
+
+
+@test_utils.test(arch=ti.metal)
+def test_exceed_max_64():
+    N = 64
+
+    @ti.kernel
+    def foo1(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+    assert foo1(A) == 64
+
+    N = 65
+
+    @ti.kernel
+    def foo2(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+
+    with pytest.raises(
+            ti.TaichiRuntimeError,
+            match=
+            "The number of elements in kernel arguments is too big! Do not exceed 64 on metal backend."
+    ):
+        foo2(A)
+
+
+@test_utils.test(arch=ti.vulkan)
+def test_exceed_max_64():
+    N = 64
+
+    @ti.kernel
+    def foo1(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+    assert foo1(A) == 64
+
+    N = 65
+
+    @ti.kernel
+    def foo2(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+
+    with pytest.raises(
+            ti.TaichiRuntimeError,
+            match=
+            "The number of elements in kernel arguments is too big! Do not exceed 64 on vulkan backend."
+    ):
+        foo2(A)
+
+
+@test_utils.test(arch=ti.x64)
+def test_exceed_max_64():
+    N = 64
+
+    @ti.kernel
+    def foo1(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+    assert foo1(A) == 64
+
+    N = 65
+
+    @ti.kernel
+    def foo2(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+
+    with pytest.raises(
+            ti.TaichiRuntimeError,
+            match=
+            "The number of elements in kernel arguments is too big! Do not exceed 64 on x64 backend."
+    ):
+        foo2(A)
+
+
+@test_utils.test(arch=ti.arm64)
+def test_exceed_max_64():
+    N = 64
+
+    @ti.kernel
+    def foo1(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+    assert foo1(A) == 64
+
+    N = 65
+
+    @ti.kernel
+    def foo2(a: ti.types.vector(N, ti.i32)) -> ti.i32:
+        return a.sum()
+
+    A = ti.Vector([1] * N)
+
+    with pytest.raises(
+            ti.TaichiRuntimeError,
+            match=
+            "The number of elements in kernel arguments is too big! Do not exceed 64 on arm64 backend."
+    ):
+        foo2(A)