「经验分享」基于PyTorch框架TBE

智聪说说网
智聪说说网
智聪说说网
34608
文章
0
评论
2022-12-2816:28:16 评论 13

  什么样的人需要进行TBE自定义算子开发:

  昇腾AI处理器不支持开发者网络中的算子开发者想要自己开发算子来提高计算性能开发者想要修改现有算子中的计算逻辑读者对象需要掌握什么:

  具备Python/C++/C语言程序开发能力

  理解数学表达式对机器学习、深度学习有一定的了解对TVM及开源TensorFlow/PyTorch框架有一定的了解大家可以从官方渠道了解更多内容

  MindStudio官方网址

  MindStudio官方论坛

  TVM的诞生

  随着深度学习的广泛应用,大量的深度学习框架及深度学习硬件平台应运而生,但不同平台的神经网络模型难以在其他硬件平台便捷的运行,无法充分利用新平台的运算性能。TVM(Tensor Virtual Machine)的诞生解决了以上问题,它是一个开源深度学习编译栈,它通过统一的中间表达(Intermediate Representation)堆栈连接深度学习模型和后端硬件平台,通过统一的结构优化Schedule,可以支持CPU、GPU和特定的加速器平台和语言。

  TVM的架构详细介绍请参考https://tvm.apache.org/。

  TBE的诞生

  TBE(Tensor Boost Engine)提供了基于TVM框架的自定义算子开发能力,通过TBE提供的API和自定义算子编程开发界面可以完成相应神经网络算子的开发。

  TBE的逻辑架构如下图所示。

  TBE工具给用户提供了多层灵活的算子开发方式,用户可以根据对硬件的理解程度自由选择,利用工具的优化和代码生成能力,生成昇腾AI处理器的高性能可执行算子。

  前端框架:包含第三方开源框架TensorFlow(TF:Google机器学习开源框架)、PyTorch(快速特征嵌入的卷积结构)。图引擎(Graph Engine:GE):Graph Engine是华为基于昇腾AI处理器软件栈对不同的机器学习框架提供统一的IR接口,对接上层网络模型框架,例如TensorFlow、PyTorch,GE的主要功能包括图准备、图拆分、图优化、图编译、图加载、图执行和图管理等(此处图指网络模型拓扑图)。融合引擎(Fusion Engine:FE):FE负责对接GE和TBE算子,具备算子信息库的加载与管理、融合规则管理、原图融合和子图优化的能力。GE在子图优化阶段将子图传递给FE,FE根据算子信息库以及FE融合优化进行预编译,例如修改数据类型、插入转换算子等,该子图将再次传递给GE进行子图合并及子图优化。张量加速引擎(TBE):TBE通过IR定义为GE的图推导提供必要的算子信息,通过算子信息库和融合规则为FE提供子图优化信息和TBE算子调用信息,TBE生成的二进制对接昇腾AI处理器,最终生成网络在昇腾AI处理器上的执行任务。TBE内部包含了特性域语言(Domain-Specific Language,DSL)模块、调度(Schedule)模块、中间表示(Intermediate Representation,IR)模块、编译优化(Pass)模块以及代码生成(CodeGen)模块如下图所示。

  算子逻辑描述:面向开发者,提供算子逻辑的编写的接口(Compute接口),使用接口来编写算子。Schedule模块:用于描述指定shape下算子如何在昇腾AI处理器上进行切分,包括Cube类算子的切分、Vector类算子的切分,它们仍然使用的是社区提供的调度原语来描述。IR模块:借用社区的IR来表示的,包括IR变形、AST树的维护等功能。编译优化(Pass):对生成的IR进行编译优化,优化的方式有双缓冲(Double Buffer)、流水线(Pipeline)同步、内存分配管理、指令映射、分块适配矩阵计算单元等。代码生成模块(CodeGen):CodeGen生成类C代码的临时文件,这个临时代码文件可以通过编译器生成算子的实现文件,可被网络模型直接加载调用。一个完整的TBE算子包含四部分:算子原型定义、对应开源框架的算子适配插件、算子信息库定义和算子实现。

  算子开发完成后在昇腾AI处理器硬件平台上的编译运行的架构如下图所示。

  上图中圆柱形图示为开发人员在自定义算子开发时需要实现的交付件。

  加载TBE算子进行模型转换的详细流程如下图所示。

  将原始第三方网络模型(TensorFlow/PyTorch)下发给GE。(注:这里网络模型的拓扑图后续简称为图。)GE调用算子插件,将TensorFlow/PyTorch网络模型中的算子映射为适配昇腾AI处理器的算子,从而将原始TensorFlow/PyTorch图解析为适配昇腾AI处理器的图。调用算子原型库校验接口进行基本参数的校验,校验通过后,会根据原型库中的推导函数推导每个节点的输出shape与dtype,进行输出tensor的静态内存的分配。GE向FE发送图优化请求,并将图下发给FE,FE匹配融合规则进行图融合,并根据fe.ini中的配置进行算子选择,选择优先级最高的算子类型进行算子匹配(默认自定义算子优先级最高),最后将优化后的整图返回给GE。GE根据图中数据将图拆分为子图并下发给FE,FE首先在子图内部插入转换算子,然后按照当前子图流程进行TBE算子预编译,对TBE算子进行UB融合(算子可以在UB中根据UB融合规则自动与其他算子的计算进行组装),并根据算子信息库中算子信息找到算子实现将其编译成算子kernel(算子的*.o与*.json文件),最后将优化后子图返回给GE。GE进行图编译,包含内存分配、流资源分配等,并向FE发送tasking请求,FE返回算子的taskinfo信息给GE,图编译完成之后生成适配昇腾AI处理器的离线模型文件(*.om)。昇腾AI软件栈提供了TBE(Tensor Boost Engine:张量加速引擎)算子开发框架,开发者可以基于此框架使用Python语言开发自定义算子,通过TBE进行算子开发有以下几种方式:

  DSL(Domain-Specific Language)开发

  为了方便开发者进行自定义算子开发,TBE DSL接口已高度封装,开发者仅需要使用DSL接口完成计算过程的表达,后续的Schedule创建、优化及编译都可通过已有接口一键式完成,适合初级开发者。DSL开发的算子性能可能较低。

  TIK开发

  TIK(Tensor Iterator Kernel)是一种基于Python语言的动态编程框架,呈现为一个Python模块。开发者可以通过调用TIK提供的API基于Python语言编写自定义算子,即TIK DSL,然后TIK编译器会将TIK DSL编译为昇腾AI处理器应用程序的二进制文件。基于TIK的自定义算子开发,提供了对Buffer的管理和数据自动同步机制,但需要开发者手动计算数据的分片和索引,需要开发者对Davinci架构有一定的了解,入门难度更高。TIK对矩阵的操作更加灵活,性能会更优。

  整个计算过程可以表现为多个输入张量经过一个计算节点得到多个张量的过程。TIK方式、DSL的开发流程本质上是一样的,只不过开发的抽象层次不一样而已。

  什么是TIK

  TIK(Tensor Iterator Kernel)是一种基于Python语言的动态编程框架,呈现为一个Python模块,运行于Host CPU上。开发者可以通过调用TIK提供的API基于Python语言编写自定义算子,然后TIK编译器会编译为昇腾AI处理器应用程序的二进制文件。

  TIK的优势

  TIK算子开发方式是一种灵活的开发方式。TIK代码在算子开发效率和算子性能自动优化上有着一定的优势:

  并行化:提供程序员串行化编程体系,方便编写算子,TIK工具自动对计算过程并行化,实现高性能。自动内存管理:程序员在编写算子的时候不用感知和管理地址,编译器会做好内存分配。灵活性:通过手动调度可以更加精确的控制数据搬运和计算流程,从而实现更高的性能,将昇腾AI处理器的能力发挥到极致。易调试:区别于其他形式算子验证的黑盒模式,开发人员没有办法一步一步地去定位算子的问题,tik的Debug模式可以帮助用户快速的定位功能问题,极大的缩短了开发调测时间。TIK编程模型

  下图是使用TIK进行编程的过程示意图,用户调用TIK API编写算子对应的Python程序后,TIK会将其转化为TIK DSL(TIK DSL是一种DSL语言,它可以在比CCE更高的抽象层次上定义CCEC程序的行为),经过编译器编译后生成CCEC文件(CCEC代码目前对于TIK编程人员无法感知),再经过CCE编译器编译后生成可运行在昇腾AI处理器上的应用程序。

  下图是基于TIK API编写Python程序的通用步骤

  TIK解决的问题

  TIK提升编程效率的办法:

  TIK工具自动处理Buffer的分配和管理TIK工具提供自动数据同步TIK工具提供硬件单元自动并行化方法TIK不能解决的问题:

  TIK不能自动进行数据切分进行算子开发前,开发者应首先进行算子分析,算子分析包含:明确算子的功能及数学表达式,选择算子开发方式(DSL方式或者TIK方式),最后细化并明确算子规格,分析流程如下所示:

  分析算子算法原理,提取出算子的数学表达式。明确算子开发方式及使用的计算接口。明确详细算子规格:明确使用的算子实现方式及算子实现接口后,需要进一步明确算子输入输出支持的数据类型,形态以及数据排布格式,明确算子实现文件名称、算子实现函数名称以及算子的类型(OpType)。操作步骤:

  1. 进入算子工程创建界面。

  1. 首次登录MindStudio:在MindStudio欢迎界面中单击“Create new project”,进入创建工程界面。

  2. 非首次登录MindStudio:在顶部菜单栏中选择“File > New > Project...”,进入创建工程界面。

  2. 创建算子工程

  1. 在左侧导航栏选择“Ascend Operator”,在右侧配置算子工程信息

  2. 单击“Next”,在弹出的页面中配置算子相关信息

  3. 单击“Finish”,完成算子工程的创建。(若工作窗口已打开其他工程,会出现如下图所示提示。

  选择“This Window”,则直接在当前工作窗口打开新创建的工程。选择“New Window”,则新建一个工作窗口打开新创建的工程。

  代码示例:

  该样例实现了从Global Memory中的A,B两处分别读取128个float16类型的数值,相加,并将结果写入Global Memory地址C中。

  算子插件的实现包含昇腾AI处理器中算子类型的注册、原始框架中算子类型的注册以及原始框架中算子属性到昇腾AI处理器中算子属性的映射,算子的映射通过Parser模块完成。插件在整网络运行场景下的实现流程如下图所示。

  首先GE接收到第三方框架的原始网络模型,并进行初始化,网络模型的拓扑图我们简称为图。GE从Register注册模块中加载TBE算子插件生成的.so文件,在OPP安装路径下的“opp/framework/”目录中。读取算子插件.so中的算子相关信息,并将其注册到算子插件的map文件中(所有算子插件的相关信息都会以map的形式存储到一个文件中)。GE向Parser模块发送调用Parser方法的请求。Parser模块根据算子类型(OpType)从算子插件的map文件中取出对应的Parser函数,并返回实现函数PaserParamsxx给Parser模块,Parser模块根据实现函数将第三方网络算子中的属性映射到昇腾AI处理器中的算子属性,即算子原型中的属性定义,完成第三方网络中算子到昇腾AI处理器中算子的映射。完成第三方网络图到适配昇腾AI处理器的图的转换后,后续会进行一些图融合、图优化的操作,然后最终将算子编译生成二进制文件。框架管理器(Framework)提供REGISTER_CUSTOM_OP宏,按照指定的算子名称完成算子的注册。

  原始框架为TensorFlow的自定义算子注册代码如下所示:

  在代码实现文件顶部使用预编译命令“#include”将插件实现函数相关的头文件包含到插件实现源文件中。register.h存储在ATC安装路径下的“include/register/” 下,包含该头文件,可使用算子注册相关类,调用相关的接口。REGISTER_CUSTOM_OP:注册自定义算子,OpType为注册到GE中的算子类型。FrameworkType:TENSORFLOW代表原始框架为TensorFlow。OriginOpType:算子在原始框架中的类型。ParseParamsByOperatorFn:用来注册解析算子属性的函数。算子原型定义规定了在昇腾AI处理器上可运行算子的约束,主要体现算子的数学含义,包含定义算子输入、输出和属性信息,基本参数的校验和shape的推导,原型定义的信息会被注册到GE的算子原型库中。网络模型生成时,GE会调用算子原型库的校验接口进行基本参数的校验,校验通过后,会根据原型库中的推导函数推导每个节点的输出shape与dtype,进行输出tensor的静态内存的分配。

  算子原型库在整个网络模型生成流程的作用如下图所示。

  其中算子注册包括OP注册、InferShape与Verify的注册,注册的时候算子类型(OpType)作为Key。

  首先GE接收到第三方框架的原始网络模型,并进行初始化,网络模型的拓扑图简称为图。算子原型库管理模块OPP算子库的“opp/op_proto/”对应目录下加载算子原型库的so文件。算子原型库管理模块根据so文件中的信息在OperatorFactory中进行算子信息注册,包括算子基本信息注册、InferShape注册、Verify注册,这三部分分别以算子类型(OpType)作为key注册到三个map文件中进行保存。图准备阶段,GE向Graph发送调用InferShape函数与Verify函数的请求。Graph会遍历图中所有节点。

  每个节点都会向OpDesc发送调用InferShape与Verify函数的请求。OpDesc从OperatorFactory中根据OpType取出对应的InferShape函数与Verify函数。OpDesc执行Verify函数进行校验,如果校验成功,则继续往下执行;如果校验失败,则直接返回。OpDesc执行InferShape函数,进行输出tensor的shape推导。OpDesc向GE返回InferShape的结果,GE后续根据InferShape结果分配输出tensor的静态内存。GE进行其他处理。算子信息库作为算子开发的交付件之一,主要体现算子在昇腾AI处理器上的具体实现规格,包括算子支持输入输出dtype、format以及输入shape等信息。网络运行时,FE会根据算子信息库中的算子信息做基本校验,选择dtype,format等信息,并根据算子信息库中信息找到对应的算子实现文件进行编译,用于生成算子二进制文件。

qq名片赞快速秒赞网址,空间说说点赞网站最低价 - 全网刷快手最便宜网站

  算子开发者需要通过配置算子信息库文件,将算子在昇腾AI处理器上相关实现信息注册到算子信息库文件中。

  算子信息库文件路径:自定义算子工程目录下的/tbe/op_info_cfg/ai_core/${soc_version}/xx.ini。 注:${soc_version}:昇腾AI处理器的版本,可从ATC安装路径的“data/platform_config”目录下查看,".ini"文件的文件名(将大写字母转换为小写)即为对应的${soc_version}。

  自定义算子开发完成后,需要对算子工程进行编译,编译出可直接安装的自定义算子run包;然后进行run包的安装,将自定义算子部署到opp算子库。

  算子工程编译的具体内容为:将算子插件实现文件、算子原型定义文件、算子信息定义文件分别编译成算子插件、算子原型库、算子信息库。算子包部署指执行自定义算子包的安装,自定义算子交付件会自动部署到opp算子库的对应目录下。

  2.8 网络运行验证

  算子部署到算子库(OPP)后,可以将其加载到网络模型中进行整网的推理验证,验证自定义算子在网络中运行结果是否正确。网络测试会覆盖算子开发的所有交付件,包含实现文件,算子原型定义、算子信息库以及算子适配插件。本节仅对验证的方法做介绍。

  基于MindStudio进行算子开发的场景下,算子网络测试的主要流程如下所示:

  使用ATC模型转换工具将包含自定义算子的原始网络模型转换为适配昇腾AI处理器的离线模型。开发AscendCL应用程序,加载转换好的离线模型文件进行模型推理,验证推理结果是否正确。基于命令行环境进行算子开发的场景下,算子网络测试的主要流程如下所示:

  使用ATC模型转换工具将包含自定义算子的原始网络模型转换为适配昇腾AI处理器的离线模型。开发AscendCL应用程序,加载转换好的离线模型文件进行模型推理,验证推理结果是否正确。基本数据类型指用于数据定义和运算的每一个数的数值类型,包括:int8, uint8, int16, uint16, int32, uint32, float16, float32,uint1(bool)等。

  不同的TIK API支持不同的数据类型。TIK为强类型语言,即不同的数据类型之间无法进行计算。TIK对象的标准数据类型,指在内存中的对象的数据类型,表示用户能以什么格式创建和读写数据,包括两种:Scalar标量数据、Tensor张量数据。Scalar标量数据对应于存储在寄存器或者Scalar Buffer中的数据,TIK提供了Scalar接口用于定义Scalar数据,接口原型为:

  dtype指定Scalar对象的数据类型,取值:int8,uint8,int16, uint16,float16,int32,uint32,float32, int64, uint64 默认值:int64name定义Scalar名字,支持string类型。名字支持数字0-9,A-Z,a-z及下划线组成的字符串,默认值:reg_buf$(COUNT), COUNT从零开始计数。init_value为初始化值: 可以是立即数(支持int, float类型)、Scalar变量、Tensor的某个值、Expr:包括Scalar变量、立即数组成的Expr。Tensor数据对应于存储Buffer中的数据,TIK提供了Tensor接口用于定义Tensor数据,用户一般只需关注张量定义的数据类型(dtype)、形状(shape)与数据存储空间(scope)即可,函数原型为:

  TIK会自动为每个申请的Tensor对象分配空间,并且避免各个数据块之间的地址冲突。且TIK会自动检查数据之间的数据依赖性,实现数据的同步。

  用户定义的Tensor在内存分配时会对起始地址进行对齐,不同scope的对齐要求请参见通用约束。用户需特别注意因地址对齐导致的超出对应内存类型的总容量时,会引起编译报错。

  对于Vector计算,一般是用Unified Buffer去存放数据,再进行计算,所以整体数据流应该是从Global Memory>Unified Buffer>Global Memory。TIK提供了data_move接口实现Global Memory和Unified Buffer间的数据搬运,函数原型为:

  在data_move的函数原型中,用户需要着重关注dst、src、nburst、burst、src_stride、dst_stride等6个参数。其中,dst、src分别表示目的操作数与源操作数,也是数据搬运的起始地址;nburst、burst分别表示待搬运数据包含的数据片段数与每个连续片段的长度(单位32 bytes);src_stride、dst_stride则分别代表相邻数据片段的间隔(即前 burst 尾与后 burst 头的间隔)。通过以上6个参数,data_move支持连续地址与间隔地址两种搬运模式。

  vec_a_gm = tik_instance.Tensor("float16" , (1,1024),"vec_a_gm", tik.scope_gm)

  vec_a_ub = tik_instance.Tensor("float16" , (1,1024),"vec_a_ub", tik.scope_ub)

  tik_instance.data_move (vec_a_ub[0, 0], vec_a_gm, sid = 0, n_burst = 1, burst = 1024/16, src_stride = 0, dst_stride = 0)

  这里n_burst = 1表示数据只搬运一次;

  burst的单位是32byte也即是block,对应float16就是16个数,因为一个float16占2个byte;因此,一个32byte可传16个float16数;有1024个数要传,所以burst的长度为1024/16,需要这么多的32byte。大家在实际计算过程中要注意计算burst这个字段时需要向上取整。

  标量计算接口主要分为单操作数和双操作数两类

  矢量计算接口中常用的主要包括单目以及双目操作接口:

  归约接口:

  数据填充接口:

  功能调试接口:

  打开MindStuido后点击左上角的File->New->Project。

  接着就能够进入工程创建部分,在MindStudio中我们可以创建很多工程,不仅仅是TBE自定义算子开发,MindX SDK、模型训练和推理也能在其中完成。

  在上图中,Name为本次TBE算子的名称,Description为本算子的描述,可不填,CANN Version为本算子使用的CANN包的版本,Project Location为本算子工程的存放位置。

  点击NEXT之后即可进入算子TBE算子创建页面。

  在上图中我们可以在New Operator中选择创建算子的格式,新上手的同学可以参考我们这里选择Sample Template,该格式能够提供一个到多个样例给我们选择。

  我们选择PyTorch框架下的Sort算子作为示例。

  点击Finish之后,选择窗口就可以完成工程的创建工作。进入算子的工程文档。

  其中在TBE算子开发过程中需要修改的文档在MindStudio中进行了标记,framework——算子实现插件存放目录;op_proto——算子IR定义文件目录;tbe——TBE算子文件目录;tbe/impl——算子实现文件目录;tbe/op_info_cfg——算子信息库文件目录;

  在framework中我们需要关注算子插件实现文件。Pytorch框架下无需修改这部分的内容。

  在op_proto中我们需要关注的是算子IR定义文件。

  tbe/impl——算子实现文件目录;tbe/op_info_cfg——算子信息库文件目录;文件在此就不做一一展示了。

  在testcase一节中有已经编写好的ut测试代码,在运行py文件前,我们需要添加python解释器。可以点击左上角File->Project Structure进行设置。

  点击之后即可进入相应界面。

  在Dependencies中的Module SDK中选择相应的Python版本后就可以点击右下加的Apply,即完成添加Python解释器的步骤。

  之后可以右键点击testcase.ut.add中test_sort_impl.py,点击Run TBE Operator ‘add’ UT Impl with coverage就可以进行UT测试。

  在进行UT测试前还需要配置一些芯片信息和用例信息。

  点击Run即可得到运行结果,并且生成反馈覆盖率的html文件。

  代码下载链接(7ukc)

  在算子实现过程中我们用到了以下四个版本的代码,大家可以拿去自行实验。

  在测试过程中我们主要使用的是coverage命令,大家将上述文件中的代码放到mul.py文件中,并在同目录下新建test_mul_impl.py。以下test文件可以复制进入test_mul_impl.py。使用命令coverage run test_mul_impl.py就可以得到debug运行的结果。

  在mul_v1代码中48行是算子提供给外部调用的接口,其中第五个参数isdebug可以忽略掉,第六个参数block_num为本次运行代码使用多少个核进行优化。在32行的cv_mul_compute函数主要负责算子计算逻辑的实现。在14行的_init()_函数主要负责参数以及Tensor的定义。

  mul_v1是能进行debug的最少需要的代码,大家可以根据这一版本的代码进行接口功能的测试。大家可以多多进行尝试,例如data_move,vec_conv,vec_add等接口,实验越多对代码的理解就能更深入。

  mul_v2和mul_v3旨在解决一个问题:UB的存储空间不足的问题。我们这边提供的解决方案就是尽可能的将更多的GM上的数据搬入到大小约为256KB的UB上进行计算。其中mul_v3中29至45行的参数对应非最后一个核的计算过程,47至57行的参数对应最后一个核的计算过程。这些参数都已经饱经风霜,大家可以复制过去用。如果想要深入理解TIK开发方式,需要手推和验证这些参数的含义。

  mul_v4是一个接近完成的代码,我们给其中各个函数做的总结如下:

  def mul:提供给外部调用的接口;

  def compute:多核开启的地方,并且向下屏蔽核的差异;

  def compute_each_core:计算分片开启的地方,并且向下屏蔽计算接口与相关参数的差异,最后一个核的代码因参数不同单独处理;

  def compute_each_times:数据类型转换开启的地方,并且向下屏蔽数据类型的差异;

  大家如果想要套用本框架,快速上手TIK的话可以修改mul_v4中的157和160行,例如add算子的功能实现可以修改接口vec_mul为vec_add就可以。至此,我们完成了本课程的介绍,非常感谢大家的学习观看,TIK的开发学习过程是漫长的,希望大家能够砥砺前行。

郑重声明:本文版权归原作者所有,转载文章仅为传播更多信息之目的,如作者信息标记有误,请第一时候联系我们修改或删除,多谢。

标签:「经验分享」基于PyTorch框架TBE

智聪说说网
  • 本文由 发表于 2022-12-2816:28:16
  • 转载请务必保留本文链接:https://www.zhicongwang.com/88476.html