Skip to content

Commit

Permalink
[Doc] Remove the limit on kernel argument size (#8408)
Browse files Browse the repository at this point in the history
Issue: #

### Brief Summary

<!--
copilot:summary
-->
### <samp>🤖[[deprecated]](https://githubnext.com/copilot-for-prs-sunset)
Generated by Copilot at b075fef</samp>

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

<!--
copilot:walkthrough
-->
### <samp>🤖[[deprecated]](https://githubnext.com/copilot-for-prs-sunset)
Generated by Copilot at b075fef</samp>

* 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>
  • Loading branch information
lin-hitonami and pre-commit-ci[bot] authored Nov 17, 2023
1 parent 5ef84bf commit bcd623f
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 10 deletions.
34 changes: 24 additions & 10 deletions docs/lang/articles/kernels/kernel_function.md
Original file line number Diff line number Diff line change
Expand Up @@ -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()`.
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand All @@ -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

Expand All @@ -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 | <ul><li>Scalar</li><li>`ti.types.matrix()`</li><li>`ti.types.vector()`</li><li>`ti.types.struct()`(Only on LLVM-based backends)</li></ul> | <ul><li>Scalar</li><li>`ti.types.matrix()`</li><li>`ti.types.vector()`</li><li>`ti.types.struct()`</li><li>...</li></ul> | <ul><li>Scalar</li><li>`ti.types.matrix()`</li><li>`ti.types.vector()`</li><li>`ti.types.struct()`</li><li>...</li></ul> |
| Maximum number of elements in arguments | <ul><li>Unlimited (CPU and CUDA)</li><li>32 (OpenGL)</li><li>64 (otherwise)</li></ul> | Unlimited | Unlimited |
| Maximum number of return statements | 1 | 1 | Unlimited |


Expand Down
16 changes: 16 additions & 0 deletions tests/python/test_argument.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
16 changes: 16 additions & 0 deletions tests/python/test_return.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit bcd623f

Please sign in to comment.