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)