（二）Taich Python部分生命周期之初始化

首先需要说明一下，此处虽说是讲Python部分的生命周期，但在初始化时taichi仍然调用了一部分pybind封装的c++代码，其中大部分都是用于初始化配置，我们也会顺带讲解一下，我仍认为这一部分属于Python部分的生命周期，等真正进入JIT编译时，才算是进入到C++部分的生命周期。

1 示例代码

import taichi as ti

ti.init(arch=ti.cpu)

x = ti.field(ti.i32, shape=(3, 4))
y = ti.ndarray(int, shape=16)
z = int(100)


@ti.kernel
def play(arg0: int, arg1: int):
    tmp = 0

    if arg1:
        tmp = tmp + 4

    for i in range(10):
        tmp += 1

    print(tmp)

    assert z == 100
    print(z)
    
    for i in x:
        x[i] = 0
    
    for i in y:
        ti.atomic_add(y[i], 5)


play(10, True)

下面我们以上面那段代码为例，taichi程序会分为两个部分，一部分是属于Python scope中的，而另一部分被装饰器所修饰的代码是属于taichi scope中的，可以看到taichi scope中的也会调用到Python全局中的field变量，这一部分我们暂时不考虑，taichi中的全局field变量会通过特殊处理链接到taichi程序中去，我们分开两部分来看，我们首先先看不含这些全局变量的taichi程序的生命周期是如何的，即我们这次调试的代码为下面的情况：

import taichi as ti

ti.init(arch=ti.cpu)

# x = ti.field(ti.i32, shape=(3, 4))
# y = ti.ndarray(int, shape=16)
# z = int(100)


@ti.kernel
def play(arg0: int, arg1: int):
    tmp = 0

    if arg1:
        tmp = tmp + 4

    for i in range(10):
        tmp += 1

    print(tmp)

    # assert z == 100
    # print(z)
    #
    # for i in x:
    #     x[i] = 0
    #
    # for i in y:
    #     ti.atomic_add(y[i], 5)


play(10, True)

2 Tachi Init

可以看到要想加载taichi程序，我们会首先需要初始化taichi包，我这里使用的cpu作为后端进行运行。

def init(arch=None,
         default_fp=None,
         default_ip=None,
         _test_mode=False,
         enable_fallback=True,
         require_version=None,
         **kwargs):
    # Check version for users every 7 days if not disabled by users.
    _version_check.start_version_check_thread()

在运行init函数之前，taichi包中有很多全局变量，下面一个是会伴随我们整个生命周期的一个Python全局变量PyTaichi，他在taichi.lang包下的impl模块中被定义

class PyTaichi:
    def __init__(self, kernels=None):
        self.materialized = False
        self.prog = None
        self.compiled_functions = {}
        self.src_info_stack = []
        self.inside_kernel = False
        self.current_kernel = None
        self.global_vars = []
        self.grad_vars = []
        self.dual_vars = []
        self.matrix_fields = []
        self.default_fp = f32
        self.default_ip = i32
        self.default_up = u32
        self.target_tape = None
        self.fwd_mode_manager = None
        self.grad_replaced = False
        self.kernels = kernels or []
        self._signal_handler_registry = None

PyTaichi的初始化只是单纯的对成员函数赋初值。

而我们之前调用的init函数是一大串初始化配置的代码，在开始时会进行版本的检查

def start_version_check_thread():
    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()

此处开启多线程进行版本检查，实际运行函数为try_check_version

def try_check_version():
    try:
        os.makedirs(_ti_core.get_repo_dir(), exist_ok=True)
        version_info_path = os.path.join(_ti_core.get_repo_dir(),
                                         'version_info')
        cur_date = datetime.date.today()
        if os.path.exists(version_info_path):
            with open(version_info_path, 'r') as f:
                version_info_file = f.readlines()
                last_time = version_info_file[0].rstrip()
                cur_uuid = version_info_file[2].rstrip()
            if cur_date.strftime('%Y-%m-%d') > last_time:
                response = check_version(cur_uuid)
                write_version_info(response, cur_uuid, version_info_path,
                                   cur_date)
        else:
            cur_uuid = str(uuid.uuid4())
            write_version_info({'status': 0}, cur_uuid, version_info_path,
                               cur_date)
            response = check_version(cur_uuid)
            write_version_info(response, cur_uuid, version_info_path, cur_date)
    # Wildcard exception to catch potential file writing errors.
    except:
        pass

在这个函数中，第一次运行taichi程序时会创建了一个taichi_cache文件夹用于放置临时缓存，windows默认在C盘根目录。同样在第一次运行时会在taichi_cache文件夹中创建一个version_info文件，这个文件内容是一传版本字符串，由时期+uuid组成，如果当前日期大于最后一次更新日期，需要对版本进行更新，日期判断只精确到天。

之后回到init函数中。

current_dir = os.getcwd()

if require_version is not None:
    check_require_version(require_version)

