def materialize(self, key=None, args=None, arg_features=None):
if key is None:
key = (self.func, 0, self.autodiff_mode)
self.runtime.materialize()
if key in self.runtime.compiled_functions:
return
grad_suffix = ""
if self.autodiff_mode == AutodiffMode.FORWARD:
grad_suffix = "_forward_grad"
elif self.autodiff_mode == AutodiffMode.REVERSE:
grad_suffix = "_reverse_grad"
kernel_name = f"{self.func.__name__}_c{self.kernel_counter}_{key[1]}{grad_suffix}"
_logging.trace(f"Compiling kernel {kernel_name}...")
tree, ctx = _get_tree_and_ctx(
self,
args=args,
excluded_parameters=self.template_slot_locations,
arg_features=arg_features)
if self.autodiff_mode != AutodiffMode.NONE:
KernelSimplicityASTChecker(self.func).visit(tree)
if impl.current_cfg().use_mesh:
taichi.lang.Mesh.update_relation(tree, ctx)
# Do not change the name of 'taichi_ast_generator'
# The warning system needs this identifier to remove unnecessary messages
def taichi_ast_generator(kernel_cxx):
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
self.runtime.current_kernel = self
try:
ctx.ast_builder = kernel_cxx.ast_builder()
transform_tree(tree, ctx)
if not ctx.is_real_function:
if self.return_type and ctx.returned != ReturnStatus.ReturnedValue:
raise TaichiSyntaxError(
"Kernel has a return type but does not have a return statement"
)
finally:
self.runtime.inside_kernel = False
self.runtime.current_kernel = None
taichi_kernel = impl.get_runtime().prog.create_kernel(
taichi_ast_generator, kernel_name, self.autodiff_mode)
self.kernel_cpp = taichi_kernel
assert key not in self.runtime.compiled_functions
self.runtime.compiled_functions[key] = self.get_function_body(
taichi_kernel)
self.compiled_kernels[key] = taichi_kernel
def __init__(self, arr, indices_first, indices_second):
self.ndarr = arr
self.arr = arr.arr
if arr.layout == Layout.SOA:
self.indices = indices_second + indices_first
else:
self.indices = indices_first + indices_second
if impl.get_runtime().ndarray_use_torch:
def getter():
return self.arr[self.indices]
def setter(value):
self.arr[self.indices] = value
else:
def getter():
self.ndarr.initialize_host_accessor()
return self.ndarr.host_accessor.getter(
*self.ndarr.pad_key(self.indices))
def setter(value):
self.ndarr.initialize_host_accessor()
self.ndarr.host_accessor.setter(
value, *self.ndarr.pad_key(self.indices))
self.getter = getter
self.setter = setter
def eig(A, dt=None):
"""Compute the eigenvalues and right eigenvectors of a real matrix.
Parameters
----------
A: ti.Matrix(n, n)
2D Matrix for which the eigenvalues and right eigenvectors will be computed.
dt: Optional[DataType]
The datatype for the eigenvalues and right eigenvectors
Returns
-------
eigenvalues: ti.Matrix(n, 2)
The eigenvalues in complex form. Each row stores one eigenvalue. The first number
of the eigenvalue represents the real part and the second number represents the
imaginary part.
eigenvectors: ti.Matrix(n*2, n)
The eigenvectors in complex form. Each column stores one eigenvector. Each eigenvector
consists of n entries, each of which is represented by two numbers for its real part
and imaginary part.
"""
if dt is None:
dt = impl.get_runtime().default_fp
from taichi.lang import linalg
if A.n == 2:
return linalg.eig2x2(A, dt)
raise Exception("Eigen solver only supports 2D matrices.")
def fixed(frac, signed=True, range=1.0, compute=None, scale=None): # pylint: disable=W0622
"""Generates a quantized type for fixed-point real numbers.
Args:
frac (int): Number of bits.
signed (bool): Signed or unsigned.
range (float): Range of the number.
compute (DataType): Type for computation.
scale (float): Scaling factor. The argument is prioritized over range.
Returns:
DataType: The specified type.
"""
if compute is None:
compute = impl.get_runtime().default_fp
if isinstance(compute, _ti_core.DataType):
compute = compute.get_ptr()
# TODO: handle cases with frac > 32
frac_type = int(bits=frac, signed=signed, compute=i32)
if scale is None:
if signed:
scale = range / 2**(frac - 1)
else:
scale = range / 2**frac
return _type_factory.get_quant_fixed_type(frac_type, compute, scale)
def _loop_range(self):
"""Gets the taichi_python.Expr wrapping the taichi_python.GlobalVariableExpression corresponding to `self` to serve as loop range.
Returns:
taichi_python.Expr: See above.
"""
return impl.get_runtime().prog.global_var_expr_from_snode(self.ptr)
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(ctx.ast_builder.make_id_expr(''))
ctx.create_variable(loop_name, loop_var)
begin = expr.Expr(0)
end = node.iter.ptr.size
ctx.ast_builder.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(impl.get_runtime().prog.config)
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)
ctx.ast_builder.end_frontend_range_for()
return None
def visit(self, tree, *arguments):
if impl.get_runtime().experimental_ast_refactor:
self.print_ast(tree, 'Initial AST')
self.rename_module.visit(tree)
self.print_ast(tree, 'AST with module renamed')
ctx = IRBuilderContext(
func=self.func,
excluded_parameters=self.excluded_parameters,
is_kernel=self.is_kernel,
arg_features=self.arg_features,
globals=self.globals,
argument_data=arguments)
# Convert Python AST to Python code that generates Taichi C++ AST.
tree = IRBuilder()(ctx, tree)
ast.fix_missing_locations(tree)
self.print_ast(tree, 'Preprocessed')
self.pass_checks.visit(tree) # does not modify the AST
return ctx.return_data
self.print_ast(tree, 'Initial AST')
self.rename_module.visit(tree)
self.print_ast(tree, 'AST with module renamed')
ctx = BuilderContext(func=self.func,
excluded_parameters=self.excluded_parameters,
is_kernel=self.is_kernel,
arg_features=self.arg_features)
# Convert Python AST to Python code that generates Taichi C++ AST.
tree = build_stmt(ctx, tree)
ast.fix_missing_locations(tree)
self.print_ast(tree, 'Preprocessed')
self.pass_checks.visit(tree) # does not modify the AST
def solve(A, b, dt=None):
"""Solve a matrix using Gauss elimination method.
Args:
A (ti.Matrix(n, n)): input nxn matrix `A`.
b (ti.Vector(n, 1)): input nx1 vector `b`.
dt (DataType): The datatype for the `A` and `b`.
Returns:
x (ti.Vector(n, 1)): the solution of Ax=b.
"""
assert A.n == A.m, "Only square matrix is supported"
assert A.n >= 2 and A.n <= 3, "Only 2D and 3D matrices are supported"
assert A.m == b.n, "Matrix and Vector dimension dismatch"
if dt is None:
dt = impl.get_runtime().default_fp
nrow, ncol = static(A.n, A.n + 1)
Ab = expr_init(Matrix.zero(dt, nrow, ncol))
lhs = tuple([e.ptr for e in A.entries])
rhs = tuple([e.ptr for e in b.entries])
for i in range(nrow):
for j in range(nrow):
Ab(i, j)._assign(lhs[nrow * i + j])
for i in range(nrow):
Ab(i, nrow)._assign(rhs[i])
if A.n == 2:
return _gauss_elimination_2x2(Ab, dt)
if A.n == 3:
return _gauss_elimination_3x3(Ab, dt)
raise Exception("Solver only supports 2D and 3D matrices.")
def __init__(self, loss=None, clear_gradients=True):
"""A context manager for reverse mode autodiff :class:`~taichi.ad.Tape`. The
context manager would catching all of the callings of functions that
decorated by :func:`~taichi.lang.kernel_impl.kernel` or
:func:`~taichi.ad.grad_replaced` under `with` statement, and calculate
all the partial gradients of a given loss variable by calling all of the
gradient function of the callings caught in reverse order while `with`
statement ended.
See also :func:`~taichi.lang.kernel_impl.kernel` and
:func:`~taichi.ad.grad_replaced` for gradient functions.
Args:
loss(:class:`~taichi.lang.expr.Expr`): The loss field, which shape should be ().
clear_gradients(Bool): Before `with` body start, clear all gradients or not.
Example::
>>> @ti.kernel
>>> def sum(a: ti.float32):
>>> for I in ti.grouped(x):
>>> y[None] += x[I] ** a
>>>
>>> with ti.Tape(loss = y):
>>> sum(2)
"""
self.calls = []
self.entered = False
self.gradient_evaluated = False
self.clear_gradients = clear_gradients
self.runtime = impl.get_runtime()
self.eval_on_exit = loss is not None
self.loss = loss
def __init__(self, name, res, vsync=False, show_window=True):
package_path = str(pathlib.Path(__file__).parent.parent)
ti_arch = default_cfg().arch
is_packed = default_cfg().packed
super().__init__(get_runtime().prog, name, res, vsync, show_window,
package_path, ti_arch, is_packed)
def block_local(*args):
"""Hints Taichi to cache the fields and to enable the BLS optimization.
Please visit https://docs.taichi-lang.org/docs/performance
for how BLS is used.
Args:
*args (List[Field]): A list of sparse Taichi fields.
"""
if impl.current_cfg().opt_level == 0:
_logging.warn("""opt_level = 1 is enforced to enable bls analysis.""")
impl.current_cfg().opt_level = 1
for a in args:
for v in a._get_field_members():
get_runtime().prog.current_ast_builder().insert_snode_access_flag(
_ti_core.SNodeAccessFlag.block_local, v.ptr)
def print_ast(tree, title=None):
if not impl.get_runtime().print_preprocessed:
return
if title is not None:
print(f'{title}:')
import astor
print(astor.to_source(tree.body[0], indent_with=' '))
def __init__(self, *args, tb=None, dtype=None):
self.tb = tb
if len(args) == 1:
if isinstance(args[0], _ti_core.Expr):
self.ptr = args[0]
elif isinstance(args[0], Expr):
self.ptr = args[0].ptr
self.tb = args[0].tb
elif is_taichi_class(args[0]):
raise TaichiTypeError(
'Cannot initialize scalar expression from '
f'taichi class: {type(args[0])}')
else:
# assume to be constant
arg = args[0]
if isinstance(arg, np.ndarray):
if arg.shape:
raise TaichiTypeError(
"Only 0-dimensional numpy array can be used to initialize a scalar expression"
)
arg = arg.dtype.type(arg)
self.ptr = make_constant_expr(arg, dtype).ptr
else:
assert False
if self.tb:
self.ptr.set_tb(self.tb)
self.ptr.type_check(impl.get_runtime().prog.config)
def Tape(loss, clear_gradients=True):
impl.get_runtime().materialize()
if len(loss.shape) != 0:
raise RuntimeError(
'The loss of `Tape` must be a 0-D field, i.e. scalar')
if not loss.snode.ptr.has_grad():
raise RuntimeError(
'Gradients of loss are not allocated, please use ti.field(..., needs_grad=True)'
' for all fields that are required by autodiff.')
if clear_gradients:
clear_all_gradients()
from taichi.lang.meta import clear_loss
clear_loss(loss)
return runtime.get_tape(loss)
def decl_sparse_matrix(dtype):
value_type = cook_dtype(dtype)
ptr_type = cook_dtype(u64)
# Treat the sparse matrix argument as a scalar since we only need to pass in the base pointer
arg_id = impl.get_runtime().prog.decl_arg(ptr_type, False)
return SparseMatrixProxy(
_ti_core.make_arg_load_expr(arg_id, ptr_type, False), value_type)
def _make_table_header(self, mode):
header_str = f'Kernel Profiler({mode}, {self._profiling_toolkit})'
arch_name = f' @ {_ti_core.arch_name(impl.current_cfg().arch).upper()}'
device_name = impl.get_runtime().prog.get_kernel_profiler_device_name()
if len(device_name) > 1: # default device_name = ' '
device_name = ' on ' + device_name
return header_str + arch_name + device_name
def __init__(self, dtype, arr_shape):
super().__init__()
self.dtype = cook_dtype(dtype)
self.arr = impl.get_runtime().prog.create_ndarray(
self.dtype, arr_shape)
self.shape = tuple(self.arr.shape)
self.element_type = dtype
def _make_table_header(mode):
header_str = f'Kernel Profiler({mode})'
arch_name = f' @ {_ti_core.arch_name(ti.cfg.arch).upper()}'
device_name = impl.get_runtime().prog.get_kernel_profiler_device_name()
if len(device_name) > 1: # default device_name = ' '
device_name = ' on ' + device_name
return header_str + arch_name + device_name
def __init__(self, dtype, arr_shape):
self.host_accessor = None
self.dtype = cook_dtype(dtype)
self.ndarray_use_torch = impl.get_runtime().ndarray_use_torch
if self.ndarray_use_torch:
assert has_pytorch(
), "PyTorch must be available if you want to create a Taichi ndarray with PyTorch as its underlying storage."
# pylint: disable=E1101
self.arr = torch.zeros(arr_shape,
dtype=to_pytorch_type(cook_dtype(dtype)))
if impl.current_cfg().arch == _ti_core.Arch.cuda:
self.arr = self.arr.cuda()
else:
self.arr = _ti_core.Ndarray(impl.get_runtime().prog,
cook_dtype(dtype), arr_shape)
def isnan(x):
"""Determines whether the parameter is a number, element-wise.
Args:
x (:mod:`~taichi.types.primitive_types`, :class:`taichi.Matrix`): The input.
Example:
>>> @ti.kernel
>>> def test():
>>> x = vec4(nan, -nan, inf, 1)
>>> ti.math.isnan(x)
>>>
>>> test()
[1, 1, 0, 0]
Returns:
For each element i of the result, returns 1 if x[i] is posititve or negative floating point NaN (Not a Number) and 0 otherwise.
"""
ftype = impl.get_runtime().default_fp
fx = ti.cast(x, ftype)
if ti.static(ftype == ti.f64):
y = ti.bit_cast(fx, ti.u64)
return (ti.cast(y >> 32, ti.u32)
& 0x7fffffff) + (ti.cast(y, ti.u32) != 0) > 0x7ff00000
y = ti.bit_cast(fx, ti.u32)
return (y & 0x7fffffff) > 0x7f800000
def mesh_patch_idx():
"""Returns the internal mesh patch id of this running thread,
only available for backends supporting `ti.extension.mesh` and to use within mesh-for loop.
Related to https://github.com/taichi-dev/taichi/issues/3608
"""
return impl.get_runtime().prog.current_ast_builder().insert_patch_idx_expr(
)
def make_var_list(size, ast_builder=None):
exprs = []
for _ in range(size):
if ast_builder is None:
exprs.append(impl.get_runtime().prog.make_id_expr(''))
else:
exprs.append(ast_builder.make_id_expr(''))
return exprs
def decl_scalar_arg(dtype):
is_ref = False
if isinstance(dtype, RefType):
is_ref = True
dtype = dtype.tp
dtype = cook_dtype(dtype)
arg_id = impl.get_runtime().prog.decl_arg(dtype, False)
return Expr(_ti_core.make_arg_load_expr(arg_id, dtype, is_ref))
def test_default_fp_ndarray(dtype):
arch = ti.lang.impl.current_cfg().arch
ti.reset()
ti.init(arch=arch, default_fp=dtype)
x = ti.Vector.ndarray(2, float, ())
assert x.dtype == impl.get_runtime().default_fp
def to_numpy(self):
"""Converts ndarray to a numpy array.
Returns:
numpy.ndarray: The result numpy array.
"""
if impl.current_cfg().ndarray_use_torch:
return self.arr.cpu().numpy()
else:
import numpy as np # pylint: disable=C0415
arr = np.zeros(shape=self.arr.shape,
dtype=to_numpy_type(self.dtype))
from taichi.lang.meta import \
ndarray_to_ext_arr # pylint: disable=C0415
ndarray_to_ext_arr(self, arr)
impl.get_runtime().sync()
return arr
def subscript(self, *indices):
assert len(indices) == 1
entry_expr = self.mesh.get_relation_access(self.from_index,
self.to_element_type,
impl.Expr(indices[0]).ptr)
entry_expr.type_check(impl.get_runtime().prog.config)
return MeshElementFieldProxy(self.mesh, self.to_element_type,
entry_expr)
def query_info(self, name):
"""For docsting of this function, see :func:`~taichi.lang.query_kernel_profile_info`."""
if self._check_not_turned_on_with_warning_message():
return None
self._update_records() # kernel records
self._count_statistics() # statistics results
# TODO : query self.StatisticalResult in python scope
return impl.get_runtime().prog.query_kernel_profile_info(name)
def __init__(self, name, res, vsync=False, show_window=True):
check_ggui_availability()
package_path = str(pathlib.Path(__file__).parent.parent)
ti_arch = default_cfg().arch
is_packed = default_cfg().packed
self.window = _ti_core.PyWindow(get_runtime().prog, name, res, vsync,
show_window, package_path, ti_arch,
is_packed)
def shape(self):
"""Gets the number of elements from root in each axis of `self`.
Returns:
Tuple[int]: The number of elements from root in each axis of `self`.
"""
impl.get_runtime().materialize()
dim = self.ptr.num_active_indices()
ret = [self.ptr.get_shape_along_axis(i) for i in range(dim)]
class callable_tuple(tuple):
@deprecated('x.shape()', 'x.shape')
def __call__(self):
return self
ret = callable_tuple(ret)
return ret
def reset():
"""Resets Taichi to its initial state.
This would destroy all the fields and kernels.
"""
impl.reset()
global runtime
runtime = impl.get_runtime()
