在AITemplate项目中添加自定义算子的完整指南

在AITemplate项目中添加自定义算子的完整指南

AITemplate AITemplate is a Python framework which renders neural network into high performance CUDA/HIP C++ code. Specialized for FP16 TensorCore (NVIDIA GPU) and MatrixCore (AMD GPU) inference. 项目地址: https://gitcode.com/gh_mirrors/ai/AITemplate

AITemplate是一个高效的深度学习模型编译框架，它允许开发者通过定义算子图来生成高性能的GPU代码。本文将详细介绍如何在AITemplate中添加一个新的自定义算子，从算子定义到代码生成再到验证的完整流程。

准备工作

在开始添加算子之前，我们需要导入一些必要的Python模块：

from typing import Any, Dict, List
import jinja2
import torch
from aitemplate import backend
from aitemplate.backend import registry
from aitemplate.backend.backend_spec import CUDASpec, ROCMSpec
from aitemplate.compiler import compile_model
from aitemplate.compiler.base import IntVar, Operator, Tensor
from aitemplate.testing import detect_target

这些模块提供了我们需要的核心功能，包括模板渲染、算子定义、模型编译等。

定义算子图节点

在AITemplate中，每个算子都是一个继承自Operator基类的Python类。下面我们定义一个简单的add_one算子，它会对输入张量的每个元素加1。

class add_one(Operator):
    def __init__(self):
        super().__init__()
        self._attrs["op"] = "add_one"  # 算子唯一标识
        self._attrs["has_profiler"] = False
        self._attrs["nop"] = False

    def __call__(self, x: Tensor) -> Tensor:
        self._attrs["inputs"] = [x]  # 记录输入张量
        self._set_depth()  # 设置算子深度
        output_shape = self._infer_shape(x)  # 推断输出形状
        output = Tensor(output_shape, src_ops={self})  # 创建输出张量
        self._attrs["outputs"] = [output]  # 记录输出张量
        return output

    def _infer_shape(self, x) -> List[IntVar]:
        return x.shape()  # 输出形状与输入相同

    def gen_function(self) -> str:
        target = backend.target.Target.current()
        func_key = f"{target.name()}.{self._attrs['op']}.gen_function"
        func = registry.get(func_key)  # 从注册表获取对应的代码生成函数
        return func(self._attrs)

关键点说明：

1. _attrs字典是算子最重要的数据结构，存储了算子的各种属性和元信息
2. _attrs["op"]必须是唯一的，它用于在后端查找对应的代码生成函数
3. _infer_shape方法负责推断输出张量的形状

定义代码生成模板

AITemplate使用Jinja2模板来生成C++代码。我们需要定义四种核心模板：

FUNC_TEMPLATE = jinja2.Template("""
{{header_files}}
namespace {
{{kernel}}
}  // namespace
{{func_signature}}
{
    invoke_add_one(output, input, num_elements, stream);
}
""")

FUNC_SIGNATURE = jinja2.Template("""
void {{func_name}}(half* output,
                  const half* input,
                  const int64_t num_elements,
                  {{prefix}}Stream_t stream)
""")

FUNC_DECL = jinja2.Template("""
{{func_signature}};
""")

FUNC_CALL_TEMPLATE = jinja2.Template("""
{{indent}}int64_t num_elements = 1;
{% for dim_name in dim_names %}
{{indent}}num_elements *= {{dim_name}};
{% endfor %}
{{indent}}{{func_name}}(
{{indent}}   {{output}}, {{input}}, num_elements, stream /* default stream */
{{indent}});
""")

这些模板分别用于：

1. 生成完整的算子函数实现
2. 定义函数签名
3. 声明函数原型
4. 生成函数调用代码

编写GPU内核

接下来我们需要编写实际的GPU计算内核。这里我们实现一个简单的逐元素加1操作：

KERNEL_TEMPLATE = jinja2.Template("""
__global__ void add_one(half* output, const half* input, const int64_t num_elements) {
  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < num_elements) {
    output[idx] = input[idx] + half(1.0);
  }
}
void invoke_add_one(half* output, const half* input, int64_t num_elements, {{prefix}}Stream_t stream) {
  if (num_elements < 1024) {
    dim3 grid(1);
    dim3 block(num_elements);
    add_one<<<grid, block, 0, stream>>>(output, input, num_elements);
  } else {
    dim3 grid((num_elements + 1024 - 1) / 1024);
    dim3 block(1024);
    add_one<<<grid, block, 0, stream>>>(output, input, num_elements);
  }
}
""")

这个内核实现了自适应线程块分配，根据输入大小自动选择最优的线程块配置。

实现代码生成函数

代码生成函数负责将模板渲染为实际的C++代码：

