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 materialize(self, key=None, args=None, arg_features=None):
if key is None:
key = (self.func, 0)
self.runtime.materialize()
if key in self.compiled_functions:
return
grad_suffix = ""
if self.is_grad:
grad_suffix = "_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.is_grad:
KernelSimplicityASTChecker(self.func).visit(tree)
# Do not change the name of 'taichi_ast_generator'
# The warning system needs this identifier to remove unnecessary messages
def taichi_ast_generator():
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:
transform_tree(tree, ctx)
if not impl.get_runtime().experimental_real_function:
if self.return_type and not ctx.returned:
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 = _ti_core.create_kernel(taichi_ast_generator,
kernel_name, self.is_grad)
self.kernel_cpp = taichi_kernel
assert key not in self.compiled_functions
self.compiled_functions[key] = self.get_function_body(taichi_kernel)
def init(arch=None,
default_fp=None,
default_ip=None,
_test_mode=False,
enable_fallback=True,
require_version=None,
**kwargs):
"""Initializes the Taichi runtime.
This should always be the entry point of your Taichi program. Most
importantly, it sets the backend used throughout the program.
Args:
arch: Backend to use. This is usually :const:`~taichi.lang.cpu` or :const:`~taichi.lang.gpu`.
default_fp (Optional[type]): Default floating-point type.
default_ip (Optional[type]): Default integral type.
require_version (Optional[string]): A version string.
**kwargs: Taichi provides highly customizable compilation through
``kwargs``, which allows for fine grained control of Taichi compiler
behavior. Below we list some of the most frequently used ones. For a
complete list, please check out
https://github.com/taichi-dev/taichi/blob/master/taichi/program/compile_config.h.
* ``cpu_max_num_threads`` (int): Sets the number of threads used by the CPU thread pool.
* ``debug`` (bool): Enables the debug mode, under which Taichi does a few more things like boundary checks.
* ``print_ir`` (bool): Prints the CHI IR of the Taichi kernels.
* ``packed`` (bool): Enables the packed memory layout. See https://docs.taichi-lang.org/docs/layout.
"""
# Check version for users every 7 days if not disabled by users.
_version_check.start_version_check_thread()
# FIXME(https://github.com/taichi-dev/taichi/issues/4811): save the current working directory since it may be
# changed by the Vulkan backend initialization on OS X.
current_dir = os.getcwd()
cfg = impl.default_cfg()
# Check if installed version meets the requirements.
if require_version is not None:
check_require_version(require_version)
# Make a deepcopy in case these args reference to items from ti.cfg, which are
# actually references. If no copy is made and the args are indeed references,
# ti.reset() could override the args to their default values.
default_fp = _deepcopy(default_fp)
default_ip = _deepcopy(default_ip)
kwargs = _deepcopy(kwargs)
reset()
spec_cfg = _SpecialConfig()
env_comp = _EnvironmentConfigurator(kwargs, cfg)
env_spec = _EnvironmentConfigurator(kwargs, spec_cfg)
# configure default_fp/ip:
# TODO: move these stuff to _SpecialConfig too:
env_default_fp = os.environ.get("TI_DEFAULT_FP")
if env_default_fp:
if default_fp is not None:
_ti_core.warn(
f'ti.init argument "default_fp" overridden by environment variable TI_DEFAULT_FP={env_default_fp}'
)
if env_default_fp == '32':
default_fp = f32
elif env_default_fp == '64':
default_fp = f64
elif env_default_fp is not None:
raise ValueError(
f'Invalid TI_DEFAULT_FP={env_default_fp}, should be 32 or 64')
env_default_ip = os.environ.get("TI_DEFAULT_IP")
if env_default_ip:
if default_ip is not None:
_ti_core.warn(
f'ti.init argument "default_ip" overridden by environment variable TI_DEFAULT_IP={env_default_ip}'
)
if env_default_ip == '32':
default_ip = i32
elif env_default_ip == '64':
default_ip = i64
elif env_default_ip is not None:
raise ValueError(
f'Invalid TI_DEFAULT_IP={env_default_ip}, should be 32 or 64')
if default_fp is not None:
impl.get_runtime().set_default_fp(default_fp)
if default_ip is not None:
impl.get_runtime().set_default_ip(default_ip)
# submodule configurations (spec_cfg):
env_spec.add('log_level', str)
env_spec.add('gdb_trigger')
env_spec.add('short_circuit_operators')
# compiler configurations (ti.cfg):
for key in dir(cfg):
if key in ['arch', 'default_fp', 'default_ip']:
continue
_cast = type(getattr(cfg, key))
if _cast is bool:
_cast = None
env_comp.add(key, _cast)
unexpected_keys = kwargs.keys()
if len(unexpected_keys):
raise KeyError(
f'Unrecognized keyword argument(s) for ti.init: {", ".join(unexpected_keys)}'
)
# dispatch configurations that are not in ti.cfg:
if not _test_mode:
_ti_core.set_core_trigger_gdb_when_crash(spec_cfg.gdb_trigger)
impl.get_runtime().short_circuit_operators = \
spec_cfg.short_circuit_operators
_logging.set_logging_level(spec_cfg.log_level.lower())
# select arch (backend):
env_arch = os.environ.get('TI_ARCH')
if env_arch is not None:
_logging.info(f'Following TI_ARCH setting up for arch={env_arch}')
arch = _ti_core.arch_from_name(env_arch)
cfg.arch = adaptive_arch_select(arch, enable_fallback, cfg.use_gles)
if cfg.arch == cc:
_ti_core.set_tmp_dir(locale_encode(prepare_sandbox()))
print(f'[Taichi] Starting on arch={_ti_core.arch_name(cfg.arch)}')
# user selected visible device
visible_device = os.environ.get("TI_VISIBLE_DEVICE")
if visible_device and (cfg.arch == vulkan or _ti_core.GGUI_AVAILABLE):
_ti_core.set_vulkan_visible_device(visible_device)
if _test_mode:
return spec_cfg
get_default_kernel_profiler().set_kernel_profiler_mode(cfg.kernel_profiler)
# create a new program:
impl.get_runtime().create_program()
_logging.trace('Materializing runtime...')
impl.get_runtime().prog.materialize_runtime()
impl._root_fb = _snode.FieldsBuilder()
if not os.environ.get("TI_DISABLE_SIGNAL_HANDLERS", False):
impl.get_runtime()._register_signal_handlers()
# Recover the current working directory (https://github.com/taichi-dev/taichi/issues/4811)
os.chdir(current_dir)
return None
def init(arch=None,
default_fp=None,
default_ip=None,
_test_mode=False,
enable_fallback=True,
**kwargs):
"""Initializes the Taichi runtime.
This should always be the entry point of your Taichi program. Most
importantly, it sets the backend used throughout the program.
Args:
arch: Backend to use. This is usually :const:`~taichi.lang.cpu` or :const:`~taichi.lang.gpu`.
default_fp (Optional[type]): Default floating-point type.
default_ip (Optional[type]): Default integral type.
**kwargs: Taichi provides highly customizable compilation through
``kwargs``, which allows for fine grained control of Taichi compiler
behavior. Below we list some of the most frequently used ones. For a
complete list, please check out
https://github.com/taichi-dev/taichi/blob/master/taichi/program/compile_config.h.
* ``cpu_max_num_threads`` (int): Sets the number of threads used by the CPU thread pool.
* ``debug`` (bool): Enables the debug mode, under which Taichi does a few more things like boundary checks.
* ``print_ir`` (bool): Prints the CHI IR of the Taichi kernels.
* ``packed`` (bool): Enables the packed memory layout. See https://docs.taichi.graphics/lang/articles/advanced/layout.
"""
# Check version for users every 7 days if not disabled by users.
skip = os.environ.get("TI_SKIP_VERSION_CHECK")
if skip != 'ON':
# We don't join this thread because we do not wish to block users.
check_version_thread = threading.Thread(target=try_check_version,
daemon=True)
check_version_thread.start()
# Make a deepcopy in case these args reference to items from ti.cfg, which are
# actually references. If no copy is made and the args are indeed references,
# ti.reset() could override the args to their default values.
default_fp = _deepcopy(default_fp)
default_ip = _deepcopy(default_ip)
kwargs = _deepcopy(kwargs)
reset()
spec_cfg = _SpecialConfig()
env_comp = _EnvironmentConfigurator(kwargs, cfg)
env_spec = _EnvironmentConfigurator(kwargs, spec_cfg)
# configure default_fp/ip:
# TODO: move these stuff to _SpecialConfig too:
env_default_fp = os.environ.get("TI_DEFAULT_FP")
if env_default_fp:
if default_fp is not None:
_ti_core.warn(
f'ti.init argument "default_fp" overridden by environment variable TI_DEFAULT_FP={env_default_fp}'
)
if env_default_fp == '32':
default_fp = f32
elif env_default_fp == '64':
default_fp = f64
elif env_default_fp is not None:
raise ValueError(
f'Invalid TI_DEFAULT_FP={env_default_fp}, should be 32 or 64')
env_default_ip = os.environ.get("TI_DEFAULT_IP")
if env_default_ip:
if default_ip is not None:
_ti_core.warn(
f'ti.init argument "default_ip" overridden by environment variable TI_DEFAULT_IP={env_default_ip}'
)
if env_default_ip == '32':
default_ip = i32
elif env_default_ip == '64':
default_ip = i64
elif env_default_ip is not None:
raise ValueError(
f'Invalid TI_DEFAULT_IP={env_default_ip}, should be 32 or 64')
if default_fp is not None:
impl.get_runtime().set_default_fp(default_fp)
if default_ip is not None:
impl.get_runtime().set_default_ip(default_ip)
# submodule configurations (spec_cfg):
env_spec.add('log_level', str)
env_spec.add('gdb_trigger')
env_spec.add('experimental_real_function')
env_spec.add('short_circuit_operators')
env_spec.add('ndarray_use_torch')
# compiler configurations (ti.cfg):
for key in dir(cfg):
if key in ['arch', 'default_fp', 'default_ip']:
continue
_cast = type(getattr(cfg, key))
if _cast is bool:
_cast = None
env_comp.add(key, _cast)
unexpected_keys = kwargs.keys()
if len(unexpected_keys):
raise KeyError(
f'Unrecognized keyword argument(s) for ti.init: {", ".join(unexpected_keys)}'
)
# dispatch configurations that are not in ti.cfg:
if not _test_mode:
set_gdb_trigger(spec_cfg.gdb_trigger)
impl.get_runtime().experimental_real_function = \
spec_cfg.experimental_real_function
impl.get_runtime().short_circuit_operators = \
spec_cfg.short_circuit_operators
impl.get_runtime().ndarray_use_torch = \
spec_cfg.ndarray_use_torch
_logging.set_logging_level(spec_cfg.log_level.lower())
# select arch (backend):
env_arch = os.environ.get('TI_ARCH')
if env_arch is not None:
_logging.info(f'Following TI_ARCH setting up for arch={env_arch}')
arch = _ti_core.arch_from_name(env_arch)
cfg.arch = adaptive_arch_select(arch, enable_fallback, cfg.use_gles)
if cfg.arch == cc:
_ti_core.set_tmp_dir(locale_encode(prepare_sandbox()))
print(f'[Taichi] Starting on arch={_ti_core.arch_name(cfg.arch)}')
# Torch based ndarray on opengl backend allocates memory on host instead of opengl backend.
# So it won't work.
if cfg.arch == opengl and spec_cfg.ndarray_use_torch:
_logging.warn(
'Opengl backend doesn\'t support torch based ndarray. Setting ndarray_use_torch to False.'
)
impl.get_runtime().ndarray_use_torch = False
if _test_mode:
return spec_cfg
get_default_kernel_profiler().set_kernel_profiler_mode(cfg.kernel_profiler)
# create a new program:
impl.get_runtime().create_program()
_logging.trace('Materializing runtime...')
impl.get_runtime().prog.materialize_runtime()
impl._root_fb = FieldsBuilder()
if not os.environ.get("TI_DISABLE_SIGNAL_HANDLERS", False):
impl.get_runtime()._register_signal_handlers()
return None