def build_assign_annotated(ctx, target, value, is_static_assign,
annotation):
"""Build an annotated assginment like this: target: annotation = value.
Args:
ctx (ast_builder_utils.BuilderContext): The builder context.
target (ast.Name): A variable name. `target.id` holds the name as
a string.
annotation: A type we hope to assign to the target
value: A node representing the value.
is_static_assign: A boolean value indicating whether this is a static assignment
"""
is_local = isinstance(target, ast.Name)
anno = ti.expr_init(annotation)
if is_static_assign:
raise TaichiSyntaxError(
"Static assign cannot be used on annotated assignment")
if is_local and not ctx.is_var_declared(target.id):
var = ti.cast(value, anno)
var = ti.expr_init(var)
ctx.create_variable(target.id, var)
else:
var = target.ptr
if var.ptr.get_ret_type() != anno:
raise TaichiSyntaxError(
"Static assign cannot have type overloading")
var.assign(value)
return var
def build_Return(ctx, node):
if not impl.get_runtime().experimental_real_function:
if ctx.is_in_non_static_control_flow():
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')
ctx.ast_builder.create_kernel_exprgroup_return(
expr.make_expr_group(
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 __call__(self, *args):
_taichi_skip_traceback = 1
if not impl.inside_kernel():
if not self.pyfunc:
raise TaichiSyntaxError(
"Taichi functions cannot be called from Python-scope."
" Use @ti.pyfunc if you wish to call Taichi functions "
"from both Python-scope and Taichi-scope.")
return self.func(*args)
if impl.get_runtime().experimental_real_function:
if impl.get_runtime().current_kernel.is_grad:
raise TaichiSyntaxError(
"Real function in gradient kernels unsupported.")
instance_id, arg_features = self.mapper.lookup(args)
key = _ti_core.FunctionKey(self.func.__name__, self.func_id,
instance_id)
if self.compiled is None:
self.compiled = {}
if key.instance_id not in self.compiled:
self.do_compile(key=key, args=args)
return self.func_call_rvalue(key=key, args=args)
else:
if self.compiled is None:
self.do_compile(key=None, args=args)
ret = self.compiled(*args)
return ret
def build_assign_basic(ctx, target, value, is_static_assign):
"""Build basic assginment like this: target = value.
Args:
ctx (ast_builder_utils.BuilderContext): The builder context.
target (ast.Name): A variable name. `target.id` holds the name as
a string.
value: A node representing the value.
is_static_assign: A boolean value indicating whether this is a static assignment
"""
is_local = isinstance(target, ast.Name)
if is_static_assign:
if not is_local:
raise TaichiSyntaxError(
"Static assign cannot be used on elements in arrays")
ctx.create_variable(target.id, value)
var = value
elif is_local and not ctx.is_var_declared(target.id):
var = impl.expr_init(value)
ctx.create_variable(target.id, var)
else:
var = build_stmt(ctx, target)
try:
var._assign(value)
except AttributeError:
raise TaichiSyntaxError(
f"Variable '{unparse(target).strip()}' cannot be assigned. Maybe it is not a Taichi object?"
)
return var
def __call__(self, *args, **kwargs):
args = _process_args(self, args, kwargs)
if not impl.inside_kernel():
if not self.pyfunc:
raise TaichiSyntaxError(
"Taichi functions cannot be called from Python-scope.")
return self.func(*args)
if self.is_real_function:
if impl.get_runtime().current_kernel.is_grad:
raise TaichiSyntaxError(
"Real function in gradient kernels unsupported.")
instance_id, _ = self.mapper.lookup(args)
key = _ti_core.FunctionKey(self.func.__name__, self.func_id,
instance_id)
if self.compiled is None:
self.compiled = {}
if key.instance_id not in self.compiled:
self.do_compile(key=key, args=args)
return self.func_call_rvalue(key=key, args=args)
tree, ctx = _get_tree_and_ctx(
self,
is_kernel=False,
args=args,
ast_builder=impl.get_runtime().prog.current_ast_builder(),
is_real_function=self.is_real_function)
ret = transform_tree(tree, ctx)
if not self.is_real_function:
if self.return_type and ctx.returned != ReturnStatus.ReturnedValue:
raise TaichiSyntaxError(
"Function has a return type but does not have a return statement"
)
return ret
def _process_args(self, args, kwargs):
ret = [argument.default for argument in self.arguments]
len_args = len(args)
if len_args > len(ret):
raise TaichiSyntaxError("Too many arguments.")
for i, arg in enumerate(args):
ret[i] = arg
for key, value in kwargs.items():
found = False
for i, arg in enumerate(self.arguments):
if key == arg.name:
if i < len_args:
raise TaichiSyntaxError(
f"Multiple values for argument '{key}'.")
ret[i] = value
found = True
break
if not found:
raise TaichiSyntaxError(f"Unexpected argument '{key}'.")
for i, arg in enumerate(ret):
if arg is inspect.Parameter.empty:
raise TaichiSyntaxError(
f"Parameter '{self.arguments[i].name}' missing.")
return ret
def __call__(self, *args):
_taichi_skip_traceback = 1
if not impl.inside_kernel():
if not self.pyfunc:
raise TaichiSyntaxError(
"Taichi functions cannot be called from Python-scope."
" Use @ti.pyfunc if you wish to call Taichi functions "
"from both Python-scope and Taichi-scope.")
return self.func(*args)
if impl.get_runtime().experimental_real_function:
if impl.get_runtime().current_kernel.is_grad:
raise TaichiSyntaxError(
"Real function in gradient kernels unsupported.")
instance_id, _ = self.mapper.lookup(args)
key = _ti_core.FunctionKey(self.func.__name__, self.func_id,
instance_id)
if self.compiled is None:
self.compiled = {}
if key.instance_id not in self.compiled:
self.do_compile(key=key, args=args)
return self.func_call_rvalue(key=key, args=args)
tree, global_vars = _get_tree_and_global_vars(self, args)
visitor = ASTTransformerTotal(is_kernel=False,
func=self,
global_vars=global_vars)
return visitor.visit(tree, *args)
def build_assign_unpack(ctx, node_target, values, is_static_assign):
"""Build the unpack assignments like this: (target1, target2) = (value1, value2).
The function should be called only if the node target is a tuple.
Args:
ctx (ast_builder_utils.BuilderContext): The builder context.
node_target (ast.Tuple): A list or tuple object. `node_target.elts` holds a
list of nodes representing the elements.
values: A node/list representing the values.
is_static_assign: A boolean value indicating whether this is a static assignment
"""
if not isinstance(node_target, ast.Tuple):
return ASTTransformer.build_assign_basic(ctx, node_target, values,
is_static_assign)
targets = node_target.elts
if isinstance(values, matrix.Matrix):
if not values.m == 1:
raise ValueError(
'Matrices with more than one columns cannot be unpacked')
values = values.entries
if not isinstance(values, collections.abc.Sequence):
raise TaichiSyntaxError(f'Cannot unpack type: {type(values)}')
if len(values) != len(targets):
raise TaichiSyntaxError(
"The number of targets is not equal to value length")
for i, target in enumerate(targets):
ASTTransformer.build_assign_basic(ctx, target, values[i],
is_static_assign)
return None
def field(cls,
members,
shape=None,
name="<Struct>",
offset=None,
needs_grad=False,
layout=Layout.AOS):
if shape is None and offset is not None:
raise TaichiSyntaxError(
"shape cannot be None when offset is being set")
field_dict = {}
for key, dtype in members.items():
field_name = name + '.' + key
if isinstance(dtype, CompoundType):
field_dict[key] = dtype.field(shape=None,
name=field_name,
offset=offset,
needs_grad=needs_grad)
else:
field_dict[key] = impl.field(dtype,
shape=None,
name=field_name,
offset=offset,
needs_grad=needs_grad)
if shape is not None:
if isinstance(shape, numbers.Number):
shape = (shape, )
if isinstance(offset, numbers.Number):
offset = (offset, )
if offset is not None and len(shape) != len(offset):
raise TaichiSyntaxError(
f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})'
)
dim = len(shape)
if layout == Layout.SOA:
for e in field_dict.values():
ti.root.dense(impl.index_nd(dim),
shape).place(e, offset=offset)
if needs_grad:
for e in field_dict.values():
ti.root.dense(impl.index_nd(dim),
shape).place(e.grad, offset=offset)
else:
ti.root.dense(impl.index_nd(dim),
shape).place(*tuple(field_dict.values()),
offset=offset)
if needs_grad:
grads = tuple(e.grad for e in field_dict.values())
ti.root.dense(impl.index_nd(dim),
shape).place(*grads, offset=offset)
return StructField(field_dict, name=name)
def build_For(ctx, node):
if node.orelse:
raise TaichiSyntaxError(
"'else' clause for 'for' not supported in Taichi kernels")
decorator = ASTTransformer.get_decorator(ctx, node.iter)
double_decorator = ''
if decorator != '' and len(node.iter.args) == 1:
double_decorator = ASTTransformer.get_decorator(
ctx, node.iter.args[0])
if decorator == 'static':
if double_decorator == 'static':
raise TaichiSyntaxError("'ti.static' cannot be nested")
with ctx.loop_scope_guard(is_static=True):
return ASTTransformer.build_static_for(
ctx, node, double_decorator == 'grouped')
with ctx.loop_scope_guard():
if decorator == 'ndrange':
if double_decorator != '':
raise TaichiSyntaxError(
"No decorator is allowed inside 'ti.ndrange")
return ASTTransformer.build_ndrange_for(ctx, node)
if decorator == 'grouped':
if double_decorator == 'static':
raise TaichiSyntaxError(
"'ti.static' is not allowed inside 'ti.grouped'")
elif double_decorator == 'ndrange':
return ASTTransformer.build_grouped_ndrange_for(ctx, node)
elif double_decorator == 'grouped':
raise TaichiSyntaxError("'ti.grouped' cannot be nested")
else:
return ASTTransformer.build_struct_for(ctx,
node,
is_grouped=True)
elif isinstance(node.iter, ast.Call) and isinstance(
node.iter.func, ast.Name) and node.iter.func.id == 'range':
return ASTTransformer.build_range_for(ctx, node)
else:
build_stmt(ctx, node.iter)
if isinstance(node.iter.ptr, mesh.MeshElementField):
if not _ti_core.is_extension_supported(
impl.default_cfg().arch, _ti_core.Extension.mesh):
raise Exception(
'Backend ' + str(impl.default_cfg().arch) +
' doesn\'t support MeshTaichi extension')
return ASTTransformer.build_mesh_for(ctx, node)
if isinstance(node.iter.ptr, mesh.MeshRelationAccessProxy):
return ASTTransformer.build_nested_mesh_for(ctx, node)
# Struct for
return ASTTransformer.build_struct_for(ctx,
node,
is_grouped=False)
def build_For(ctx, node):
if node.orelse:
raise TaichiSyntaxError(
"'else' clause for 'for' not supported in Taichi kernels")
with ctx.control_scope_guard():
decorator = IRBuilder.get_decorator(ctx, node.iter)
double_decorator = ''
if decorator != '' and len(node.iter.args) == 1:
double_decorator = IRBuilder.get_decorator(
ctx, node.iter.args[0])
ast.fix_missing_locations(node)
if decorator == 'static':
if double_decorator == 'static':
raise TaichiSyntaxError("'ti.static' cannot be nested")
return IRBuilder.build_static_for(
ctx, node, double_decorator == 'grouped')
if decorator == 'ndrange':
if double_decorator != '':
raise TaichiSyntaxError(
"No decorator is allowed inside 'ti.ndrange")
return IRBuilder.build_ndrange_for(ctx, node)
if decorator == 'grouped':
if double_decorator == 'static':
raise TaichiSyntaxError(
"'ti.static' is not allowed inside 'ti.grouped'")
elif double_decorator == 'ndrange':
return IRBuilder.build_grouped_ndrange_for(ctx, node)
elif double_decorator == 'grouped':
raise TaichiSyntaxError("'ti.grouped' cannot be nested")
else:
return IRBuilder.build_struct_for(ctx,
node,
is_grouped=True)
elif isinstance(node.iter, ast.Call) and isinstance(
node.iter.func, ast.Name) and node.iter.func.id == 'range':
return IRBuilder.build_range_for(ctx, node)
elif isinstance(node.iter, ast.Attribute) and isinstance(
build_stmt(ctx, node.iter).value.ptr, impl.MeshInstance):
if not ti.is_extension_supported(ti.cfg.arch,
ti.extension.mesh):
raise Exception('Backend ' + str(ti.cfg.arch) +
' doesn\'t support MeshTaichi extension')
return IRBuilder.build_mesh_for(ctx, node)
else: # Struct for
if not isinstance(node.iter, ast.Attribute):
build_stmt(ctx, node.iter)
return IRBuilder.build_struct_for(ctx, node, is_grouped=False)
def visit_While(self, node):
if node.orelse:
raise TaichiSyntaxError(
"'else' clause for 'while' not supported in Taichi kernels")
with self.control_scope():
self.current_control_scope().append('while')
template = '''
if 1:
ti.core.begin_frontend_while(ti.Expr(1).ptr)
__while_cond = 0
if __while_cond:
pass
else:
break
ti.core.pop_scope()
'''
cond = node.test
t = ast.parse(template).body[0]
t.body[1].value = cond
t.body = t.body[:3] + node.body + t.body[3:]
self.generic_visit(t, ['body'])
return ast.copy_location(t, node)
def element_wise_writeback_binary(self, foo, other):
ret = self.empty_copy()
if isinstance(other, (dict)):
other = Struct(other)
if is_taichi_class(other):
other = other.variable()
if foo.__name__ == 'assign' and not isinstance(other, Struct):
raise TaichiSyntaxError(
'cannot assign scalar expr to '
f'taichi class {type(self)}, maybe you want to use `a.fill(b)` instead?'
)
if isinstance(other, Struct):
if self.entries.keys() != other.entries.keys():
raise TypeError(
f"Member mismatch between structs {self.keys}, {other.keys}"
)
for k, v in self.items:
if isinstance(v, expr.Expr):
ret.entries[k] = foo(v, other.entries[k])
else:
ret.entries[k] = v.element_wise_binary(
foo, other.entries[k])
else: # assumed to be scalar
for k, v in self.items:
if isinstance(v, expr.Expr):
ret.entries[k] = foo(v, other)
else:
ret.entries[k] = v.element_wise_binary(foo, other)
return ret
def build_Assert(ctx, node):
extra_args = []
if node.msg is not None:
if ASTTransformer._is_string_mod_args(node.msg):
msg, extra_args = ASTTransformer._handle_string_mod_args(
ctx, node.msg)
else:
msg = build_stmt(ctx, node.msg)
if isinstance(node.msg, ast.Constant):
msg = str(msg)
elif isinstance(node.msg, ast.Str):
pass
elif isinstance(msg, collections.abc.Sequence) and len(
msg) > 0 and msg[0] == "__ti_format__":
msg, extra_args = ASTTransformer.ti_format_list_to_assert_msg(
msg)
else:
raise TaichiSyntaxError(
f"assert info must be constant or formatted string, not {type(msg)}"
)
else:
msg = unparse(node.test)
test = build_stmt(ctx, node.test)
impl.ti_assert(test, msg.strip(), extra_args)
return None
def build_struct_for(ctx, node, is_grouped):
# for i, j in x
# for I in ti.grouped(x)
targets = IRBuilder.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).ptr
loop_indices = ti.lang.expr.make_var_list(
size=len(loop_var.shape))
expr_group = ti.lang.expr.make_expr_group(loop_indices)
ti.begin_frontend_struct_for(expr_group, loop_var)
ctx.create_variable(target, ti.Vector(loop_indices, dt=ti.i32))
node.body = build_stmts(ctx, node.body)
ti.core.end_frontend_range_for()
else:
_vars = []
for name in targets:
var = ti.Expr(ti.core.make_id_expr(""))
_vars.append(var)
ctx.create_variable(name, var)
loop_var = node.iter.ptr
expr_group = ti.lang.expr.make_expr_group(*_vars)
ti.begin_frontend_struct_for(expr_group, loop_var)
node.body = build_stmts(ctx, node.body)
ti.core.end_frontend_range_for()
return node
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 build_While(ctx, node):
if node.orelse:
raise TaichiSyntaxError(
"'else' clause for 'while' not supported in Taichi kernels")
with ctx.control_scope_guard():
ctx.current_control_scope().append('while')
template = '''
if 1:
ti.core.begin_frontend_while(ti.Expr(1).ptr)
__while_cond = 0
if __while_cond:
pass
else:
break
ti.core.pop_scope()
'''
cond = node.test
t = ast.parse(template).body[0]
t.body[1].value = cond
t.body = t.body[:3] + node.body + t.body[3:]
t.body = build_stmts(ctx, t.body)
return ast.copy_location(t, node)
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 __init__(self, *args):
args = list(args)
for i, arg in enumerate(args):
if not isinstance(arg, collections.abc.Sequence):
args[i] = (0, arg)
if len(args[i]) != 2:
raise TaichiSyntaxError(
"Every argument of ndrange should be a scalar or a tuple/list like (begin, end)"
)
args[i] = (args[i][0], ops.max(args[i][0], args[i][1]))
for arg in args:
for bound in arg:
if not isinstance(bound, int) and not (
isinstance(bound, Expr)
and is_integral(bound.ptr.get_ret_type())):
raise TaichiTypeError(
"Every argument of ndrange should be an integer scalar or a tuple/list of (int, int)"
)
self.bounds = args
self.dimensions = [None] * len(args)
for i, bound in enumerate(self.bounds):
self.dimensions[i] = bound[1] - bound[0]
self.acc_dimensions = self.dimensions.copy()
for i in reversed(range(len(self.bounds) - 1)):
self.acc_dimensions[
i] = self.acc_dimensions[i] * self.acc_dimensions[i + 1]
if len(self.acc_dimensions
) == 0: # for the empty case, e.g. ti.ndrange()
self.acc_dimensions = [1]
def maybe_transform_ti_func_call_to_stmt(ti_func, *args, **kwargs):
_taichi_skip_traceback = 1
if '_sitebuiltins' == getattr(ti_func, '__module__', '') and getattr(
getattr(ti_func, '__class__', ''), '__name__', '') == 'Quitter':
raise TaichiSyntaxError(f'exit or quit not supported in Taichi-scope')
if getattr(ti_func, '__module__',
'') == '__main__' and not getattr(ti_func, '__wrapped__', ''):
warnings.warn(
f'Calling into non-Taichi function {ti_func.__name__}.'
' This means that scope inside that function will not be processed'
' by the Taichi transformer. Proceed with caution! '
' Maybe you want to decorate it with @ti.func?',
UserWarning,
stacklevel=2)
is_taichi_function = getattr(ti_func, '_is_taichi_function', False)
# If is_taichi_function is true: call a decorated Taichi function
# in a Taichi kernel/function.
if is_taichi_function and get_runtime().experimental_real_function:
# Compiles the function here.
# Invokes Func.__call__.
func_call_result = ti_func(*args, **kwargs)
return _ti_core.insert_expr_stmt(func_call_result.ptr)
else:
return ti_func(*args, **kwargs)
def linearize_entry_id(self, *args):
assert 1 <= len(args) <= 2
if len(args) == 1 and isinstance(args[0], (list, tuple)):
args = args[0]
if len(args) == 1:
args = args + (0, )
_taichi_skip_traceback = 1
# TODO(#1004): See if it's possible to support indexing at runtime
for i, a in enumerate(args):
if not isinstance(a, int):
raise TaichiSyntaxError(
f'The {i}-th index of a Matrix/Vector must be a compile-time constant '
f'integer, got {type(a)}.\n'
'This is because matrix operations will be **unrolled** at compile-time '
'for performance reason.\n'
'If you want to *iterate through matrix elements*, use a static range:\n'
' for i in ti.static(range(3)):\n'
' print(i, "-th component is", vec[i])\n'
'See https://taichi.readthedocs.io/en/stable/meta.html#when-to-use-for-loops-with-ti-static for more details.'
)
assert 0 <= args[0] < self.n, \
f"The 0-th matrix index is out of range: 0 <= {args[0]} < {self.n}"
assert 0 <= args[1] < self.m, \
f"The 1-th matrix index is out of range: 0 <= {args[1]} < {self.m}"
return args[0] * self.m + args[1]
def build_mesh_for(ctx, node):
targets = IRBuilder.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():
element_dict = {
'verts': ti.core.MeshElementType.Vertex,
'edges': ti.core.MeshElementType.Edge,
'faces': ti.core.MeshElementType.Face,
'cells': ti.core.MeshElementType.Cell
}
var = ti.Expr(ti.core.make_id_expr(""))
ctx.mesh = node.iter.value.ptr
assert isinstance(ctx.mesh, impl.MeshInstance)
mesh_idx = ti.MeshElementFieldProxy(ctx.mesh,
element_dict[node.iter.attr],
var.ptr)
ctx.create_variable(target, mesh_idx)
ti.core.begin_frontend_mesh_for(mesh_idx.ptr, ctx.mesh.mesh_ptr,
element_dict[node.iter.attr])
node.body = build_stmts(ctx, node.body)
ctx.mesh = None
ti.core.end_frontend_range_for()
return node
def place(
self,
members,
reorder=False,
needs_grad=False,
):
"""Declares mesh attributes for the mesh element in current mesh builder.
Args:
members (Dict[str, Union[PrimitiveType, VectorType, MatrixType]]): \
names and types for element attributes.
reorder: True if reorders the internal memory for coalesced data access within mesh-for loop.
needs_grad: True if needs to record grad.
Example::
>>> vec3 = ti.types.vector(3, ti.f32)
>>> mesh = ti.TriMesh()
>>> mesh.faces.place({'area' : ti.f32}) # declares a mesh attribute `area` for each face element.
>>> mesh.verts.place({'pos' : vec3}, reorder=True) # declares a mesh attribute `pos` for each vertex element, and reorder it in memory.
"""
self.builder.elements.add(self._type)
for key, dtype in members.items():
if key in {'verts', 'edges', 'faces', 'cells'}:
raise TaichiSyntaxError(
f"'{key}' cannot use as attribute name. It has been reserved as ti.Mesh's keyword."
)
self.attr_dict[key] = MeshAttrType(key, dtype, reorder, needs_grad)
def build_grouped_ndrange_for(ctx, node):
with ctx.variable_scope_guard():
ndrange_var = ti.expr_init(build_stmt(ctx, node.iter.args[0]).ptr)
ndrange_begin = ti.cast(ti.Expr(0), ti.i32)
ndrange_end = ti.cast(
ti.Expr(ti.subscript(ndrange_var.acc_dimensions, 0)), ti.i32)
ndrange_loop_var = ti.Expr(ti.core.make_id_expr(''))
ti.core.begin_frontend_range_for(ndrange_loop_var.ptr,
ndrange_begin.ptr,
ndrange_end.ptr)
targets = IRBuilder.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 = ti.expr_init(
ti.Vector([0] * len(ndrange_var.dimensions), dt=ti.i32))
ctx.create_variable(target, target_var)
I = ti.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
ti.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])
node.body = build_stmts(ctx, node.body)
ti.core.end_frontend_range_for()
return node
def wrapped(*args, **kwargs):
# If we reach here (we should never), it means the class is not decorated
# with @ti.data_oriented, otherwise getattr would have intercepted the call.
clsobj = type(args[0])
assert not hasattr(clsobj, '_data_oriented')
raise TaichiSyntaxError(
f'Please decorate class {clsobj.__name__} with @ti.data_oriented'
)
def wrapped(a, b):
if is_taichi_class(a):
return a._element_wise_writeback_binary(imp_foo, b)
if is_taichi_class(b):
raise TaichiSyntaxError(
f'cannot augassign taichi class {type(b)} to scalar expr')
else:
return imp_foo(a, b)
def build_Compare(ctx, node):
operands = build_exprs(ctx, [node.left] + list(node.comparators))
operators = []
for i in range(len(node.ops)):
if isinstance(node.ops[i], ast.Lt):
op_str = 'Lt'
elif isinstance(node.ops[i], ast.LtE):
op_str = 'LtE'
elif isinstance(node.ops[i], ast.Gt):
op_str = 'Gt'
elif isinstance(node.ops[i], ast.GtE):
op_str = 'GtE'
elif isinstance(node.ops[i], ast.Eq):
op_str = 'Eq'
elif isinstance(node.ops[i], ast.NotEq):
op_str = 'NotEq'
elif isinstance(node.ops[i], ast.In):
raise TaichiSyntaxError(
'"in" is not supported in Taichi kernels.')
elif isinstance(node.ops[i], ast.NotIn):
raise TaichiSyntaxError(
'"not in" is not supported in Taichi kernels.')
elif isinstance(node.ops[i], ast.Is):
raise TaichiSyntaxError(
'"is" is not supported in Taichi kernels.')
elif isinstance(node.ops[i], ast.IsNot):
raise TaichiSyntaxError(
'"is not" is not supported in Taichi kernels.')
else:
raise Exception(f'Unknown operator {node.ops[i]}')
operators += [
ast.copy_location(ast.Str(s=op_str, kind=None), node)
]
call = ast.Call(
func=parse_expr('ti.chain_compare'),
args=[
ast.copy_location(ast.List(elts=operands, ctx=ast.Load()),
node),
ast.copy_location(ast.List(elts=operators, ctx=ast.Load()),
node)
],
keywords=[])
call = ast.copy_location(call, node)
return call
def taichi_ast_generator():
_taichi_skip_traceback = 1
if self.runtime.inside_kernel:
raise TaichiSyntaxError(
"Kernels cannot call other kernels. I.e., nested kernels are not allowed. Please check if you have direct/indirect invocation of kernels within kernels. Note that some methods provided by the Taichi standard library may invoke kernels, and please move their invocations to Python-scope."
)
self.runtime.inside_kernel = True
compiled()
self.runtime.inside_kernel = False
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 build_Compare(ctx, node):
build_stmt(ctx, node.left)
build_stmts(ctx, node.comparators)
ops = {
ast.Eq: lambda l, r: l == r,
ast.NotEq: lambda l, r: l != r,
ast.Lt: lambda l, r: l < r,
ast.LtE: lambda l, r: l <= r,
ast.Gt: lambda l, r: l > r,
ast.GtE: lambda l, r: l >= r,
}
ops_static = {
ast.In: lambda l, r: l in r,
ast.NotIn: lambda l, r: l not in r,
ast.Is: lambda l, r: l is r,
ast.IsNot: lambda l, r: l is not r,
}
if ctx.is_in_static_scope():
ops = {**ops, **ops_static}
operands = [node.left.ptr
] + [comparator.ptr for comparator in node.comparators]
val = True
for i, node_op in enumerate(node.ops):
l = operands[i]
r = operands[i + 1]
op = ops.get(type(node_op))
if isinstance(node_op, (ast.Is, ast.IsNot)):
name = "is" if isinstance(node_op, ast.Is) else "is not"
warnings.warn_explicit(
f'Operator "{name}" in Taichi scope is deprecated. Please avoid using it.',
DeprecationWarning, ctx.file,
node.lineno + ctx.lineno_offset)
if op is None:
if type(node_op) in ops_static:
raise TaichiSyntaxError(
f'"{type(node_op).__name__}" is only supported inside `ti.static`.'
)
else:
raise TaichiSyntaxError(
f'"{type(node_op).__name__}" is not supported in Taichi kernels.'
)
val = ti_ops.logical_and(val, op(l, r))
node.ptr = val
return node.ptr
