From c8ef7741e33d4cab26decb6e8db4b907b68a7223 Mon Sep 17 00:00:00 2001 From: Lin Jiang Date: Mon, 20 Nov 2023 10:12:16 +0800 Subject: [PATCH] [Doc] Update the doc of real functions (#8412) 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 32b8bfa Add a new section on Taichi real functions and update the existing section on real function arguments and return values in `docs/lang/articles/kernels/kernel_function.md`. Fix the section title and warn about a possible bug on older NVIDIA GPUs. ### Walkthrough ### 🤖[[deprecated]](https://githubnext.com/copilot-for-prs-sunset) Generated by Copilot at 32b8bfa * Add a new section on Taichi real functions, explaining their concept, usage, and benefits ([link](https://github.com/taichi-dev/taichi/pull/8412/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL299-R362)) * Update the section on Taichi real function arguments and return values, showing the new syntax and semantics for passing scalars, vectors, and matrices ([link](https://github.com/taichi-dev/taichi/pull/8412/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL299-R362)) * Add a warning about a possible bug when passing scalars by reference on older NVIDIA GPUs, and suggest a workaround ([link](https://github.com/taichi-dev/taichi/pull/8412/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL299-R362)) * Add an example of using recursion in a real function, demonstrating its power and flexibility ([link](https://github.com/taichi-dev/taichi/pull/8412/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL299-R362)) * Modify the section title to match the style of the other sections, using "Kernel Functions" instead of "Kernel Function" ([link](https://github.com/taichi-dev/taichi/pull/8412/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL299-R362)) --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/lang/articles/kernels/kernel_function.md | 59 ++++++++++++++++++- 1 file changed, 58 insertions(+), 1 deletion(-) diff --git a/docs/lang/articles/kernels/kernel_function.md b/docs/lang/articles/kernels/kernel_function.md index 26382c8c98267..4ea02eee3c268 100644 --- a/docs/lang/articles/kernels/kernel_function.md +++ b/docs/lang/articles/kernels/kernel_function.md @@ -296,18 +296,75 @@ Return values of a Taichi inline function can be scalars, `ti.types.matrix()`, ` ## Taichi real function +Taichi real functions are Taichi functions that are compiled into separate functions (like the device functions in CUDA) and can be called recursively at runtime. +The code inside the Taichi real function are executed serially, which means that you cannot write parallel loop inside it. +However, if the real function is called inside a parallel loop, the real function will be executed in parallel along with other parts of the parallel loop. + +If you want to do deep runtime recursion on CUDA, you may need to increase the stack size by passing `cuda_stack_limit` to `ti.init()`. + +Taichi real functions are only supported on the LLVM-based backends (CPU and CUDA backends). + ### 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. +The scalar, `ti.types.matrix()`, `ti.types.vector()`, and `ti.types.struct()` arguments are passed by value, while the `ti.types.ndarray()`, `ti.field()`, and `ti.template()` arguments are passed by reference. + Note that you must type hint the function arguments. +#### Passing a scalar by reference +The Taichi real function also supports passing a scalar by reference. To do this, you need to wrap the type hint with `ti.ref()`. + +Here is an example of passing an integer by reference: + +```python +@ti.real_func +def add_one(a: ti.ref(ti.i32)): + a += 1 + +@ti.kernel +def foo(): + a = 1 + add_one(a) + print(a) + +foo() # Prints 2 +``` + +:::caution WARNING + +Passing scalars by reference may be buggy on NVIDIA GPUs with Pascal or older architecture (for example GTX 1080 Ti). +We recommend using the latest NVIDIA GPUs (at least 20-series) if you want to pass a scalar by reference. + +::: + ### Return values -Return values of a Taichi inline function can be scalars, `ti.types.matrix()`, `ti.types.vector()`, `ti.types.struct()`, or other types. Note the following: +Return values of a Taichi real function can be scalars, `ti.types.matrix()`, `ti.types.vector()`, `ti.types.struct()`, or other types. Note the following: - You must type hint the return values of a Taichi real function. - A Taichi real function *can* have more than one `return` statement. +The example below calls the real function `sum_func` recursively to calculate the sum of `1` to `n`. +Inside the real function, there are two `return` statements, and the recursion depth is not a constant number. +The cuda stack limit is set to 32kB to allow deep runtime recursion. + +```python skip-ci:ToyDemo +ti.init(arch=ti.cuda, cuda_stack_limit=32768) + +@ti.real_func +def sum_func(n: ti.i32) -> ti.i32: + if n == 0: + return 0 + return sum_func(n - 1) + n + +@ti.kernel +def sum(n: ti.i32) -> ti.i32: + return sum_func(n) + +print(sum(100)) # 5050 +``` + +You can find more examples of the real function in the [repository](https://github.com/taichi-dev/taichi/tree/master/python/taichi/examples/real_func). ## A recap: Taichi kernel vs. Taichi inline function vs. Taichi real function