wehub-resource-sync
140 lines
5.0 KiB
C++
140 lines
5.0 KiB
C++
/* Copyright (c) 2021 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 <stdint.h>
|
|
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/extensions.h"
|
|
#include "unsupported/Eigen/CXX11/Tensor"
|
|
|
|
namespace phi {
|
|
|
|
// EigenDim converts DDim into Eigen::DSizes.
|
|
template <int D>
|
|
struct EigenDim {
|
|
using Type = Eigen::DSizes<int64_t, D>;
|
|
|
|
static Type From(const DDim& dims) {
|
|
PADDLE_ENFORCE_EQ(arity(dims),
|
|
D,
|
|
common::errors::InvalidArgument(
|
|
"Input dimension size should be equal to %d, but "
|
|
"received dimension size is %d.",
|
|
arity(dims),
|
|
D));
|
|
Type ret;
|
|
for (int64_t d = 0; d < arity(dims); d++) {
|
|
ret[d] = dims[d];
|
|
}
|
|
return ret;
|
|
}
|
|
};
|
|
|
|
// Interpret phi::Tensor as EigenTensor and EigenConstTensor.
|
|
template <typename T, size_t D, int MajorType = Eigen::RowMajor>
|
|
struct EigenTensor {
|
|
// TODO(qijun) Now, default type in unaligned, and we will make a benchmark on
|
|
// the speed of aligned and unaligned version in future.
|
|
using Type = Eigen::TensorMap<Eigen::Tensor<T, D, MajorType, int64_t>>;
|
|
|
|
using ConstType =
|
|
Eigen::TensorMap<Eigen::Tensor<const T, D, MajorType, int64_t>>;
|
|
|
|
static Type From(DenseTensor& tensor, DDim dims) { // NOLINT
|
|
// why tensor.data<T>() not work?
|
|
// return Type(const_cast<T*>(reinterpret_cast<const T*>(tensor.data())),
|
|
// EigenDim<D>::From(dims));
|
|
return Type(const_cast<T*>(tensor.data<T>()), EigenDim<D>::From(dims));
|
|
}
|
|
|
|
static Type From(DenseTensor& tensor) { // NOLINT
|
|
return From(tensor, tensor.dims());
|
|
} // NOLINT
|
|
|
|
static ConstType From(const DenseTensor& tensor, DDim dims) {
|
|
// return ConstType(reinterpret_cast<const T*>(tensor.data()),
|
|
// EigenDim<D>::From(dims));
|
|
return ConstType(tensor.data<T>(), EigenDim<D>::From(dims));
|
|
}
|
|
|
|
static ConstType From(const DenseTensor& tensor) {
|
|
return From(tensor, tensor.dims());
|
|
}
|
|
};
|
|
|
|
template <typename T, int MajorType = Eigen::RowMajor>
|
|
struct EigenMatrix : public EigenTensor<T, 2, MajorType> {
|
|
static typename EigenMatrix::Type Reshape(DenseTensor& tensor, // NOLINT
|
|
int num_col_dims) {
|
|
int rank = tensor.dims().size();
|
|
PADDLE_ENFORCE_EQ((num_col_dims > 0 && num_col_dims < rank),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Input dimension number(num_col_dims) must be "
|
|
"between 0 and %d, but received number is %d.",
|
|
rank,
|
|
num_col_dims));
|
|
return EigenMatrix::From(tensor,
|
|
flatten_to_2d(tensor.dims(), num_col_dims));
|
|
}
|
|
|
|
static typename EigenMatrix::ConstType Reshape(const DenseTensor& tensor,
|
|
int num_col_dims) {
|
|
int rank = tensor.dims().size();
|
|
PADDLE_ENFORCE_EQ((num_col_dims > 0 && num_col_dims < rank),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Input dimension number(num_col_dims) must be "
|
|
"between 0 and %d, but received number is %d.",
|
|
rank,
|
|
num_col_dims));
|
|
return EigenMatrix::From(tensor,
|
|
flatten_to_2d(tensor.dims(), num_col_dims));
|
|
}
|
|
};
|
|
|
|
template <typename T, int MajorType = Eigen::RowMajor>
|
|
struct EigenVector : public EigenTensor<T, 1, MajorType> {
|
|
// Flatten reshapes a Tensor into an EigenVector.
|
|
static typename EigenVector::Type Flatten(DenseTensor& tensor) { // NOLINT
|
|
return EigenVector::From(tensor, {product(tensor.dims())});
|
|
}
|
|
|
|
static typename EigenVector::ConstType Flatten(
|
|
const DenseTensor& tensor) { // NOLINT
|
|
return EigenVector::From(tensor, {product(tensor.dims())});
|
|
}
|
|
};
|
|
|
|
template <typename T, int MajorType = Eigen::RowMajor>
|
|
struct EigenScalar {
|
|
// Scalar tensor (implemented as a rank-0 tensor) of scalar type T.
|
|
using Type = Eigen::TensorMap<
|
|
Eigen::TensorFixedSize<T, Eigen::Sizes<>, MajorType, int64_t>>;
|
|
using ConstType = Eigen::TensorMap<
|
|
Eigen::TensorFixedSize<const T, Eigen::Sizes<>, MajorType, int64_t>>;
|
|
|
|
static Type From(DenseTensor& tensor) { // NOLINT
|
|
return Type(const_cast<T*>(tensor.data<T>()));
|
|
}
|
|
|
|
static ConstType From(const DenseTensor& tensor) {
|
|
return ConstType(tensor.data<T>());
|
|
}
|
|
};
|
|
|
|
} // namespace phi
|