Add Transpose function (#91)

* Add Transpose function

* csrcs->csrc

* Add transpose unittest

* Add reduce_max_large_dim unittest

csrc/fastdeploy/function/reduce.cc

@@ -12,11 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "fastdeploy/function/reduce.h"
+
 #include <set>
 
 #include "fastdeploy/function/eigen.h"
-#include "fastdeploy/function/reduce.h"
 #include "fastdeploy/function/reduce_functor.h"
+#include "fastdeploy/function/transpose.h"
 #include "fastdeploy/utils/utils.h"
 
 namespace fastdeploy {
@@ -71,7 +73,7 @@ void ReduceFunctor(const FDTensor& input, FDTensor* output,
 inline void GetShuffledDim(const std::vector<int64_t>& src_dims,
                            std::vector<int64_t>* dst_dims,
                            const std::vector<int64_t>& reduced_dims,
-                           std::vector<int>* perm_axis) {
+                           std::vector<int64_t>* perm_axis) {
   // check if it's a reduced dim
   std::vector<bool> src_dims_check(src_dims.size(), false);
   size_t src_size = src_dims.size();
@@ -104,19 +106,33 @@ template <typename OutT>
 void GetShuffledInput(const FDTensor& input, FDTensor* shuffled_input,
                       const std::vector<int64_t>& dims) {
   auto shuffled_dims = input.shape;
-  std::vector<int> perm_axis(input.shape.size());
+  std::vector<int64_t> perm_axis(input.shape.size());
   GetShuffledDim(input.shape, &shuffled_dims, dims, &perm_axis);
 
   shuffled_input->Allocate(shuffled_dims, input.dtype);
-  // TODO(zhoushunjie) : Need to implement trans function
-  // phi::funcs::TransposeNormal<DeviceContext, OutT> trans;
-  // trans(dev_ctx, input, shuffled_input, perm_axis);
+  Transpose(input, shuffled_input, perm_axis);
 }
 
 //////////////// HandleLargeDim
 template <typename OutT, typename Functor>
 void HandleLargeDim(const FDTensor& input, FDTensor* output,
                     const std::vector<int64_t>& dims, bool keep_dim) {
+  auto out_dims = input.shape;
+  std::vector<int64_t> dims_ref = dims;
+  auto x_rank = input.shape.size();
+  for (size_t i = 0; i < dims_ref.size(); ++i) {
+    if (dims_ref[i] < 0) dims_ref[i] = x_rank + dims_ref[i];
+    out_dims[dims_ref[i]] = 1;
+  }
+  if (!keep_dim) {
+    const int kDelFlag = -2;
+    for (size_t i = 0; i < dims_ref.size(); ++i) {
+      out_dims[dims_ref[i]] = kDelFlag;
+    }
+    out_dims.erase(remove(out_dims.begin(), out_dims.end(), kDelFlag),
+                   out_dims.end());
+  }
+  output->Allocate(out_dims, TypeToDataType<OutT>::dtype);
   //  shuffle the reduced dim to the end
   FDTensor shuffled_input;
   GetShuffledInput<OutT>(input, &shuffled_input, dims);
@@ -126,11 +142,9 @@ void HandleLargeDim(const FDTensor& input, FDTensor* output,
   const int64_t reduced = shuffled_input.Numel() / unreduced;
   shuffled_input.Allocate({unreduced, reduced}, TypeToDataType<OutT>::dtype);
 
-  auto output_dim = output->shape;
-  output->Allocate({unreduced}, TypeToDataType<OutT>::dtype);
-
+  output->shape = {unreduced};
   ReduceFunctor<OutT, 2, 1, Functor>(shuffled_input, output, {1}, keep_dim);
-  output->shape = output_dim;
+  output->shape = out_dims;
 }
 
 ////////////// ReduceKernel
@@ -152,7 +166,7 @@ void ReduceKernelImpl(const FDTensor& input, FDTensor* output,
   } else {
     int ndim = input.shape.size();
     int rdim = dims.size();
-    if (ndim > 3) {
+    if (ndim > 4) {
       HandleLargeDim<OutT, Functor>(input, output, dims, keep_dim);
     } else {
       HANDLE_REDUCE_DIM(4, 3);

csrc/fastdeploy/function/transpose.cc

@@ -0,0 +1,115 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "fastdeploy/function/transpose.h"
+#include "fastdeploy/function/eigen.h"
+#include "fastdeploy/utils/utils.h"
+
+namespace fastdeploy {
+#ifdef ENABLE_FDTENSOR_FUNC
+
+template <typename T>
+struct TransposeNormalKernel {
+  void operator()(const FDTensor& in, FDTensor* out,
+                  const std::vector<int64_t>& axis) {
+    const int rank = axis.size();
+    auto in_stride = GetStride(in.shape);
+    auto out_stride = GetStride(out->shape);
+    const T* in_ptr = reinterpret_cast<const T*>(in.Data());
+    T* out_ptr = reinterpret_cast<T*>(out->Data());
+
+    auto transpose_helper = [&](int64_t beg, int64_t end) {
+      for (int64_t out_idx = beg; out_idx < end; ++out_idx) {
+        int64_t in_idx = 0;
+        int64_t tmp_idx = out_idx;
+        // calculate the input index
+        for (int i = 0; i < rank; ++i) {
+          const int64_t coordinate = tmp_idx / out_stride[i];
+          tmp_idx -= coordinate * out_stride[i];
+          in_idx += coordinate * in_stride[axis[i]];
+        }
+        out_ptr[out_idx] = in_ptr[in_idx];
+      }
+    };
+    transpose_helper(0, out->Numel());
+  }
+};
+
+template <typename T, int Rank>
+struct TransposeKernelImpl {
+  void operator()(const FDTensor& in, FDTensor* out,
+                  const std::vector<int64_t>& axis) {
+    Eigen::array<int, Rank> permute;
+    for (int i = 0; i < Rank; i++) {
+      permute[i] = axis[i];
+    }
+
+    auto& place = *EigenDeviceWrapper::GetInstance()->GetDevice();
+    auto eigen_in = EigenTensor<T, Rank>::From(in);
+    auto eigen_out = EigenTensor<T, Rank>::From(*out);
+    eigen_out.device(place) = eigen_in.shuffle(permute);
+  }
+};
+
+template <typename T>
+void TransposeKernel(const FDTensor& x, FDTensor* out,
+                     const std::vector<int64_t>& axis) {
+  int rank = axis.size();
+  switch (rank) {
+    case 1:
+      TransposeKernelImpl<T, 1> trans1;
+      trans1(x, out, axis);
+      break;
+    case 2:
+      TransposeKernelImpl<T, 2> trans2;
+      trans2(x, out, axis);
+      break;
+    case 3:
+      TransposeKernelImpl<T, 3> trans3;
+      trans3(x, out, axis);
+      break;
+    case 4:
+      TransposeKernelImpl<T, 4> trans4;
+      trans4(x, out, axis);
+      break;
+    default:
+      // for rank >= 4 situation
+      TransposeNormalKernel<T> trans_normal;
+      trans_normal(x, out, axis);
+  }
+}
+
+void Transpose(const FDTensor& x, FDTensor* out,
+               const std::vector<int64_t>& dims) {
+  size_t dims_size = dims.size();
+  FDASSERT(dims_size == x.shape.size(),
+           "The input tensor's dimension should be equal to the dims's size.");
+  std::vector<int> count(dims_size, 0);
+  for (size_t i = 0; i < dims_size; i++) {
+    FDASSERT(dims[i] >= 0, "The dims should be greater than or equal to 0.");
+    FDASSERT(dims[i] < static_cast<int>(dims_size) && ++count[dims[i]] == 1,
+             "Each element of Attribute axis should be a unique value range "
+             "from 0 to (dims - 1), where the dims is the axis's size, unique "
+             "value means this axis value can appear only once. ");
+  }
+  std::vector<int64_t> out_dims(dims_size);
+  for (size_t i = 0; i < dims_size; i++) {
+    out_dims[i] = x.shape[dims[i]];
+  }
+  out->Allocate(out_dims, x.dtype);
+  FD_VISIT_ALL_TYPES(x.dtype, "TransposeKernel",
+                     ([&] { TransposeKernel<data_t>(x, out, dims); }));
+}
+#endif
+}  // namespace fastdeploy

csrc/fastdeploy/function/transpose.h

@@ -0,0 +1,29 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "fastdeploy/core/fd_tensor.h"
+
+namespace fastdeploy {
+#ifdef ENABLE_FDTENSOR_FUNC
+/** Excute the transpose operation for input FDTensor along given dims.
+    @param x The input tensor.
+    @param out The output tensor which stores the result.
+    @param dims The vector of axis which the input tensor will transpose.
+*/
+FASTDEPLOY_DECL void Transpose(const FDTensor& x, FDTensor* out,
+                               const std::vector<int64_t>& dims);
+#endif
+}  // namespace fastdeploy