| 33 | |
| 34 | template <typename T, bool conjugate> |
| 35 | void TransposeSimple(const CPUDevice& device, const Tensor& in, |
| 36 | const gtl::ArraySlice<int32> perm, Tensor* out) { |
| 37 | const int ndims = in.dims(); |
| 38 | gtl::InlinedVector<int64, 8> in_strides = ComputeStride<int64>(in.shape()); |
| 39 | gtl::InlinedVector<int64, 8> out_strides = ComputeStride<int64>(out->shape()); |
| 40 | const T* p = reinterpret_cast<const T*>(in.tensor_data().data()); |
| 41 | T* q = reinterpret_cast<T*>(const_cast<char*>((out->tensor_data().data()))); |
| 42 | auto transpose_fn = [=, &in_strides, &out_strides, &perm](int64 begin, |
| 43 | int64 end) { |
| 44 | for (int64 o_idx = begin; o_idx < end; ++o_idx) { |
| 45 | int64 i_idx = 0; |
| 46 | int64 t = o_idx; |
| 47 | for (int i = 0; i < ndims; ++i) { |
| 48 | const int64 ratio = t / out_strides[i]; |
| 49 | t -= ratio * out_strides[i]; |
| 50 | i_idx += ratio * in_strides[perm[i]]; |
| 51 | } |
| 52 | if (conjugate) { |
| 53 | q[o_idx] = Eigen::numext::conj(p[i_idx]); |
| 54 | } else { |
| 55 | q[o_idx] = p[i_idx]; |
| 56 | } |
| 57 | } |
| 58 | }; |
| 59 | double cycles_per_element = |
| 60 | (conjugate ? 1 : 0) + ndims * (Eigen::TensorOpCost::DivCost<int64>() + |
| 61 | 2 * Eigen::TensorOpCost::MulCost<int64>() + |
| 62 | 2 * Eigen::TensorOpCost::AddCost<int64>()); |
| 63 | Eigen::TensorOpCost cost(/*bytes_loaded=*/sizeof(T), |
| 64 | /*bytes_stored=*/sizeof(T), cycles_per_element); |
| 65 | device.parallelFor(in.NumElements(), cost, std::move(transpose_fn)); |
| 66 | } |
| 67 | |
| 68 | } // namespace |
| 69 |
nothing calls this directly
no test coverage detected