def build_range_for(ctx, node):
with ctx.variable_scope_guard():
loop_name = node.target.id
ctx.check_loop_var(loop_name)
loop_var = expr.Expr(_ti_core.make_id_expr(''))
ctx.create_variable(loop_name, loop_var)
if len(node.iter.args) not in [1, 2]:
raise TaichiSyntaxError(
f"Range should have 1 or 2 arguments, found {len(node.iter.args)}"
)
if len(node.iter.args) == 2:
begin = ti_ops.cast(
expr.Expr(build_stmt(ctx, node.iter.args[0])),
primitive_types.i32)
end = ti_ops.cast(
expr.Expr(build_stmt(ctx, node.iter.args[1])),
primitive_types.i32)
else:
begin = ti_ops.cast(expr.Expr(0), primitive_types.i32)
end = ti_ops.cast(
expr.Expr(build_stmt(ctx, node.iter.args[0])),
primitive_types.i32)
_ti_core.begin_frontend_range_for(loop_var.ptr, begin.ptr, end.ptr)
build_stmts(ctx, node.body)
_ti_core.end_frontend_range_for()
return None
def bit_cast(obj, dtype):
"""Copy and cast a scalar to a specified data type with its underlying
bits preserved. Must be called in taichi scope.
This function is equivalent to `reinterpret_cast` in C++.
Args:
obj (:mod:`~taichi.types.primitive_types`): Input scalar.
dtype (:mod:`~taichi.types.primitive_types`): Target data type, must have \
the same precision bits as the input (hence `f32` -> `f64` is not allowed).
Returns:
A copy of `obj`, casted to the specified data type `dtype`.
Example::
>>> @ti.kernel
>>> def test():
>>> x = 3.14
>>> y = ti.bit_cast(x, ti.i32)
>>> print(y) # 1078523331
>>>
>>> z = ti.bit_cast(y, ti.f32)
>>> print(z) # 3.14
"""
dtype = cook_dtype(dtype)
if is_taichi_class(obj):
raise ValueError('Cannot apply bit_cast on Taichi classes')
else:
return expr.Expr(_ti_core.bits_cast(expr.Expr(obj).ptr, dtype))
def build_ndrange_for(ctx, node):
with ctx.variable_scope_guard():
ndrange_var = impl.expr_init(build_stmt(ctx, node.iter))
ndrange_begin = ti_ops.cast(expr.Expr(0), primitive_types.i32)
ndrange_end = ti_ops.cast(
expr.Expr(impl.subscript(ndrange_var.acc_dimensions, 0)),
primitive_types.i32)
ndrange_loop_var = expr.Expr(_ti_core.make_id_expr(''))
_ti_core.begin_frontend_range_for(ndrange_loop_var.ptr,
ndrange_begin.ptr,
ndrange_end.ptr)
I = impl.expr_init(ndrange_loop_var)
targets = ASTTransformer.get_for_loop_targets(node)
for i, target in enumerate(targets):
if i + 1 < len(targets):
target_tmp = impl.expr_init(
I // ndrange_var.acc_dimensions[i + 1])
else:
target_tmp = impl.expr_init(I)
ctx.create_variable(
target,
impl.expr_init(target_tmp + impl.subscript(
impl.subscript(ndrange_var.bounds, i), 0)))
if i + 1 < len(targets):
I.assign(I -
target_tmp * ndrange_var.acc_dimensions[i + 1])
build_stmts(ctx, node.body)
_ti_core.end_frontend_range_for()
return None
def build_nested_mesh_for(ctx, node):
targets = ASTTransformer.get_for_loop_targets(node)
if len(targets) != 1:
raise TaichiSyntaxError(
"Nested-mesh for should have 1 loop target, found {len(targets)}"
)
target = targets[0]
with ctx.variable_scope_guard():
ctx.mesh = node.iter.ptr.mesh
assert isinstance(ctx.mesh, impl.MeshInstance)
loop_name = node.target.id + '_index__'
loop_var = expr.Expr(_ti_core.make_id_expr(''))
ctx.create_variable(loop_name, loop_var)
begin = expr.Expr(0)
end = node.iter.ptr.size
_ti_core.begin_frontend_range_for(loop_var.ptr, begin.ptr, end.ptr)
entry_expr = _ti_core.get_relation_access(
ctx.mesh.mesh_ptr, node.iter.ptr.from_index.ptr,
node.iter.ptr.to_element_type, loop_var.ptr)
entry_expr.type_check()
mesh_idx = mesh.MeshElementFieldProxy(
ctx.mesh, node.iter.ptr.to_element_type, entry_expr)
ctx.create_variable(target, mesh_idx)
build_stmts(ctx, node.body)
_ti_core.end_frontend_range_for()
return None
def cast(obj, dtype):
"""Copy and cast a scalar or a matrix to a specified data type.
Must be called in Taichi scope.
Args:
obj (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
Input scalar or matrix.
dtype (:mod:`~taichi.types.primitive_types`): A primitive type defined in :mod:`~taichi.types.primitive_types`.
Returns:
A copy of `obj`, casted to the specified data type `dtype`.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Matrix([0, 1, 2], ti.i32)
>>> y = ti.cast(x, ti.f32)
>>> print(y)
>>>
>>> test()
[0.0, 1.0, 2.0]
"""
dtype = cook_dtype(dtype)
if is_taichi_class(obj):
# TODO: unify with element_wise_unary
return obj.cast(dtype)
return expr.Expr(_ti_core.value_cast(expr.Expr(obj).ptr, dtype))
def _handle_string_mod_args(ctx, node):
msg = build_stmt(ctx, node.left)
args = build_stmt(ctx, node.right)
if not isinstance(args, collections.abc.Sequence):
args = (args, )
args = [expr.Expr(x).ptr for x in args]
return msg, args
def shfl_up_f32(mask, val, offset):
return expr.Expr(
_ti_core.insert_internal_func_call(
"cuda_shfl_up_sync_f32",
# lane offset is 0 for warp size 32
expr.make_expr_group(mask, val, offset, 0),
False))
def shfl_sync_i32(mask, val, offset):
return expr.Expr(
_ti_core.insert_internal_func_call(
# lane offset is 31 for warp size 32
"cuda_shfl_sync_i32",
expr.make_expr_group(mask, val, offset, 31),
False))
def atomic_max(x, y):
"""Atomically compute the maximum of `x` and `y`, element-wise.
Store the result in `x`, and return the old value of `x`.
`x` must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
The input.
Returns:
The old value of `x`.
Example::
>>> @ti.kernel
>>> def test():
>>> x = 1
>>> y = 2
>>> z = ti.atomic_max(x, y)
>>> print(x) # 2 the new value of x
>>> print(z) # 1, the old value of x
>>>
>>> ti.atomic_max(1, x) # will raise TaichiSyntaxError
"""
return impl.expr_init(
expr.Expr(_ti_core.expr_atomic_max(x.ptr, y.ptr), tb=stack_info()))
def shfl_down_i32(mask, val, offset):
# Here we use 31 as the last argument since 32 (warp size) does not work
# for some reason. Using 31 leads to the desired behavior.
return expr.Expr(
_ti_core.insert_internal_func_call(
"cuda_shfl_down_sync_i32",
expr.make_expr_group(mask, val, offset, 31), False))
def atomic_xor(x, y):
"""Atomically compute the bit-wise XOR of `x` and `y`, element-wise.
Store the result in `x`, and return the old value of `x`.
`x` must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
The input. When both are matrices they must have the same shape.
Returns:
The old value of `x`.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([-1, 0, 1])
>>> y = ti.Vector([1, 2, 3])
>>> z = ti.atomic_xor(x, y)
>>> print(x) # [-2, 2, 2] the new value of x
>>> print(z) # [-1, 0, 1], the old value of x
>>>
>>> ti.atomic_xor(1, x) # will raise TaichiSyntaxError
"""
return impl.expr_init(
expr.Expr(_ti_core.expr_atomic_bit_xor(x.ptr, y.ptr), tb=stack_info()))
def build_Return(ctx, node):
if not impl.get_runtime().experimental_real_function:
if ctx.is_in_non_static():
raise TaichiSyntaxError(
"Return inside non-static if/for is not supported")
build_stmt(ctx, node.value)
if ctx.is_kernel or impl.get_runtime().experimental_real_function:
# TODO: check if it's at the end of a kernel, throw TaichiSyntaxError if not
if node.value is not None:
if ctx.func.return_type is None:
raise TaichiSyntaxError(
f'A {"kernel" if ctx.is_kernel else "function"} '
'with a return value must be annotated '
'with a return type, e.g. def func() -> ti.f32')
_ti_core.create_kernel_return(
ti_ops.cast(expr.Expr(node.value.ptr),
ctx.func.return_type).ptr)
# For args[0], it is an ast.Attribute, because it loads the
# attribute, |ptr|, of the expression |ret_expr|. Therefore we
# only need to replace the object part, i.e. args[0].value
else:
ctx.return_data = node.value.ptr
if not impl.get_runtime().experimental_real_function:
ctx.returned = True
return None
def atomic_sub(x, y):
"""Atomically subtract `x` by `y`, store the result in `x`,
and return the old value of `x`.
`x` must be a writable target, constant expressions or scalars
are not allowed.
Args:
x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \
The input.
Returns:
The old value of `x`.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.Vector([0, 0, 0])
>>> y = ti.Vector([1, 2, 3])
>>> z = ti.atomic_sub(x, y)
>>> print(x) # [-1, -2, -3] the new value of x
>>> print(z) # [0, 0, 0], the old value of x
>>>
>>> ti.atomic_sub(1, x) # will raise TaichiSyntaxError
"""
return impl.expr_init(
expr.Expr(_ti_core.expr_atomic_sub(x.ptr, y.ptr), tb=stack_info()))
def build_call_if_is_type(ctx, node, args, keywords):
func = node.func.ptr
if id(func) in primitive_types.type_ids:
if len(args) != 1 or keywords or isinstance(args[0], expr.Expr):
raise TaichiSyntaxError(
"Type annotation can only be given to a single literal.")
node.ptr = expr.Expr(args[0], dtype=func)
return True
return False
def append(node, indices, val):
"""Append a value `val` to a SNode `node` at index `indices`.
Args:
node (:class:`~taichi.SNode`): Input SNode.
indices (Union[int, :class:`~taichi.Vector`]): the indices to visit.
val (:mod:`~taichi.types`): the data to be appended.
"""
a = impl.expr_init(
_ti_core.insert_append(node._snode.ptr, expr.make_expr_group(indices),
expr.Expr(val).ptr))
return a
def length(node, indices):
"""Return the length of the dynamic SNode `node` at index `indices`.
Args:
node (:class:`~taichi.SNode`): a dynamic SNode.
indices (Union[int, :class:`~taichi.Vector`]): the indices to query.
Returns:
int: the length of cell `node[indices]`.
"""
return expr.Expr(
_ti_core.insert_len(node._snode.ptr, expr.make_expr_group(indices)))
def random(dtype=float):
"""The random function.
Args:
dtype (DataType): Type of the random variable.
Returns:
A random variable whose type is `dtype`.
"""
dtype = cook_dtype(dtype)
x = expr.Expr(_ti_core.make_rand_expr(dtype))
return impl.expr_init(x)
def 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.
Args:
f (Union[:class:`~taichi.Field`, :class:`~taichi.MatrixField`]): Input taichi field for memory address query.
indices (Union[int, :class:`~taichi.Vector`]): The specified field indices of the query.
Returns:
ti.u64: The memory address of `f[indices]`.
"""
return expr.Expr(
_ti_core.expr_get_addr(f._snode.ptr, expr.make_expr_group(indices)))
def is_active(node, indices):
"""Explicitly query whether a cell in a SNode `node` at location
`indices` is active or not.
Args:
node (:class:`~taichi.SNode`): Must be a pointer, hash or bitmasked node.
indices (Union[int, list, :class:`~taichi.Vector`]): the indices to visit.
Returns:
bool: the cell `node[indices]` is active or not.
"""
return expr.Expr(
_ti_core.insert_is_active(node._snode.ptr,
expr.make_expr_group(indices)))
def build_grouped_ndrange_for(ctx, node):
with ctx.variable_scope_guard():
ndrange_var = impl.expr_init(build_stmt(ctx, node.iter.args[0]))
ndrange_begin = ti_ops.cast(expr.Expr(0), primitive_types.i32)
ndrange_end = ti_ops.cast(
expr.Expr(impl.subscript(ndrange_var.acc_dimensions, 0)),
primitive_types.i32)
ndrange_loop_var = expr.Expr(_ti_core.make_id_expr(''))
_ti_core.begin_frontend_range_for(ndrange_loop_var.ptr,
ndrange_begin.ptr,
ndrange_end.ptr)
targets = ASTTransformer.get_for_loop_targets(node)
if len(targets) != 1:
raise TaichiSyntaxError(
f"Group for should have 1 loop target, found {len(targets)}"
)
target = targets[0]
target_var = impl.expr_init(
matrix.Vector([0] * len(ndrange_var.dimensions),
dt=primitive_types.i32))
ctx.create_variable(target, target_var)
I = impl.expr_init(ndrange_loop_var)
for i in range(len(ndrange_var.dimensions)):
if i + 1 < len(ndrange_var.dimensions):
target_tmp = I // ndrange_var.acc_dimensions[i + 1]
else:
target_tmp = I
impl.subscript(target_var,
i).assign(target_tmp + ndrange_var.bounds[i][0])
if i + 1 < len(ndrange_var.dimensions):
I.assign(I -
target_tmp * ndrange_var.acc_dimensions[i + 1])
build_stmts(ctx, node.body)
_ti_core.end_frontend_range_for()
return None
def build_While(ctx, node):
if node.orelse:
raise TaichiSyntaxError(
"'else' clause for 'while' not supported in Taichi kernels")
with ctx.loop_scope_guard():
_ti_core.begin_frontend_while(expr.Expr(1).ptr)
while_cond = build_stmt(ctx, node.test)
impl.begin_frontend_if(while_cond)
_ti_core.begin_frontend_if_true()
_ti_core.pop_scope()
_ti_core.begin_frontend_if_false()
_ti_core.insert_break_stmt()
_ti_core.pop_scope()
build_stmts(ctx, node.body)
_ti_core.pop_scope()
return None
def build_mesh_for(ctx, node):
targets = ASTTransformer.get_for_loop_targets(node)
if len(targets) != 1:
raise TaichiSyntaxError(
"Mesh for should have 1 loop target, found {len(targets)}")
target = targets[0]
with ctx.variable_scope_guard():
var = expr.Expr(_ti_core.make_id_expr(""))
ctx.mesh = node.iter.ptr.mesh
assert isinstance(ctx.mesh, impl.MeshInstance)
mesh_idx = mesh.MeshElementFieldProxy(ctx.mesh,
node.iter.ptr._type, var.ptr)
ctx.create_variable(target, mesh_idx)
_ti_core.begin_frontend_mesh_for(mesh_idx.ptr, ctx.mesh.mesh_ptr,
node.iter.ptr._type)
build_stmts(ctx, node.body)
ctx.mesh = None
_ti_core.end_frontend_range_for()
return None
def build_Return(ctx, node):
if not ctx.is_real_function:
if ctx.is_in_non_static_control_flow():
raise TaichiSyntaxError(
"Return inside non-static if/for is not supported")
if node.value is not None:
build_stmt(ctx, node.value)
if node.value is None or node.value.ptr is None:
if not ctx.is_real_function:
ctx.returned = ReturnStatus.ReturnedVoid
return None
if ctx.is_kernel or ctx.is_real_function:
# TODO: check if it's at the end of a kernel, throw TaichiSyntaxError if not
if ctx.func.return_type is None:
raise TaichiSyntaxError(
f'A {"kernel" if ctx.is_kernel else "function"} '
'with a return value must be annotated '
'with a return type, e.g. def func() -> ti.f32')
if id(ctx.func.return_type) in primitive_types.type_ids:
ctx.ast_builder.create_kernel_exprgroup_return(
expr.make_expr_group(
ti_ops.cast(expr.Expr(node.value.ptr),
ctx.func.return_type).ptr))
elif isinstance(ctx.func.return_type, MatrixType):
ctx.ast_builder.create_kernel_exprgroup_return(
expr.make_expr_group([
ti_ops.cast(exp, ctx.func.return_type.dtype) for exp in
itertools.chain.from_iterable(node.value.ptr.to_list())
]))
else:
raise TaichiSyntaxError(
"The return type is not supported now!")
# For args[0], it is an ast.Attribute, because it loads the
# attribute, |ptr|, of the expression |ret_expr|. Therefore we
# only need to replace the object part, i.e. args[0].value
else:
ctx.return_data = node.value.ptr
if not ctx.is_real_function:
ctx.returned = ReturnStatus.ReturnedValue
return None
def build_struct_for(ctx, node, is_grouped):
# for i, j in x
# for I in ti.grouped(x)
targets = ASTTransformer.get_for_loop_targets(node)
for target in targets:
ctx.check_loop_var(target)
with ctx.variable_scope_guard():
if is_grouped:
if len(targets) != 1:
raise TaichiSyntaxError(
f"Group for should have 1 loop target, found {len(targets)}"
)
target = targets[0]
loop_var = build_stmt(ctx, node.iter)
loop_indices = expr.make_var_list(size=len(loop_var.shape),
ast_builder=ctx.ast_builder)
expr_group = expr.make_expr_group(loop_indices)
impl.begin_frontend_struct_for(ctx.ast_builder, expr_group,
loop_var)
ctx.create_variable(
target, matrix.Vector(loop_indices,
dt=primitive_types.i32))
build_stmts(ctx, node.body)
ctx.ast_builder.end_frontend_struct_for()
else:
_vars = []
for name in targets:
var = expr.Expr(ctx.ast_builder.make_id_expr(""))
_vars.append(var)
ctx.create_variable(name, var)
loop_var = node.iter.ptr
expr_group = expr.make_expr_group(*_vars)
impl.begin_frontend_struct_for(ctx.ast_builder, expr_group,
loop_var)
build_stmts(ctx, node.body)
ctx.ast_builder.end_frontend_struct_for()
return None
def random(dtype=float):
"""Return a single random float/integer according to the specified data type.
Must be called in taichi scope.
If the required `dtype` is float type, this function returns a random number
sampled from the uniform distribution in the half-open interval [0, 1).
For integer types this function returns a random integer in the
half-open interval [0, 2^32) if a 32-bit integer is required,
or a random integer in the half-open interval [0, 2^64) if a
64-bit integer is required.
Args:
dtype (:mod:`~taichi.types.primitive_types`): Type of the required random value.
Returns:
A random value with type `dtype`.
Example::
>>> @ti.kernel
>>> def test():
>>> x = ti.random(float)
>>> print(x) # 0.090257
>>>
>>> y = ti.random(ti.f64)
>>> print(y) # 0.716101627301
>>>
>>> i = ti.random(ti.i32)
>>> print(i) # -963722261
>>>
>>> j = ti.random(ti.i64)
>>> print(j) # 73412986184350777
"""
dtype = cook_dtype(dtype)
x = expr.Expr(_ti_core.make_rand_expr(dtype))
return impl.expr_init(x)
def length(l, indices):
return expr.Expr(
_ti_core.insert_len(l.snode.ptr, expr.make_expr_group(indices)))
def is_active(l, indices):
return expr.Expr(
_ti_core.insert_is_active(l.snode.ptr, expr.make_expr_group(indices)))
def atomic_xor(a, b):
return impl.expr_init(
expr.Expr(_ti_core.expr_atomic_bit_xor(a.ptr, b.ptr), tb=stack_info()))
def append(l, indices, val):
a = impl.expr_init(
_ti_core.insert_append(l.snode.ptr, expr.make_expr_group(indices),
expr.Expr(val).ptr))
return a
def expr_python_mod(a, b):
# a % b = a - (a // b) * b
quotient = expr.Expr(_ti_core.expr_floordiv(a, b))
multiply = expr.Expr(_ti_core.expr_mul(b, quotient.ptr))
return _ti_core.expr_sub(a, multiply.ptr)
