本文分享自华为云社区《CANN算子：利用迭代器高效实现Tensor数据切割分块处理》，作者： CatherineWang 。
任务场景及目标
在 CANN aicpu 算子开发实现中，经常需要对 n 维 Tensor 进行切片（slice）、切块（dice）、转置（transpose）、交换指定维度数据（shuffle）等操作。上述操作实质上是按照指定规律依次进行数据读取，并将读取到的数据写入新的数据地址中。
本文以 Diagonal 算子为例，介绍并详细讲解如何利用迭代器对 n 维 Tensor 进行基于位置坐标的大批量数据读取工作。
Diagonal 算子希望对指定两个维度的数据进行对角元素的提取，最终返回张量的对角线元素。本质上该算子通过属性 dim1 和 dim2 确定一个矩阵，返回该矩阵的对角元素（存在偏移量 offset），并将其放置在最后一维。非 dim1 和 dim2 的维度，将会被当成 batch 维度处理。
常规方案：
方案一：将 shape 为 s，元素个数为 numel 的 输入 Tensor：x 转化为 Eigen::Tensor：eigen_x；对 eigen_x 进行 shuffle 操作，将 dim1 和 dim2 换至倒数第二和倒数第一维；通过 reshape 操作将 eigen_x 变化为一个三维 Eigen::Tensor：reshape_x，shape=(numel/ s[dim1]/s[dim2]，s[dim1]，s[dim2])；对后两维数据取对角元素，并将最终数据赋值给输出数据地址。注意：由于 Eigen::Tensor<typename T, int NumIndices_>不能够动态设置维度，即 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(=60*1+20*2+5*1+3)位对应值。
2.迭代器
定义迭代器 PositionIterator，包含私有成员 pos_和 shape_，其中 pos_为初始位置，shape_为标准形状。通过重载++符号，对 pos_进行修改，实现迭代器的自增操作。基于上述迭代器，可以实现对给定的 shape 依次取位操作。如给定对于给定的 shape=(d_1,d_2,…,d_n)，从初始位置(0,0,…,0)开始，依次取(0,0,…,0,0), (0,0,…,0,1),…,(0,0,…,0,d_n-1), (0,0,…,1,0), (0,0,…,1,1),…, (d_1 - 1,d_2 - 1,…,d_{n-1}-1,d_{n}-1)。
事实上，可以将上述迭代器理解为一种进制，对于给定的标准形状 shape_=(d_1,d_2,…,d_n)，第 i 位运算时便是逢 d_i 进 1。同时通过 PositionIterator .End()控制迭代器的结束。具体实现如下：
template <typename T>
class PositionIterator {
public:
PositionIterator(){};
~PositionIterator(){};
PositionIterator(std::vector<T> stt, std::vector<T> sh) {
if (stt.size() != sh.size()) {
PositionIterator();
} else {
for (unsigned int i = 0; i < sh.size(); i++) {
if (stt[i] >= sh[i]) {
PositionIterator();
}
}
pos_ = stt;
shape_ = sh;
}
}
PositionIterator operator++() {
pos_[shape_.size() - 1] += 1;
for (unsigned int i = shape_.size() - 1; i > 0; i--) {
if (pos_[i] / shape_[i] != 0) {
pos_[i - 1] += pos_[i] / shape_[i];
pos_[i] = pos_[i] % shape_[i];
}
}
return *this;
}
bool End() {
if (pos_[0] != shape_[0]) {
return false;
}
return true;
}
std::vector<T> GetPos() { return pos_; }
std::vector<T> GetShape() { return shape_; }
private:
std::vector<T> pos_;
std::vector<T> shape_;
};
复制代码
Diagonal 算子的实现
利用迭代器，在一般情况下，我们只需要两层 for 循环，便可以实现 Diagonal 算子的计算过程。第一层 for 循环用于确定除 dim1 和 dim2 维度的位置坐标，第二层 for 循环用于对 dim1 和 dim2 对应维度确定对角元素位置，通过这样的两层 for 循环，便可将对角元素位置确定。通过这样的取值处理，相较于 Eigen 实现思路，计算速度有着明显的提升，且无维度限制，st 测试结果对比如下：
具体实现可参见如下代码：
template <typename T>
uint32_t DiagonalCpuKernel::DoComputeType(CpuKernelContext &ctx,
const int64_t &offset,
const int64_t &dim1,
const int64_t &dim2) {
// Get the inuput and output
Tensor *input_x = ctx.Input(0);
Tensor *y = ctx.Output(0);
// Get some information of input
auto x_shape = input_x->GetTensorShape();
std::vector<int64_t> x_shape_ = x_shape->GetDimSizes();
const int64_t x_dim = x_shape->GetDims();
auto dataptr = reinterpret_cast<T *>(ctx.Input(0)->GetData());
auto y_dataptr = reinterpret_cast<T *>(y->GetData());
// Compute
// 首先计算出对角线元素个数
int64_t dsize = OffsetSize(offset, dim1, dim2, x_shape_);
// 生成输入Tensor的步长向量x_stride
std::vector<int64_t> x_stride = ConstructStride<int64_t>(x_shape_);
// 分情况讨论，2维和大于2维的情况
if (x_dim != N2) {
//set the vx_shape and vx_stride
// 生成x_shape和x_stride中除去dim1和dim2对应值的vx_shape与vx_stride
std::vector<int64_t> vx_shape, vx_stride;
for (unsigned int tmp_dim = 0; tmp_dim < x_shape_.size(); tmp_dim++) {
if (tmp_dim != dim1 && tmp_dim != dim2) {
vx_shape.push_back(x_shape_[tmp_dim]);
vx_stride.push_back(x_stride[tmp_dim]);
}
}
// set the y_shape, y_stride, vy_stride
// 生成输出Tensor的形状及步长向量：y_shape和y_stride
std::vector<int64_t> y_shape = vx_shape;
y_shape.push_back(dsize);
std::vector<int64_t> y_stride =
ConstructStride<int64_t>(y_shape);
// 生成输出Tensor的出去最后一维的步长向量：vy_stride
std::vector<int64_t> vy_stride = y_stride;
vy_stride.pop_back();
// 读取对角数据
std::vector<int64_t> v_start(vx_shape.size(), 0);
for (PositionIterator<int64_t> myiter(v_start, vx_shape); !myiter.End();
++myiter) {
// 利用迭代器确定除dim1和dim2维度的位置坐标
auto p = myiter.GetPos();
// 通过步长向量和位置坐标计算出输入和输出的基础位置值base_pos1和outbase_pos
int64_t base_pos1 = MulSum<int64_t>(p, vx_stride);
int64_t outbase_pos = MulSum<int64_t>(p, vy_stride);
for (int i = 0; i < dsize; i++) {
// 结合前面计算出的基础位置值，对dim1和dim2对应维度确定对角元素位置，并赋值给输出数据地址（get_data涉及对上对角还是下对角取元素，不影响对迭代器作用的理解）
int64_t base_pos2 = i * (x_stride[dim1] + x_stride[dim2]);
int64_t arr[N2] = {x_stride[dim1], x_stride[dim2]};
y_dataptr[outbase_pos + i] =
get_data(base_pos1 + base_pos2, offset, arr, dataptr);
}
}
} else {
for (int i = 0; i < dsize; i++) {
int64_t base_pos = i * (x_stride[dim1] + x_stride[dim2]);
int64_t arr[N2] = {x_stride[dim1], x_stride[dim2]};
y_dataptr[i] = get_data(base_pos, offset, arr, dataptr);
}
}
return KERNEL_STATUS_OK;
}
复制代码
迭代器的其他用法
1、数据切条：如 Sort 算子中，用迭代器对 Tensor 数据关于 tmp_axis 维度进行取条，以进行后续的排序操作。
for (position_iterator<int64_t> mit(v_start, v_shape); !mit.end(); ++mit) {
auto p = mit.get_pos();
int axis_len = input_shape_[tmp_axis];
std::vector<ValueIndex<T>> data_(axis_len);
int base_pos = mul_sum<int64_t>(p, v_stride);
for (int32_t i = 0; i < axis_len; i++) {
data_[i].value = x_dataptr[base_pos + i * input_stride[tmp_axis]];
data_[i].index = i;
}
复制代码
2、数据切块：切块处理可以利用两个迭代器循环叠加，也可以利用一个迭代器和两个坐标位置 for 循环
3、关于指定维度 dim，对 Tensor 降维拆分为 N 子 Tensor：如 UniqueConsecutive 算子中，首先需要关于属性 axis 维，将原本的 Tensor 数据拆分为 input_shape[axis]个子 Tensor（此处用 vector 存储 Tensor 中的数据）。
std::vector<std::vector<T1>> data_;
for (int64_t i = 0; i < dim0; i++) {
std::vector<T1> tmp_v1;
for (PositionIterator<int64_t> mit(v_start, v_shape); !mit.End(); ++mit) {
auto pos = mit.GetPos();
tmp_v1.push_back(
x_dataptr[MulSum<int64_t>(pos, v_stride) + i * input_stride[axis]]);
}
data_.push_back(tmp_v1);
}
复制代码
点击关注，第一时间了解华为云新鲜技术~
评论