CANN運算元:利用迭代器高效實現Tensor資料切割分塊處理

華為雲開發者聯盟 發表於 2022-07-04
摘要:本文以Diagonal運算元為例,介紹並詳細講解如何利用迭代器對n維Tensor進行基於位置座標的大批量資料讀取工作。

本文分享自華為雲社群《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測試結果對比如下:

CANN運算元:利用迭代器高效實現Tensor資料切割分塊處理

具體實現可參見如下程式碼:

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);
  }

 

點選關注,第一時間瞭解華為雲新鮮技術~

相關文章