[Lang] Customized struct support

docs/lang/api/struct.md

@@ -0,0 +1,108 @@
+---
+sidebar_position: 3
+---
+
+# Structs
+
+Mixed-data-type records can be created in Taichi as custom structs. A struct in Taichi can have two forms. Similar to vectors and matrices, structs have both local variable and global field forms.
+
+## Declaration
+
+### As global vector fields
+
+::: {.function}
+ti.struct.field(members, shape = None, offset = None)
+
+parameter members
+
+: (Dict[str, DataType]) name and data type for each struct member. The data type of a member can be either a primitive type (numbers) or a compound type (vectors, matrices, structs).
+
+parameter shape
+
+: (optional, scalar or tuple) shape of the struct field, see
+`tensor`{.interpreted-text role="ref"}
+
+parameter offset
+
+: (optional, scalar or tuple) see `offset`{.interpreted-text
+role="ref"}
+
+For example, this creates a Struct field of the two float members `a` and `b`: :
+
+    # Python-scope
+    x = ti.Struct.field({'a': ti.f32, 'b': fi.f32}, shape=(5, 4))
+
+A struct field with vector, matrix, or struct components can be created with compound types: :
+
+    # Python-scope
+    vec3 = ti.types.vector(3, float)
+    x = ti.Struct.field({'a': ti.f32, 'b': vec3}, shape=(5, 4))
+
+:::
+
+### As a temporary local variable
+
+A local Struct variable can be created with *either* a dictionary *or* keyword arguments.
+
+::: {.function}
+ti.Struct(members, **kwargs)
+
+parameter members
+
+: (Dict) The dictionary containing struct members.
+
+parameter **kwargs
+
+: The keyword arguments to specify struct members.
+
+Lists and nested dictionaries in the member dictionary or keyword arguments will be converted into local `ti.Matrix` and `ti.Struct`, respectively.
+
+For example, this creates a struct with a float member `a` and vector member `b`:
+
+    # Taichi-scope
+    x = ti.Struct({'a': 1.0, 'b': [1.0, 1.0, 1.0]})
+    # or
+    x = ti.Struct(a=1.0, b=ti.Vector([1.0, 1.0, 1.0]))
+
+:::
+
+## Accessing components
+
+### As global struct fields
+
+Global struct field members are accessed as object attributes. For example, this extracts the member `a` of struct `x[6, 3]`: :
+
+    a = x[6, 3].a
+
+    # or
+    s = x[6, 3]
+    a = s.a
+
+In contrast to vector and matrix fields, struct members in a global struct field can be accessed in both attribute-first and index-first manners:
+
+    a = x[6, 3].a
+    # is equivalent to
+    a = x.a[6, 3]
+
+This allows for all field elements for a given struct field member to be extracted as a whole as a scalar, vector, or matrix field by accessing the member attributes of *the global struct field`:
+
+    # Python-scope
+    vec3 = ti.types.vector(3, float)
+    x = ti.Struct.field({'a': ti.f32, 'b': vec3}, shape=(5, 4))
+
+    x.a.fill(1.0) # x.a is equivalent to ti.field(ti.f32, shape=(5, 4))
+    x.b.fill(2.0) # x.b is equivalent to ti.Vector.field(3, ti.f32, shape=(5, 4))
+    a = x.a.to_numpy()
+    b = x.b.to_numpy()
+
+### As a temporary local variable
+
+Members of a local struct can be accessed using both attributes (object-like) or keys (dict-like):
+
+    x = ti.Struct(a=1.0, b=ti.Vector([1.0, 1.0, 1.0]))
+    a = x['a'] # a = 1.0
+    x.b = ti.Vector([1.0, 2.0, 3.0])
+
+## Element-wise operations (WIP)
+
+TODO: add element wise operations docs

docs/lang/articles/basic/external.md

@@ -83,7 +83,7 @@ field.from_numpy(array)  # the input array must be of shape (233, 666, 3)
 ```
 
 - For matrix fields, if the matrix is `n*m`, then **the shape of NumPy
-  array should be** `(*field_shape, matrix_n, matrix_m)`:
+array should be** `(*field_shape, matrix_n, matrix_m)`:
 
 ```python
 field = ti.Matrix.field(3, 4, ti.i32, shape=(233, 666))
@@ -97,6 +97,20 @@ array.shape  # (233, 666, 3, 4)
 field.from_numpy(array)  # the input array must be of shape (233, 666, 3, 4)
 ```
 
+- For struct fields, the external array will be exported as **a dictionary of arrays** with the keys being struct member names and values being struct member arrays. Nested structs will be exported as nested dictionaries:
+
+```python
+field = ti.Struct.field({'a': ti.i32, 'b': ti.types.vector(float, 3)} shape=(233, 666))
+field.shape # (233, 666)
+
+array_dict = field.to_numpy()
+array_dict.keys() # dict_keys(['a', 'b'])
+array_dict['a'].shape # (233, 666)
+array_dict['b'].shape # (233, 666, 3)
+
+field.from_numpy(array_dict) # the input array must have the same keys as the field=
+```
+
 ## Using external arrays as Taichi kernel arguments
 
 Use the type hint `ti.ext_arr()` for passing external arrays as kernel

