diff --git a/paddle/fluid/operators/math/matrix_inverse.cc b/paddle/fluid/operators/math/matrix_inverse.cc index 25bc5d725e1fd0b3d896ec44a603ce2834b70b44..60481491cb4b406809993c69503bff90bada8092 100644 --- a/paddle/fluid/operators/math/matrix_inverse.cc +++ b/paddle/fluid/operators/math/matrix_inverse.cc @@ -23,34 +23,10 @@ namespace math { template class MatrixInverseFunctor { - using Matrix = - Eigen::Matrix; - using EigenMatrixMap = Eigen::Map; - using ConstEigenMatrixMap = Eigen::Map; - public: void operator()(const platform::CPUDeviceContext& context, const framework::Tensor& a, framework::Tensor* a_inv) { - const auto& mat_dims = a.dims(); - const int rank = mat_dims.size(); - int n = mat_dims[rank - 1]; - int batch_size = rank > 2 ? a.numel() / (n * n) : 1; - - const T* a_ptr = a.data(); - T* a_inv_ptr = a_inv->mutable_data(context.GetPlace()); - - for (int i = 0; i < batch_size; ++i) { - ConstEigenMatrixMap mat(a_ptr + i * n * n, n, n); - EigenMatrixMap mat_inv(a_inv_ptr + i * n * n, n, n); - Eigen::PartialPivLU lu; - lu.compute(mat); - - const T min_abs_pivot = lu.matrixLU().diagonal().cwiseAbs().minCoeff(); - PADDLE_ENFORCE_GT( - min_abs_pivot, static_cast(0), - platform::errors::InvalidArgument("Input is not invertible.")); - mat_inv.noalias() = lu.inverse(); - } + compute_inverse_eigen(context, a, a_inv); } }; diff --git a/paddle/fluid/operators/math/matrix_inverse.cu.cc b/paddle/fluid/operators/math/matrix_inverse.cu.cc index 7f5df114680552b96053418aba7ba6cf3f3a75a3..5deedf084c6970fb15f1d6c1eb5ab8457aff6b7f 100644 --- a/paddle/fluid/operators/math/matrix_inverse.cu.cc +++ b/paddle/fluid/operators/math/matrix_inverse.cu.cc @@ -12,6 +12,7 @@ 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 "paddle/fluid/operators/math/matrix_inverse.h" #include "paddle/fluid/operators/math/blas.h" namespace paddle { @@ -32,6 +33,7 @@ class MatrixInverseFunctor { public: void operator()(const platform::CUDADeviceContext& context, const framework::Tensor& a, framework::Tensor* a_inv) { +#ifndef PADDLE_WITH_HIP const auto& mat_dims = a.dims(); const int rank = mat_dims.size(); int n = mat_dims[rank - 1]; @@ -111,6 +113,9 @@ class MatrixInverseFunctor { "non-singular matrix", i, info[i], info[i])); } +#else + compute_inverse_eigen(context, a, a_inv); +#endif } }; diff --git a/paddle/fluid/operators/math/matrix_inverse.h b/paddle/fluid/operators/math/matrix_inverse.h index f0baf0b250e757315da6e45cdb9d92c1336e9c26..fb58b483666526c0b7e745d1e49e308235871d57 100644 --- a/paddle/fluid/operators/math/matrix_inverse.h +++ b/paddle/fluid/operators/math/matrix_inverse.h @@ -15,6 +15,8 @@ limitations under the License. */ #pragma once #include +#include "Eigen/Core" +#include "Eigen/LU" #include "paddle/fluid/framework/tensor.h" #include "paddle/fluid/platform/device_context.h" @@ -22,6 +24,36 @@ namespace paddle { namespace operators { namespace math { +template +void compute_inverse_eigen(const DeviceContext& context, + const framework::Tensor& a, + framework::Tensor* a_inv) { + using Matrix = + Eigen::Matrix; + using EigenMatrixMap = Eigen::Map; + using ConstEigenMatrixMap = Eigen::Map; + const auto& mat_dims = a.dims(); + const int rank = mat_dims.size(); + int n = mat_dims[rank - 1]; + int batch_size = rank > 2 ? a.numel() / (n * n) : 1; + + const T* a_ptr = a.data(); + T* a_inv_ptr = a_inv->mutable_data(context.GetPlace()); + + for (int i = 0; i < batch_size; ++i) { + ConstEigenMatrixMap mat(a_ptr + i * n * n, n, n); + EigenMatrixMap mat_inv(a_inv_ptr + i * n * n, n, n); + Eigen::PartialPivLU lu; + lu.compute(mat); + + const T min_abs_pivot = lu.matrixLU().diagonal().cwiseAbs().minCoeff(); + PADDLE_ENFORCE_GT( + min_abs_pivot, static_cast(0), + platform::errors::InvalidArgument("Input is not invertible.")); + mat_inv.noalias() = lu.inverse(); + } +} + template class MatrixInverseFunctor { public: