def rescale_index(a, b, I): """Rescales the index 'I' of field (or SNode) 'a' to match the shape of SNode 'b' Parameters ---------- a: ti.field(), ti.Vector.field, ti.Matrix.field() input taichi field or snode b: ti.field(), ti.Vector.field, ti.Matrix.field() output taichi field or snode I: ti.Vector() grouped loop index Returns ------- Ib: ti.Vector() rescaled grouped loop index """ assert isinstance( a, (Field, SNode)), "The first argument must be a field or an SNode" assert isinstance( b, (Field, SNode)), "The second argument must be a field or an SNode" if isinstance(I, list): I = matrix.Vector(I) else: assert isinstance( I, matrix.Matrix ), "The third argument must be an index (list or ti.Vector)" entries = [I(i) for i in range(I.n)] for n in range(min(I.n, min(len(a.shape), len(b.shape)))): if a.shape[n] > b.shape[n]: entries[n] = I(n) // (a.shape[n] // b.shape[n]) if a.shape[n] < b.shape[n]: entries[n] = I(n) * (b.shape[n] // a.shape[n]) return matrix.Vector(entries)
def rescale_index(a, b, I): """Rescales the index 'I' of field (or SNode) 'a' to match the shape of SNode 'b'. Args: a, b (Union[:class:`~taichi.Field`, :class:`~taichi.MatrixField`): Input taichi fields or snodes. I (Union[list, :class:`~taichi.Vector`]): grouped loop index. Returns: Ib (:class:`~taichi.Vector`): rescaled grouped loop index """ assert isinstance( a, (Field, SNode)), "The first argument must be a field or an SNode" assert isinstance( b, (Field, SNode)), "The second argument must be a field or an SNode" if isinstance(I, list): I = matrix.Vector(I) else: assert isinstance( I, matrix.Matrix ), "The third argument must be an index (list or ti.Vector)" entries = [I(i) for i in range(I.n)] for n in range(min(I.n, min(len(a.shape), len(b.shape)))): if a.shape[n] > b.shape[n]: entries[n] = I(n) // (a.shape[n] // b.shape[n]) if a.shape[n] < b.shape[n]: entries[n] = I(n) * (b.shape[n] // a.shape[n]) return matrix.Vector(entries)
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 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