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 bit_cast(obj, dtype): """Copy and cast a scalar to a specified data type with its underlying bits preserved. Must be called in taichi scope. This function is equivalent to `reinterpret_cast` in C++. Args: obj (:mod:`~taichi.types.primitive_types`): Input scalar. dtype (:mod:`~taichi.types.primitive_types`): Target data type, must have \ the same precision bits as the input (hence `f32` -> `f64` is not allowed). Returns: A copy of `obj`, casted to the specified data type `dtype`. Example:: >>> @ti.kernel >>> def test(): >>> x = 3.14 >>> y = ti.bit_cast(x, ti.i32) >>> print(y) # 1078523331 >>> >>> z = ti.bit_cast(y, ti.f32) >>> print(z) # 3.14 """ dtype = cook_dtype(dtype) if is_taichi_class(obj): raise ValueError('Cannot apply bit_cast on Taichi classes') else: return expr.Expr(_ti_core.bits_cast(expr.Expr(obj).ptr, dtype))
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_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 cast(obj, dtype): """Copy and cast a scalar or a matrix to a specified data type. Must be called in Taichi scope. Args: obj (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ Input scalar or matrix. dtype (:mod:`~taichi.types.primitive_types`): A primitive type defined in :mod:`~taichi.types.primitive_types`. Returns: A copy of `obj`, casted to the specified data type `dtype`. Example:: >>> @ti.kernel >>> def test(): >>> x = ti.Matrix([0, 1, 2], ti.i32) >>> y = ti.cast(x, ti.f32) >>> print(y) >>> >>> test() [0.0, 1.0, 2.0] """ dtype = cook_dtype(dtype) if is_taichi_class(obj): # TODO: unify with element_wise_unary return obj.cast(dtype) return expr.Expr(_ti_core.value_cast(expr.Expr(obj).ptr, dtype))
def _handle_string_mod_args(ctx, node): msg = build_stmt(ctx, node.left) args = build_stmt(ctx, node.right) if not isinstance(args, collections.abc.Sequence): args = (args, ) args = [expr.Expr(x).ptr for x in args] return msg, args
def shfl_up_f32(mask, val, offset): return expr.Expr( _ti_core.insert_internal_func_call( "cuda_shfl_up_sync_f32", # lane offset is 0 for warp size 32 expr.make_expr_group(mask, val, offset, 0), False))
def shfl_sync_i32(mask, val, offset): return expr.Expr( _ti_core.insert_internal_func_call( # lane offset is 31 for warp size 32 "cuda_shfl_sync_i32", expr.make_expr_group(mask, val, offset, 31), False))
def atomic_max(x, y): """Atomically compute the maximum of `x` and `y`, element-wise. Store the result in `x`, and return the old value of `x`. `x` must be a writable target, constant expressions or scalars are not allowed. Args: x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ The input. Returns: The old value of `x`. Example:: >>> @ti.kernel >>> def test(): >>> x = 1 >>> y = 2 >>> z = ti.atomic_max(x, y) >>> print(x) # 2 the new value of x >>> print(z) # 1, the old value of x >>> >>> ti.atomic_max(1, x) # will raise TaichiSyntaxError """ return impl.expr_init( expr.Expr(_ti_core.expr_atomic_max(x.ptr, y.ptr), tb=stack_info()))
def shfl_down_i32(mask, val, offset): # Here we use 31 as the last argument since 32 (warp size) does not work # for some reason. Using 31 leads to the desired behavior. return expr.Expr( _ti_core.insert_internal_func_call( "cuda_shfl_down_sync_i32", expr.make_expr_group(mask, val, offset, 31), False))
def atomic_xor(x, y): """Atomically compute the bit-wise XOR of `x` and `y`, element-wise. Store the result in `x`, and return the old value of `x`. `x` must be a writable target, constant expressions or scalars are not allowed. Args: x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ The input. When both are matrices they must have the same shape. Returns: The old value of `x`. Example:: >>> @ti.kernel >>> def test(): >>> x = ti.Vector([-1, 0, 1]) >>> y = ti.Vector([1, 2, 3]) >>> z = ti.atomic_xor(x, y) >>> print(x) # [-2, 2, 2] the new value of x >>> print(z) # [-1, 0, 1], the old value of x >>> >>> ti.atomic_xor(1, x) # will raise TaichiSyntaxError """ return impl.expr_init( expr.Expr(_ti_core.expr_atomic_bit_xor(x.ptr, y.ptr), tb=stack_info()))
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 atomic_sub(x, y): """Atomically subtract `x` by `y`, store the result in `x`, and return the old value of `x`. `x` must be a writable target, constant expressions or scalars are not allowed. Args: x, y (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): \ The input. Returns: The old value of `x`. Example:: >>> @ti.kernel >>> def test(): >>> x = ti.Vector([0, 0, 0]) >>> y = ti.Vector([1, 2, 3]) >>> z = ti.atomic_sub(x, y) >>> print(x) # [-1, -2, -3] the new value of x >>> print(z) # [0, 0, 0], the old value of x >>> >>> ti.atomic_sub(1, x) # will raise TaichiSyntaxError """ return impl.expr_init( expr.Expr(_ti_core.expr_atomic_sub(x.ptr, y.ptr), tb=stack_info()))
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 append(node, indices, val): """Append a value `val` to a SNode `node` at index `indices`. Args: node (:class:`~taichi.SNode`): Input SNode. indices (Union[int, :class:`~taichi.Vector`]): the indices to visit. val (:mod:`~taichi.types`): the data to be appended. """ a = impl.expr_init( _ti_core.insert_append(node._snode.ptr, expr.make_expr_group(indices), expr.Expr(val).ptr)) return a
def length(node, indices): """Return the length of the dynamic SNode `node` at index `indices`. Args: node (:class:`~taichi.SNode`): a dynamic SNode. indices (Union[int, :class:`~taichi.Vector`]): the indices to query. Returns: int: the length of cell `node[indices]`. """ return expr.Expr( _ti_core.insert_len(node._snode.ptr, expr.make_expr_group(indices)))
def random(dtype=float): """The random function. Args: dtype (DataType): Type of the random variable. Returns: A random variable whose type is `dtype`. """ dtype = cook_dtype(dtype) x = expr.Expr(_ti_core.make_rand_expr(dtype)) return impl.expr_init(x)
def get_addr(f, indices): """Query the memory address (on CUDA/x64) of field `f` at index `indices`. Currently, this function can only be called inside a taichi kernel. Args: f (Union[:class:`~taichi.Field`, :class:`~taichi.MatrixField`]): Input taichi field for memory address query. indices (Union[int, :class:`~taichi.Vector`]): The specified field indices of the query. Returns: ti.u64: The memory address of `f[indices]`. """ return expr.Expr( _ti_core.expr_get_addr(f._snode.ptr, expr.make_expr_group(indices)))
def is_active(node, indices): """Explicitly query whether a cell in a SNode `node` at location `indices` is active or not. Args: node (:class:`~taichi.SNode`): Must be a pointer, hash or bitmasked node. indices (Union[int, list, :class:`~taichi.Vector`]): the indices to visit. Returns: bool: the cell `node[indices]` is active or not. """ return expr.Expr( _ti_core.insert_is_active(node._snode.ptr, expr.make_expr_group(indices)))
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 build_While(ctx, node): if node.orelse: raise TaichiSyntaxError( "'else' clause for 'while' not supported in Taichi kernels") with ctx.loop_scope_guard(): _ti_core.begin_frontend_while(expr.Expr(1).ptr) while_cond = build_stmt(ctx, node.test) impl.begin_frontend_if(while_cond) _ti_core.begin_frontend_if_true() _ti_core.pop_scope() _ti_core.begin_frontend_if_false() _ti_core.insert_break_stmt() _ti_core.pop_scope() build_stmts(ctx, node.body) _ti_core.pop_scope() return None
def build_mesh_for(ctx, node): targets = ASTTransformer.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(): var = expr.Expr(_ti_core.make_id_expr("")) ctx.mesh = node.iter.ptr.mesh assert isinstance(ctx.mesh, impl.MeshInstance) mesh_idx = mesh.MeshElementFieldProxy(ctx.mesh, node.iter.ptr._type, var.ptr) ctx.create_variable(target, mesh_idx) _ti_core.begin_frontend_mesh_for(mesh_idx.ptr, ctx.mesh.mesh_ptr, node.iter.ptr._type) build_stmts(ctx, node.body) ctx.mesh = None _ti_core.end_frontend_range_for() return None
def build_Return(ctx, node): if not ctx.is_real_function: if ctx.is_in_non_static_control_flow(): raise TaichiSyntaxError( "Return inside non-static if/for is not supported") if node.value is not None: build_stmt(ctx, node.value) if node.value is None or node.value.ptr is None: if not ctx.is_real_function: ctx.returned = ReturnStatus.ReturnedVoid return None if ctx.is_kernel or ctx.is_real_function: # TODO: check if it's at the end of a kernel, throw TaichiSyntaxError if not 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') if id(ctx.func.return_type) in primitive_types.type_ids: ctx.ast_builder.create_kernel_exprgroup_return( expr.make_expr_group( ti_ops.cast(expr.Expr(node.value.ptr), ctx.func.return_type).ptr)) elif isinstance(ctx.func.return_type, MatrixType): ctx.ast_builder.create_kernel_exprgroup_return( expr.make_expr_group([ ti_ops.cast(exp, ctx.func.return_type.dtype) for exp in itertools.chain.from_iterable(node.value.ptr.to_list()) ])) else: raise TaichiSyntaxError( "The return type is not supported now!") # 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 ctx.is_real_function: ctx.returned = ReturnStatus.ReturnedValue return None
def build_struct_for(ctx, node, is_grouped): # for i, j in x # for I in ti.grouped(x) targets = ASTTransformer.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) loop_indices = expr.make_var_list(size=len(loop_var.shape), ast_builder=ctx.ast_builder) expr_group = expr.make_expr_group(loop_indices) impl.begin_frontend_struct_for(ctx.ast_builder, expr_group, loop_var) ctx.create_variable( target, matrix.Vector(loop_indices, dt=primitive_types.i32)) build_stmts(ctx, node.body) ctx.ast_builder.end_frontend_struct_for() else: _vars = [] for name in targets: var = expr.Expr(ctx.ast_builder.make_id_expr("")) _vars.append(var) ctx.create_variable(name, var) loop_var = node.iter.ptr expr_group = expr.make_expr_group(*_vars) impl.begin_frontend_struct_for(ctx.ast_builder, expr_group, loop_var) build_stmts(ctx, node.body) ctx.ast_builder.end_frontend_struct_for() return None
def random(dtype=float): """Return a single random float/integer according to the specified data type. Must be called in taichi scope. If the required `dtype` is float type, this function returns a random number sampled from the uniform distribution in the half-open interval [0, 1). For integer types this function returns a random integer in the half-open interval [0, 2^32) if a 32-bit integer is required, or a random integer in the half-open interval [0, 2^64) if a 64-bit integer is required. Args: dtype (:mod:`~taichi.types.primitive_types`): Type of the required random value. Returns: A random value with type `dtype`. Example:: >>> @ti.kernel >>> def test(): >>> x = ti.random(float) >>> print(x) # 0.090257 >>> >>> y = ti.random(ti.f64) >>> print(y) # 0.716101627301 >>> >>> i = ti.random(ti.i32) >>> print(i) # -963722261 >>> >>> j = ti.random(ti.i64) >>> print(j) # 73412986184350777 """ dtype = cook_dtype(dtype) x = expr.Expr(_ti_core.make_rand_expr(dtype)) return impl.expr_init(x)
def length(l, indices): return expr.Expr( _ti_core.insert_len(l.snode.ptr, expr.make_expr_group(indices)))
def is_active(l, indices): return expr.Expr( _ti_core.insert_is_active(l.snode.ptr, expr.make_expr_group(indices)))
def atomic_xor(a, b): return impl.expr_init( expr.Expr(_ti_core.expr_atomic_bit_xor(a.ptr, b.ptr), tb=stack_info()))
def append(l, indices, val): a = impl.expr_init( _ti_core.insert_append(l.snode.ptr, expr.make_expr_group(indices), expr.Expr(val).ptr)) return a
def expr_python_mod(a, b): # a % b = a - (a // b) * b quotient = expr.Expr(_ti_core.expr_floordiv(a, b)) multiply = expr.Expr(_ti_core.expr_mul(b, quotient.ptr)) return _ti_core.expr_sub(a, multiply.ptr)