docs/lang/articles/basic/field.md

@@ -20,7 +20,7 @@ A simple example might help you understand scalar fields. Assume you have a rect
 heat field on the wok:
 
 ``` python
-heat_field = taichi.field(dtype=ti.f32, shape=(width_wok, height_wok))
+heat_field = ti.field(dtype=ti.f32, shape=(width_wok, height_wok))
 ```
 
 - Every global variable is an N-dimensional field.
@@ -40,7 +40,7 @@ heat_field = taichi.field(dtype=ti.f32, shape=(width_wok, height_wok))
 ## Vector fields
 We are all live in a gravitational field which is a vector field. At each position of the 3D space, there is a gravity force vector. The gravitational field could be represent with:
 ```python
-gravitational_field = taichi.Vector.field(n = 3,dtype=ti.f32,shape=(x,y,z))
+gravitational_field = ti.Vector.field(n = 3,dtype=ti.f32,shape=(x,y,z))
 ```
 `x,y,z` are the sizes of each dimension of the 3D space respectively.  `n` is the number of elements of the gravity force vector.
 
@@ -51,7 +51,7 @@ gravitational_field = taichi.Vector.field(n = 3,dtype=ti.f32,shape=(x,y,z))
 Field elements can also be matrices. In continuum mechanics, each
 infinitesimal point in a material exists a strain and a stress tensor. The strain and stress tensor is a 3 by 3 matrix in the 3D space. To represent this tensor field we could use:
 ```python
-strain_tensor_field = taichi.Matrix.field(n = 3,m = 3, dtype=ti.f32, shape=(x,y,z))
+strain_tensor_field = ti.Matrix.field(n = 3,m = 3, dtype=ti.f32, shape=(x,y,z))
 ```
 
 `x,y,z` are the sizes of each dimension of the 3D material respectively. `n, m` are the dimensions of the strain tensor.
@@ -87,3 +87,33 @@ declare a field of size `64`. E.g., instead of declaring
 `ti.Matrix.field(64, 32, dtype=ti.f32, shape=(3, 2))`, declare
 `ti.Matrix.field(3, 2, dtype=ti.f32, shape=(64, 32))`. Try to put large
 dimensions to fields instead of matrices.
+
+## Struct fields
+In addition to vectors and matrices, field elements can be user-defined structs. A struct variable may contain scalars, vectors/matrices, or other structs as its members. A struct field is created by providing a dictionary of name and data type of each member. For example, a 1D field of particles with position, velocity, acceleration, and mass for each particle can be represented as:
+```python
+particle_field = ti.Struct.field({
+    "pos": ti.types.vector(3, ti.f32),
+    "vel": ti.types.vector(3, ti.f32),
+    "acc": ti.types.vector(3, ti.f32),
+    "mass": ti.f32,
+  }, shape=(n,))
+```
+[Compound types](type.md#compound-types) (`ti.types.vector`, `ti.types.matrix`, and `ti.types.struct`) need to be used to create vectors, matrices, or structs as field members. Apart from using `ti.Struct.field`, the above particle field can be alternatively created using field creation from compound types as:
+```python
+vec3f = ti.types.vector(3, ti.f32)
+particle = ti.types.struct(
+  pos=vec3f, vel=vec3f, acc=vec3f, mass=ti.f32,
+)
+particle_field = particle.field(shape=(n,))
+```
+Members of a struct field can be accessed either locally (i.e., member of a struct field element) or globally (i.e., member field of a struct field):
+```python
+# set the position of the first particle to origin
+particle_field[0] # local ti.Struct
+particle_field[0].pos = ti.Vector([0.0, 0.0, 0.0])
+
+# make the mass of all particles be 1
+particle_field.mass # global ti.Vector.field
+particle_field.mass.fill(1.0)
+```
+- See [Structs](../../api/struct.md) for more on structs.

docs/lang/articles/basic/type.md

@@ -4,6 +4,10 @@ sidebar_position: 2
 
 # Type system
 
+Data types in Taichi consist of Primitive Types and Compound Types. Primitive Types are the numerical data types used by backends, while Compound Types are user-defined types of data records composed of multiple members.
+
+## Primitive types
+
 Taichi supports common numerical data types. Each type is denoted as a
 character indicating its _category_ and a number of _precision bits_,
 e.g., `i32` and `f64`.
@@ -29,9 +33,9 @@ For example, the two most commonly used types:
 - `i32` represents a 32-bit signed integer.
 - `f32` represents a 32-bit floating pointer number.
 
-## Supported types
+## Supported primitive types
 
-Currently, supported basic types in Taichi are
+Currently, supported primitive types in Taichi are
 
 - int8 `ti.i8`
 - int16 `ti.i16`
@@ -184,3 +188,38 @@ the same width as the the old type. For example, bit-casting `i32` to
 :::note
 For people from C++, `ti.bit_cast` is equivalent to `reinterpret_cast`.
 :::
+
+## Compound types
+
+User-defined compound types are created using the `ti.types` module. Supported compound types include vectors, matrices, and structs:
+
+```python
+vec2i = ti.types.vector(2, ti.i32)
+vec3f = ti.types.vector(3, float)
+mat2f = ti.types.matrix(2, 2, float)
+ray = ti.types.struct(ro=vec3f, rd=vec3f, l=ti.f32)
+```
+
+### Creating fields
+
+Fields of a given compound type can be created with the `.field()` method of a Compound Type:
+
+```python
+# ti.Vector.field(2, dtype=ti.i32, shape=(233, 666))
+x = vec2i.field(shape=(233, 666))
+
+# ti.Matrix.field(2, 2, dtype=ti.i32, shape=(233, 666))
+x = mat2f.field(shape=(233, 666))
+
+# ti.Struct.field({'ro': vec3f, 'rd': vec3f, 'l': ti.f32}, shape=(233, 666))
+x = ray.field(shape=(233, 666))
+```
+
+### Creating local variables
+Compound types can be directly called to create matrix or struct instances. Vectors and matrices can be created using GLSL-like broadcast syntax since the shape of the vector or matrix is already known:
+```python
+ray = ray3f(0.0) # ti.Struct(ro=[0.0, 0.0, 0.0], rd=[0.0, 0.0, 0.0], l=0.0)
+ro = vec3f(0.0) # ti.Vector([0.0, 0.0, 0.0])
+rd = vec3f(vec2i(0), 1) # ti.Vector([0.0, 0.0, 1.0]), will perform implicit cast
+ray2 = ray3f(ro=ro, rd=rd, l=1.0)
+```

python/taichi/core/primitive_types.py

@@ -57,4 +57,6 @@
     'u32',
     'uint64',
     'u64',
+    'real_types',
+    'integer_types',
 ]

