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)