From 7e5b4d69dde992bdb61754128f73c2350d0c06df Mon Sep 17 00:00:00 2001 From: luobingchun <2974874275@qq.com> Date: Wed, 3 Aug 2022 18:34:58 +0800 Subject: [PATCH] [assistant][ops] Add sparse operator SparseAddmm --- .../cpu/kernel/sparse_addmm_cpu_kernel.cc | 370 ++++++++++++++++++ .../cpu/kernel/sparse_addmm_cpu_kernel.h | 58 +++ mindspore/core/ops/core_ops.h | 2 + mindspore/core/ops/sparse_addmm.cc | 124 ++++++ mindspore/core/ops/sparse_addmm.h | 41 ++ .../mindspore/ops/_op_impl/aicpu/__init__.py | 1 + .../ops/_op_impl/aicpu/sparseaddmm.py | 87 ++++ .../mindspore/ops/operations/sparse_ops.py | 54 ++- .../utils/block_util.py | 3 + tests/ut/python/ops/test_ops.py | 11 + 10 files changed, 750 insertions(+), 1 deletion(-) create mode 100644 mindspore/ccsrc/plugin/device/cpu/kernel/sparse_addmm_cpu_kernel.cc create mode 100644 mindspore/ccsrc/plugin/device/cpu/kernel/sparse_addmm_cpu_kernel.h create mode 100644 mindspore/core/ops/sparse_addmm.cc create mode 100644 mindspore/core/ops/sparse_addmm.h create mode 100644 mindspore/python/mindspore/ops/_op_impl/aicpu/sparseaddmm.py diff --git a/mindspore/ccsrc/plugin/device/cpu/kernel/sparse_addmm_cpu_kernel.cc b/mindspore/ccsrc/plugin/device/cpu/kernel/sparse_addmm_cpu_kernel.cc new file mode 100644 index 00000000000..3d849373d22 --- /dev/null +++ b/mindspore/ccsrc/plugin/device/cpu/kernel/sparse_addmm_cpu_kernel.cc @@ -0,0 +1,370 @@ +/** + * Copyright 2022 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 "plugin/device/cpu/kernel/sparse_addmm_cpu_kernel.h" +#include +#include +#include + +namespace mindspore { +namespace kernel { +namespace { +constexpr size_t kSparseAddmmInputsNum = 7; +constexpr size_t kSparseAddmmOutputsNum = 1; +constexpr size_t kSparseAddmmOutputShapeSize = 2; +constexpr size_t kSparseAddmmDenseShapeSize = 2; +constexpr size_t kIndicesSizeNum = 2; +constexpr size_t kIndices2rdDimNum = 2; +constexpr size_t kShapeValue = 0; +constexpr size_t kIndex0 = 0; +constexpr size_t kIndex1 = 1; +constexpr size_t kIndex2 = 2; +constexpr size_t kIndex3 = 3; +constexpr size_t kIndex4 = 4; +constexpr size_t kIndex5 = 5; +constexpr size_t kIndex6 = 6; +} // namespace + +void SparseAddmmCpuKernelMod::InitKernel(const CNodePtr &kernel_node) { + MS_EXCEPTION_IF_NULL(kernel_node); + kernel_name_ = common::AnfAlgo::GetCNodeName(kernel_node); + auto indices_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, INDICES); + if (indices_shape.size() != kIndicesSizeNum && indices_shape[1] != kIndices2rdDimNum) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ + << "', it requires 'indices' should be a 2-D Tensor and the second dimension length " + "should be 2, but got 'indices' shape: " + << Vector2Str(indices_shape); + } + auto values_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, VALUES); + if (values_shape.size() != 1 || values_shape[0] != indices_shape[0]) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ + << "', it requires 'values' should be a 1-D Tensor and the first dimension length " + " should be equal to the first dimension length of 'indices', but got 'values' shape: " + << Vector2Str(values_shape) << " and 'indices' shape: " << Vector2Str(indices_shape); + } + output_shape_ = Convert2SizeT(common::AnfAlgo::GetOutputInferShape(kernel_node, 0)); + values_size_ = LongToSize(values_shape[0]); + b_shape_ = Convert2SizeT(common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, DENSE)); + if (b_shape_.size() != kSparseAddmmDenseShapeSize) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the dimension of 'dense' should be " + << kSparseAddmmDenseShapeSize << "-D, but got " << b_shape_.size() << "-D"; + } + if (output_shape_.size() != kSparseAddmmOutputShapeSize) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the dimension of output should be " + << kSparseAddmmOutputShapeSize << "-D, but got " << output_shape_.size() << "-D"; + } + + auto kernel_attr = GetKernelAttrFromNode(kernel_node); + auto [is_match, index] = MatchKernelAttr(kernel_attr, GetOpSupport()); + if (!is_match) { + MS_LOG(EXCEPTION) << "SparseAddmm does not support this kernel data type: " << kernel_attr; + } + kernel_func_ = func_list_[index].second; +} + +template +bool SparseAddmmCpuKernelMod::LaunchKernel(const std::vector &inputs, + const std::vector &outputs) { + CHECK_KERNEL_INPUTS_NUM(inputs.size(), kSparseAddmmInputsNum, kernel_name_); + CHECK_KERNEL_OUTPUTS_NUM(outputs.size(), kSparseAddmmOutputsNum, kernel_name_); + auto ret = memset_s(outputs[0]->addr, outputs[0]->size, 0, outputs[0]->size); + if (ret != EOK) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', memset output failed. Error no: " << ret; + } + + auto *a_indices = reinterpret_cast(inputs[kIndex0]->addr); + auto *a_values = reinterpret_cast(inputs[kIndex1]->addr); + auto *x1_shape = reinterpret_cast(inputs[kIndex2]->addr); + auto *b = reinterpret_cast(inputs[kIndex3]->addr); + auto *c = reinterpret_cast(inputs[kIndex4]->addr); + auto *alpha = reinterpret_cast(inputs[kIndex5]->addr); + auto *beta = reinterpret_cast(inputs[kIndex6]->addr); + auto *out = reinterpret_cast(outputs[kIndex0]->addr); + + const size_t indices_length = inputs[kIndex0]->size / sizeof(I); + const size_t values_length = inputs[kIndex1]->size / sizeof(T); + const size_t b_length = inputs[kIndex3]->size / sizeof(T); + + const size_t dim_num = 2; + const size_t out_dim_0 = output_shape_[0]; + const size_t out_dim_1 = output_shape_[1]; + const size_t b_dim_0 = b_shape_[0]; + const size_t b_dim_1 = b_shape_[1]; + const size_t same_dim = b_dim_0; + + const I x1_shape_0 = x1_shape[0]; + const I x1_shape_1 = x1_shape[1]; + + const size_t x1_shape_0_s = IntToSize(x1_shape_0); + const size_t x1_shape_1_s = IntToSize(x1_shape_1); + if (x1_shape_0_s <= kShapeValue || x1_shape_1_s <= kShapeValue) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the value of 'x1_shape' should be greater than 0."; + } + if (x1_shape_1_s != b_dim_0) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ + << "', the col of 'x1_shape' should be equal to the row of 'x2_dense'," + " but got col: " + << x1_shape_1_s << ", row: " << b_dim_0; + } + + for (size_t i = 0; i < values_size_; ++i) { + if (i * dim_num + 1 >= indices_length) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the index of 'indices' out of bounds."; + } + if (i >= values_length) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the index of 'values' out of bounds."; + } + + const int row = a_indices[i * dim_num]; + const int col = a_indices[i * dim_num + 1]; + if (row >= SizeToInt(out_dim_0) || row < 0 || col >= SizeToInt(same_dim) || col < 0) { + MS_EXCEPTION(ValueError) << "The indices including out of bounds index, row range: [0, " << out_dim_0 + << "), col range: [0, " << same_dim << "), but got row: " << row << ", col: " << col; + } + + const size_t row_s = IntToSize(row); + const size_t col_s = IntToSize(col); + for (size_t n = 0; n < out_dim_1; ++n) { + if (col_s * b_dim_1 + n >= b_length) { + MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the index of 'b' out of bounds."; + } + const T b_value = b[col_s * b_dim_1 + n]; + out[row_s * out_dim_1 + n] += *(alpha)*a_values[i] * b_value; + } + } + + for (size_t i = 0; i < out_dim_0; ++i) { + for (size_t j = 0; j < out_dim_1; ++j) { + const T c_value = c[i * out_dim_1 + j]; + out[i * out_dim_1 + j] += *(beta)*c_value; + } + } + + return true; +} + +std::vector> SparseAddmmCpuKernelMod::func_list_ = { + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt8) + .AddOutputAttr(kNumberTypeInt8), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt8) + .AddOutputAttr(kNumberTypeInt8), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt16) + .AddOutputAttr(kNumberTypeInt16), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt16) + .AddOutputAttr(kNumberTypeInt16), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeInt32), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeInt32), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeInt64), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeInt64), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeFloat64) + .AddOutputAttr(kNumberTypeFloat64), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeFloat64) + .AddOutputAttr(kNumberTypeFloat64), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeUInt8) + .AddOutputAttr(kNumberTypeUInt8), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeUInt8) + .AddOutputAttr(kNumberTypeUInt8), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeUInt16) + .AddOutputAttr(kNumberTypeUInt16), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeUInt16) + .AddOutputAttr(kNumberTypeUInt16), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddOutputAttr(kNumberTypeUInt32), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddOutputAttr(kNumberTypeUInt32), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeUInt64) + .AddOutputAttr(kNumberTypeUInt64), + &SparseAddmmCpuKernelMod::LaunchKernel}, + {KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeUInt64) + .AddOutputAttr(kNumberTypeUInt64), + &SparseAddmmCpuKernelMod::LaunchKernel}}; + +std::vector SparseAddmmCpuKernelMod::GetOpSupport() { + std::vector support_list; + (void)std::transform(func_list_.begin(), func_list_.end(), std::back_inserter(support_list), + [](const std::pair &pair) { return pair.first; }); + return support_list; +} + +MS_KERNEL_FACTORY_REG(NativeCpuKernelMod, SparseAddmm, SparseAddmmCpuKernelMod); +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/plugin/device/cpu/kernel/sparse_addmm_cpu_kernel.h b/mindspore/ccsrc/plugin/device/cpu/kernel/sparse_addmm_cpu_kernel.h new file mode 100644 index 00000000000..ba8a85bda8a --- /dev/null +++ b/mindspore/ccsrc/plugin/device/cpu/kernel/sparse_addmm_cpu_kernel.h @@ -0,0 +1,58 @@ +/** + * Copyright 2022 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 MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSE_ADDMM_CPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSE_ADDMM_CPU_KERNEL_H_ + +#include +#include +#include "plugin/device/cpu/kernel/cpu_kernel.h" +#include "plugin/factory/ms_factory.h" + +namespace mindspore { +namespace kernel { +class SparseAddmmCpuKernelMod : public DeprecatedNativeCpuKernelMod { + public: + SparseAddmmCpuKernelMod() = default; + ~SparseAddmmCpuKernelMod() override = default; + + void InitKernel(const CNodePtr &kernel_node) override; + + bool Launch(const std::vector &inputs, const std::vector &workspace, + const std::vector &outputs) override { + return kernel_func_(this, inputs, outputs); + } + + protected: + std::vector GetOpSupport() override; + + private: + template + bool LaunchKernel(const std::vector &inputs, const std::vector &outputs); + using SparseAddmmFunc = std::function &, + const std::vector &)>; + static std::vector> func_list_; + SparseAddmmFunc kernel_func_; + + std::vector output_shape_; + std::vector b_shape_; + size_t output_size_{0}; + size_t values_size_{0}; + enum input_list_ { INDICES, VALUES, SPARSE_SHAPE, DENSE, DENSE_X2, ALPHA, BETA }; +}; +} // namespace kernel +} // namespace mindspore +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_RMSPROP_CPU_KERNEL_H_ diff --git a/mindspore/core/ops/core_ops.h b/mindspore/core/ops/core_ops.h index c98068884c0..f672983eb82 100644 --- a/mindspore/core/ops/core_ops.h +++ b/mindspore/core/ops/core_ops.h @@ -288,6 +288,7 @@ constexpr auto kCOOTensorDenseMatmul = "COOTensorDenseMatmul"; // Sparse ops constexpr auto kSparseTensorDenseMatmul = "SparseTensorDenseMatmul"; +constexpr auto kSparseAddmm = "SparseAddmm"; constexpr auto kCSRReduceSum = "CSRReduceSum"; constexpr auto kCSRMV = "CSRMV"; constexpr auto kCSRMM = "CSRMM"; @@ -923,6 +924,7 @@ GVAR_DEF(PrimitivePtr, kPrimCSRTensorGetDenseShape, std::make_shared( // Sparse ops GVAR_DEF(PrimitivePtr, kPrimSparseTensorDenseMatmul, std::make_shared(kSparseTensorDenseMatmul)); +GVAR_DEF(PrimitivePtr, kPrimSparseAddmm, std::make_shared(kSparseAddmm)); GVAR_DEF(PrimitivePtr, kPrimCOOTensorDenseMatmul, std::make_shared(kCOOTensorDenseMatmul)); GVAR_DEF(PrimitivePtr, kPrimCSRReduceSum, std::make_shared(kCSRReduceSum)); GVAR_DEF(PrimitivePtr, kPrimCSRMV, std::make_shared(kCSRMV)); diff --git a/mindspore/core/ops/sparse_addmm.cc b/mindspore/core/ops/sparse_addmm.cc new file mode 100644 index 00000000000..2876fdff9f1 --- /dev/null +++ b/mindspore/core/ops/sparse_addmm.cc @@ -0,0 +1,124 @@ +/** + * Copyright 2022 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 +#include +#include +#include +#include +#include "ops/sparse_addmm.h" +#include "ops/op_utils.h" +#include "utils/check_convert_utils.h" +#include "abstract/ops/primitive_infer_map.h" +#include "mindapi/src/helper.h" + +namespace mindspore { +namespace ops { +namespace { +abstract::ShapePtr SparseAddmmInferShape(const PrimitivePtr &primitive, + const std::vector &input_args) { + auto indices_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex0]->BuildShape())[kShape]; + auto values_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex1]->BuildShape())[kShape]; + auto shape_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex2]->BuildShape())[kShape]; + auto x2_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex3]->BuildShape())[kShape]; + auto x3_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex4]->BuildShape())[kShape]; + auto alpha_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex5]->BuildShape())[kShape]; + auto beta_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex6]->BuildShape())[kShape]; + const int kDimensionTwo = 2; + const int kDimensionOne = 1; + if (indices_shape.size() != kDimensionTwo) { + MS_EXCEPTION(ValueError) << "For '" << primitive->name() << "', the input indices should " + << "have rank 2, but got " << indices_shape.size() << "."; + } + if (indices_shape[1] != kDimensionTwo) { + MS_EXCEPTION(ValueError) << "For '" << primitive->name() << "', the 2nd dimension of indices " + << "should be 2, but got " << indices_shape[1] << "."; + } + if (values_shape.size() != kDimensionOne) { + MS_EXCEPTION(ValueError) << "For '" << primitive->name() << "', the input value should " + << "have rank 1, but got " << values_shape.size() << "."; + } + if (shape_shape.size() != kDimensionOne) { + MS_EXCEPTION(ValueError) << "For '" << primitive->name() << "', the input shape should " + << "have rank 1, but got " << shape_shape.size() << "."; + } + if (shape_shape[0] != kDimensionTwo) { + MS_EXCEPTION(ValueError) << "For '" << primitive->name() << "', the 1st dimension of input shape " + << "should be 2, but got " << shape_shape[0] << "."; + } + if (x2_shape.size() != kDimensionTwo) { + MS_EXCEPTION(ValueError) << "For '" << primitive->name() << "', the shape of input dense " + << "should be [2], but got [" << x2_shape.size() << "]."; + } + if (x3_shape.size() != kDimensionTwo) { + MS_EXCEPTION(ValueError) << "For '" << primitive->name() << "', the shape of input dense " + << "should be [2], but got [" << x3_shape.size() << "]."; + } + if (alpha_shape.size() != kDimensionOne) { + MS_EXCEPTION(ValueError) << "For '" << primitive->name() << "', the input shape should " + << "have rank 1, but got " << alpha_shape.size() << "."; + } + if (beta_shape.size() != kDimensionOne) { + MS_EXCEPTION(ValueError) << "For '" << primitive->name() << "', the input shape should " + << "have rank 1, but got " << beta_shape.size() << "."; + } + return std::make_shared(x3_shape); +} +TypePtr SparseAddmmInferType(const PrimitivePtr &primitive, const std::vector &input_args) { + std::map types; + std::set valid_types = {kFloat32, kFloat64, kInt32, kInt64, kInt16, + kInt8, kUInt32, kUInt64, kUInt16, kUInt8}; + TypePtr indices_type = input_args[kInputIndex0]->BuildType(); + TypePtr values_type = input_args[kInputIndex1]->BuildType(); + TypePtr shape_type = input_args[kInputIndex2]->BuildType(); + TypePtr x2_type = input_args[kInputIndex3]->BuildType(); + TypePtr x3_type = input_args[kInputIndex4]->BuildType(); + TypePtr alpha_type = input_args[kInputIndex5]->BuildType(); + TypePtr beta_type = input_args[kInputIndex6]->BuildType(); + auto prim_name = primitive->name(); + + (void)types.emplace("x1_values", values_type); + (void)types.emplace("x2", x2_type); + (void)types.emplace("x3", x3_type); + (void)types.emplace("alpha", alpha_type); + (void)types.emplace("beta", beta_type); + (void)CheckAndConvertUtils::CheckTensorTypeSame(types, valid_types, primitive->name()); + + const std::set valid_type = {kInt64, kInt32}; + (void)CheckAndConvertUtils::CheckTensorTypeValid("indices", indices_type, valid_type, prim_name); + + (void)CheckAndConvertUtils::CheckTensorTypeValid("sparse_shape", shape_type, valid_type, prim_name); + + auto tensor_type = x3_type->cast(); + auto tensor_element = tensor_type->element(); + + return tensor_element; +} +} // namespace +AbstractBasePtr SparseAddmmInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, + const std::vector &input_args) { + MS_EXCEPTION_IF_NULL(primitive); + const int64_t input_num = 7; + CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, primitive->name()); + // infer type + auto infer_type = SparseAddmmInferType(primitive, input_args); + // infer shape + auto infer_shape = SparseAddmmInferShape(primitive, input_args); + return std::make_shared(infer_type, infer_shape); +} +MIND_API_OPERATOR_IMPL(SparseAddmm, BaseOperator); +REGISTER_PRIMITIVE_EVAL_IMPL(SparseAddmm, prim::kPrimSparseAddmm, SparseAddmmInfer, nullptr, true); +} // namespace ops +} // namespace mindspore diff --git a/mindspore/core/ops/sparse_addmm.h b/mindspore/core/ops/sparse_addmm.h new file mode 100644 index 00000000000..01945fce641 --- /dev/null +++ b/mindspore/core/ops/sparse_addmm.h @@ -0,0 +1,41 @@ +/** + * Copyright 2022 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 MINDSPORE_CORE_OPS_SPARSE_ADDMM_H_ +#define MINDSPORE_CORE_OPS_SPARSE_ADDMM_H_ +#include +#include +#include "ops/base_operator.h" +#include "mindapi/base/types.h" + +namespace mindspore { +namespace ops { +constexpr auto kNameSparseAddmm = "SparseAddmm"; +class MIND_API SparseAddmm : public BaseOperator { + public: + MIND_API_BASE_MEMBER(SparseAddmm); + SparseAddmm() : BaseOperator(kNameSparseAddmm) { + InitIOName({"indices", "values", "sparse_shape", "x2_dense", "x3_dense", "alpha", "beta"}, {"output"}); + } + void Init() const {} +}; +abstract::AbstractBasePtr SparseAddmmInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive, + const std::vector &input_args); +using PrimSparseAddmmPtr = std::shared_ptr; +} // namespace ops +} // namespace mindspore + +#endif diff --git a/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py b/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py index e424ff4d155..2d0c0ef9079 100644 --- a/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py +++ b/mindspore/python/mindspore/ops/_op_impl/aicpu/__init__.py @@ -98,6 +98,7 @@ from .random_poisson import _random_poisson_aicpu from .random_choice_with_mask import _random_choice_with_mask_aicpu from .rsqrt import _rsqrt_aicpu from .rsqrt_grad import _rsqrt_grad_aicpu +from .sparseaddmm import _sparse_addmm_aicpu from .search_sorted import _search_sorted_aicpu from .stack import _stack_aicpu from .unstack import _unstack_aicpu diff --git a/mindspore/python/mindspore/ops/_op_impl/aicpu/sparseaddmm.py b/mindspore/python/mindspore/ops/_op_impl/aicpu/sparseaddmm.py new file mode 100644 index 00000000000..485812ec441 --- /dev/null +++ b/mindspore/python/mindspore/ops/_op_impl/aicpu/sparseaddmm.py @@ -0,0 +1,87 @@ +# Copyright 2022 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. +# ============================================================================ + +"""SparseAddmm op""" +from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType + +sparseaddmm_op_info = AiCPURegOp("SparseAddmm") \ + .fusion_type("OPAQUE") \ + .input(0, "x1_indices", "required") \ + .input(1, "x1_values", "required") \ + .input(2, "x1_shape", "required") \ + .input(3, "x2", "required") \ + .input(4, "x3", "required") \ + .input(5, "alpha", "required") \ + .input(6, "beta", "required") \ + .output(0, "y", "required") \ + .dtype_format(DataType.I32_Default, DataType.I8_Default, DataType.I32_Default, DataType.I8_Default, + DataType.I8_Default, DataType.I8_Default, DataType.I8_Default, DataType.I8_Default) \ + .dtype_format(DataType.I64_Default, DataType.I8_Default, DataType.I64_Default, DataType.I8_Default, + DataType.I8_Default, DataType.I8_Default, DataType.I8_Default, DataType.I8_Default) \ + .dtype_format(DataType.I32_Default, DataType.I16_Default, DataType.I32_Default, DataType.I16_Default, + DataType.I16_Default, DataType.I16_Default, DataType.I16_Default, DataType.I16_Default) \ + .dtype_format(DataType.I64_Default, DataType.I16_Default, DataType.I64_Default, DataType.I16_Default, + DataType.I16_Default, DataType.I16_Default, DataType.I16_Default, DataType.I16_Default) \ + .dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default, + DataType.I32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \ + .dtype_format(DataType.I64_Default, DataType.I32_Default, DataType.I64_Default, DataType.I32_Default, + DataType.I32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \ + .dtype_format(DataType.I32_Default, DataType.I64_Default, DataType.I32_Default, DataType.I64_Default, + DataType.I64_Default, DataType.I64_Default, DataType.I64_Default, DataType.I64_Default) \ + .dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.I64_Default, DataType.I64_Default, + DataType.I64_Default, DataType.I64_Default, DataType.I64_Default, DataType.I64_Default) \ + .dtype_format(DataType.I32_Default, DataType.U8_Default, DataType.I32_Default, DataType.U8_Default, + DataType.U8_Default, DataType.U8_Default, DataType.U8_Default, DataType.U8_Default) \ + .dtype_format(DataType.I64_Default, DataType.U8_Default, DataType.I64_Default, DataType.U8_Default, + DataType.U8_Default, DataType.U8_Default, DataType.U8_Default, DataType.U8_Default) \ + .dtype_format(DataType.I32_Default, DataType.U16_Default, DataType.I32_Default, DataType.U16_Default, + DataType.U16_Default, DataType.U16_Default, DataType.U16_Default, DataType.U16_Default) \ + .dtype_format(DataType.I64_Default, DataType.U16_Default, DataType.I64_Default, DataType.U16_Default, + DataType.U16_Default, DataType.U16_Default, DataType.U16_Default, DataType.U16_Default) \ + .dtype_format(DataType.I32_Default, DataType.U32_Default, DataType.I32_Default, DataType.U32_Default, + DataType.U32_Default, DataType.U32_Default, DataType.U32_Default, DataType.U32_Default) \ + .dtype_format(DataType.I64_Default, DataType.U32_Default, DataType.I64_Default, DataType.U32_Default, + DataType.U32_Default, DataType.U32_Default, DataType.U32_Default, DataType.U32_Default) \ + .dtype_format(DataType.I32_Default, DataType.U64_Default, DataType.I32_Default, DataType.U64_Default, + DataType.U64_Default, DataType.U64_Default, DataType.U64_Default, DataType.U64_Default) \ + .dtype_format(DataType.I64_Default, DataType.U64_Default, DataType.I64_Default, DataType.U64_Default, + DataType.U64_Default, DataType.U64_Default, DataType.U64_Default, DataType.U64_Default) \ + .dtype_format(DataType.I32_Default, DataType.F16_Default, DataType.I32_Default, DataType.F16_Default, + DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \ + .dtype_format(DataType.I64_Default, DataType.F16_Default, DataType.I64_Default, DataType.F16_Default, + DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.F16_Default) \ + .dtype_format(DataType.I32_Default, DataType.F32_Default, DataType.I32_Default, DataType.F32_Default, + DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ + .dtype_format(DataType.I64_Default, DataType.F32_Default, DataType.I64_Default, DataType.F32_Default, + DataType.F32_Default, DataType.F32_Default, DataType.F32_Default, DataType.F32_Default) \ + .dtype_format(DataType.I32_Default, DataType.F64_Default, DataType.I32_Default, DataType.F64_Default, + DataType.F64_Default, DataType.F64_Default, DataType.F64_Default, DataType.F64_Default) \ + .dtype_format(DataType.I64_Default, DataType.F64_Default, DataType.I64_Default, DataType.F64_Default, + DataType.F64_Default, DataType.F64_Default, DataType.F64_Default, DataType.F64_Default) \ + .dtype_format(DataType.I32_Default, DataType.C64_Default, DataType.I32_Default, DataType.C64_Default, + DataType.C64_Default, DataType.C64_Default, DataType.C64_Default, DataType.C64_Default) \ + .dtype_format(DataType.I64_Default, DataType.C64_Default, DataType.I64_Default, DataType.C64_Default, + DataType.C64_Default, DataType.C64_Default, DataType.C64_Default, DataType.C64_Default) \ + .dtype_format(DataType.I32_Default, DataType.C128_Default, DataType.I32_Default, DataType.C128_Default, + DataType.C128_Default, DataType.C128_Default, DataType.C128_Default, DataType.C128_Default) \ + .dtype_format(DataType.I64_Default, DataType.C128_Default, DataType.I64_Default, DataType.C128_Default, + DataType.C128_Default, DataType.C128_Default, DataType.C128_Default, DataType.C128_Default) \ + .get_op_info() + + +@op_info_register(sparseaddmm_op_info) +def _sparse_addmm_aicpu(): + """SparseAddmm AiCPU register""" + return diff --git a/mindspore/python/mindspore/ops/operations/sparse_ops.py b/mindspore/python/mindspore/ops/operations/sparse_ops.py index 4ad84df26cf..0f8063e7e72 100644 --- a/mindspore/python/mindspore/ops/operations/sparse_ops.py +++ b/mindspore/python/mindspore/ops/operations/sparse_ops.py @@ -1,6 +1,6 @@ # coding: utf-8 -# Copyright 2020-2021 Huawei Technologies Co., Ltd +# Copyright 2020-2022 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. @@ -531,6 +531,58 @@ class Sspaddmm(Primitive): 'x3_dense', 'alpha', 'beta'], outputs=['y_indices', 'y_values', 'y_shape']) +class SparseAddmm(Primitive): + """ + Multiplies sparse matrix `A` by dense matrix `B` * `alpha` and add dense matrix `C` * `beta`. + The rank of sparse matrix and dense matrix must equal to `2`. + + Inputs: + - **indices** (Tensor) - A 2-D Tensor, represents the position of the element in the sparse tensor. + Support int32, int64, each element value should be a non-negative int number. The shape is :math:`(n, 2)`. + - **values** (Tensor) - A 1-D Tensor, represents the value corresponding to the position in the `indices`. + Support float32, float64, int8, int16, int32, int64, uint8, uint16, uint32, uint64. + The shape should be :math:`(n,)`. + - **sparse_shape** (Tensor) - A positive int tuple which specifies the shape of sparse tensor. + Support int32, int64, should have 2 elements, represent sparse tensor shape is :math:`(N, C)`. + - **x2_dense** (Tensor) - A 2-D Tensor, the dtype is same as `values`. + - **x3_dense** (Tensor) - A 2-D Tensor, the dtype is same as `values`. + - **alpha** (Tensor) - A 1-D Tensor, the dtype is same as `values`. + - **beta** (Tensor) - A 1-D Tensor, the dtype is same as `values`. + + Outputs: + Tensor, the dtype is the same as `values`. + + Raises: + TypeError: If dtype of `indices`, dtype of `values` and dtype of `dense` don't meet the parameter description. + ValueError: If `sparse_shape`, shape of `indices, shape of `values`, and shape of `dense` don't meet the + parameter description. + + Supported Platforms: + ``Ascend`` ``CPU`` + + Examples: + >>> indices = Tensor([[0, 1], [1, 2]], dtype=ms.int32) + >>> values = Tensor([1, 2], dtype=ms.float32) + >>> sparse_shape = Tensor([1, 2], dtype=ms.int32) + >>> x2_dense = Tensor([[1,1], [2,2], [3,3], [4,4]], dtype=ms.float32) + >>> x3_dense = Tensor([[2,2], [6,6], [0,0]], dtype=ms.float32) + >>> alpha = Tensor([1], dtype=ms.float32) + >>> beta = Tensor([1], dtype=ms.float32) + >>> sparse_addmm = ops.SparseAddmm() + >>> out = sparse_addmm(indices, values, sparse_shape, x2_dense, x3_dense, alpha, beta) + >>> print(out) + [[4 4] + [12 12] + [0 0]] + """ + + @prim_attr_register + def __init__(self): + """Initialize SparseAddmm""" + self.init_prim_io_names(inputs=['indices', 'values', 'sparse_shape', 'x2_dense', 'x3_dense', 'alpha', 'beta'], + outputs=['output']) + + class SparseConcat(Primitive): """ concatenates the input SparseTensor(COO format) along the specified dimension. demo API now diff --git a/tests/mindspore_test_framework/utils/block_util.py b/tests/mindspore_test_framework/utils/block_util.py index ba900cd4c33..f8ef8420cea 100644 --- a/tests/mindspore_test_framework/utils/block_util.py +++ b/tests/mindspore_test_framework/utils/block_util.py @@ -224,6 +224,9 @@ class InputOpNet(nn.Cell): x = self.op(x1, x2, x3, x4, x5, self.c1, self.c2, self.c3, self.c4) return x + def construct7_c0(self, x1, x2, x3, x4, x5, x6, x7): + x = self.op(x1, x2, x3, x4, x5, x6, x7) + return x def construct9_c0(self, x1, x2, x3, x4, x5, x6, x7, x8, x9): x = self.op(x1, x2, x3, x4, x5, x6, x7, x8, x9) diff --git a/tests/ut/python/ops/test_ops.py b/tests/ut/python/ops/test_ops.py index 5b705624aa6..232aee5d18a 100755 --- a/tests/ut/python/ops/test_ops.py +++ b/tests/ut/python/ops/test_ops.py @@ -125,6 +125,7 @@ from mindspore.ops.operations.sparse_ops import SparseMatrixNNZ from mindspore.ops.operations.sparse_ops import SparseTensorDenseAdd from mindspore.ops.operations.sparse_ops import SparseMatrixTranspose from mindspore.ops.operations.sparse_ops import CSRSparseMatrixToSparseTensor +from mindspore.ops.operations.sparse_ops import SparseAddmm from mindspore.ops.operations.sparse_ops import SparseTensorToCSRSparseMatrix from mindspore.ops.operations.sparse_ops import SparseSparseMinimum from mindspore.ops.operations.sparse_ops import SparseSegmentSqrtN @@ -3227,6 +3228,16 @@ test_case_array_ops = [ 'block': UnravelIndex(), 'desc_inputs': [Tensor(np.array([5, 5]).astype(np.int64)), Tensor(np.array([3, 3]).astype(np.int64))], 'skip': ['backward']}), + ('SparseAddmm', { + 'block': SparseAddmm(), + 'desc_inputs': [Tensor(np.array([[0, 1], [1, 2]]).astype(np.int32)), + Tensor(np.array([1, 2]).astype(np.int32)), + Tensor(np.array([3, 4]).astype(np.int32)), + Tensor(np.array([[1, 1], [2, 2], [3, 3], [4, 4]]).astype(np.int32)), + Tensor(np.array([[2, 2], [6, 6], [0, 0]]).astype(np.int32)), + Tensor(np.array([1]).astype(np.int32)), + Tensor(np.array([1]).astype(np.int32))], + 'skip': ['backward']}), ('SpaceToDepth', { 'block': P.SpaceToDepth(2), 'desc_inputs': [[1, 3, 2, 2]],