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