Skip to content

Commit

Permalink
[Doc] Add doc about struct arguments (#7959)
Browse files Browse the repository at this point in the history
Issue: #

### Brief Summary

<!--
copilot:summary
-->
### <samp>🤖 Generated by Copilot at adf95e4</samp>

This pull request improves the documentation of `ti.kernel` and
`ti.func` by using the new `ti.types` module and showing how to work
with structs and matrices.

### Walkthrough

<!--
copilot:walkthrough
-->
### <samp>🤖 Generated by Copilot at adf95e4</samp>

* Update the argument and return value types of kernels and functions to
use the `ti.types` module
([link](https://github.com/taichi-dev/taichi/pull/7959/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL96-R98),
[link](https://github.com/taichi-dev/taichi/pull/7959/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL138-R158),
[link](https://github.com/taichi-dev/taichi/pull/7959/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL279-R298),
[link](https://github.com/taichi-dev/taichi/pull/7959/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL287-R306),
[link](https://github.com/taichi-dev/taichi/pull/7959/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cL300-R319))
* Add an example of passing a struct with a matrix to a kernel by value
in `docs/lang/articles/kernels/kernel_function.md`
([link](https://github.com/taichi-dev/taichi/pull/7959/files?diff=unified&w=0#diff-305d66d1c231266cb14a0d2b166363840f782874eb6c5b3d5d4e15e08f3cb04cR112-R130))

---------

Co-authored-by: Zhao Liang <mathzhaoliang@gmail.com>
  • Loading branch information
lin-hitonami and neozhaoliang authored May 10, 2023
1 parent dd5762d commit 4a24872
Showing 1 changed file with 22 additions and 7 deletions.
29 changes: 22 additions & 7 deletions docs/lang/articles/kernels/kernel_function.md
Original file line number Diff line number Diff line change
Expand Up @@ -93,9 +93,9 @@ Kernels in Taichi can be called either directly or from inside a native Python f

A kernel can accept multiple arguments. However, it's important to note that you can't pass arbitrary Python objects to a kernel. This is because Python objects can be dynamic and may contain data that the Taichi compiler cannot recognize.

The kernel can accept various argument types, including scalars, `ti.Matrix`, `ti.Vector`, `ti.types.ndarray()`, and `ti.template()`. These argument types make it easy to pass data from the Python scope to the Taichi scope. You can find the supported types in the `ti.types` module. For more information on this, see the [Type System](../type_system/type.md).
The kernel can accept various argument types, including scalars, `ti.types.matrix()`, `ti.types.vector()`, `ti.types.struct()`, `ti.types.ndarray()`, and `ti.template()`. These argument types make it easy to pass data from the Python scope to the Taichi scope. You can find the supported types in the `ti.types` module. For more information on this, see the [Type System](../type_system/type.md).

Scalars, `ti.Vector` and `ti.Matrix` are passed by value, which means that the kernel receives a copy of the argument. However, `ti.types.ndarray()` and `ti.template()` are passed by reference, which means that any changes made to the argument inside the kernel will affect the original value as well.
Scalars, `ti.types.matrix()`, `ti.types.vector()`, and `ti.types.struct()` are passed by value, which means that the kernel receives a copy of the argument. However, `ti.types.ndarray()` and `ti.template()` are passed by reference, which means that any changes made to the argument inside the kernel will affect the original value as well.

Note that we won't cover `ti.template()` here as it is a more advanced topic and is discussed in [Metaprogramming](../advanced/meta.md#template-metaprogramming).

Expand All @@ -109,6 +109,21 @@ def my_kernel(x: int, y: float):
my_kernel(1, 1.0) # Prints 2.0
```

Here is another example of passing a nested struct argument with a matrix to a kernel by value, in which we created a struct type `transform_type` that contains two members: a rotation matrix `R` and a translation vector `T`. We then created another struct type `pos_type` that has `transform_type` as its member and passed an instance of `pos_type` to a kernel.

```python
transform_type = ti.types.struct(R=ti.math.mat3, T=ti.math.vec3)
pos_type = ti.types.struct(x=ti.math.vec3, trans=transform_type)

@ti.kernel
def kernel_with_nested_struct_arg(p: pos_type) -> ti.math.vec3:
return p.trans.R @ p.x + p.trans.T

trans = transform_type(ti.math.mat3(1), [1, 1, 1])
p = pos_type(x=[1, 1, 1], trans=trans)
print(kernel_with_nested_struct_arg(p)) # [4., 4., 4.]
```

You can use `ti.types.ndarray()` as a type hint to pass a `ndarray` from NumPy or a `tensor` from PyTorch to a kernel. Taichi recognizes the shape and data type of these data structures, which allows you to access their attributes in a kernel.

In the example below, `x` is updated after `my_kernel()` is called since it is passed by reference:
Expand All @@ -135,8 +150,8 @@ print(x) # Prints [5, 7, 9]

### Return value

In Taichi, a kernel is allowed to have a maximum of one return value, which could either be a scalar, `ti.Matrix`, or `ti.Vector`.
Moreover, in the LLVM-based backends (CPU and CUDA backends), a return value could also be a `ti.Struct`.
In Taichi, a kernel is allowed to have a maximum of one return value, which could either be a scalar, `ti.types.matrix()`, or `ti.types.vector()`.
Moreover, in the LLVM-based backends (CPU and CUDA backends), a return value could also be a `ti.types.struct()`.

Here is an example of a kernel that returns a ti.Struct:

Expand Down Expand Up @@ -276,15 +291,15 @@ All Taichi functions are force-inlined. This means that if you call a Taichi fun

### Arguments

A Taichi function can accept multiple arguments, which may include scalar, `ti.Matrix`, `ti.Vector`, `ti.types.ndarray()`, `ti.template()`, `ti.field`, and `ti.Struct` types. Note that some of the restrictions on kernel arguments do not apply to Taichi functions:
A Taichi 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.


### Return values

Return values of a Taichi function can be scalars, `ti.Matrix`, `ti.Vector`, `ti.Struct`, or other types. Note the following:
Return values of a Taichi function can be scalars, `ti.types.matrix()`, `ti.types.vector()`, `ti.types.struct()`, or other types. Note the following:

- Unlike a kernel, a Taichi function can have multiple return values.
- It is *not* required (but recommended) to type hint the return values of a Taichi function.
Expand All @@ -297,7 +312,7 @@ Return values of a Taichi function can be scalars, `ti.Matrix`, `ti.Vector`, `ti
| Call scope | Python scope | Taichi scope |
| Type hint arguments | Mandatory | Recommended |
| Type hint return values | Mandatory | Recommended |
| Return type | <ul><li>Scalar</li><li>`ti.Vector`</li><li>`ti.Matrix`</li><li>`ti.Struct`(Only on LLVM-based backends)</li></ul> | <ul><li>Scalar</li><li>`ti.Vector`</li><li>`ti.Matrix`</li><li>`ti.Struct`</li><li>...</li></ul> |
| 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> |
| Maximum number of elements in arguments | <ul><li>32 (OpenGL)</li><li>64 (otherwise)</li></ul> | Unlimited |
| Maximum number of return values in a return statement | 1 | Unlimited |

Expand Down

0 comments on commit 4a24872

Please sign in to comment.