def build_assign_basic(ctx, target, value, is_static_assign): """Build basic assginment like this: target = value. Args: ctx (ast_builder_utils.BuilderContext): The builder context. target (ast.Name): A variable name. `target.id` holds the name as a string. value: A node representing the value. is_static_assign: A boolean value indicating whether this is a static assignment """ is_local = isinstance(target, ast.Name) if is_static_assign: if not is_local: raise TaichiSyntaxError( "Static assign cannot be used on elements in arrays") ctx.create_variable(target.id, value) var = value elif is_local and not ctx.is_var_declared(target.id): var = impl.expr_init(value) ctx.create_variable(target.id, var) else: var = build_stmt(ctx, target) try: var.assign(value) except AttributeError: raise TaichiSyntaxError( f"Variable '{unparse(target).strip()}' cannot be assigned. Maybe it is not a Taichi object?" ) return var
def _process_args(self, args, kwargs): ret = [argument.default for argument in self.arguments] len_args = len(args) if len_args > len(ret): raise TaichiSyntaxError("Too many arguments.") for i, arg in enumerate(args): ret[i] = arg for key, value in kwargs.items(): found = False for i, arg in enumerate(self.arguments): if key == arg.name: if i < len_args: raise TaichiSyntaxError( f"Multiple values for argument '{key}'.") ret[i] = value found = True break if not found: raise TaichiSyntaxError(f"Unexpected argument '{key}'.") for i, arg in enumerate(ret): if arg is inspect.Parameter.empty: raise TaichiSyntaxError( f"Parameter '{self.arguments[i].name}' missing.") return ret
def __call__(self, *args, **kwargs): args = _process_args(self, args, kwargs) if not impl.inside_kernel(): if not self.pyfunc: raise TaichiSyntaxError( "Taichi functions cannot be called from Python-scope.") return self.func(*args) if self.is_real_function: if impl.get_runtime( ).current_kernel.autodiff_mode != AutodiffMode.NONE: raise TaichiSyntaxError( "Real function in gradient kernels unsupported.") instance_id, _ = self.mapper.lookup(args) key = _ti_core.FunctionKey(self.func.__name__, self.func_id, instance_id) if self.compiled is None: self.compiled = {} if key.instance_id not in self.compiled: self.do_compile(key=key, args=args) return self.func_call_rvalue(key=key, args=args) tree, ctx = _get_tree_and_ctx( self, is_kernel=False, args=args, ast_builder=impl.get_runtime().prog.current_ast_builder(), is_real_function=self.is_real_function) ret = transform_tree(tree, ctx) if not self.is_real_function: if self.return_type and ctx.returned != ReturnStatus.ReturnedValue: raise TaichiSyntaxError( "Function has a return type but does not have a return statement" ) return ret
def __call__(self, *args): _taichi_skip_traceback = 1 if not impl.inside_kernel(): if not self.pyfunc: raise TaichiSyntaxError( "Taichi functions cannot be called from Python-scope." " Use @ti.pyfunc if you wish to call Taichi functions " "from both Python-scope and Taichi-scope.") return self.func(*args) if impl.get_runtime().experimental_ast_refactor: tree, global_vars = _get_tree_and_global_vars(self, args) visitor = ASTTransformerTotal(is_kernel=False, func=self, globals=global_vars) return visitor.visit(tree, *args) if impl.get_runtime().experimental_real_function: if impl.get_runtime().current_kernel.is_grad: raise TaichiSyntaxError( "Real function in gradient kernels unsupported.") instance_id, arg_features = self.mapper.lookup(args) key = _ti_core.FunctionKey(self.func.__name__, self.func_id, instance_id) if self.compiled is None: self.compiled = {} if key.instance_id not in self.compiled: self.do_compile(key=key, args=args) return self.func_call_rvalue(key=key, args=args) else: if self.compiled is None: self.do_compile(key=None, args=args) ret = self.compiled(*args) return ret
def build_assign_unpack(ctx, node_target, values, is_static_assign): """Build the unpack assignments like this: (target1, target2) = (value1, value2). The function should be called only if the node target is a tuple. Args: ctx (ast_builder_utils.BuilderContext): The builder context. node_target (ast.Tuple): A list or tuple object. `node_target.elts` holds a list of nodes representing the elements. values: A node/list representing the values. is_static_assign: A boolean value indicating whether this is a static assignment """ if not isinstance(node_target, ast.Tuple): return ASTTransformer.build_assign_basic(ctx, node_target, values, is_static_assign) targets = node_target.elts if isinstance(values, matrix.Matrix): if not values.m == 1: raise ValueError( 'Matrices with more than one columns cannot be unpacked') values = values.entries if not isinstance(values, collections.abc.Sequence): raise TaichiSyntaxError(f'Cannot unpack type: {type(values)}') if len(values) != len(targets): raise TaichiSyntaxError( "The number of targets is not equal to value length") for i, target in enumerate(targets): ASTTransformer.build_assign_basic(ctx, target, values[i], is_static_assign) return None
def __call__(self, *args): if not impl.inside_kernel(): if not self.pyfunc: raise TaichiSyntaxError( "Taichi functions cannot be called from Python-scope." " Use @ti.pyfunc if you wish to call Taichi functions " "from both Python-scope and Taichi-scope.") return self.func(*args) if impl.get_runtime().experimental_real_function: if impl.get_runtime().current_kernel.is_grad: raise TaichiSyntaxError( "Real function in gradient kernels unsupported.") instance_id, _ = self.mapper.lookup(args) key = _ti_core.FunctionKey(self.func.__name__, self.func_id, instance_id) if self.compiled is None: self.compiled = {} if key.instance_id not in self.compiled: self.do_compile(key=key, args=args) return self.func_call_rvalue(key=key, args=args) tree, ctx = _get_tree_and_ctx( self, is_kernel=False, args=args, ast_builder=impl.get_runtime().prog.current_ast_builder()) ret = transform_tree(tree, ctx) if not impl.get_runtime().experimental_real_function: if self.return_type and not ctx.returned: raise TaichiSyntaxError( "Function has a return type but does not have a return statement" ) return ret
def build_assign_annotated(ctx, target, value, is_static_assign, annotation): """Build an annotated assginment like this: target: annotation = value. Args: ctx (ast_builder_utils.BuilderContext): The builder context. target (ast.Name): A variable name. `target.id` holds the name as a string. annotation: A type we hope to assign to the target value: A node representing the value. is_static_assign: A boolean value indicating whether this is a static assignment """ is_local = isinstance(target, ast.Name) anno = impl.expr_init(annotation) if is_static_assign: raise TaichiSyntaxError( "Static assign cannot be used on annotated assignment") if is_local and not ctx.is_var_declared(target.id): var = ti_ops.cast(value, anno) var = impl.expr_init(var) ctx.create_variable(target.id, var) else: var = build_stmt(ctx, target) if var.ptr.get_ret_type() != anno: raise TaichiSyntaxError( "Static assign cannot have type overloading") var.assign(value) return var
def build_Return(ctx, node): if not impl.get_runtime().experimental_real_function: if ctx.is_in_non_static(): raise TaichiSyntaxError( "Return inside non-static if/for is not supported") build_stmt(ctx, node.value) if ctx.is_kernel or impl.get_runtime().experimental_real_function: # TODO: check if it's at the end of a kernel, throw TaichiSyntaxError if not if node.value is not None: if ctx.func.return_type is None: raise TaichiSyntaxError( f'A {"kernel" if ctx.is_kernel else "function"} ' 'with a return value must be annotated ' 'with a return type, e.g. def func() -> ti.f32') _ti_core.create_kernel_return( ti_ops.cast(expr.Expr(node.value.ptr), ctx.func.return_type).ptr) # For args[0], it is an ast.Attribute, because it loads the # attribute, |ptr|, of the expression |ret_expr|. Therefore we # only need to replace the object part, i.e. args[0].value else: ctx.return_data = node.value.ptr if not impl.get_runtime().experimental_real_function: ctx.returned = True return None
def field(cls, members, shape=None, name="<Struct>", offset=None, needs_grad=False, layout=Layout.AOS): if shape is None and offset is not None: raise TaichiSyntaxError( "shape cannot be None when offset is being set") field_dict = {} for key, dtype in members.items(): field_name = name + '.' + key if isinstance(dtype, CompoundType): field_dict[key] = dtype.field(shape=None, name=field_name, offset=offset, needs_grad=needs_grad) else: field_dict[key] = impl.field(dtype, shape=None, name=field_name, offset=offset, needs_grad=needs_grad) if shape is not None: if isinstance(shape, numbers.Number): shape = (shape, ) if isinstance(offset, numbers.Number): offset = (offset, ) if offset is not None and len(shape) != len(offset): raise TaichiSyntaxError( f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})' ) dim = len(shape) if layout == Layout.SOA: for e in field_dict.values(): ti.root.dense(impl.index_nd(dim), shape).place(e, offset=offset) if needs_grad: for e in field_dict.values(): ti.root.dense(impl.index_nd(dim), shape).place(e.grad, offset=offset) else: ti.root.dense(impl.index_nd(dim), shape).place(*tuple(field_dict.values()), offset=offset) if needs_grad: grads = tuple(e.grad for e in field_dict.values()) ti.root.dense(impl.index_nd(dim), shape).place(*grads, offset=offset) return StructField(field_dict, name=name)
def build_For(ctx, node): if node.orelse: raise TaichiSyntaxError( "'else' clause for 'for' not supported in Taichi kernels") decorator = ASTTransformer.get_decorator(ctx, node.iter) double_decorator = '' if decorator != '' and len(node.iter.args) == 1: double_decorator = ASTTransformer.get_decorator( ctx, node.iter.args[0]) if decorator == 'static': if double_decorator == 'static': raise TaichiSyntaxError("'ti.static' cannot be nested") with ctx.loop_scope_guard(is_static=True): return ASTTransformer.build_static_for( ctx, node, double_decorator == 'grouped') with ctx.loop_scope_guard(): if decorator == 'ndrange': if double_decorator != '': raise TaichiSyntaxError( "No decorator is allowed inside 'ti.ndrange") return ASTTransformer.build_ndrange_for(ctx, node) if decorator == 'grouped': if double_decorator == 'static': raise TaichiSyntaxError( "'ti.static' is not allowed inside 'ti.grouped'") elif double_decorator == 'ndrange': return ASTTransformer.build_grouped_ndrange_for(ctx, node) elif double_decorator == 'grouped': raise TaichiSyntaxError("'ti.grouped' cannot be nested") else: return ASTTransformer.build_struct_for(ctx, node, is_grouped=True) elif isinstance(node.iter, ast.Call) and isinstance( node.iter.func, ast.Name) and node.iter.func.id == 'range': return ASTTransformer.build_range_for(ctx, node) else: build_stmt(ctx, node.iter) if isinstance(node.iter.ptr, mesh.MeshElementField): if not _ti_core.is_extension_supported( impl.default_cfg().arch, _ti_core.Extension.mesh): raise Exception( 'Backend ' + str(impl.default_cfg().arch) + ' doesn\'t support MeshTaichi extension') return ASTTransformer.build_mesh_for(ctx, node) if isinstance(node.iter.ptr, mesh.MeshRelationAccessProxy): return ASTTransformer.build_nested_mesh_for(ctx, node) # Struct for return ASTTransformer.build_struct_for(ctx, node, is_grouped=False)
def build_For(ctx, node): if node.orelse: raise TaichiSyntaxError( "'else' clause for 'for' not supported in Taichi kernels") with ctx.control_scope_guard(): decorator = IRBuilder.get_decorator(ctx, node.iter) double_decorator = '' if decorator != '' and len(node.iter.args) == 1: double_decorator = IRBuilder.get_decorator( ctx, node.iter.args[0]) ast.fix_missing_locations(node) if decorator == 'static': if double_decorator == 'static': raise TaichiSyntaxError("'ti.static' cannot be nested") return IRBuilder.build_static_for( ctx, node, double_decorator == 'grouped') if decorator == 'ndrange': if double_decorator != '': raise TaichiSyntaxError( "No decorator is allowed inside 'ti.ndrange") return IRBuilder.build_ndrange_for(ctx, node) if decorator == 'grouped': if double_decorator == 'static': raise TaichiSyntaxError( "'ti.static' is not allowed inside 'ti.grouped'") elif double_decorator == 'ndrange': return IRBuilder.build_grouped_ndrange_for(ctx, node) elif double_decorator == 'grouped': raise TaichiSyntaxError("'ti.grouped' cannot be nested") else: return IRBuilder.build_struct_for(ctx, node, is_grouped=True) elif isinstance(node.iter, ast.Call) and isinstance( node.iter.func, ast.Name) and node.iter.func.id == 'range': return IRBuilder.build_range_for(ctx, node) elif isinstance(node.iter, ast.Attribute) and isinstance( build_stmt(ctx, node.iter).value.ptr, impl.MeshInstance): if not ti.is_extension_supported(ti.cfg.arch, ti.extension.mesh): raise Exception('Backend ' + str(ti.cfg.arch) + ' doesn\'t support MeshTaichi extension') return IRBuilder.build_mesh_for(ctx, node) else: # Struct for if not isinstance(node.iter, ast.Attribute): build_stmt(ctx, node.iter) return IRBuilder.build_struct_for(ctx, node, is_grouped=False)
def build_While(ctx, node): if node.orelse: raise TaichiSyntaxError( "'else' clause for 'while' not supported in Taichi kernels") with ctx.control_scope_guard(): ctx.current_control_scope().append('while') template = ''' if 1: ti.core.begin_frontend_while(ti.Expr(1).ptr) __while_cond = 0 if __while_cond: pass else: break ti.core.pop_scope() ''' cond = node.test t = ast.parse(template).body[0] t.body[1].value = cond t.body = t.body[:3] + node.body + t.body[3:] t.body = build_stmts(ctx, t.body) return ast.copy_location(t, node)
def build_struct_for(ctx, node, is_grouped): # for i, j in x # for I in ti.grouped(x) targets = IRBuilder.get_for_loop_targets(node) for target in targets: ctx.check_loop_var(target) with ctx.variable_scope_guard(): if is_grouped: if len(targets) != 1: raise TaichiSyntaxError( f"Group for should have 1 loop target, found {len(targets)}" ) target = targets[0] loop_var = build_stmt(ctx, node.iter).ptr loop_indices = ti.lang.expr.make_var_list( size=len(loop_var.shape)) expr_group = ti.lang.expr.make_expr_group(loop_indices) ti.begin_frontend_struct_for(expr_group, loop_var) ctx.create_variable(target, ti.Vector(loop_indices, dt=ti.i32)) node.body = build_stmts(ctx, node.body) ti.core.end_frontend_range_for() else: _vars = [] for name in targets: var = ti.Expr(ti.core.make_id_expr("")) _vars.append(var) ctx.create_variable(name, var) loop_var = node.iter.ptr expr_group = ti.lang.expr.make_expr_group(*_vars) ti.begin_frontend_struct_for(expr_group, loop_var) node.body = build_stmts(ctx, node.body) ti.core.end_frontend_range_for() return node
def __init__(self, *args): args = list(args) for i, arg in enumerate(args): if not isinstance(arg, collections.abc.Sequence): args[i] = (0, arg) if len(args[i]) != 2: raise TaichiSyntaxError( "Every argument of ndrange should be a scalar or a tuple/list like (begin, end)" ) args[i] = (args[i][0], ops.max(args[i][0], args[i][1])) for arg in args: for bound in arg: if not isinstance(bound, int) and not ( isinstance(bound, Expr) and is_integral(bound.ptr.get_ret_type())): raise TaichiTypeError( "Every argument of ndrange should be an integer scalar or a tuple/list of (int, int)" ) self.bounds = args self.dimensions = [None] * len(args) for i, bound in enumerate(self.bounds): self.dimensions[i] = bound[1] - bound[0] self.acc_dimensions = self.dimensions.copy() for i in reversed(range(len(self.bounds) - 1)): self.acc_dimensions[ i] = self.acc_dimensions[i] * self.acc_dimensions[i + 1] if len(self.acc_dimensions ) == 0: # for the empty case, e.g. ti.ndrange() self.acc_dimensions = [1]
def linearize_entry_id(self, *args): assert 1 <= len(args) <= 2 if len(args) == 1 and isinstance(args[0], (list, tuple)): args = args[0] if len(args) == 1: args = args + (0, ) _taichi_skip_traceback = 1 # TODO(#1004): See if it's possible to support indexing at runtime for i, a in enumerate(args): if not isinstance(a, int): raise TaichiSyntaxError( f'The {i}-th index of a Matrix/Vector must be a compile-time constant ' f'integer, got {type(a)}.\n' 'This is because matrix operations will be **unrolled** at compile-time ' 'for performance reason.\n' 'If you want to *iterate through matrix elements*, use a static range:\n' ' for i in ti.static(range(3)):\n' ' print(i, "-th component is", vec[i])\n' 'See https://taichi.readthedocs.io/en/stable/meta.html#when-to-use-for-loops-with-ti-static for more details.' ) assert 0 <= args[0] < self.n, \ f"The 0-th matrix index is out of range: 0 <= {args[0]} < {self.n}" assert 0 <= args[1] < self.m, \ f"The 1-th matrix index is out of range: 0 <= {args[1]} < {self.m}" return args[0] * self.m + args[1]
def build_range_for(ctx, node): with ctx.variable_scope_guard(): loop_name = node.target.id ctx.check_loop_var(loop_name) loop_var = expr.Expr(_ti_core.make_id_expr('')) ctx.create_variable(loop_name, loop_var) if len(node.iter.args) not in [1, 2]: raise TaichiSyntaxError( f"Range should have 1 or 2 arguments, found {len(node.iter.args)}" ) if len(node.iter.args) == 2: begin = ti_ops.cast( expr.Expr(build_stmt(ctx, node.iter.args[0])), primitive_types.i32) end = ti_ops.cast( expr.Expr(build_stmt(ctx, node.iter.args[1])), primitive_types.i32) else: begin = ti_ops.cast(expr.Expr(0), primitive_types.i32) end = ti_ops.cast( expr.Expr(build_stmt(ctx, node.iter.args[0])), primitive_types.i32) _ti_core.begin_frontend_range_for(loop_var.ptr, begin.ptr, end.ptr) build_stmts(ctx, node.body) _ti_core.end_frontend_range_for() return None
def build_Assert(ctx, node): extra_args = [] if node.msg is not None: if ASTTransformer._is_string_mod_args(node.msg): msg, extra_args = ASTTransformer._handle_string_mod_args( ctx, node.msg) else: msg = build_stmt(ctx, node.msg) if isinstance(node.msg, ast.Constant): msg = str(msg) elif isinstance(node.msg, ast.Str): pass elif isinstance(msg, collections.abc.Sequence) and len( msg) > 0 and msg[0] == "__ti_format__": msg, extra_args = ASTTransformer.ti_format_list_to_assert_msg( msg) else: raise TaichiSyntaxError( f"assert info must be constant or formatted string, not {type(msg)}" ) else: msg = unparse(node.test) test = build_stmt(ctx, node.test) impl.ti_assert(test, msg.strip(), extra_args) return None
def build_mesh_for(ctx, node): targets = IRBuilder.get_for_loop_targets(node) if len(targets) != 1: raise TaichiSyntaxError( "Mesh for should have 1 loop target, found {len(targets)}") target = targets[0] with ctx.variable_scope_guard(): element_dict = { 'verts': ti.core.MeshElementType.Vertex, 'edges': ti.core.MeshElementType.Edge, 'faces': ti.core.MeshElementType.Face, 'cells': ti.core.MeshElementType.Cell } var = ti.Expr(ti.core.make_id_expr("")) ctx.mesh = node.iter.value.ptr assert isinstance(ctx.mesh, impl.MeshInstance) mesh_idx = ti.MeshElementFieldProxy(ctx.mesh, element_dict[node.iter.attr], var.ptr) ctx.create_variable(target, mesh_idx) ti.core.begin_frontend_mesh_for(mesh_idx.ptr, ctx.mesh.mesh_ptr, element_dict[node.iter.attr]) node.body = build_stmts(ctx, node.body) ctx.mesh = None ti.core.end_frontend_range_for() return node
def place( self, members, reorder=False, needs_grad=False, ): """Declares mesh attributes for the mesh element in current mesh builder. Args: members (Dict[str, Union[PrimitiveType, VectorType, MatrixType]]): \ names and types for element attributes. reorder: True if reorders the internal memory for coalesced data access within mesh-for loop. needs_grad: True if needs to record grad. Example:: >>> vec3 = ti.types.vector(3, ti.f32) >>> mesh = ti.TriMesh() >>> mesh.faces.place({'area' : ti.f32}) # declares a mesh attribute `area` for each face element. >>> mesh.verts.place({'pos' : vec3}, reorder=True) # declares a mesh attribute `pos` for each vertex element, and reorder it in memory. """ self.builder.elements.add(self._type) for key, dtype in members.items(): if key in {'verts', 'edges', 'faces', 'cells'}: raise TaichiSyntaxError( f"'{key}' cannot use as attribute name. It has been reserved as ti.Mesh's keyword." ) self.attr_dict[key] = MeshAttrType(key, dtype, reorder, needs_grad)
def element_wise_writeback_binary(self, foo, other): ret = self.empty_copy() if isinstance(other, (dict)): other = Struct(other) if is_taichi_class(other): other = other.variable() if foo.__name__ == 'assign' and not isinstance(other, Struct): raise TaichiSyntaxError( 'cannot assign scalar expr to ' f'taichi class {type(self)}, maybe you want to use `a.fill(b)` instead?' ) if isinstance(other, Struct): if self.entries.keys() != other.entries.keys(): raise TypeError( f"Member mismatch between structs {self.keys}, {other.keys}" ) for k, v in self.items: if isinstance(v, expr.Expr): ret.entries[k] = foo(v, other.entries[k]) else: ret.entries[k] = v.element_wise_binary( foo, other.entries[k]) else: # assumed to be scalar for k, v in self.items: if isinstance(v, expr.Expr): ret.entries[k] = foo(v, other) else: ret.entries[k] = v.element_wise_binary(foo, other) return ret
def build_nested_mesh_for(ctx, node): targets = ASTTransformer.get_for_loop_targets(node) if len(targets) != 1: raise TaichiSyntaxError( "Nested-mesh for should have 1 loop target, found {len(targets)}" ) target = targets[0] with ctx.variable_scope_guard(): ctx.mesh = node.iter.ptr.mesh assert isinstance(ctx.mesh, impl.MeshInstance) loop_name = node.target.id + '_index__' loop_var = expr.Expr(_ti_core.make_id_expr('')) ctx.create_variable(loop_name, loop_var) begin = expr.Expr(0) end = node.iter.ptr.size _ti_core.begin_frontend_range_for(loop_var.ptr, begin.ptr, end.ptr) entry_expr = _ti_core.get_relation_access( ctx.mesh.mesh_ptr, node.iter.ptr.from_index.ptr, node.iter.ptr.to_element_type, loop_var.ptr) entry_expr.type_check() mesh_idx = mesh.MeshElementFieldProxy( ctx.mesh, node.iter.ptr.to_element_type, entry_expr) ctx.create_variable(target, mesh_idx) build_stmts(ctx, node.body) _ti_core.end_frontend_range_for() return None
def maybe_transform_ti_func_call_to_stmt(ti_func, *args, **kwargs): _taichi_skip_traceback = 1 if '_sitebuiltins' == getattr(ti_func, '__module__', '') and getattr( getattr(ti_func, '__class__', ''), '__name__', '') == 'Quitter': raise TaichiSyntaxError(f'exit or quit not supported in Taichi-scope') if getattr(ti_func, '__module__', '') == '__main__' and not getattr(ti_func, '__wrapped__', ''): warnings.warn( f'Calling into non-Taichi function {ti_func.__name__}.' ' This means that scope inside that function will not be processed' ' by the Taichi transformer. Proceed with caution! ' ' Maybe you want to decorate it with @ti.func?', UserWarning, stacklevel=2) is_taichi_function = getattr(ti_func, '_is_taichi_function', False) # If is_taichi_function is true: call a decorated Taichi function # in a Taichi kernel/function. if is_taichi_function and get_runtime().experimental_real_function: # Compiles the function here. # Invokes Func.__call__. func_call_result = ti_func(*args, **kwargs) return _ti_core.insert_expr_stmt(func_call_result.ptr) else: return ti_func(*args, **kwargs)
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 visit_While(self, node): if node.orelse: raise TaichiSyntaxError( "'else' clause for 'while' not supported in Taichi kernels") with self.control_scope(): self.current_control_scope().append('while') template = ''' if 1: ti.core.begin_frontend_while(ti.Expr(1).ptr) __while_cond = 0 if __while_cond: pass else: break ti.core.pop_scope() ''' cond = node.test t = ast.parse(template).body[0] t.body[1].value = cond t.body = t.body[:3] + node.body + t.body[3:] self.generic_visit(t, ['body']) return ast.copy_location(t, node)
def wrapped(*args, **kwargs): # If we reach here (we should never), it means the class is not decorated # with @ti.data_oriented, otherwise getattr would have intercepted the call. clsobj = type(args[0]) assert not hasattr(clsobj, '_data_oriented') raise TaichiSyntaxError( f'Please decorate class {clsobj.__name__} with @ti.data_oriented' )
def wrapped(a, b): if is_taichi_class(a): return a._element_wise_writeback_binary(imp_foo, b) if is_taichi_class(b): raise TaichiSyntaxError( f'cannot augassign taichi class {type(b)} to scalar expr') else: return imp_foo(a, b)
def build_Compare(ctx, node): operands = build_exprs(ctx, [node.left] + list(node.comparators)) operators = [] for i in range(len(node.ops)): if isinstance(node.ops[i], ast.Lt): op_str = 'Lt' elif isinstance(node.ops[i], ast.LtE): op_str = 'LtE' elif isinstance(node.ops[i], ast.Gt): op_str = 'Gt' elif isinstance(node.ops[i], ast.GtE): op_str = 'GtE' elif isinstance(node.ops[i], ast.Eq): op_str = 'Eq' elif isinstance(node.ops[i], ast.NotEq): op_str = 'NotEq' elif isinstance(node.ops[i], ast.In): raise TaichiSyntaxError( '"in" is not supported in Taichi kernels.') elif isinstance(node.ops[i], ast.NotIn): raise TaichiSyntaxError( '"not in" is not supported in Taichi kernels.') elif isinstance(node.ops[i], ast.Is): raise TaichiSyntaxError( '"is" is not supported in Taichi kernels.') elif isinstance(node.ops[i], ast.IsNot): raise TaichiSyntaxError( '"is not" is not supported in Taichi kernels.') else: raise Exception(f'Unknown operator {node.ops[i]}') operators += [ ast.copy_location(ast.Str(s=op_str, kind=None), node) ] call = ast.Call( func=parse_expr('ti.chain_compare'), args=[ ast.copy_location(ast.List(elts=operands, ctx=ast.Load()), node), ast.copy_location(ast.List(elts=operators, ctx=ast.Load()), node) ], keywords=[]) call = ast.copy_location(call, node) return call
def taichi_ast_generator(): _taichi_skip_traceback = 1 if self.runtime.inside_kernel: raise TaichiSyntaxError( "Kernels cannot call other kernels. I.e., nested kernels are not allowed. Please check if you have direct/indirect invocation of kernels within kernels. Note that some methods provided by the Taichi standard library may invoke kernels, and please move their invocations to Python-scope." ) self.runtime.inside_kernel = True compiled() self.runtime.inside_kernel = False
def build_call_if_is_type(ctx, node, args, keywords): func = node.func.ptr if id(func) in primitive_types.type_ids: if len(args) != 1 or keywords or isinstance(args[0], expr.Expr): raise TaichiSyntaxError( "Type annotation can only be given to a single literal.") node.ptr = expr.Expr(args[0], dtype=func) return True return False
def build_Compare(ctx, node): build_stmt(ctx, node.left) build_stmts(ctx, node.comparators) ops = { ast.Eq: lambda l, r: l == r, ast.NotEq: lambda l, r: l != r, ast.Lt: lambda l, r: l < r, ast.LtE: lambda l, r: l <= r, ast.Gt: lambda l, r: l > r, ast.GtE: lambda l, r: l >= r, } ops_static = { ast.In: lambda l, r: l in r, ast.NotIn: lambda l, r: l not in r, ast.Is: lambda l, r: l is r, ast.IsNot: lambda l, r: l is not r, } if ctx.is_in_static_scope(): ops = {**ops, **ops_static} operands = [node.left.ptr ] + [comparator.ptr for comparator in node.comparators] val = True for i, node_op in enumerate(node.ops): l = operands[i] r = operands[i + 1] op = ops.get(type(node_op)) if isinstance(node_op, (ast.Is, ast.IsNot)): name = "is" if isinstance(node_op, ast.Is) else "is not" warnings.warn_explicit( f'Operator "{name}" in Taichi scope is deprecated. Please avoid using it.', DeprecationWarning, ctx.file, node.lineno + ctx.lineno_offset) if op is None: if type(node_op) in ops_static: raise TaichiSyntaxError( f'"{type(node_op).__name__}" is only supported inside `ti.static`.' ) else: raise TaichiSyntaxError( f'"{type(node_op).__name__}" is not supported in Taichi kernels.' ) val = ti_ops.logical_and(val, op(l, r)) node.ptr = val return node.ptr