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_Subscript(ctx, node):
build_stmt(ctx, node.value)
build_stmt(ctx, node.slice)
if not ASTTransformer.is_tuple(node.slice):
node.slice.ptr = [node.slice.ptr]
node.ptr = impl.subscript(node.value.ptr, *node.slice.ptr)
return node.ptr
def subscript(self, *indices):
_taichi_skip_traceback = 1
if self.is_global():
ret = self.empty_copy()
for i, e in enumerate(self.entries):
ret.entries[i] = impl.subscript(e, *indices)
return ret
else:
assert len(indices) in [1, 2]
i = indices[0]
j = 0 if len(indices) == 1 else indices[1]
return self(i, j)
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
def subscript(self, *indices):
_taichi_skip_traceback = 1
if self.is_global():
ret = self.empty_copy()
for i, e in enumerate(self.entries):
ret.entries[i] = impl.subscript(e, *indices)
return ret
else:
assert len(indices) in [1, 2]
i = indices[0]
j = 0 if len(indices) == 1 else indices[1]
# ptr.is_global_ptr() will check whether it's an element in the field (which is different from ptr.is_global_var()).
if isinstance(self.entries[0],
ti.Expr) and self.entries[0].ptr.is_global_ptr(
) and ti.is_extension_supported(
ti.cfg.arch, ti.extension.dynamic_index):
return ti.subscript_with_offset(self.entries[0], (i, j),
self.m, True)
else:
return self(i, j)
def foo(n):
return ti.atomic_add(impl.subscript(b, None), n)
def fibonacci(x):
return impl.subscript(bar(x), 1, 0)
