def is_arch_supported(arch): """Checks whether an arch is supported on the machine. Args: arch (taichi_core.Arch): Specified arch. Returns: bool: Whether `arch` is supported on the machine. """ arch_table = { cuda: _ti_core.with_cuda, metal: _ti_core.with_metal, opengl: _ti_core.with_opengl, cc: _ti_core.with_cc, vulkan: lambda: _ti_core.with_vulkan(), 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 is_arch_supported(arch): arch_table = { cuda: _ti_core.with_cuda, metal: _ti_core.with_metal, opengl: _ti_core.with_opengl, cc: _ti_core.with_cc, 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 add `export TI_WITH_{arch.upper()}=0` " f" to environment variables to depress this warning message.") return False
def add(self, key, cast=None): cast = cast or self.bool_int self.keys.append(key) # TI_ASYNC= : no effect # TI_ASYNC=0 : False # TI_ASYNC=1 : True name = 'TI_' + key.upper() value = os.environ.get(name, '') if len(value): self[key] = cast(value) if key in self.kwargs: _ti_core.warn( f'ti.init argument "{key}" overridden by environment variable {name}={value}' ) del self.kwargs[key] # mark as recognized elif key in self.kwargs: self[key] = self.kwargs[key] del self.kwargs[key] # mark as recognized
def init(arch=None, default_fp=None, default_ip=None, _test_mode=False, **kwargs): # 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') # 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 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) print(f'[Taichi] Starting on arch={_ti_core.arch_name(ti.cfg.arch)}') if _test_mode: return spec_cfg # create a new program: impl.get_runtime().create_program() impl._root_fb = FieldsBuilder()
def init(arch=None, default_fp=None, default_ip=None, _test_mode=False, **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_fp (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. """ # 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('experimental_ast_refactor') # 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 'use_unified_memory' in unexpected_keys: _ti_core.warn( f'"use_unified_memory" is a deprecated option, as taichi no longer have the option of using unified memory.' ) del kwargs['use_unified_memory'] 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( ).experimental_ast_refactor = spec_cfg.experimental_ast_refactor 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) 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)}') 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()