PyPTO通过函数定义在NPU硬件上构建可编译的计算图结构,并利用@pypto.frontend.jit装饰器实现即时编译(JIT),从而充分发挥NPU的并行计算能力,从而提升算子执行效率。
在执行JIT编译前,需要定义Kernel函数,获取输入输出Tensor、配置Tiling信息并实现计算逻辑。
基础函数定义:
[object Object]多输入输出函数定义:
[object Object]
当通过PyPTO函数完成kernel的计算流及数据流的编写,可以加上pypto.frontend.jit的装饰器,标记该函数为JIT编译目标,触发PyPTO的编译流程。
[object Object]
JIT编译流程为:
- 首次调用时,函数以“记录模式”执行,操作被记录并优化为计算图,经编译器生成针对NPU的优化代码后缓存二进制文件;
- 后续调用时,函数直接调用缓存的二进制文件在NPU上执行,无需重复编译。
JIT装饰器支持参数配置,可根据配置支持不同的条件编译:
[object Object]
JIT配置选项说明如下:
- codegen_options:代码生成设置。
- host_options:主机端选项。
- pass_options:编译器PASS传递选项。
- runtime_options:运行时执行选项。
- verify_options:精度检验工具选项。
- debug_options:性能数据采集功能配置选项。
除了使用JIT装饰器启用不同配置外,还可以直接在代码中调用pypto.set_codegen_options、pypto.set_host_options、pypto.set_pass_options、pypto.set_runtime_options,pypto.set_verify_options, pypto.set_debug_options接口进行配置,例如:
[object Object]
建议优先使用JIT入参配置各类选项,因为JIT配置选项提供了配置的便利性,同时避免在计算函数内部出现与数据流和计算不相关的代码。
您可以定义多个JIT函数并将它们一起使用:
[object Object]