运行在昇腾AI处理器的AI Core上,鉴于AI Core的强大算力,主要负责执行矩阵、向量、标量相关的计算密集型算子。
其中,TBE(Tensor Boost Engine)提供了基于张量虚拟机(Tensor Virtual Machine,TVM)框架的自定义算子开发能力,提供了用户开发自定义算子所需工具。TBE工具给用户提供了两种算子开发方式 — DSL与TIK,开发者可以根据需求自由选择,两种开发方式的区别如下:
DSL接口已高度封装,用户仅需要使用DSL接口完成计算过程的表达,后续的算子调度、算子优化及编译都可通过已有的接口一键式完成,适合初级开发用户。
开发者可以通过调用TIK提供的API基于Python语言编写自定义算子,然后TIK编译器会将其编译为适配昇腾AI处理器应用程序的二进制文件。但TIK需要用户手工控制数据搬运和计算流程,入门较高,但开发方式比较灵活,在性能上有一定的优势。
运行在昇腾AI处理器的AI CPU上。以下几种场景下,可使用AI CPU方式实现自定义算子。
例如,Dump、profiling等控制算子,Queue、Stack等资源状态类算子,TopK、Where等检索类算子。
若存在相同算子类型的TBE算子与AI CPU算子,TBE算子的执行优先级高于AI CPU算子。