Skip to content

7.9。模板注释

原文： http://numba.pydata.org/numba-doc/latest/developer/stencil.html

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

7.9.1。模板装饰

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

7.9.2。处理三种模式

如上所述，Numba 支持从@jit编译函数内部或外部调用模板，有或没有 parallel = True 选项。

7.9.2.1。外部 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的调用在给定结果和参数类型的情况下生成模板函数，最后执行生成的模板函数并返回其结果。

7.9.2.2。没有parallel=True 的 Jit

构造时，StencilFunc将自身插入到键入上下文的用户函数集中，并提供_type_me回调。通过这种方式，标准 Numba 编译器能够确定StencilFunc的输出类型和签名。每个StencilFunc都维护一个先前看到的输入参数类型和关键字类型组合的缓存。如果先前看到，StencilFunc返回计算的签名。如果以前没有计算过，StencilFunc通过在模板内核上运行 Numba 编译器前端，然后在 Numba IR （IR）上执行类型推断来获得标量返回类型，从而计算模板的返回类型内核。从那里，构造一个 Numpy 数组类型，其元素类型与标量返回类型匹配。

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

因此，在这种模式下，生成的模板函数是一个独立的函数，称为 jitted 代码中的正常函数。

7.9.2.3。 Jit 与parallel=True

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

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

7.9.3。创建模板功能

从概念上讲，通过在内核周围添加循环代码，根据循环索引将相对内核索引转换为绝对数组索引，并使用要分配的语句替换内核的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中，以便对同一模板的其他调用不需要再次执行此过程。

7.9.4。例外情况

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

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

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

回到顶部