（三）Python部分生命周期之Kernel_python itk setkernel-优快云博客

本文链接：https://blog.youkuaiyun.com/weixin_44153058/article/details/128909508

文章详细介绍了Taichi编程框架中@kernel装饰器的工作原理，包括如何判断kernel是否在类中定义，以及Kernel类的初始化和编译过程。在初始化Kernel时，获取函数参数并进行类型检查，然后在第一次调用时进行编译。编译涉及AST树的生成和C++内核的创建。此外，文章还提到了自动微分模式和多线程中的GIL管理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1. kernel装饰器

之后我们看一下@ti.kernel装饰器

def _kernel_impl(_func, level_of_class_stackframe, verbose=False):
    is_classkernel = _inside_class(level_of_class_stackframe + 1)
    if verbose:
        print(f'kernel={
     _func.__name__} is_classkernel={
     is_classkernel}')
    primal = Kernel(_func,
                    autodiff_mode=AutodiffMode.NONE,
                    _classkernel=is_classkernel)
    adjoint = Kernel(_func,
                     autodiff_mode=AutodiffMode.REVERSE,
                     _classkernel=is_classkernel)
    primal.grad = adjoint
    if is_classkernel:
        @functools.wraps(_func)
        def wrapped(*args, **kwargs):
            clsobj = type(args[0])
            assert not hasattr(clsobj, '_data_oriented')
            raise TaichiSyntaxError(
                f'Please decorate class {
     clsobj.__name__} with @ti.data_oriented'
            )
    else:
        @functools.wraps(_func)
        def wrapped(*args, **kwargs):
            try:
                return primal(*args, **kwargs)
            except (TaichiCompilationError, TaichiRuntimeError) as e:
                raise type(e)('\n' + str(e)) from None
        wrapped.grad = adjoint
    wrapped._is_wrapped_kernel = True
    wrapped._is_classkernel = is_classkernel
    wrapped._primal = primal
    wrapped._adjoint = adjoint
    return wrapped

def kernel(fn):
    return _kernel_impl(fn, level_of_class_stackframe=3)

kernel是一个装饰器函数，其内部调用了_kernel_impl作为具体实现。这里首先判断了当前的kernel是否在一个类中，是的，taichi允许你在一个类中定义kernel，这给了开发者极大的便利，同时taichi用了非常巧妙的方式来处理在类中的Kernel，这里我们暂时不讨论，我先关注于在类外面的kernel函数。而_inside_class方法具体是如何判断是否在类中的，则用了我们的老朋友，inspect模块来获取解释器堆栈上的语句信息，通过正则匹配检测class关键字实现的。

接下来实例化了两个Kernel类，这个Kernel类是taichi程序的编译核心，也是执行入口，具体编译过程在Kernel的__call__魔法函数中。

之后是一个内部闭包wrapped，functools.wrap装饰器的作用是不改变调用对象，单纯改变所装饰对象的一些属性，也就是这里最后返回的是wrapped函数，但是属性则依然是传入的play函数，所以我们打印play函数的__name__和__doc__属性就会发现没有发生改变，但是具体调用逻辑则是调用Kernel的__call__函数。

那我们接下来就很好看一下Kernel这个类：

class Kernel:
    counter = 0

    def __init__(self, _func, autodiff_mode, _classkernel=False):
        self.func = _func
        self.kernel_counter = Kernel.counter
        Kernel.counter += 1
        assert autodiff_mode in (AutodiffMode.NONE, AutodiffMode.VALIDATION,
                                 AutodiffMode.FORWARD, AutodiffMode.REVERSE)
        self.autodiff_mode = autodiff_mode
        self.grad = None
        self.arguments = []
        self.return_type