if "packed" in kwargs:
    if kwargs["packed"] is True:
        warnings.warn(
            "Currently packed=True is the default setting and the switch will be removed in v1.4.0.",
            DeprecationWarning)
    else:
        warnings.warn(
            "The automatic padding mode (packed=False) will no longer exist in v1.4.0. The switch will "
            "also be removed then. Make sure your code doesn't rely on it.",
            DeprecationWarning)

if "default_up" in kwargs:
    raise KeyError(
        "'default_up' is always the unsigned type of 'default_ip'. Please set 'default_ip' instead."
    )
default_fp = deepcopy(default_fp)
default_ip = deepcopy(default_ip)
kwargs = deepcopy(kwargs)

之后获取到了当前运行的py文件所在的文件夹，对参数中进行了检测，如果存在不合法会进行警告和报错，对合法参数进行深拷贝复制。

def reset():
    global pytaichi
    old_kernels = pytaichi.kernels
    pytaichi.clear()
    pytaichi = PyTaichi(old_kernels)
    for k in old_kernels:
        k.reset()

之后进行reset操作，reset的具体实现在Impl中，清空当前全局的pytaichi变量中的kernel集合，并对每一个kernel进行复位。

cfg = impl.default_cfg()
cfg.offline_cache = True  # Enable offline cache in frontend instead of C++ side

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 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)

接下来就是进行配置设置，上面代码全为配置设置，获取的Config类是由C++ Pybind绑定的，如下：

struct CompileConfig {
  Arch arch;
  bool debug;
  bool cfg_optimization;
  bool check_out_of_bound;
  bool validate_autodiff;
  int simd_width;
  int opt_level;
  int external_optimization_level;
  int max_vector_width;
  bool packed;
  bool print_preprocessed_ir;
	 ......

  CompileConfig();
};

默认的配置都是对这些成员变量进行初始化，接下来的是一些特殊环境配置，我初始化时并没有进行设置，这些部分都是默认值，大部分的控制流都不会进入，之后是加入日志等级和gdb等信息。

# 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()

逐级遍历Config中的属性，将属性名作为键值，类型作为value存到env.comp中。

unexpected_keys = kwargs.keys()

if len(unexpected_keys):
    raise KeyError(
        f'Unrecognized keyword argument(s) for ti.init: {", ".join(unexpected_keys)}'
    )

之后是查看是否有不需要的参数，有则报错

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 cfg.debug:
    impl.get_runtime()._register_signal_handlers()

os.chdir(current_dir)
return None

接下来就是初始化一些核心类用于创建和编译Kernel。这里首先会创建一个Program，这个Program可以视为一个整体的Taichi程序，会存放全部的filed等全局变量和所有的kernel。Program是一个c++类，通过Pybind绑定成Python对象，Program本身并不提供任何实现，具体实现有其子类提供，默认使用的是LLVM版本的Program，这段代码可以在taichi项目中的c++源码文件夹taichi下面的program文件夹中的program.cpp中找到：

Program::Program(Arch desired_arch) : snode_rw_accessors_bank_(this) {
  TI_TRACE("Program initializing...");
  .......
  main_thread_id_ = std::this_thread::get_id();
  ......

  profiler = make_profiler(config.arch, config.kernel_profiler);
  if (arch_uses_llvm(config.arch)) {
#ifdef TI_WITH_LLVM
    if (config.arch != Arch::dx12) {
      program_impl_ = std::make_unique<LlvmProgramImpl>(config, profiler.get());
    } else {
      // NOTE: use Dx12ProgramImpl to avoid using LlvmRuntimeExecutor for dx12.
#ifdef TI_WITH_DX12
      TI_ASSERT(directx12::is_dx12_api_available());
      program_impl_ = std::make_unique<Dx12ProgramImpl>(config);
#else
      TI_ERROR("This taichi is not compiled with DX12");
#endif
    }
#else
    TI_ERROR("This taichi is not compiled with LLVM");
#endif
  } else if (config.arch == Arch::metal) {
#ifdef TI_WITH_METAL
    TI_ASSERT(metal::is_metal_api_available());
    program_impl_ = std::make_unique<MetalProgramImpl>(config);
#else
    TI_ERROR("This taichi is not compiled with Metal")
#endif

我们下面的也都以默认的LLVM版本的Program实现进行讲解，后续有时间我会讲解其他backend实现。

让我们回到Python的init函数中去，在创建完Program后。调用了materialize_runtime方法来创建运行环境，我们具体来看一下LLVMProgramImpl的实现：

void materialize_runtime(MemoryPool *memory_pool,
                         KernelProfilerBase *profiler,
                         uint64 **result_buffer_ptr) override {
  runtime_exec_->materialize_runtime(memory_pool, profiler,
                                     result_buffer_ptr);
}

额，又包了一层，这里调用的是其执行器的初始化方法，继续往下走可以看到这里是对jit进行一些配置。到此为止，超长的Init函数就看完了。