python/taichi/lang/__init__.py

@@ -4,7 +4,7 @@
 
 from taichi.core.util import locale_encode
 from taichi.core.util import ti_core as _ti_core
-from taichi.lang import impl
+from taichi.lang import impl, types
 from taichi.lang.enums import Layout
 from taichi.lang.exception import InvalidOperationError
 from taichi.lang.impl import *
@@ -16,6 +16,7 @@
 from taichi.lang.ops import *
 from taichi.lang.quant_impl import quant
 from taichi.lang.runtime_ops import async_flush, sync
+from taichi.lang.struct import Struct
 from taichi.lang.transformer import TaichiSyntaxError
 from taichi.lang.type_factory_impl import type_factory
 from taichi.lang.util import (has_pytorch, is_taichi_class, python_scope,

python/taichi/lang/field.py

@@ -206,6 +206,9 @@ def initialize_host_accessors(self):
             SNodeHostAccessor(e.ptr.snode()) for e in self.vars
         ]
 
+    def host_access(self, key):
+        return [SNodeHostAccess(e, key) for e in self.host_accessors]
+
 
 class ScalarField(Field):
     """Taichi scalar field with SNode implementation.

python/taichi/lang/impl.py

@@ -1,6 +1,6 @@
 import numbers
-import types
 import warnings
+from types import FunctionType, MethodType
 
 import numpy as np
 from taichi.core.util import ti_core as _ti_core
@@ -11,6 +11,7 @@
 from taichi.lang.matrix import MatrixField
 from taichi.lang.ndarray import ScalarNdarray
 from taichi.lang.snode import SNode
+from taichi.lang.struct import StructField
 from taichi.lang.tape import TapeImpl
 from taichi.lang.util import (cook_dtype, has_pytorch, is_taichi_class,
                               python_scope, taichi_scope, to_pytorch_type)
@@ -156,6 +157,10 @@ def subscript(value, *indices):
                 Expr(_ti_core.subscript(e.ptr, indices_expr_group))
                 for e in value.get_field_members()
             ])
+        elif isinstance(value, StructField):
+            return ti.Struct(
+                {k: subscript(v, *indices)
+                 for k, v in value.items})
         else:
             return Expr(_ti_core.subscript(var, indices_expr_group))
     elif isinstance(value, AnyArray):
@@ -819,7 +824,7 @@ def static(x, *xs):
         return x
     elif isinstance(x, Field):
         return x
-    elif isinstance(x, (types.FunctionType, types.MethodType)):
+    elif isinstance(x, (FunctionType, MethodType)):
         return x
     else:
         raise ValueError(