关于模板的说明

Numba 提供了 @stencil 装饰器来表示模板计算。本文档解释了此功能在 Numba 提供的几种不同模式中是如何实现的。目前，支持从非 JIT 编译代码调用模板，也支持从 JIT 编译代码调用，无论是否带有 parallel=True 选项。

模板装饰器

模板装饰器本身只返回一个 StencilFunc 对象。此对象封装了程序中指定的原始模板内核函数以及传递给模板装饰器的选项。另外值得注意的是，在模板的第一次编译之后，计算出的模板邻域会存储在 StencilFunc 对象的 neighborhood 属性中。

处理三种模式

如上所述，Numba 支持在 @jit 编译函数内部或外部调用模板，无论是否带有 parallel=True 选项。

JIT 上下文之外

StencilFunc 覆盖了 __call__ 方法，因此对 StencilFunc 对象的调用会执行模板。

def __call__(self, *args, **kwargs):
    result = kwargs.get('out')

    new_stencil_func = self._stencil_wrapper(result, None, *args)

    if result is None:
        return new_stencil_func.entry_point(*args)
    else:
        return new_stencil_func.entry_point(*args, result)

首先，检查可选的 out 参数是否存在。如果存在，则输出数组存储在 result 中。然后，对 _stencil_wrapper 的调用根据结果和参数类型生成模板函数，最后执行生成的模板函数并返回其结果。

不带 `parallel=True` 的 JIT 模式

构建时，StencilFunc 会将自身插入到类型化上下文的用户函数集中，并提供 _type_me 回调。通过这种方式，标准 Numba 编译器能够确定 StencilFunc 的输出类型和签名。每个 StencilFunc 维护一个先前见过的输入参数类型和关键字类型组合的缓存。如果以前见过，StencilFunc 返回计算出的签名。如果以前未计算过，StencilFunc 通过对模板内核运行 Numba 编译器前端，然后对 Numba IR (IR) 执行类型推断以获取内核的标量返回类型来计算模板的返回类型。由此，构造一个元素类型与该标量返回类型匹配的 Numpy 数组类型。

在为先前未见的输入和关键字类型组合计算模板签名后，StencilFunc 随后创建模板函数本身。StencilFunc 然后将新的模板函数的定义安装到目标上下文中，以便 JIT 编译的代码能够调用它。

因此，在此模式下，生成的模板函数是一个独立的函数，可以在 JIT 编译代码中像普通函数一样被调用。

带 `parallel=True` 的 JIT 模式

当从带有 parallel=True 的 JIT 编译上下文调用 StencilFunc 时，不会使用由创建模板函数生成的单独模板函数。相反，会在当前函数中创建 parfors（阶段 5b：执行自动并行化），该函数实现模板。此代码再次从模板内核开始，并执行类似的内核大小计算，但随后不是使用标准的 Python 循环语法，而是创建相应的 parfors，以便模板的执行将并行进行。

模板到 parfor 的转换也可以通过设置 parallel={'stencil': False} 来选择性禁用，以及阶段 5b：执行自动并行化中描述的其他子选项。

创建模板函数

从概念上讲，模板函数是通过在用户指定的模板内核周围添加循环代码，将相对内核索引转换为基于循环索引的绝对数组索引，并将内核的 return 语句替换为将计算值赋给输出数组的语句来创建的。

为了完成此转换，首先创建模板内核 IR 的副本，以便对不同模板签名的后续 IR 修改不会相互影响。

然后，采用类似于为 parfors 创建 GUFunc 的方法。在文本缓冲区中，创建一个具有唯一名称的 Python 函数。将输入数组参数添加到函数定义中，如果 out 参数类型存在，则将 out 参数添加到模板函数定义中。如果 out 参数不存在，则首先使用 numpy.zeros 创建一个与输入数组形状相同的输出数组。

然后分析内核以计算模板大小和边界的形状（如果存在，则使用 neighborhood 模板装饰器参数）。然后，为输入数组的每个维度向模板函数定义添加一个 for 循环。每个循环的范围由先前计算的模板内核大小控制，以便输出图像的边界不会被修改，而是保持原样。最内层 for 循环的主体是一个单独的 sentinel 语句，该语句在 IR 中很容易识别。使用带文本缓冲区的 exec 调用强制生成模板函数，并使用 eval 访问相应的函数，然后使用 run_frontend 获取模板函数 IR。

对模板函数 IR 和内核 IR 执行各种重命名和重新标记，以便两者可以无冲突地组合。内核 IR 中的相对索引（即 getitem 调用）被替换为表达式，其中相应的循环索引变量添加到相对索引中。内核 IR 中的 return 语句被替换为输出数组中相应元素的 setitem。然后扫描模板函数 IR 以查找哨兵，并将哨兵替换为修改后的内核 IR。

接下来，使用 compile_ir 编译组合的模板函数 IR。生成的编译结果会缓存在 StencilFunc 中，以便对同一模板的其他调用无需再次执行此过程。

引发的异常

在模板编译期间会执行各种检查，以确保用户指定的选项彼此不冲突或与其他运行时参数不冲突。例如，如果用户已手动为模板装饰器指定 neighborhood，则该邻域的长度必须与输入数组的维度匹配。如果情况并非如此，则会引发 ValueError。

如果未指定邻域，则必须推断它，并且推断内核的要求是所有索引都是常量整数。如果不是，则会引发 ValueError，指示内核索引不能为非常量。

最后，模板实现通过对模板内核运行 Numba 类型推断来检测输出数组类型。如果此内核的返回类型与传递给 cval 模板装饰器选项的值的类型不匹配，则会引发 ValueError。