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