taichi.lang.snode

Module Contents

Classes

SNode

A Python-side SNode wrapper.

Functions

rescale_index(a, b, I)	Rescales the index 'I' of field (or SNode) 'a' to match the shape of SNode 'b'
append(l, indices, val)
is_active(l, indices)
activate(l, indices)
deactivate(l, indices)
length(l, indices)
get_addr(f, indices)	Query the memory address (on CUDA/x64) of field f at index indices.

class taichi.lang.snode.SNode(ptr)

A Python-side SNode wrapper.

For more information on Taichi’s SNode system, please check out these references:

• https://docs.taichi.graphics/lang/articles/advanced/sparse

• https://yuanming.taichi.graphics/publication/2019-taichi/taichi-lang.pdf

Arg:

ptr (pointer): The C++ side SNode pointer.

dense(self, axes, dimensions)

Adds a dense SNode as a child component of self.

Parameters

• axes (List[Axis]) – Axes to activate.

• dimensions (Union[List[int], int]) – Shape of each axis.

Returns

The added SNode instance.

pointer(self, axes, dimensions)

Adds a pointer SNode as a child component of self.

Parameters

• axes (List[Axis]) – Axes to activate.

• dimensions (Union[List[int], int]) – Shape of each axis.

Returns

The added SNode instance.

static hash(axes, dimensions)

Not supported.

dynamic(self, axis, dimension, chunk_size=None)

Adds a dynamic SNode as a child component of self.

Parameters

• axis (List[Axis]) – Axis to activate, must be 1.

• dimension (int) – Shape of the axis.

• chunk_size (int) – Chunk size.

Returns

The added SNode instance.

bitmasked(self, axes, dimensions)

Adds a bitmasked SNode as a child component of self.

Parameters

• axes (List[Axis]) – Axes to activate.

• dimensions (Union[List[int], int]) – Shape of each axis.

Returns

The added SNode instance.

bit_struct(self, num_bits: int)

Adds a bit_struct SNode as a child component of self.

Parameters

num_bits – Number of bits to use.

Returns

The added SNode instance.

bit_array(self, axes, dimensions, num_bits)

Adds a bit_array SNode as a child component of self.

Parameters

• axes (List[Axis]) – Axes to activate.

• dimensions (Union[List[int], int]) – Shape of each axis.

• num_bits (int) – Number of bits to use.

Returns

The added SNode instance.

place(self, *args, offset=None, shared_exponent=False)

Places a list of Taichi fields under the self container.

Parameters

• *args (List[ti.field]) – A list of Taichi fields to place.

• offset (Union[Number, tuple[Number]]) – Offset of the field domain.

• shared_exponent (bool) – Only useful for quant types.

Returns

The self container.

lazy_grad(self)

Automatically place the adjoint fields following the layout of their primal fields.

Users don’t need to specify needs_grad when they define scalar/vector/matrix fields (primal fields) using autodiff. When all the primal fields are defined, using taichi.root.lazy_grad() could automatically generate their corresponding adjoint fields (gradient field).

To know more details about primal, adjoint fields and lazy_grad(), please see Page 4 and Page 13-14 of DiffTaichi Paper: https://arxiv.org/pdf/1910.00935.pdf

parent(self, n=1)

Gets an ancestor of self in the SNode tree.

Parameters

n (int) – the number of levels going up from self.

Returns

The n-th parent of self.

Return type

Union[None, _Root, SNode]

path_from_root(self)

Gets the path from root to self in the SNode tree.

Returns

The list of SNodes on the path from root to self.

Return type

List[Union[_Root, SNode]]

property dtype(self)

Gets the data type of self.

Returns

The data type of self.

Return type

DataType

property id(self)

Gets the id of self.

Returns

The id of self.

Return type

int

property shape(self)

Gets the number of elements from root in each axis of self.

Returns

The number of elements from root in each axis of self.

Return type

Tuple[int]

loop_range(self)

Gets the taichi_core.Expr wrapping the taichi_core.GlobalVariableExpression corresponding to self to serve as loop range.

Returns

See above.

Return type

taichi_core.Expr

property name(self)

Gets the name of self.

Returns

The name of self.

Return type

str

property snode(self)

Gets self.

Returns

self.

Return type

SNode

property needs_grad(self)

Checks whether self has a corresponding gradient SNode.

Returns

Whether self has a corresponding gradient SNode.

Return type

bool

get_children(self)

Gets all children components of self.

Returns

All children components of self.

Return type

List[SNode]

property num_dynamically_allocated(self)

property cell_size_bytes(self)

property offset_bytes_in_parent_cell(self)

deactivate_all(self)

Recursively deactivate all children components of self.

physical_index_position(self)

Gets mappings from virtual axes to physical axes.

Returns

Mappings from virtual axes to physical axes.

Return type

Dict[int, int]

taichi.lang.snode.rescale_index(a, b, I)

Rescales the index ‘I’ of field (or SNode) ‘a’ to match the shape of SNode ‘b’

Parameters

• a (ti.field(), ti.Vector.field, ti.Matrix.field()) – input taichi field or snode

• b (ti.field(), ti.Vector.field, ti.Matrix.field()) – output taichi field or snode

• I (ti.Vector()) – grouped loop index

Returns

Ib – rescaled grouped loop index

Return type

ti.Vector()

taichi.lang.snode.append(l, indices, val)

taichi.lang.snode.is_active(l, indices)

taichi.lang.snode.activate(l, indices)

taichi.lang.snode.deactivate(l, indices)

taichi.lang.snode.length(l, indices)

taichi.lang.snode.get_addr(f, indices)

Query the memory address (on CUDA/x64) of field f at index indices.

Currently, this function can only be called inside a taichi kernel.

Parameters

• f (Union[ti.field, ti.Vector.field, ti.Matrix.field]) – Input taichi field for memory address query.

• indices (Union[int, ti.Vector()]) – The specified field indices of the query.

Returns

The memory address of f[indices].

Return type

ti.u64