def is_arch_supported(arch, use_gles=False): """Checks whether an arch is supported on the machine. Args: arch (taichi_core.Arch): Specified arch. use_gles (bool): If True, check is GLES is available otherwise check if GLSL is available. Only effective when `arch` is `ti.opengl`. Default is `False`. Returns: bool: Whether `arch` is supported on the machine. """ arch_table = { cuda: _ti_core.with_cuda, metal: _ti_core.with_metal, opengl: functools.partial(_ti_core.with_opengl, use_gles), cc: _ti_core.with_cc, vulkan: _ti_core.with_vulkan, dx11: _ti_core.with_dx11, wasm: lambda: True, cpu: lambda: True, } with_arch = arch_table.get(arch, lambda: False) try: return with_arch() except Exception as e: arch = _ti_core.arch_name(arch) _ti_core.warn( f"{e.__class__.__name__}: '{e}' occurred when detecting " f"{arch}, consider adding `TI_ENABLE_{arch.upper()}=0` " f" to environment variables to suppress this warning message.") return False
def _check_not_turned_on_with_warning_message(self): if self._profiling_mode is False: _ti_core.warn( 'use \'ti.init(kernel_profiler = True)\' to turn on KernelProfiler.' ) return True return False
def get_predefined_cupti_metrics(name=''): if name not in predefined_cupti_metrics: _ti_core.warn("Valid Taichi predefined metrics list (str):") for key in predefined_cupti_metrics: _ti_core.warn(f" '{key}'") return None return predefined_cupti_metrics[name]
def set_toolkit(self, toolkit_name='default'): if self._check_not_turned_on_with_warning_message(): return False status = impl.get_runtime().prog.set_kernel_profiler_toolkit( toolkit_name) if status is True: self._profiling_toolkit = toolkit_name else: _ti_core.warn( f'Failed to set kernel profiler toolkit ({toolkit_name}) , keep using ({self._profiling_toolkit}).' ) return status
def get_predefined_cupti_metrics(name=''): """Returns the specified cupti metric. Accepted arguments are 'global_access', 'shared_access', 'atomic_access', 'cache_hit_rate', 'device_utilization'. Args: name (str): cupti metri name. """ if name not in predefined_cupti_metrics: _ti_core.warn("Valid Taichi predefined metrics list (str):") for key in predefined_cupti_metrics: _ti_core.warn(f" '{key}'") return None return predefined_cupti_metrics[name]
def add(self, key, _cast=None): _cast = _cast or self.bool_int self.keys.append(key) # TI_OFFLINE_CACHE= : no effect # TI_OFFLINE_CACHE=0 : False # TI_OFFLINE_CACHE=1 : True name = 'TI_' + key.upper() value = os.environ.get(name, '') if key in self.kwargs: self[key] = self.kwargs[key] if value: _ti_core.warn( f'Environment variable {name}={value} overridden by ti.init argument "{key}"' ) del self.kwargs[key] # mark as recognized elif value: self[key] = _cast(value)
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': try_check_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) ti.reset() spec_cfg = _SpecialConfig() env_comp = _EnvironmentConfigurator(kwargs, ti.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 = ti.f32 elif env_default_fp == '64': default_fp = ti.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 = ti.i32 elif env_default_ip == '64': default_ip = ti.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('print_preprocessed') env_spec.add('log_level', str) env_spec.add('gdb_trigger') env_spec.add('excepthook') env_spec.add('experimental_real_function') env_spec.add('short_circuit_operators') # compiler configurations (ti.cfg): for key in dir(ti.cfg): if key in ['arch', 'default_fp', 'default_ip']: continue _cast = type(getattr(ti.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.set_gdb_trigger(spec_cfg.gdb_trigger) impl.get_runtime().print_preprocessed = spec_cfg.print_preprocessed impl.get_runtime().experimental_real_function = \ spec_cfg.experimental_real_function impl.get_runtime().short_circuit_operators = \ spec_cfg.short_circuit_operators ti.set_logging_level(spec_cfg.log_level.lower()) if spec_cfg.excepthook: # TODO(#1405): add a way to restore old excepthook ti.enable_excepthook() # select arch (backend): env_arch = os.environ.get('TI_ARCH') if env_arch is not None: ti.info(f'Following TI_ARCH setting up for arch={env_arch}') arch = _ti_core.arch_from_name(env_arch) ti.cfg.arch = adaptive_arch_select(arch, enable_fallback, ti.cfg.use_gles) if ti.cfg.arch == cc: _ti_core.set_tmp_dir(locale_encode(prepare_sandbox())) print(f'[Taichi] Starting on arch={_ti_core.arch_name(ti.cfg.arch)}') # Torch based ndarray on opengl backend allocates memory on host instead of opengl backend. # So it won't work. if ti.cfg.arch == opengl and ti.cfg.ndarray_use_torch: ti.warn( 'Opengl backend doesn\'t support torch based ndarray. Setting ndarray_use_torch to False.' ) ti.cfg.ndarray_use_torch = False if _test_mode: return spec_cfg get_default_kernel_profiler().set_kernel_profiler_mode( ti.cfg.kernel_profiler) # create a new program: impl.get_runtime().create_program() ti.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