CANN算子利用迭代器高效实现Tensor数据切割分块处理

　　摘要：本文以Diagonal算子为例，介绍并详细讲解如何利用迭代器对n维Tensor进行基于位置坐标的大批量数据读取工作。
　　本文分享自华为云社区《CANN算子：利用迭代器高效实现Tensor数据切割分块处理云社区华为云》，作者：CatherineWang。任务场景及目标
　　在CANNaicpu算子开发实现中，经常需要对n维Tensor进行切片（slice）、切块（dice）、转置（transpose）、交换指定维度数据（shuffle）等操作。上述操作实质上是按照指定规律依次进行数据读取，并将读取到的数据写入新的数据地址中。
　　本文以Diagonal算子为例，介绍并详细讲解如何利用迭代器对n维Tensor进行基于位置坐标的大批量数据读取工作。
　　Diagonal算子希望对指定两个维度的数据进行对角元素的提取，最终返回张量的对角线元素。本质上该算子通过属性dim1和dim2确定一个矩阵，返回该矩阵的对角元素（存在偏移量offset），并将其放置在最后一维。非dim1和dim2的维度，将会被当成batch维度处理。常规方案：
　　方案一：将shape为s，元素个数为numel的输入Tensor：x转化为Eigen：：Tensor：eigenx；对eigenx进行shuffle操作，将dim1和dim2换至倒数第二和倒数第一维；通过reshape操作将eigenx变化为一个三维Eigen：：Tensor：reshapex，shape（numels〔dim1〕s〔dim2〕，s〔dim1〕，s〔dim2〕）；对后两维数据取对角元素，并将最终数据赋值给输出数据地址。注意：由于Eigen：：Tensor不能够动态设置维度，即NumIndices项必须是一个具体的值，因此需要提前定义对应维度的Eigen：：Tensor备用。
　　方案二：对于一个n维的Tensor，利用n层for循环进行数据的定位读取，并取对角值。
　　可以看出上述两个方案对动态大小的输入计算实现处理都较为繁琐，需要提前分情况设置对应维度的Eigen：：Tensor或是for循环逻辑结构，即存在维数限制。准备知识及分析
　　我们知道再AICPU中，对于一个Tensor，我们能够通过GetTensorShape、GetData等函数获得Tensor形状大小、具体数据地址等信息。但我们不能通过位置坐标的形式直接获得指定位置的数据值。1。步长
　　首先介绍步长（stride）这一概念（对这部分知识已掌握的可以直接跳转下一部分内容）。stride是在指定维度dim中从一个元素跳到下一个元素所必需的步长。例如，对于一个shape（2，3，4，5）的Tensor，其stride（60，20，5，1）。因此如果想要获取到上述Tensor中位置坐标为〔1，2，1，3〕的数据，只需要找到数据地址中第108（601202513）位对应值。2。迭代器
　　定义迭代器PositionIterator，包含私有成员pos和shape，其中pos为初始位置，shape为标准形状。通过重载符号，对pos进行修改，实现迭代器的自增操作。基于上述迭代器，可以实现对给定的shape依次取位操作。如给定对于给定的shape（d1，d2，，dn），从初始位置（0，0，，0）开始，依次取（0，0，，0，0），（0，0，，0，1），，（0，0，，0，dn1），（0，0，，1，0），（0，0，，1，1），，（d11，d21，，d｛n1｝1，d｛n｝1）。
　　事实上，可以将上述迭代器理解为一种进制，对于给定的标准形状shape（d1，d2，，dn），第i位运算时便是逢di进1。同时通过PositionIterator。End（）控制迭代器的结束。具体实现如下：templatetypenameTclassPositionIterator｛public：PositionIterator（）｛｝；PositionIterator（）｛｝；PositionIterator（std：：vectorTstt，std：：vectorTsh）｛if（stt。size（）！sh。size（））｛PositionIterator（）；｝else｛for（unsignedinti0；ish。size（）；i）｛if（stt〔i〕sh〔i〕）｛PositionIterator（）；｝｝posstt；shapesh；｝｝PositionIteratoroperator（）｛pos〔shape。size（）1〕1；for（unsignedintishape。size（）1；i0；i）｛if（pos〔i〕shape〔i〕！0）｛pos〔i1〕pos〔i〕shape〔i〕；pos〔i〕pos〔i〕shape〔i〕；｝｝returnthis；｝boolEnd（）｛if（pos〔0〕！shape〔0〕）｛returnfalse；｝returntrue；｝std：：vectorTGetPos（）｛returnpos；｝std：：vectorTGetShape（）｛returnshape；｝private：std：：vectorTpos；std：：vectorTshape；｝；Diagonal算子的实现
　　利用迭代器，在一般情况下，我们只需要两层for循环，便可以实现Diagonal算子的计算过程。第一层for循环用于确定除dim1和dim2维度的位置坐标，第二层for循环用于对dim1和dim2对应维度确定对角元素位置，通过这样的两层for循环，便可将对角元素位置确定。通过这样的取值处理，相较于Eigen实现思路，计算速度有着明显的提升，且无维度限制，st测试结果对比如下：
　　具体实现可参见如下代码：templatetypenameTuint32tDiagonalCpuKernel：：DoComputeType（CpuKernelContextctx，constint64toffset，constint64tdim1，constint64tdim2）｛GettheinuputandoutputTensorinputxctx。Input（0）；Tensoryctx。Output（0）；GetsomeinformationofinputautoxshapeinputxGetTensorShape（）；std：：vectorint64txshapexshapeGetDimSizes（）；constint64txdimxshapeGetDims（）；autodataptrreinterpretcastT（ctx。Input（0）GetData（））；autoydataptrreinterpretcastT（yGetData（））；Compute首先计算出对角线元素个数int64tdsizeOffsetSize（offset，dim1，dim2，xshape）；生成输入Tensor的步长向量xstridestd：：vectorint64txstrideConstructStrideint64t（xshape）；分情况讨论，2维和大于2维的情况if（xdim！N2）｛setthevxshapeandvxstride生成xshape和xstride中除去dim1和dim2对应值的vxshape与vxstridestd：：vectorint64tvxshape，vxstride；for（unsignedinttmpdim0；tmpdimxshape。size（）；tmpdim）｛if（tmpdim！dim1tmpdim！dim2）｛vxshape。pushback（xshape〔tmpdim〕）；vxstride。pushback（xstride〔tmpdim〕）；｝｝settheyshape，ystride，vystride生成输出Tensor的形状及步长向量：yshape和ystridestd：：vectorint64tyshapevxshape；yshape。pushback（dsize）；std：：vectorint64tystrideConstructStrideint64t（yshape）；生成输出Tensor的出去最后一维的步长向量：vystridestd：：vectorint64tvystrideystride；vystride。popback（）；读取对角数据std：：vectorint64tvstart（vxshape。size（），0）；for（PositionIteratorint64tmyiter（vstart，vxshape）；！myiter。End（）；myiter）｛利用迭代器确定除dim1和dim2维度的位置坐标autopmyiter。GetPos（）；通过步长向量和位置坐标计算出输入和输出的基础位置值basepos1和outbaseposint64tbasepos1MulSumint64t（p，vxstride）；int64toutbaseposMulSumint64t（p，vystride）；for（inti0；idsize；i）｛结合前面计算出的基础位置值，对dim1和dim2对应维度确定对角元素位置，并赋值给输出数据地址（getdata涉及对上对角还是下对角取元素，不影响对迭代器作用的理解）int64tbasepos2i（xstride〔dim1〕xstride〔dim2〕）；int64tarr〔N2〕｛xstride〔dim1〕，xstride〔dim2〕｝；ydataptr〔outbaseposi〕getdata（basepos1basepos2，offset，arr，dataptr）；｝｝｝else｛for（inti0；idsize；i）｛int64tbaseposi（xstride〔dim1〕xstride〔dim2〕）；int64tarr〔N2〕｛xstride〔dim1〕，xstride〔dim2〕｝；ydataptr〔i〕getdata（basepos，offset，arr，dataptr）；｝｝returnKERNELSTATUSOK；｝迭代器的其他用法
　　1、数据切条：如Sort算子中，用迭代器对Tensor数据关于tmpaxis维度进行取条，以进行后续的排序操作。for（positioniteratorint64tmit（vstart，vshape）；！mit。end（）；mit）｛autopmit。getpos（）；intaxisleninputshape〔tmpaxis〕；std：：vectorValueIndexTdata（axislen）；intbaseposmulsumint64t（p，vstride）；for（int32ti0；iaxislen；i）｛data〔i〕。valuexdataptr〔baseposiinputstride〔tmpaxis〕〕；data〔i〕。indexi；｝
　　2、数据切块：切块处理可以利用两个迭代器循环叠加，也可以利用一个迭代器和两个坐标位置for循环
　　3、关于指定维度dim，对Tensor降维拆分为N子Tensor：如UniqueConsecutive算子中，首先需要关于属性axis维，将原本的Tensor数据拆分为inputshape〔axis〕个子Tensor（此处用vector存储Tensor中的数据）。std：：vectorstd：：vectorT1data；for（int64ti0；idim0；i）｛std：：vectorT1tmpv1；for（PositionIteratorint64tmit（vstart，vshape）；！mit。End（）；mit）｛autoposmit。GetPos（）；tmpv1。pushback（xdataptr〔MulSumint64t（pos，vstride）iinputstride〔axis〕〕）；｝data。pushback（tmpv1）；｝
　　点击下方，第一时间了解华为云新鲜技术
　　华为云博客大数据博客AI博客云计算博客开发者中心华为云