def gen_function_call(func_attrs: Dict[str, Any], indent="  ", is_cuda=False) -> str:
    output_name = FUNC_CALL_FP16_PARAM_TEMPLATE.render(
        name=func_attrs["outputs"][0]._attrs["name"], is_cuda=is_cuda
    )
    input_name = FUNC_CALL_FP16_PARAM_TEMPLATE.render(
        name=func_attrs["inputs"][0]._attrs["name"], is_cuda=is_cuda
    )
    dim_names = [dim._attrs["name"] for dim in func_attrs["inputs"][0].shape()]
    return FUNC_CALL_TEMPLATE.render(
          func_name=func_attrs["name"],
          output=output_name,
          input=input_name,
          dim_names=dim_names,
          indent=indent,
    )

def gen_function(func_attrs: Dict[str, Any], header_files: str, backend_spec) -> str:
    prefix = backend_spec.prefix
    return FUNC_TEMPLATE.render(
          header_files=header_files,
          kernel=KERNEL_TEMPLATE.render(prefix=prefix),
          func_signature=FUNC_SIGNATURE.render(
              func_name=func_attrs["name"], prefix=prefix
          ),
    )

def gen_function_decl(func_attrs: Dict[str, Any], backend_spec) -> str:
    return FUNC_DECL.render(
          func_signature=FUNC_SIGNATURE.render(
              func_name=func_attrs["name"],
              prefix=backend_spec.prefix,
          ).strip()
    )

注册后端实现

最后，我们需要将代码生成函数注册到对应的后端（CUDA和ROCm）：

# CUDA后端注册
CUDA_HEADER_FILES = """
#include <cuda_fp16.h>
"""

@registry.reg("cuda.add_one.gen_function")
def cuda_add_one_gen_function(func_attrs: Dict[str, Any]) -> str:
    return gen_function(func_attrs, CUDA_HEADER_FILES, CUDASpec())

@registry.reg("cuda.add_one.func_decl")
def cuda_add_one_gen_function_decl(func_attrs: Dict[str, Any]) -> str:
    return gen_function_decl(func_attrs, CUDASpec())

@registry.reg("cuda.add_one.func_call")
def cuda_add_one_gen_function_call(func_attrs: Dict[str, Any], indent="  ") -> str:
    return gen_function_call(func_attrs, indent, is_cuda=True)

# ROCm后端注册（可选）
HIP_HEADER_FILES = """
#include <hip/hip_fp16.h>
#include <hip/hip_runtime.h>
"""

@registry.reg("rocm.add_one.gen_function")
def rocm_add_one_gen_function(func_attrs: Dict[str, Any]) -> str:
    return gen_function(func_attrs, HIP_HEADER_FILES, ROCMSpec())

@registry.reg("rocm.add_one.func_decl")
def rocm_add_one_gen_function_decl(func_attrs: Dict[str, Any]) -> str:
    return gen_function_decl(func_attrs, ROCMSpec())

@registry.reg("rocm.add_one.func_call")
def rocm_add_one_gen_function_call(func_attrs: Dict[str, Any], indent="  ") -> str:
    return gen_function_call(func_attrs, indent, is_cuda=False)

验证算子功能

完成算子实现后，我们需要验证它的正确性：

def create_ait_model(shapes):
    X = Tensor(
          shape=shapes,
          dtype="float16",
          name="X",
          is_input=True,
    )
    Y = add_one()(X)
    Y._attrs["is_output"] = True
    Y._attrs["name"] = "Y"
    return Y

def verify_add_one():
    shapes = [16, 512]
    x = torch.randn(shapes).cuda().half()
    y_pt = x + 1.0  # PyTorch参考实现

    Y = create_ait_model([16, 512])
    target = detect_target()
    with compile_model(Y, target, "./tmp", "add_one") as module:
        y = torch.empty(shapes).cuda().half()
        inputs = {"X": x}
        outputs = {"Y": y}
        module.run_with_tensors(inputs, outputs)
        print(torch.allclose(y, y_pt, atol=1e-2, rtol=1e-2))  # 比较结果

总结

通过本文，我们详细介绍了在AITemplate中添加自定义算子的完整流程：

1. 定义算子图节点类
2. 实现形状推断逻辑
3. 编写Jinja2模板
4. 实现GPU计算内核
5. 编写代码生成函数
6. 注册到对应后端
7. 验证算子正确性

这个过程展示了AITemplate强大的可扩展性，开发者可以根据需要添加各种自定义算子，充分利用GPU的计算能力。

AITemplate AITemplate is a Python framework which renders neural network into high performance CUDA/HIP C++ code. Specialized for FP16 TensorCore (NVIDIA GPU) and MatrixCore (AMD GPU) inference. 项目地址: https://gitcode.com/gh_mirrors/ai/AITemplate