From bcd623fddccde980a01c2a9bb4facb56862cc26e Mon Sep 17 00:00:00 2001 From: Lin Jiang Date: Fri, 17 Nov 2023 11:18:17 +0800 Subject: [PATCH] [Doc] Remove the limit on kernel argument size (#8408) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Issue: # ### Brief Summary ### 🤖[[deprecated]](https://githubnext.com/copilot-for-prs-sunset) Generated by Copilot at b075fef This pull request improves the documentation of various types of functions in Taichi by explaining the new features and limitations introduced in v1.7.0. It also warns the users about the possible performance and memory issues of using large arguments or return values in `kernel_function.md`. ### Walkthrough ### 🤖[[deprecated]](https://githubnext.com/copilot-for-prs-sunset) Generated by Copilot at b075fef * Warn users about the potential issues of passing or returning large values in kernel functions, and provide suggestions on how to avoid them ([link](https://github.com/taichi-dev/taichi/pull/8408/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cR157-R166), [link](https://github.com/taichi-dev/taichi/pull/8408/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cR201-R210)) * Simplify the documentation of Taichi inline and real functions by removing redundant information on the arguments ([link](https://github.com/taichi-dev/taichi/pull/8408/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL267-R289), [link](https://github.com/taichi-dev/taichi/pull/8408/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL284-R303)) * Update the table that compares Taichi kernel, Taichi inline function, and Taichi real function by removing the row on the maximum number of elements in arguments, since this limit has been removed in Taichi v1.7.0 ([link](https://github.com/taichi-dev/taichi/pull/8408/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL305)) --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/lang/articles/kernels/kernel_function.md | 34 +++++++++++++------ tests/python/test_argument.py | 16 +++++++++ tests/python/test_return.py | 16 +++++++++ 3 files changed, 56 insertions(+), 10 deletions(-) diff --git a/docs/lang/articles/kernels/kernel_function.md b/docs/lang/articles/kernels/kernel_function.md index 3bf2ff39f7023..26382c8c98267 100644 --- a/docs/lang/articles/kernels/kernel_function.md +++ b/docs/lang/articles/kernels/kernel_function.md @@ -154,6 +154,16 @@ You can also use argument packs if you want to pass many arguments to a kernel. When defining the arguments of a kernel in Taichi, please make sure that each of the arguments has type hint. +:::caution WARNING + +We have removed the limit on the size of the argument in Taichi v1.7.0. +However, please keep in mind that the size of arguments in a kernel should be small. +When you pass a large argument to a kernel, the compile time will increase significantly. +If you find yourself passing a large argument to a kernel, you may want to consider using a `ti.field()` or a `ti.types.ndarray()` instead. + +We have not tested arguments with a very large size (>4KB), and we do not guarantee that it will work properly. +::: + ### Return value In Taichi, a kernel can have multiple return values, and the return values can either be a scalar, `ti.types.matrix()`, or `ti.types.vector()`. @@ -188,6 +198,16 @@ When defining the return value of a kernel in Taichi, it is important to follow - Use type hint to specify the return value of a kernel. - Make sure that you have at most one return statement in a kernel. +:::caution WARNING + +We have removed the limit on the size of the return values in Taichi v1.7.0. +However, please keep in mind that the size of return values in a kernel should be small. +When the return value of the kernel is very large, the compile time will increase significantly. +If you find your return value is very large, you may want to consider using a `ti.field()` or a `ti.types.ndarray()` instead. + +We have not tested return values with a very large size (>4KB), and we do not guarantee that it will work properly. +::: + #### Automatic type cast In the following code snippet, the return value is automatically cast into the hinted type: @@ -264,11 +284,8 @@ All Taichi inline functions are force-inlined. This means that if you call a Tai ### Arguments -A Taichi inline function can accept multiple arguments, which may include scalar, `ti.types.matrix()`, `ti.types.vector()`, `ti.types.struct()`, `ti.types.ndarray()`, `ti.field()`, and `ti.template()` types. Note that some of the restrictions on kernel arguments do not apply to Taichi functions: - -- It is not strictly required to type hint the function arguments (but it is still recommended). -- You can pass an unlimited number of elements in the function arguments. - +A Taichi inline function can accept multiple arguments, which may include scalar, `ti.types.matrix()`, `ti.types.vector()`, `ti.types.struct()`, `ti.types.ndarray()`, `ti.field()`, and `ti.template()` types. +Note that unlike Taichi kernels, it is not strictly required to type hint the function arguments (but it is still recommended). ### Return values @@ -281,10 +298,8 @@ Return values of a Taichi inline function can be scalars, `ti.types.matrix()`, ` ### Arguments -A Taichi real function can accept multiple arguments, which may include scalar, `ti.types.matrix()`, `ti.types.vector()`, `ti.types.struct()`, `ti.types.ndarray()`, `ti.field()`, and `ti.template()` types. Note the following: - -- You must type hint the function arguments. -- You can pass an unlimited number of elements in the function arguments. +A Taichi real function can accept multiple arguments, which may include scalar, `ti.types.matrix()`, `ti.types.vector()`, `ti.types.struct()`, `ti.types.ndarray()`, `ti.field()`, and `ti.template()` types. +Note that you must type hint the function arguments. ### Return values @@ -302,7 +317,6 @@ Return values of a Taichi inline function can be scalars, `ti.types.matrix()`, ` | Type hint arguments | Mandatory | Recommended | Mandatory | | Type hint return values | Mandatory | Recommended | Mandatory | | Return type | | | | -| Maximum number of elements in arguments | | Unlimited | Unlimited | | Maximum number of return statements | 1 | 1 | Unlimited | diff --git a/tests/python/test_argument.py b/tests/python/test_argument.py index a957765375fc1..6c7bf5a77eef7 100644 --- a/tests/python/test_argument.py +++ b/tests/python/test_argument.py @@ -274,3 +274,19 @@ def foo(a: ti.f32) -> ti.f32: return bar(a) assert foo(1.5) == 1.0 + + +@test_utils.test(exclude=[ti.amdgpu]) +def test_arg_4k(): + vec1024 = ti.types.vector(1024, ti.i32) + + @ti.kernel + def bar(a: vec1024) -> ti.i32: + ret = 0 + for i in range(1024): + ret += a[i] + + return ret + + a = vec1024([i for i in range(1024)]) + assert bar(a) == 523776 diff --git a/tests/python/test_return.py b/tests/python/test_return.py index 9ea1df4b02821..69e58ce9ed92b 100644 --- a/tests/python/test_return.py +++ b/tests/python/test_return.py @@ -343,3 +343,19 @@ def foo() -> tp: return bar() assert foo().a == 0 + + +@test_utils.test(exclude=[ti.amdgpu]) +def test_ret_4k(): + vec1024 = ti.types.vector(1024, ti.i32) + + @ti.kernel + def foo() -> vec1024: + ret = vec1024(0) + for i in range(1024): + ret[i] = i + return ret + + ret = foo() + for i in range(1024): + assert ret[i] == i