[vulkan] Relax number of array args for each kernel (#5689)

Author: ailzhang

docs/lang/articles/kernels/syntax.md

@@ -72,12 +72,12 @@ my_kernel(24, 3.2)  # The system prints 27.2
 
 The upper limit for element numbers is backend-specific:
 
-- 8 on OpenGL backend
-- 64 on CPU, Vulkan, CUDA, or Metal
+- 64 on CPU, Vulkan, CUDA, OpenGL or Metal
 
 :::note
 - The number of elements in a scalar argument is always 1.
 - The number of the elements in a `ti.Matrix` or in a `ti.Vector` is the actual number of scalars inside of them.
+- The upper limit of array arguments (`ti.types.ndarray()`) is 32 and they must be the first 32 among all 64 arguments.
 :::
 
 ```python

python/taichi/lang/kernel_impl.py

@@ -759,16 +759,14 @@ def func__(*args):
             ) and self.runtime.target_tape and not self.runtime.grad_replaced:
                 self.runtime.target_tape.insert(self, args)
 
-            if actual_argument_slot > 8 and (
-                    impl.current_cfg().arch == _ti_core.opengl
-                    or impl.current_cfg().arch == _ti_core.cc):
+            if actual_argument_slot > 8 and impl.current_cfg(
+            ).arch == _ti_core.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 (
-                (impl.current_cfg().arch != _ti_core.opengl
-                 and impl.current_cfg().arch != _ti_core.cc)):
+            if actual_argument_slot > 64 and impl.current_cfg(
+            ).arch != _ti_core.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."
                 )

taichi/inc/constants.h

@@ -5,10 +5,10 @@
 constexpr int taichi_max_num_indices = 8;
 // legacy: only used in cc and opengl backends
 constexpr int taichi_max_num_args = 8;
-// used in llvm backend: only the first 16 arguments can be types.ndarray
+// used in llvm backend: only the first 32 arguments can be types.ndarray
 // TODO: refine argument passing
 constexpr int taichi_max_num_args_total = 64;
-constexpr int taichi_max_num_args_extra = 16;
+constexpr int taichi_max_num_args_extra = 32;
 constexpr int taichi_max_num_snodes = 1024;
 constexpr int kMaxNumSnodeTreesLlvm = 512;
 constexpr int taichi_max_gpu_block_dim = 1024;

tests/python/test_aot.py

@@ -264,40 +264,6 @@ def init(d: ti.i32, density1: ti.types.ndarray(),
     assert (density6.to_numpy() == (np.zeros(shape=(n, n)) + 6)).all()
 
 
-@test_utils.test(arch=ti.opengl)
-def test_opengl_exceed_max_ssbo():
-    # 8 ndarrays + args > 8 (maximum allowed)
-    n = 4
-    density1 = ti.ndarray(dtype=ti.f32, shape=(n, n))
-    density2 = ti.ndarray(dtype=ti.f32, shape=(n, n))
-    density3 = ti.ndarray(dtype=ti.f32, shape=(n, n))
-    density4 = ti.ndarray(dtype=ti.f32, shape=(n, n))
-    density5 = ti.ndarray(dtype=ti.f32, shape=(n, n))
-    density6 = ti.ndarray(dtype=ti.f32, shape=(n, n))
-    density7 = ti.ndarray(dtype=ti.f32, shape=(n, n))
-    density8 = ti.ndarray(dtype=ti.f32, shape=(n, n))
-
-    @ti.kernel
-    def init(d: ti.i32, density1: ti.types.ndarray(),
-             density2: ti.types.ndarray(), density3: ti.types.ndarray(),
-             density4: ti.types.ndarray(), density5: ti.types.ndarray(),
-             density6: ti.types.ndarray(), density7: ti.types.ndarray(),
-             density8: ti.types.ndarray()):
-        for i, j in density1:
-            density1[i, j] = d + 1
-            density2[i, j] = d + 2
-            density3[i, j] = d + 3
-            density4[i, j] = d + 4
-            density5[i, j] = d + 5
-            density6[i, j] = d + 6
-            density7[i, j] = d + 7
-            density8[i, j] = d + 8
-
-    with pytest.raises(RuntimeError):
-        init(0, density1, density2, density3, density4, density5, density6,
-             density7, density8)
-
-
 @test_utils.test(arch=[ti.opengl, ti.vulkan])
 def test_mpm99_aot():
     quality = 1  # Use a larger value for higher-res simulations

tests/python/test_argument.py

@@ -4,7 +4,7 @@
 from tests import test_utils
 
 
-@test_utils.test(arch=[ti.opengl, ti.cc])
+@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,
@@ -181,3 +181,54 @@ def duplicate():
     with pytest.raises(ti.TaichiSyntaxError,
                        match="Multiple values for argument 'a'"):
         duplicate()
+
+
+@test_utils.test(exclude=[ti.cc])
+def test_args_with_many_ndarrays():
+
+    particle_num = 0
+    cluster_num = 0
+    permu_num = 0
+
+    particlePosition = ti.Vector.ndarray(3, ti.f32, shape=10)
+    outClusterPosition = ti.Vector.ndarray(3, ti.f32, shape=10)
+    outClusterOffsets = ti.ndarray(ti.i32, shape=10)
+    outClusterSizes = ti.ndarray(ti.i32, shape=10)
+    outClusterIndices = ti.ndarray(ti.i32, shape=10)
+
+    particle_pos = ti.Vector.ndarray(3, ti.f32, shape=20)
+    particle_prev_pos = ti.Vector.ndarray(3, ti.f32, shape=20)
+    particle_rest_pos = ti.Vector.ndarray(3, ti.f32, shape=20)
+    particle_index = ti.ndarray(ti.i32, shape=20)
+
+    cluster_rest_mass_center = ti.Vector.ndarray(3, ti.f32, shape=20)
+    cluster_begin = ti.ndarray(ti.i32, shape=20)
+
+    @ti.kernel
+    def ti_import_cluster_data(
+        center: ti.types.vector(3,
+                                ti.f32), particle_num: int, cluster_num: int,
+        permu_num: int, particlePosition: ti.types.ndarray(field_dim=1),
+        outClusterPosition: ti.types.ndarray(field_dim=1),
+        outClusterOffsets: ti.types.ndarray(field_dim=1),
+        outClusterSizes: ti.types.ndarray(field_dim=1),
+        outClusterIndices: ti.types.ndarray(field_dim=1),
+        particle_pos: ti.types.ndarray(field_dim=1),
+        particle_prev_pos: ti.types.ndarray(field_dim=1),
+        particle_rest_pos: ti.types.ndarray(field_dim=1),
+        cluster_rest_mass_center: ti.types.ndarray(field_dim=1),
+        cluster_begin: ti.types.ndarray(field_dim=1),
+        particle_index: ti.types.ndarray(field_dim=1)):
+
+        added_permu_num = outClusterIndices.shape[0]
+
+        for i in range(added_permu_num):
+            particle_index[i] = 1.0
+
+    center = ti.math.vec3(0, 0, 0)
+    ti_import_cluster_data(center, particle_num, cluster_num, permu_num,
+                           particlePosition, outClusterPosition,
+                           outClusterOffsets, outClusterSizes,
+                           outClusterIndices, particle_pos, particle_prev_pos,
+                           particle_rest_pos, cluster_rest_mass_center,
+                           cluster_begin, particle_index)