关于模板的使用说明

Numba 提供了 @stencil 装饰器来表示模板计算。本文档解释了此功能在 Numba 中几种不同模式下的实现方式。目前，从非 jitted 代码调用 stencil 以及从 jitted 代码调用 stencil 都是支持的，无论是否使用 parallel=True 选项。

模板装饰器

stencil装饰器本身只返回一个 StencilFunc 对象。该对象封装了程序中指定的原始stencil内核函数以及传递给stencil装饰器的选项。值得注意的是，在stencil第一次编译后，计算出的stencil邻域存储在 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 随后在目标上下文中安装新模板函数的定义，以便即时编译的代码能够调用它。

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

使用 `parallel=True` 进行 Jit

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

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。