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