def build_grouped_ndrange_for(ctx, node):
with ctx.variable_scope_guard():
ndrange_var = ti.expr_init(build_stmt(ctx, node.iter.args[0]).ptr)
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 = IRBuilder.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])
node.body = build_stmts(ctx, node.body)
ti.core.end_frontend_range_for()
return node
def build_ndrange_for(ctx, node):
with ctx.variable_scope_guard():
ndrange_var = ti.expr_init(build_stmt(ctx, node.iter).ptr)
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)
I = ti.expr_init(ndrange_loop_var)
targets = IRBuilder.get_for_loop_targets(node)
for i, target in enumerate(targets):
if i + 1 < len(targets):
target_tmp = ti.expr_init(
I // ndrange_var.acc_dimensions[i + 1])
else:
target_tmp = ti.expr_init(I)
ctx.create_variable(
target,
ti.expr_init(
target_tmp +
ti.subscript(ti.subscript(ndrange_var.bounds, i), 0)))
if i + 1 < len(targets):
I.assign(I -
target_tmp * ndrange_var.acc_dimensions[i + 1])
node.body = build_stmts(ctx, node.body)
ti.core.end_frontend_range_for()
return node
def build_Subscript(ctx, node):
node.value = build_stmt(ctx, node.value)
node.slice = build_stmt(ctx, node.slice)
if not IRBuilder.is_tuple(node.slice):
node.slice.ptr = [node.slice.ptr]
node.ptr = ti.subscript(node.value.ptr, *node.slice.ptr)
return node
def build_Subscript(ctx, node):
node.value = build_stmt(ctx, node.value)
node.slice = build_stmt(ctx, node.slice)
if not isinstance(node.slice, ast.Tuple):
node.slice.ptr = [node.slice.ptr]
node.ptr = ti.subscript(node.value.ptr, *node.slice.ptr)
return node
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 = ti.subscript(node.value.ptr, *node.slice.ptr)
return node.ptr
def iFrame(self):
'''
(TS, int32, RO) Current frame number start from 0.
'''
if not self.has_input:
raise Exception(
'Add ``self.define_input()`` to ``on_init`` if you '
'wish to use inputs')
if ti.inside_kernel():
return ti.subscript(self._iFrame, None)
else:
return self._iFrame[None]
def iTime(self):
'''
(TS, float32, RO) Current time in seconds.
'''
if not self.has_input:
raise Exception(
'Add ``self.define_input()`` to ``on_init`` if you '
'wish to use inputs')
if ti.inside_kernel():
return ti.subscript(self._iTime, None)
else:
return self._iTime[None]
def _subscript(self, *indices):
args = [ti.subscript(e, *indices) for e in self.entries]
return self.__class__(*args)
def fibonacci(x):
return tc.subscript(bar(x), 1, 0)
def subscript(self, *indices):
return ti.subscript(self.core, *indices)
def foo(n):
return ti.atomic_add(ti.subscript(b, None), n)