def ext_arr_to_matrix(arr: ext_arr(), mat: template(), as_vector: template()):
for I in grouped(mat):
for p in static(range(mat.n)):
for q in static(range(mat.m)):
if static(as_vector):
mat[I][p] = arr[I, p]
else:
mat[I][p, q] = arr[I, p, q]
def matrix_to_ext_arr(mat: template(), arr: ndarray_type.ndarray(),
as_vector: template()):
for I in grouped(mat):
for p in static(range(mat.n)):
for q in static(range(mat.m)):
if static(as_vector):
arr[I, p] = mat[I][p]
else:
arr[I, p, q] = mat[I][p, q]
def ndarray_matrix_to_ext_arr(ndarray: any_arr(), arr: ext_arr(),
as_vector: template()):
for I in grouped(ndarray):
for p in static(range(ndarray[I].n)):
for q in static(range(ndarray[I].m)):
if static(as_vector):
arr[I, p] = ndarray[I][p]
else:
arr[I, p, q] = ndarray[I][p, q]
def ndarray_matrix_to_ext_arr(ndarray: ndarray_type.ndarray(),
arr: ndarray_type.ndarray(),
layout_is_aos: template(),
as_vector: template()):
for I in grouped(ndarray):
for p in static(range(ndarray[I].n)):
for q in static(range(ndarray[I].m)):
if static(as_vector):
if static(layout_is_aos):
arr[I, p] = ndarray[I][p]
else:
arr[p, I] = ndarray[I][p]
else:
if static(layout_is_aos):
arr[I, p, q] = ndarray[I][p, q]
else:
arr[p, q, I] = ndarray[I][p, q]
def build_static_for(ctx, node, is_grouped):
if is_grouped:
assert len(node.iter.args[0].args) == 1
ndrange_arg = build_stmt(ctx, node.iter.args[0].args[0])
if not isinstance(ndrange_arg, ndrange):
raise TaichiSyntaxError(
"Only 'ti.ndrange' is allowed in 'ti.static(ti.grouped(...))'."
)
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]
for value in impl.grouped(ndrange_arg):
with ctx.variable_scope_guard():
ctx.create_variable(target, value)
build_stmts(ctx, node.body)
status = ctx.loop_status()
if status == LoopStatus.Break:
break
elif status == LoopStatus.Continue:
ctx.set_loop_status(LoopStatus.Normal)
else:
build_stmt(ctx, node.iter)
targets = ASTTransformer.get_for_loop_targets(node)
for target_values in node.iter.ptr:
if not isinstance(
target_values,
collections.abc.Sequence) or len(targets) == 1:
target_values = [target_values]
with ctx.variable_scope_guard():
for target, target_value in zip(targets, target_values):
ctx.create_variable(target, target_value)
build_stmts(ctx, node.body)
status = ctx.loop_status()
if status == LoopStatus.Break:
break
elif status == LoopStatus.Continue:
ctx.set_loop_status(LoopStatus.Normal)
return None
def fill_tensor(tensor: template(), val: template()):
for I in grouped(tensor):
tensor[I] = val
def ext_arr_to_ndarray(arr: ndarray_type.ndarray(),
ndarray: ndarray_type.ndarray()):
for I in grouped(ndarray):
ndarray[I] = arr[I]
def ndarray_to_ndarray(ndarray: any_arr(), other: any_arr()):
for I in grouped(ndarray):
ndarray[I] = other[I]
def ext_arr_to_tensor(arr: ext_arr(), tensor: template()):
for I in grouped(tensor):
tensor[I] = arr[I]
def ndarray_to_ndarray(ndarray: ndarray_type.ndarray(),
other: ndarray_type.ndarray()):
for I in grouped(ndarray):
ndarray[I] = other[I]
def tensor_to_ext_arr(tensor: template(), arr: ndarray_type.ndarray()):
for I in grouped(tensor):
arr[I] = tensor[I]
def ext_arr_to_ndarray(arr: ext_arr(), ndarray: any_arr()):
for I in grouped(ndarray):
ndarray[I] = arr[I]
def snode_deactivate(b: template()):
for I in grouped(b):
deactivate(b, I)
def fill_ndarray(ndarray: ndarray_type.ndarray(), val: template()):
for I in grouped(ndarray):
ndarray[I] = val
def clear_gradients(_vars: template()):
for I in grouped(ScalarField(Expr(_vars[0]))):
for s in static(_vars):
ScalarField(Expr(s))[I] = 0
def ndarray_to_ext_arr(ndarray: any_arr(), arr: ext_arr()):
for I in grouped(ndarray):
arr[I] = ndarray[I]
def tensor_to_ext_arr(tensor: template(), arr: ext_arr()):
for I in grouped(tensor):
arr[I] = tensor[I]
def fill_ndarray_matrix(ndarray: any_arr(), val: template()):
for I in grouped(ndarray):
ndarray[I].fill(val)
def fill_ndarray(ndarray: any_arr(), val: template()):
for I in grouped(ndarray):
ndarray[I] = val
def fill_matrix(mat: template(), vals: template()):
for I in grouped(mat):
for p in static(range(mat.n)):
for q in static(range(mat.m)):
mat[I][p, q] = vals[p][q]
def tensor_to_image(tensor: template(), arr: ext_arr()):
for I in grouped(tensor):
t = ops.cast(tensor[I], f32)
arr[I, 0] = t
arr[I, 1] = t
arr[I, 2] = t
def snode_deactivate_dynamic(b: template()):
for I in grouped(b.parent()):
deactivate(b, I)
def vector_to_image(mat: template(), arr: ndarray_type.ndarray()):
for I in grouped(mat):
for p in static(range(mat.n)):
arr[I, p] = ops.cast(mat[I][p], f32)
if static(mat.n <= 2):
arr[I, 2] = 0
def fill_ndarray_matrix(ndarray: ndarray_type.ndarray(), val: template()):
for I in grouped(ndarray):
ndarray[I].fill(val)
def tensor_to_tensor(tensor: template(), other: template()):
for I in grouped(tensor):
tensor[I] = other[I]
def ndarray_to_ext_arr(ndarray: ndarray_type.ndarray(),
arr: ndarray_type.ndarray()):
for I in grouped(ndarray):
arr[I] = ndarray[I]
def ext_arr_to_tensor(arr: ndarray_type.ndarray(), tensor: template()):
for I in grouped(tensor):
tensor[I] = arr[I]
def tensor_to_image(tensor: template(), arr: ndarray_type.ndarray()):
for I in grouped(tensor):
t = ops.cast(tensor[I], f32)
arr[I, 0] = t
arr[I, 1] = t
arr[I, 2] = t
def field_fill_taichi_scope(F: template(), val: template()):
for I in grouped(F):
F[I] = val
