趋近智
尽管机器学习 (machine learning)编译器擅长优化流行框架中常见的标准操作,但应用和研究经常需要一些不被原生支持的操作。这些操作可能包括新颖的激活函数 (activation function)、专门的数据预处理步骤、利用独特硬件功能的操作,或在主编译器流程之外开发的高度优化核函数。高级机器学习运行时系统必须提供机制来整合这些定制算子和核函数。
整合定制算子变得必要,原因有以下几点:
整合定制算子通常需要编译器前端、编译器优化遍和运行时系统之间的配合。
整合并执行定制算子的流程。编译器保留定制算子节点并生成调用运行时的代码,运行时随后查找并调用已注册的核函数实现。
主要方面包括:
MyCustomAttention、SpecialPreprocessing)。这必须与图表示中使用的标识符匹配。一个简化的注册 API 可能如下所示(C++):
// 核函数的前向声明
Status my_custom_op_cpu_kernel(KernelContext* context);
// 注册函数(通常在库加载时调用)
void register_ops() {
OpRegistry* registry = Runtime::GetGlobalOpRegistry();
registry->Register("MyCustomOp")
.Device(DeviceType::CPU)
.Implementation(my_custom_op_cpu_kernel)
.Input("input_tensor", DataType::FLOAT32) // 可选元数据
.Output("output_tensor", DataType::FLOAT32); // 可选元数据
// 也可能在此处注册 GPU 版本
// registry->Register("MyCustomOp").Device(DeviceType::GPU)...
}
实现可以静态链接到主应用程序中,也可以动态加载(例如,从共享对象 .so 或动态链接库 .dll)。动态加载提供了灵活性,允许用户在不重新编译整个运行时系统的情况下添加定制算子。
定制核函数的函数签名非常重要。运行时需要传递所有必要信息:
cudaStream_t)、临时内存分配器,以及可能的性能分析工具。void* 或类型化指针),位于正确的设备上。float32、int8)。dilation_rate、epsilon)。一个典型的核函数签名可能如下所示:
// 简化的核函数上下文结构
struct KernelContext {
void* stream; // 例如,cudaStream_t 或等效物
Allocator* temp_allocator;
// ... 其他上下文信息
};
// 简化的张量信息结构
struct TensorInfo {
void* data;
DataType dtype;
std::vector<int64_t> shape;
std::vector<int64_t> strides;
DeviceType device;
};
// 定制核函数的示例签名
Status my_custom_op_gpu_kernel(
KernelContext* context,
const std::vector<TensorInfo>& inputs,
const std::vector<TensorInfo>& outputs,
const std::map<std::string, AttributeValue>& attributes
) {
// 实现时使用 context->stream, inputs[0].data, 等。
// 检查属性、输入形状/类型。
// 启动 GPU 核函数。
// 返回 Status::OK 或错误代码。
}
运行时负责确保输入张量数据在定制核函数预期其存在的设备上可用。如果用 CPU 张量输入调用定制 GPU 核函数,运行时必须管理数据传输(可能是异步的)。类似地,在设备上产生的输出可能需要传回。
定制核函数,特别是 GPU 核函数,通常异步执行。核函数实现必须正确使用提供的执行流(示例中的 context->stream)来排队其工作。运行时需要管理依赖关系,确保定制核函数启动与先行操作同步,并且后续操作在必要时等待定制核函数完成(例如,通过记录和等待与流关联的事件)。不正确的同步是整合定制核函数时常见的错误源。
temp_allocator 可确保缓冲区在运行时的内存计划内管理;独立分配内存可能会干扰运行时的优化和跟踪。有效支持定制算子是灵活强大的机器学习 (machine learning)运行时系统的一个显著特点,它使用户能够突破性能限制,尝试新颖的模型架构,这些能力是仅靠标准编译器优化难以实现的。
这部分内容有帮助吗?
© 2026 ApX Machine LearningAI伦理与透明度•