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_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_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 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():
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])
build_stmts(ctx, node.body)
ctx.mesh = None
_ti_core.end_frontend_range_for()
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 = 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))
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)
build_stmts(ctx, node.body)
_ti_core.end_frontend_range_for()
return None
def build_grouped_ndrange_for(ctx, node):
with ctx.variable_scope_guard():
ndrange_var = ti.expr_init(build_stmt(ctx, node.iter.args[0]))
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 = 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 = 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])
build_stmts(ctx, node.body)
_ti_core.end_frontend_range_for()
return None
def create_field_member(dtype, name):
dtype = cook_dtype(dtype)
# primal
x = Expr(_ti_core.make_id_expr(""))
x.declaration_tb = get_traceback(stacklevel=4)
x.ptr = _ti_core.global_new(x.ptr, dtype)
x.ptr.set_name(name)
x.ptr.set_is_primal(True)
pytaichi.global_vars.append(x)
x_grad = None
if _ti_core.needs_grad(dtype):
# adjoint
x_grad = Expr(_ti_core.make_id_expr(""))
x_grad.ptr = _ti_core.global_new(x_grad.ptr, dtype)
x_grad.ptr.set_name(name + ".grad")
x_grad.ptr.set_is_primal(False)
x.ptr.set_grad(x_grad.ptr)
return x, x_grad
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 make_var_list(size):
exprs = []
for _ in range(size):
exprs.append(_ti_core.make_id_expr(''))
return exprs
