def build_ndrange_for(ctx, node):
with ctx.variable_scope_guard():
ndrange_var = ti.expr_init(build_stmt(ctx, node.iter).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)
I = ti.expr_init(ndrange_loop_var)
targets = IRBuilder.get_for_loop_targets(node)
for i, target in enumerate(targets):
if i + 1 < len(targets):
target_tmp = ti.expr_init(
I // ndrange_var.acc_dimensions[i + 1])
else:
target_tmp = ti.expr_init(I)
ctx.create_variable(
target,
ti.expr_init(
target_tmp +
ti.subscript(ti.subscript(ndrange_var.bounds, i), 0)))
if i + 1 < len(targets):
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 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 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 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 set_vertices(self, data: ti.ext_arr()):
self.n_vertices[None] = min(ti.Expr(self.pos.shape[0]), data.shape[0])
for i in range(self.n_vertices[None]):
for j in ti.static(range(3)):
self.pos[i][j] = data[i, 0, j]
self.tex[i][j] = data[i, 1, j]
self.nrm[i][j] = data[i, 2, j]
def shape(self):
if self.ptr.is_external_var():
import taichi as ti
dim = ti.get_external_tensor_dim(self.ptr)
ret = [
ti.Expr(ti.get_external_tensor_shape_along_axis(self.ptr, i))
for i in range(dim)
]
return ret
return self.snode.shape
def __init__(self, *args):
import taichi as ti
args = list(args)
for i in range(len(args)):
if isinstance(args[i], list):
args[i] = list(args[i])
if not isinstance(args[i], tuple):
args[i] = (0, args[i])
assert len(args[i]) == 2
args[i] = ti.Expr(args[i][0]), ti.Expr(args[i][1])
self.bounds = args
self.dimensions = [None] * len(args)
for i in range(len(self.bounds)):
self.dimensions[i] = self.bounds[i][1] - self.bounds[i][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]
def func_call_rvalue(self, key, args):
# Skip the template args, e.g., |self|
assert impl.get_runtime().experimental_real_function
non_template_args = []
for i in range(len(self.argument_annotations)):
if not isinstance(self.argument_annotations[i], template):
non_template_args.append(args[i])
non_template_args = impl.make_expr_group(non_template_args)
return ti.Expr(
_ti_core.make_func_call_expr(
self.taichi_functions[key.instance_id], non_template_args))
def build_range_for(ctx, node):
with ctx.variable_scope_guard():
loop_name = node.target.id
ctx.check_loop_var(loop_name)
loop_var = ti.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.cast(
ti.Expr(build_stmt(ctx, node.iter.args[0]).ptr), ti.i32)
end = ti.cast(ti.Expr(build_stmt(ctx, node.iter.args[1]).ptr),
ti.i32)
else:
begin = ti.cast(ti.Expr(0), ti.i32)
end = ti.cast(ti.Expr(build_stmt(ctx, node.iter.args[0]).ptr),
ti.i32)
ti.core.begin_frontend_range_for(loop_var.ptr, begin.ptr, end.ptr)
node.body = build_stmts(ctx, node.body)
ti.core.end_frontend_range_for()
return node
def build_While(ctx, node):
if node.orelse:
raise TaichiSyntaxError(
"'else' clause for 'while' not supported in Taichi kernels")
with ctx.control_scope_guard():
ti.core.begin_frontend_while(ti.Expr(1).ptr)
while_cond = build_stmt(ctx, node.test).ptr
ti.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()
node.body = build_stmts(ctx, node.body)
ti.core.pop_scope()
return node
def build_Return(ctx, node):
node.value = 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.cast(ti.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
return ast.Pass()
ctx.return_data = node.value.ptr
return node
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 loop_var in targets:
ctx.check_loop_var(loop_var)
if is_grouped:
pass
# template = '''
# if 1:
# ___loop_var = 0
# {} = ti.lang.expr.make_var_vector(size=len(___loop_var.shape))
# ___expr_group = ti.lang.expr.make_expr_group({})
# ti.begin_frontend_struct_for(___expr_group, ___loop_var)
# ti.core.end_frontend_range_for()
# '''.format(vars, vars)
# t = ast.parse(template).body[0]
# cut = 4
# t.body[0].value = node.iter
# t.body = t.body[:cut] + node.body + t.body[cut:]
else:
with ctx.variable_scope_guard():
for name in targets:
ctx.create_variable(name,
ti.Expr(ti.core.make_id_expr("")))
vars = [ctx.get_var_by_name(name) for name in targets]
node.iter = build_stmt(ctx, node.iter)
ti.begin_frontend_struct_for(
ti.lang.expr.make_expr_group(*vars), node.iter.ptr)
node.body = build_stmts_wo_scope(ctx, node.body)
ti.core.end_frontend_range_for()
return node
def clear_gradients(vars: ti.template()):
for I in ti.grouped(ti.Expr(vars[0])):
for s in ti.static(vars):
ti.Expr(s)[I] = 0
def loop_range(self):
import taichi as ti
return ti.Expr(ti.core.global_var_expr_from_snode(self.ptr))
def ti_assert(cond, msg, extra_args):
# Mostly a wrapper to help us convert from ti.Expr (defined in Python) to
# taichi_lang_core.Expr (defined in C++)
import taichi as ti
taichi_lang_core.create_assert_stmt(
ti.Expr(cond).ptr, msg, [ti.Expr(x).ptr for x in extra_args])
def set_triangles(self, data: ti.ext_arr()):
self.n_triangles[None] = min(ti.Expr(self.faces.shape[0]), data.shape[0])
for i in range(self.n_triangles[None]):
for j, k in ti.static(ti.ndrange(3, 3)):
self.faces[i][j, k] = data[i, j]
