def ndarray(dtype, shape, layout=Layout.NULL):
"""Defines a Taichi ndarray with scalar elements.
Args:
dtype (Union[DataType, MatrixType]): Data type of each element. This can be either a scalar type like ti.f32 or a compound type like ti.types.vector(3, ti.i32).
shape (Union[int, tuple[int]]): Shape of the ndarray.
layout (Layout, optional): Layout of ndarray, only applicable when element is non-scalar type. Default is Layout.AOS.
Example:
The code below shows how a Taichi ndarray with scalar elements can be declared and defined::
>>> x = ti.ndarray(ti.f32, shape=(16, 8)) # ndarray of shape (16, 8), each element is ti.f32 scalar.
>>> vec3 = ti.types.vector(3, ti.i32)
>>> y = ti.ndarray(vec3, shape=(10, 2)) # ndarray of shape (10, 2), each element is a vector of 3 ti.i32 scalars.
>>> matrix_ty = ti.types.matrix(3, 4, float)
>>> z = ti.ndarray(matrix_ty, shape=(4, 5), layout=ti.Layout.SOA) # ndarray of shape (4, 5), each element is a matrix of (3, 4) ti.float scalars.
"""
if isinstance(shape, numbers.Number):
shape = (shape, )
if dtype in all_types:
assert layout == Layout.NULL
return ScalarNdarray(dtype, shape)
if isinstance(dtype, MatrixType):
layout = Layout.AOS if layout == Layout.NULL else layout
return MatrixNdarray(dtype.n, dtype.m, dtype.dtype, shape, layout)
raise TaichiRuntimeError(
f'{dtype} is not supported as ndarray element type')
def produce_injected_args(kernel, symbolic_args=None):
injected_args = []
for i, arg in enumerate(kernel.arguments):
anno = arg.annotation
if isinstance(anno, template_types):
if not isinstance(anno, NdarrayType):
raise TaichiCompilationError(
f'Expected Ndaray type, got {anno}')
if symbolic_args is not None:
element_shape = tuple(symbolic_args[i].element_shape)
element_dim = len(element_shape)
dtype = symbolic_args[i].dtype()
else:
element_shape = anno.element_shape
element_dim = anno.element_dim
dtype = anno.dtype
if element_shape is None or anno.field_dim is None:
raise TaichiCompilationError(
'Please either specify both `element_shape` and `field_dim` '
'in the param annotation, or provide an example '
f'ndarray for param={arg.name}')
if element_dim is None or element_dim == 0:
injected_args.append(
ScalarNdarray(dtype, (2, ) * anno.field_dim))
elif element_dim == 1:
injected_args.append(
VectorNdarray(element_shape[0],
dtype=dtype,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
elif element_dim == 2:
injected_args.append(
MatrixNdarray(element_shape[0],
element_shape[1],
dtype=dtype,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
else:
raise RuntimeError('')
elif isinstance(anno, MatrixType):
if not isinstance(symbolic_args[i], list):
raise RuntimeError('Expected a symbolic arg with Matrix type.')
symbolic_mat_n = len(symbolic_args[i])
symbolic_mat_m = len(symbolic_args[i][0])
if symbolic_mat_m != anno.m or symbolic_mat_n != anno.n:
raise RuntimeError(
f'Matrix dimension mismatch, expected ({anno.n}, {anno.m}) '
f'but dispathed shape ({symbolic_mat_n}, {symbolic_mat_m}).'
)
injected_args.append(Matrix([0] * anno.n * anno.m, dt=anno.dtype))
else:
# For primitive types, we can just inject a dummy value.
injected_args.append(0)
return injected_args
def add_kernel(self, kernel_fn, example_any_arrays=None, name=None):
"""Add a taichi kernel to the AOT module.
Args:
kernel_fn (Function): the function decorated by taichi `kernel`.
example_any_arrays (Dict[int, ti.ndarray]): a dict where key is arg_id and key is example any_arr input.
name (str): Name to identify this kernel in the module. If not
provided, uses the built-in ``__name__`` attribute of `kernel_fn`.
"""
name = name or kernel_fn.__name__
kernel = kernel_fn._primal
assert isinstance(kernel, kernel_impl.Kernel)
injected_args = []
num_arr = len([
anno for anno in kernel.argument_annotations
if isinstance(anno, ArgAnyArray)
])
assert example_any_arrays is None or num_arr == len(
example_any_arrays
), f'Need {num_arr} example any_arr inputs but got {len(example_any_arrays)}'
i = 0
for anno in kernel.argument_annotations:
if isinstance(anno, ArgAnyArray):
if example_any_arrays:
injected_args.append(example_any_arrays[i])
else:
assert anno.element_shape is not None and anno.field_dim is not None, 'Please either specify element_shape & field_dim in the kernel arg annotation or provide a dict of example ndarrays.'
if anno.element_dim == 0:
injected_args.append(
ScalarNdarray(dtype=f32,
shape=(2, ) * anno.field_dim))
elif anno.element_dim == 1:
injected_args.append(
VectorNdarray(anno.element_shape[0],
dtype=f32,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
elif anno.element_dim == 2:
injected_args.append(
MatrixNdarray(anno.element_shape[0],
anno.element_shape[1],
dtype=f32,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
else:
raise RuntimeError('')
else:
# For primitive types, we can just inject a dummy value.
injected_args.append(0)
i = i + 1
kernel.ensure_compiled(*injected_args)
self._aot_builder.add(name, kernel.kernel_cpp)
# kernel AOT
self._kernels.append(kernel)
def ndarray(dtype, shape):
"""Defines a Taichi ndarray with scalar elements.
Args:
dtype (DataType): Data type of each value.
shape (Union[int, tuple[int]]): Shape of the ndarray.
Example:
The code below shows how a Taichi ndarray with scalar elements can be declared and defined::
>>> x = ti.ndarray(ti.f32, shape=(16, 8))
"""
if isinstance(shape, numbers.Number):
shape = (shape, )
return ScalarNdarray(dtype, shape)
def produce_injected_args(kernel, symbolic_args=None):
injected_args = []
for j, arg in enumerate(kernel.arguments):
anno = arg.annotation
if isinstance(anno, template_types):
if not isinstance(anno, NdarrayType):
raise TaichiCompilationError(
f'Expected Ndaray type, got {anno}')
if symbolic_args is not None:
anno.element_shape = tuple(symbolic_args[j].element_shape)
anno.element_dim = len(anno.element_shape)
if anno.element_shape is None or anno.field_dim is None:
raise TaichiCompilationError(
'Please either specify both `element_shape` and `field_dim` '
'in the param annotation, or provide an example '
f'ndarray for param={arg.name}')
if anno.element_dim == 0:
injected_args.append(ScalarNdarray(f32,
(2, ) * anno.field_dim))
elif anno.element_dim == 1:
injected_args.append(
VectorNdarray(anno.element_shape[0],
dtype=f32,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
elif anno.element_dim == 2:
injected_args.append(
MatrixNdarray(anno.element_shape[0],
anno.element_shape[1],
dtype=f32,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
else:
raise RuntimeError('')
else:
# For primitive types, we can just inject a dummy value.
injected_args.append(0)
return injected_args
def produce_injected_args(kernel, symbolic_args=None):
injected_args = []
for i, arg in enumerate(kernel.arguments):
anno = arg.annotation
if isinstance(anno, template_types):
if not isinstance(anno, NdarrayType):
raise TaichiCompilationError(
f'Expected Ndaray type, got {anno}')
if symbolic_args is not None:
element_shape = tuple(symbolic_args[i].element_shape)
element_dim = len(element_shape)
field_dim = symbolic_args[i].field_dim
dtype = symbolic_args[i].dtype()
else:
element_shape = anno.element_shape
element_dim = anno.element_dim
field_dim = anno.field_dim
dtype = anno.dtype
if element_shape is None or field_dim is None:
raise TaichiCompilationError(
'Please either specify both `element_shape` and `field_dim` '
'in the param annotation, or provide an example '
f'ndarray for param={arg.name}')
if anno.field_dim is not None and field_dim != anno.field_dim:
raise TaichiCompilationError(
f'{field_dim} from Arg {arg.name} doesn\'t match kernel\'s annotated field_dim={anno.field_dim}'
)
if anno.dtype is not None and not check_type_match(
dtype, anno.dtype):
raise TaichiCompilationError(
f' Arg {arg.name}\'s dtype {dtype.to_string()} doesn\'t match kernel\'s annotated dtype={anno.dtype.to_string()}'
)
if element_dim is None or element_dim == 0:
injected_args.append(ScalarNdarray(dtype, (2, ) * field_dim))
elif element_dim == 1:
injected_args.append(
VectorNdarray(element_shape[0],
dtype=dtype,
shape=(2, ) * field_dim,
layout=Layout.AOS))
elif element_dim == 2:
injected_args.append(
MatrixNdarray(element_shape[0],
element_shape[1],
dtype=dtype,
shape=(2, ) * field_dim,
layout=Layout.AOS))
else:
raise RuntimeError('')
elif isinstance(anno, (TextureType, RWTextureType)):
if symbolic_args is None:
raise RuntimeError(
'Texture type annotation doesn\'t have enough information for aot. Please either specify the channel_format, shape and num_channels in the graph arg declaration.'
)
texture_shape = tuple(symbolic_args[i].texture_shape)
channel_format = symbolic_args[i].channel_format()
num_channels = symbolic_args[i].num_channels
injected_args.append(
Texture(channel_format, num_channels, texture_shape))
elif isinstance(anno, MatrixType):
if not isinstance(symbolic_args[i], list):
raise RuntimeError('Expected a symbolic arg with Matrix type.')
symbolic_mat_n = len(symbolic_args[i])
symbolic_mat_m = len(symbolic_args[i][0])
if symbolic_mat_m != anno.m or symbolic_mat_n != anno.n:
raise RuntimeError(
f'Matrix dimension mismatch, expected ({anno.n}, {anno.m}) '
f'but dispatched shape ({symbolic_mat_n}, {symbolic_mat_m}).'
)
injected_args.append(Matrix([0] * anno.n * anno.m, dt=anno.dtype))
else:
if symbolic_args is not None:
dtype = symbolic_args[i].dtype()
else:
dtype = anno
if not check_type_match(dtype, anno):
raise TaichiCompilationError(
f' Arg {arg.name}\'s dtype {dtype.to_string()} doesn\'t match kernel\'s annotated dtype={anno.to_string()}'
)
# For primitive types, we can just inject a dummy value.
injected_args.append(0)
return injected_args
def add_kernel(self, kernel_fn, template_args=None, name=None):
"""Add a taichi kernel to the AOT module.
Args:
kernel_fn (Function): the function decorated by taichi `kernel`.
template_args (Dict[str, Any]): a dict where key is the template
parameter name, and value is the instantiating arg. Note that this
works for both :class:`~taichi.types.template` and for
`:class:`~taichi.types.ndarray`.
name (str): Name to identify this kernel in the module. If not
provided, uses the built-in ``__name__`` attribute of `kernel_fn`.
"""
kernel_name = name or kernel_fn.__name__
kernel = kernel_fn._primal
assert isinstance(kernel, kernel_impl.Kernel)
injected_args = []
template_types = (NdarrayType, template)
num_template_args = len([
anno for anno in kernel.argument_annotations
if isinstance(anno, template_types)
])
if template_args is not None and num_template_args != len(
template_args):
raise TaichiCompilationError(
f'Need {num_template_args} inputs to instantiate the template '
f'parameters, got {len(template_args)}')
i = 0
for arg_name, anno in zip(kernel.argument_names,
kernel.argument_annotations):
if isinstance(anno, template_types):
if template_args:
injected_args.append(template_args[arg_name])
else:
if not isinstance(anno, NdarrayType):
raise TaichiCompilationError(
f'Expected Ndaray type, got {anno}')
if anno.element_shape is None or anno.field_dim is None:
raise TaichiCompilationError(
'Please either specify both `element_shape` and `field_dim` '
'in the param annotation, or provide an example '
f'ndarray for param={name}')
if anno.element_dim == 0:
injected_args.append(
ScalarNdarray(f32, (2, ) * anno.field_dim))
elif anno.element_dim == 1:
injected_args.append(
VectorNdarray(anno.element_shape[0],
dtype=f32,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
elif anno.element_dim == 2:
injected_args.append(
MatrixNdarray(anno.element_shape[0],
anno.element_shape[1],
dtype=f32,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
else:
raise RuntimeError('')
i = i + 1
else:
# For primitive types, we can just inject a dummy value.
injected_args.append(0)
kernel.ensure_compiled(*injected_args)
self._aot_builder.add(kernel_name, kernel.kernel_cpp)
# kernel AOT
self._kernels.append(kernel)