diff --git a/mindspore/ccsrc/include/common/utils/utils.h b/mindspore/ccsrc/include/common/utils/utils.h index 754d123eb6b..bc341a24ff5 100644 --- a/mindspore/ccsrc/include/common/utils/utils.h +++ b/mindspore/ccsrc/include/common/utils/utils.h @@ -181,6 +181,7 @@ constexpr auto kCeLUOpName = "CeLU"; constexpr auto kCeluV2OpName = "CeluV2"; constexpr auto kChannelShuffleOpName = "ChannelShuffle"; constexpr auto kCheckNumericsOpName = "CheckNumerics"; +constexpr auto kCholeskyOpName = "Cholesky"; constexpr auto kCholeskyGradOpName = "CholeskyGrad"; constexpr auto kCholeskyInverseOpName = "CholeskyInverse"; constexpr auto kCholeskySolveOpName = "CholeskySolve"; diff --git a/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_ops/cpu_kernel/ms_kernel/cholesky.cc b/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_ops/cpu_kernel/ms_kernel/cholesky.cc new file mode 100644 index 00000000000..0a031deb902 --- /dev/null +++ b/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_ops/cpu_kernel/ms_kernel/cholesky.cc @@ -0,0 +1,101 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * 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 "cholesky.h" +#include +#include +#include +#include +#include +#include "utils/eigen_tensor.h" +#include "utils/kernel_util.h" +#include "cpu_kernel_utils.h" + +namespace { +const uint32_t kInputNum = 1; +const uint32_t kOutputNum = 1; +const char *Cholesky = "Cholesky"; +} // namespace + +namespace aicpu { +uint32_t CholeskyCpuKernel::Compute(CpuKernelContext &ctx) { + KERNEL_HANDLE_ERROR(NormalCheck(ctx, kInputNum, kOutputNum), "Cholesky check input and output failed."); + + Tensor *input = ctx.Input(0); + AttrValue *upper = ctx.GetAttr("upper"); + bool upperinfo = (upper == nullptr) ? false : upper->GetBool(); + auto data_type_in = input->GetDataType(); + + switch (data_type_in) { + case DT_FLOAT: + return ComputeKernel(ctx, upperinfo); + case DT_DOUBLE: + return ComputeKernel(ctx, upperinfo); + default: + KERNEL_LOG_ERROR("Cholesky kernel data type [%s] not support.", DTypeStr(data_type_in).c_str()); + return KERNEL_STATUS_PARAM_INVALID; + } +} + +REGISTER_CPU_KERNEL(Cholesky, CholeskyCpuKernel); + +template +uint32_t CholeskyCpuKernel::ComputeKernel(CpuKernelContext &ctx, const bool &upper) { + auto inputptr = reinterpret_cast(ctx.Input(0)->GetData()); + auto outputptr = reinterpret_cast(ctx.Output(0)->GetData()); + std::vector dims = ctx.Input(0)->GetTensorShape()->GetDimSizes(); + int64_t dimsnum = ctx.Input(0)->GetTensorShape()->GetDims(); + int64_t dim = ctx.Input(0)->GetTensorShape()->GetDimSize(dimsnum - 1); + int64_t n = dim; + int64_t count = 1; + int64_t no_batch = 2; + + if (dimsnum > no_batch) { + for (int64_t m = 0; m < dimsnum - no_batch; m++) { + count *= dims[m]; + } + } + + Eigen::Matrix A(n, n); + + for (int64_t k = 0; k < count; k++) { + for (int64_t i = 0; i < n * n; i++) { + A.data()[i] = inputptr[k * n * n + i]; + } + + Eigen::LLT> A_llt(A); + + if (!A.isApprox(A.transpose()) || A_llt.info() == Eigen::NumericalIssue) { + KERNEL_LOG_ERROR("There exists non semi-positive definitie matrix!"); + return KERNEL_STATUS_INNER_ERROR; + } + + Eigen::Matrix L = A_llt.matrixL(); + + Eigen::Matrix U = A_llt.matrixU(); + + if (!upper) { + for (int64_t i = 0; i < n * n; i++) { + outputptr[k * n * n + i] = L.data()[i]; + } + } else { + for (int64_t i = 0; i < n * n; i++) { + outputptr[k * n * n + i] = U.data()[i]; + } + } + } + return KERNEL_STATUS_OK; +} +} // namespace aicpu \ No newline at end of file diff --git a/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_ops/cpu_kernel/ms_kernel/cholesky.h b/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_ops/cpu_kernel/ms_kernel/cholesky.h new file mode 100644 index 00000000000..6ef6ef67e57 --- /dev/null +++ b/mindspore/ccsrc/plugin/device/ascend/kernel/aicpu/aicpu_ops/cpu_kernel/ms_kernel/cholesky.h @@ -0,0 +1,38 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * 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. + */ + +#ifndef AICPU_KERNELS_NORMALIZED_CHOLESKY_H_ +#define AICPU_KERNELS_NORMALIZED_CHOLESKY_H_ + +#include "cpu_ops_kernel.h" +#include "cpu_types.h" +#include "utils/bcast.h" + +namespace aicpu { +class CholeskyCpuKernel : public CpuKernel { + public: + CholeskyCpuKernel() = default; + ~CholeskyCpuKernel() override = default; + + protected: + uint32_t Compute(CpuKernelContext &ctx) override; + + private: + template + uint32_t ComputeKernel(CpuKernelContext &ctx, const bool &upper); +}; +} // namespace aicpu +#endif diff --git a/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.cc b/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.cc index 42846fd88ab..6a9d9af3adf 100644 --- a/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.cc +++ b/mindspore/ccsrc/plugin/device/ascend/optimizer/mindir/aicpu_lib_select.cc @@ -57,6 +57,7 @@ const AnfNodePtr AICpuLibSelectPass::Process(const FuncGraphPtr &graph, const An mindspore::kCacheSwapTableOpName, mindspore::kCauchyOpName, mindspore::kChannelShuffleOpName, + mindspore::kCholeskyOpName, mindspore::kCholeskyGradOpName, mindspore::kCholeskyInverseOpName, mindspore::kCholeskySolveOpName, diff --git a/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py b/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py index 95f98262668..09cb406f94c 100644 --- a/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py +++ b/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py @@ -342,3 +342,4 @@ from .lu_unpack_grad import _lu_unpack_grad_aicpu from .matrix_exp import _matrix_exp_aicpu from .pad_v3_grad import _pad_v3_grad_aicpu from .pad_v3 import _pad_v3_aicpu +from .cholesky import _cholesky_aicpu