forked from mindspore-Ecosystem/mindspore
!35425 [assistant][SparseTensorDenseAdd] Add new sparse operator SparseTensorDenseAdd
Merge pull request !35425 from 靳铎/SparseTensorDenseAdd
This commit is contained in:
Gitee
commit
600a7b14dc
|
|
@ -0,0 +1,300 @@
|
|||
/**
|
||||
* Copyright 2021-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_tensor_dense_add_cpu_kernel.h"
|
||||
#include <algorithm>
|
||||
#include <utility>
|
||||
#include <complex>
|
||||
#include <functional>
|
||||
#include <type_traits>
|
||||
#include "plugin/device/cpu/hal/device/cpu_device_address.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
namespace {
|
||||
constexpr size_t kIndicesShapeSize = 2;
|
||||
constexpr size_t kSparseTensorDenseAddInputsNum = 4;
|
||||
constexpr size_t kSparseTensorDenseAddOutputsNum = 1;
|
||||
using complex64 = std::complex<float>;
|
||||
using complex128 = std::complex<double>;
|
||||
} // namespace
|
||||
|
||||
void SparseTensorDenseAddCpuKernelMod::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() != kIndicesShapeSize) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', it requires 'x1_indices' must be a " << kIndicesShapeSize
|
||||
<< "-D Tensor, but got " << indices_shape.size() << "-D";
|
||||
}
|
||||
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 'x1_values' must be a 1-D Tensor and the first dimension length "
|
||||
<< "must be equal to the first dimension length of 'indices', but got 'x1_values' shape: "
|
||||
<< Vector2Str(values_shape) << " and 'x1_indices' shape: " << Vector2Str(indices_shape);
|
||||
}
|
||||
auto shape_shape_ = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, SPARSE_SHAPE);
|
||||
x2_shape_ = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, DENSE);
|
||||
size_t x1_rank = shape_shape_[0];
|
||||
size_t x2_rank = x2_shape_.size();
|
||||
if (x1_rank != x2_rank) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_
|
||||
<< "', x1 and x2 must have same ranks, but got 'x1' shape: " << Vector2Str(shape_shape_)
|
||||
<< "and 'x2' shape: " << Vector2Str(x2_shape_);
|
||||
}
|
||||
values_size_ = values_shape[0];
|
||||
output_shape_ = common::AnfAlgo::GetOutputInferShape(kernel_node, 0);
|
||||
auto kernel_attr = GetKernelAttrFromNode(kernel_node);
|
||||
auto [is_match, index] = MatchKernelAttr(kernel_attr, GetOpSupport());
|
||||
if (!is_match) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_
|
||||
<< "SparseTensorDenseAdd does not support this kernel data type: " << kernel_attr;
|
||||
}
|
||||
kernel_func_ = func_list_[index].second;
|
||||
}
|
||||
|
||||
template <typename I, typename T>
|
||||
bool SparseTensorDenseAddCpuKernelMod::LaunchKernel(const std::vector<kernel::AddressPtr> &inputs,
|
||||
const std::vector<kernel::AddressPtr> &outputs) {
|
||||
CHECK_KERNEL_INPUTS_NUM(inputs.size(), kSparseTensorDenseAddInputsNum, kernel_name_);
|
||||
CHECK_KERNEL_OUTPUTS_NUM(outputs.size(), kSparseTensorDenseAddOutputsNum, kernel_name_);
|
||||
if (outputs[0]->size == 0) {
|
||||
MS_LOG(WARNING) << "For '" << kernel_name_ << "', output memory size must be greater than 0, but got 0.";
|
||||
return true;
|
||||
}
|
||||
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;
|
||||
}
|
||||
const auto *indices_addr = reinterpret_cast<I *>(inputs[0]->addr);
|
||||
const auto *values_addr = reinterpret_cast<T *>(inputs[1]->addr);
|
||||
const auto *shape_addr = reinterpret_cast<I *>(inputs[2]->addr);
|
||||
const auto *x2_addr = reinterpret_cast<T *>(inputs[3]->addr);
|
||||
auto *output_addr = reinterpret_cast<T *>(outputs[0]->addr);
|
||||
const size_t indices_length = inputs[0]->size / sizeof(I);
|
||||
const size_t values_length = inputs[1]->size / sizeof(T);
|
||||
const size_t x2_length = inputs[3]->size / sizeof(T);
|
||||
const size_t out_length = outputs[0]->size / sizeof(T);
|
||||
size_t rank = output_shape_.size();
|
||||
for (size_t i = 0; i < x2_length; i++) {
|
||||
if (i > out_length) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the index of 'x2' out of bounds.";
|
||||
}
|
||||
output_addr[i] = x2_addr[i];
|
||||
}
|
||||
for (size_t i = 0; i < rank; i++) {
|
||||
size_t x1_shape_i = shape_addr[i];
|
||||
size_t x2_shape_i = x2_shape_[i];
|
||||
if (x1_shape_i != x2_shape_i) {
|
||||
MS_EXCEPTION(RuntimeError) << "For '" << kernel_name_ << "', Dimension [" << i << "] does not equal"
|
||||
<< "(no broadcasting is supported): x1_shape side " << x1_shape_i
|
||||
<< " vs x2_shape side " << x2_shape_i << ".";
|
||||
}
|
||||
}
|
||||
for (size_t i = 0; i < values_size_; ++i) {
|
||||
if (i >= values_length) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the index of 'x1_values' out of bounds.";
|
||||
}
|
||||
size_t out_index = 0;
|
||||
for (size_t j = 0; j < rank; j++) {
|
||||
if (i * rank + j >= indices_length) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the index of 'x1_indices' out of bounds.";
|
||||
}
|
||||
int index = indices_addr[i * rank + j];
|
||||
if (index >= SizeToInt(output_shape_[j]) || index < 0) {
|
||||
MS_EXCEPTION(ValueError) << "For '" << kernel_name_ << "', the " << i << "th x1_value in " << j
|
||||
<< "th dimension index: " << index << " of 'output' out of bounds: [0, "
|
||||
<< output_shape_[j] << ")";
|
||||
}
|
||||
size_t count = 1;
|
||||
for (size_t k = j + 1; k < rank; k++) {
|
||||
count *= output_shape_[k];
|
||||
}
|
||||
out_index += IntToSize(index) * count;
|
||||
}
|
||||
output_addr[out_index] += values_addr[i];
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<std::pair<KernelAttr, SparseTensorDenseAddCpuKernelMod::SparseTensorDenseAddFunc>>
|
||||
SparseTensorDenseAddCpuKernelMod::func_list_ = {
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt8)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt8)
|
||||
.AddOutputAttr(kNumberTypeInt8),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, int8_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt16)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt16)
|
||||
.AddOutputAttr(kNumberTypeInt16),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, int16_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddOutputAttr(kNumberTypeInt32),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, int32_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddOutputAttr(kNumberTypeInt64),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, int64_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeUInt8)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeUInt8)
|
||||
.AddOutputAttr(kNumberTypeUInt8),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, uint8_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeUInt16)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeUInt16)
|
||||
.AddOutputAttr(kNumberTypeUInt16),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, uint16_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeFloat16)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeFloat16)
|
||||
.AddOutputAttr(kNumberTypeFloat16),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, float16>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeFloat32)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeFloat32)
|
||||
.AddOutputAttr(kNumberTypeFloat32),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, float>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeFloat64)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeFloat64)
|
||||
.AddOutputAttr(kNumberTypeFloat64),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, double>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeComplex64)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeComplex64)
|
||||
.AddOutputAttr(kNumberTypeComplex64),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, complex64>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeComplex128)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeComplex128)
|
||||
.AddOutputAttr(kNumberTypeComplex128),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int32_t, complex128>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt8)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt8)
|
||||
.AddOutputAttr(kNumberTypeInt8),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, int8_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt16)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt16)
|
||||
.AddOutputAttr(kNumberTypeInt16),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, int16_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt32)
|
||||
.AddOutputAttr(kNumberTypeInt32),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, int32_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddOutputAttr(kNumberTypeInt64),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, int64_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeUInt8)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeUInt8)
|
||||
.AddOutputAttr(kNumberTypeUInt8),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, uint8_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeUInt16)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeUInt16)
|
||||
.AddOutputAttr(kNumberTypeUInt16),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, uint16_t>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeFloat16)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeFloat16)
|
||||
.AddOutputAttr(kNumberTypeFloat16),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, float16>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeFloat32)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeFloat32)
|
||||
.AddOutputAttr(kNumberTypeFloat32),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, float>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeFloat64)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeFloat64)
|
||||
.AddOutputAttr(kNumberTypeFloat64),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, double>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeComplex64)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeComplex64)
|
||||
.AddOutputAttr(kNumberTypeComplex64),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, complex64>},
|
||||
{KernelAttr()
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeComplex128)
|
||||
.AddInputAttr(kNumberTypeInt64)
|
||||
.AddInputAttr(kNumberTypeComplex128)
|
||||
.AddOutputAttr(kNumberTypeComplex128),
|
||||
&SparseTensorDenseAddCpuKernelMod::LaunchKernel<int64_t, complex128>}};
|
||||
|
||||
std::vector<KernelAttr> SparseTensorDenseAddCpuKernelMod::GetOpSupport() {
|
||||
std::vector<KernelAttr> support_list;
|
||||
(void)std::transform(func_list_.begin(), func_list_.end(), std::back_inserter(support_list),
|
||||
[](const std::pair<KernelAttr, SparseTensorDenseAddFunc> &pair) { return pair.first; });
|
||||
return support_list;
|
||||
}
|
||||
|
||||
MS_KERNEL_FACTORY_REG(NativeCpuKernelMod, SparseTensorDenseAdd, SparseTensorDenseAddCpuKernelMod);
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,63 @@
|
|||
/**
|
||||
* Copyright 2021-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_TENSOR_DENSE_ADD_CPU_KERNEL_H_
|
||||
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSE_TENSOR_DENSE_ADD_CPU_KERNEL_H_
|
||||
|
||||
#include <memory>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
#include <utility>
|
||||
#include "plugin/device/cpu/kernel/cpu_kernel.h"
|
||||
#include "plugin/factory/ms_factory.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
class SparseTensorDenseAddCpuKernelMod : public DeprecatedNativeCpuKernelMod {
|
||||
public:
|
||||
SparseTensorDenseAddCpuKernelMod() = default;
|
||||
~SparseTensorDenseAddCpuKernelMod() override = default;
|
||||
|
||||
void InitKernel(const CNodePtr &kernel_node) override;
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs) override {
|
||||
return kernel_func_(this, inputs, outputs);
|
||||
}
|
||||
|
||||
protected:
|
||||
std::vector<KernelAttr> GetOpSupport() override;
|
||||
|
||||
private:
|
||||
template <typename I, typename T>
|
||||
bool LaunchKernel(const std::vector<kernel::AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &outputs);
|
||||
template <typename I, typename T>
|
||||
bool LaunchKernelComplex(const std::vector<kernel::AddressPtr> &inputs,
|
||||
const std::vector<kernel::AddressPtr> &outputs);
|
||||
using SparseTensorDenseAddFunc =
|
||||
std::function<bool(SparseTensorDenseAddCpuKernelMod *, const std::vector<kernel::AddressPtr> &,
|
||||
const std::vector<kernel::AddressPtr> &)>;
|
||||
static std::vector<std::pair<KernelAttr, SparseTensorDenseAddFunc>> func_list_;
|
||||
SparseTensorDenseAddFunc kernel_func_;
|
||||
|
||||
ShapeVector x2_shape_;
|
||||
ShapeVector output_shape_;
|
||||
size_t values_size_{0};
|
||||
enum input_list_ { INDICES, VALUES, SPARSE_SHAPE, DENSE };
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSE_TENSOR_DENSE_ADD_CPU_KERNEL_H_
|
||||
|
|
@ -282,6 +282,7 @@ constexpr auto kSparseMatrixNNZ = "SparseMatrixNNZ";
|
|||
|
||||
// Sparse Grad ops
|
||||
constexpr auto kSparseAddGrad = "SparseAddGrad";
|
||||
constexpr auto kSparseTensorDenseAdd = "SparseTensorDenseAdd";
|
||||
|
||||
// Meta Function Graph
|
||||
constexpr auto kJ = "J";
|
||||
|
|
@ -886,6 +887,7 @@ GVAR_DEF(PrimitivePtr, kPrimSparseMatrixNNZ, std::make_shared<Primitive>(kSparse
|
|||
|
||||
// Sparse Grad ops
|
||||
GVAR_DEF(PrimitivePtr, kPrimSparseAddGrad, std::make_shared<Primitive>(kSparseAddGrad));
|
||||
GVAR_DEF(PrimitivePtr, kPrimSparseTensorDenseAdd, std::make_shared<Primitive>(kSparseTensorDenseAdd));
|
||||
|
||||
// TensorList
|
||||
GVAR_DEF(PrimitivePtr, kPrimTensorListFromTensor, std::make_shared<Primitive>("TensorListFromTensor"));
|
||||
|
|
|
|||
|
|
@ -0,0 +1,102 @@
|
|||
/**
|
||||
* 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 <set>
|
||||
#include <vector>
|
||||
#include <memory>
|
||||
#include "ops/sparse_tensor_dense_add.h"
|
||||
#include "ops/op_utils.h"
|
||||
#include "utils/check_convert_utils.h"
|
||||
#include "abstract/ops/primitive_infer_map.h"
|
||||
#include "ops/primitive_c.h"
|
||||
#include "mindapi/src/helper.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace ops {
|
||||
abstract::ShapePtr SparseTensorDenseAddInferShape(const PrimitivePtr &prim,
|
||||
const std::vector<AbstractBasePtr> &input_args) {
|
||||
auto prim_name = prim->name();
|
||||
auto x1_indices_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape())[kShape];
|
||||
auto x1_values_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[1]->BuildShape())[kShape];
|
||||
auto x1_shape_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[2]->BuildShape())[kShape];
|
||||
auto x2_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[3]->BuildShape())[kShape];
|
||||
int64_t x2_shape_size = x2_shape.size();
|
||||
const int kDimensionOne = 1;
|
||||
const int kDimensionTwo = 2;
|
||||
const int kDimensionFive = 5;
|
||||
if (x1_indices_shape.size() != kDimensionTwo) {
|
||||
MS_EXCEPTION(ValueError) << "For " << prim_name
|
||||
<< ", the 'x1_indices' should have rank 2, but got: " << x1_indices_shape.size();
|
||||
}
|
||||
if (x1_shape_shape.size() != kDimensionOne) {
|
||||
MS_EXCEPTION(ValueError) << "For " << prim_name
|
||||
<< ", the 'x1_shape' should have rank 1, but got: : " << x1_shape_shape.size();
|
||||
}
|
||||
if (x1_values_shape.size() != kDimensionOne || x1_values_shape[0] != x1_indices_shape[0]) {
|
||||
MS_EXCEPTION(ValueError) << "For '" << prim_name
|
||||
<< "', the 'x1_values' must be a 1-D tensor and the first dimension length"
|
||||
<< " must be equal to the first dimension length of 'x1_indices', but got "
|
||||
<< x1_values_shape[0] << " vs " << x1_indices_shape[0] << ".";
|
||||
}
|
||||
if (x1_shape_shape[0] != x1_indices_shape[1]) {
|
||||
MS_EXCEPTION(ValueError) << "For '" << prim_name
|
||||
<< "', the length of 'x1_shape' should be equal to the second dimension"
|
||||
<< " length of 'x1_indices', but got " << x1_shape_shape[0] << " vs "
|
||||
<< x1_indices_shape[1] << ".";
|
||||
}
|
||||
if (x1_shape_shape[0] != x2_shape_size) {
|
||||
MS_EXCEPTION(ValueError) << "For '" << prim_name
|
||||
<< "', the rank of 'x1_shape' should be equal to the rank of 'x2_shape', but got "
|
||||
<< x1_shape_shape[0] << " vs " << x2_shape_size << ".";
|
||||
}
|
||||
if (x2_shape.size() > kDimensionFive || x2_shape.size() < kDimensionOne) {
|
||||
MS_EXCEPTION(ValueError) << "For '" << prim_name
|
||||
<< "', Only tensors with ranks between 1 and 5 are currently supported. "
|
||||
<< "Tensor rank: " << x2_shape.size() << ".";
|
||||
}
|
||||
ShapeVector output_shape = x2_shape;
|
||||
return std::make_shared<abstract::Shape>(output_shape);
|
||||
}
|
||||
|
||||
TypePtr SparseTensorDenseAddInferType(const PrimitivePtr &prim, const std::vector<AbstractBasePtr> &input_args) {
|
||||
auto indices_type = input_args[kInputIndex0]->BuildType();
|
||||
auto values_type = input_args[kInputIndex1]->BuildType();
|
||||
auto shape_type = input_args[kInputIndex2]->BuildType();
|
||||
auto x2_type = input_args[kInputIndex3]->BuildType();
|
||||
const std::set<TypePtr> valid_indices_types = {kInt32, kInt64};
|
||||
const std::set<TypePtr> valid_values_types = {kInt8, kInt16, kInt32, kInt64, kUInt8, kUInt16, kUInt32,
|
||||
kUInt64, kFloat16, kFloat32, kFloat64, kComplex64, kComplex128};
|
||||
(void)CheckAndConvertUtils::CheckTensorTypeSame({{"indices", indices_type}, {"shape", shape_type}},
|
||||
valid_indices_types, prim->name());
|
||||
(void)CheckAndConvertUtils::CheckTensorTypeSame({{"values", values_type}, {"x2", x2_type}}, valid_values_types,
|
||||
prim->name());
|
||||
return x2_type;
|
||||
}
|
||||
|
||||
AbstractBasePtr SparseTensorDenseAddInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &prim,
|
||||
const std::vector<AbstractBasePtr> &input_args) {
|
||||
MS_EXCEPTION_IF_NULL(prim);
|
||||
constexpr int inputs_num = 4;
|
||||
CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, inputs_num, prim->name());
|
||||
auto infer_type = SparseTensorDenseAddInferType(prim, input_args);
|
||||
auto infer_shape = SparseTensorDenseAddInferShape(prim, input_args);
|
||||
return abstract::MakeAbstract(infer_shape, infer_type);
|
||||
}
|
||||
MIND_API_OPERATOR_IMPL(SparseTensorDenseAdd, BaseOperator);
|
||||
REGISTER_PRIMITIVE_EVAL_IMPL(SparseTensorDenseAdd, prim::kPrimSparseTensorDenseAdd, SparseTensorDenseAddInfer, nullptr,
|
||||
true);
|
||||
} // namespace ops
|
||||
} // namespace mindspore
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
/**
|
||||
* 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 MINDSPORE_CORE_OPS_SPARSE_TENSOR_DENSE_ADD_H_
|
||||
#define MINDSPORE_CORE_OPS_SPARSE_TENSOR_DENSE_ADD_H_
|
||||
#include <vector>
|
||||
#include <memory>
|
||||
#include "ops/base_operator.h"
|
||||
#include "mindapi/base/types.h"
|
||||
#include "ops/primitive_c.h"
|
||||
#include "abstract/abstract_value.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace ops {
|
||||
constexpr auto kNameSparseTensorDenseAdd = "SparseTensorDenseAdd";
|
||||
/// \brief Add a sparse tensor with a dense tensor.
|
||||
/// Refer to Python API @ref mindspore.ops.SparseTensorDenseAdd for more details.
|
||||
class MIND_API SparseTensorDenseAdd : public BaseOperator {
|
||||
public:
|
||||
MIND_API_BASE_MEMBER(SparseTensorDenseAdd);
|
||||
/// \brief Constructor.
|
||||
SparseTensorDenseAdd() : BaseOperator(kNameSparseTensorDenseAdd) {
|
||||
InitIOName({"x1_indices", "x1_values", "x1_shape", "x2"}, {"output"});
|
||||
}
|
||||
/// \brief Init.
|
||||
void Init() const {}
|
||||
};
|
||||
AbstractBasePtr SparseTensorDenseAddInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
|
||||
const std::vector<AbstractBasePtr> &input_args);
|
||||
} // namespace ops
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CORE_OPS_SPARSE_TENSOR_DENSE_ADD_H_
|
||||
|
|
@ -21,5 +21,6 @@ from . import grad_nn_ops
|
|||
from . import grad_math_ops
|
||||
from . import grad_linalg_ops
|
||||
from . import grad_image_ops
|
||||
from . import grad_sparse
|
||||
|
||||
__all__ = ['get_bprop_fn']
|
||||
|
|
|
|||
|
|
@ -0,0 +1,14 @@
|
|||
"""Define the grad rules of math related operations."""
|
||||
|
||||
from .. import functional as F
|
||||
from .._grad.grad_base import bprop_getters
|
||||
from ..composite.multitype_ops.zeros_like_impl import zeros_like
|
||||
from ..operations.sparse_ops import SparseTensorDenseAdd
|
||||
|
||||
|
||||
@bprop_getters.register(SparseTensorDenseAdd)
|
||||
def get_bprop_sparse_tensor_dense_add(self):
|
||||
"""Grad definition for `SparseTensorDenseAdd` operation."""
|
||||
def bprop(x1_indices, x1_values, x1_shape, x2, out, dout):
|
||||
return (zeros_like(x1_indices), F.gather_nd(dout, x1_indices), zeros_like(x1_shape), dout,)
|
||||
return bprop
|
||||
|
|
@ -103,6 +103,7 @@ from .reverse_sequence import _reverse_sequence_aicpu
|
|||
from .matrix_inverse import _matrix_inverse_aicpu
|
||||
from .matrix_determinant import _matrix_determinant_aicpu
|
||||
from .log_matrix_determinant import _log_matrix_determinant_aicpu
|
||||
from .sparse_tensor_dense_add import _sparse_tensor_dense_add_aicpu
|
||||
from .lstsq import _lstsq_aicpu
|
||||
from .crop_and_resize import _crop_and_resize_aicpu
|
||||
from .crop_and_resize_grad_boxes import _crop_and_resize_grad_boxes_aicpu
|
||||
|
|
|
|||
|
|
@ -0,0 +1,84 @@
|
|||
# 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.
|
||||
# ============================================================================
|
||||
|
||||
"""SparseTensorDenseAdd op"""
|
||||
from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
|
||||
|
||||
sparse_tensor_dense_add_op_info = AiCPURegOp("SparseTensorDenseAdd") \
|
||||
.fusion_type("OPAQUE") \
|
||||
.input(0, "x1_indices", "required") \
|
||||
.input(1, "x1_values", "required") \
|
||||
.input(2, "x1_shape", "required") \
|
||||
.input(3, "x2", "required") \
|
||||
.output(0, "y", "required") \
|
||||
.dtype_format(DataType.I32_Default, DataType.F16_Default, DataType.I32_Default,
|
||||
DataType.F16_Default, DataType.F16_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.F32_Default, DataType.I32_Default,
|
||||
DataType.F32_Default, DataType.F32_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.F64_Default, DataType.I32_Default,
|
||||
DataType.F64_Default, DataType.F64_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.I8_Default, DataType.I32_Default,
|
||||
DataType.I8_Default, DataType.I8_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.I16_Default, DataType.I32_Default,
|
||||
DataType.I16_Default, DataType.I16_Default) \
|
||||
.dtype_format(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) \
|
||||
.dtype_format(DataType.I32_Default, DataType.U8_Default, DataType.I32_Default,
|
||||
DataType.U8_Default, DataType.U8_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.U16_Default, DataType.I32_Default,
|
||||
DataType.U16_Default, DataType.U16_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.U32_Default, DataType.I32_Default,
|
||||
DataType.U32_Default, DataType.U32_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.C64_Default, DataType.I32_Default,
|
||||
DataType.C64_Default, DataType.C64_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.C128_Default, DataType.I32_Default,
|
||||
DataType.C128_Default, DataType.C128_Default) \
|
||||
.dtype_format(DataType.I32_Default, DataType.U64_Default, DataType.I32_Default,
|
||||
DataType.U64_Default, DataType.U64_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.F16_Default, DataType.I64_Default,
|
||||
DataType.F16_Default, DataType.F16_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.F32_Default, DataType.I64_Default,
|
||||
DataType.F32_Default, DataType.F32_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.F64_Default, DataType.I64_Default,
|
||||
DataType.F64_Default, DataType.F64_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.I8_Default, DataType.I64_Default,
|
||||
DataType.I8_Default, DataType.I8_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.I16_Default, DataType.I64_Default,
|
||||
DataType.I16_Default, DataType.I16_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.I32_Default, DataType.I64_Default,
|
||||
DataType.I32_Default, DataType.I32_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.I64_Default,
|
||||
DataType.I64_Default, DataType.I64_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.U8_Default, DataType.I64_Default,
|
||||
DataType.U8_Default, DataType.U8_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.U16_Default, DataType.I64_Default,
|
||||
DataType.U16_Default, DataType.U16_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.U32_Default, DataType.I64_Default,
|
||||
DataType.U32_Default, DataType.U32_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.U64_Default, DataType.I64_Default,
|
||||
DataType.U64_Default, DataType.U64_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.C64_Default, DataType.I64_Default,
|
||||
DataType.C64_Default, DataType.C64_Default) \
|
||||
.dtype_format(DataType.I64_Default, DataType.C128_Default, DataType.I64_Default,
|
||||
DataType.C128_Default, DataType.C128_Default) \
|
||||
.get_op_info()
|
||||
|
||||
|
||||
@op_info_register(sparse_tensor_dense_add_op_info)
|
||||
def _sparse_tensor_dense_add_aicpu():
|
||||
"""SparseTensorDenseAdd AiCPU register"""
|
||||
return
|
||||
|
|
@ -89,6 +89,53 @@ class SparseToDense(PrimitiveWithInfer):
|
|||
return out
|
||||
|
||||
|
||||
class SparseTensorDenseAdd(Primitive):
|
||||
"""
|
||||
Add a sparse tensor and a dense tensor to get a dense tensor.
|
||||
|
||||
Inputs:
|
||||
- **x1_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)`.
|
||||
- **x1_values** (Tensor) - A 1-D Tensor, represents the value corresponding to the position in the `indices`.
|
||||
The shape should be :math:`(n,)`.
|
||||
- **x1_shape** (tuple(int)) - A positive int tuple which specifies the shape of sparse tensor,
|
||||
should have 2 elements, represent sparse tensor shape is :math:`(N, C)`.
|
||||
-**x2** (Tensor)- A dense Tensor, the dtype is same as `values`.
|
||||
|
||||
Returns:
|
||||
Tensor, add result of sparse tensor and dense tensor. The dtype is same as `values`,
|
||||
and the shape is `x1_shape`.
|
||||
|
||||
Raises:
|
||||
TypeError: If the dtype of `x1_indices` and 'x1_shape' is neither int32 nor int64.
|
||||
ValueError: If `x1_shape`, shape of `x1_indices`, shape of `x1_values` and shape
|
||||
of 'x2' don't meet the parameter description.
|
||||
|
||||
Supported Platforms:
|
||||
``Ascend`` ``CPU``
|
||||
|
||||
Examples:
|
||||
>>> from mindspore import Tensor
|
||||
>>> from mindspore.ops import operations as ops
|
||||
>>> from mindspore.common import dtype as mstype
|
||||
>>> x1_indices = Tensor([[0, 0], [0, 1]], dtype=mstype.int64)
|
||||
>>> x1_values = Tensor([1, 1], dtype=mstype.float32)
|
||||
>>> x1_shape = Tensor([3, 3], dtype=mstype.int64)
|
||||
>>> x2= Tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]], dtype=mstype.float32)
|
||||
>>> sparse_tensor_dense_add = ops.SparseTensorDenseAdd()
|
||||
>>> out = sparse_tensor_dense_add(x1_indices, x1_values, x1_shape, x2)
|
||||
>>> print(out)
|
||||
[[2. 2. 1.]
|
||||
[1. 1. 1.]
|
||||
[1. 1. 1.]]
|
||||
"""
|
||||
|
||||
@prim_attr_register
|
||||
def __init__(self):
|
||||
"""Initialize SparseTensorDenseAdd."""
|
||||
self.init_prim_io_names(inputs=['x1_indices', 'x1_values', 'x1_shape', 'x2'], outputs=['y'])
|
||||
|
||||
|
||||
class SparseTensorDenseMatmul(Primitive):
|
||||
"""
|
||||
Multiplies sparse matrix `A` by dense matrix `B`.
|
||||
|
|
|
|||
|
|
@ -104,6 +104,7 @@ from mindspore.ops.operations.nn_ops import ReLUV3
|
|||
from mindspore.ops.operations.sparse_ops import DenseToCSRSparseMatrix, Sspaddmm
|
||||
from mindspore.ops.operations.sparse_ops import SparseTensorDenseMatmul
|
||||
from mindspore.ops.operations.sparse_ops import SparseMatrixNNZ
|
||||
from mindspore.ops.operations.sparse_ops import SparseTensorDenseAdd
|
||||
from mindspore.ops.operations.other_ops import BlackmanWindow
|
||||
from mindspore.ops.operations.nn_ops import SparseApplyCenteredRMSProp
|
||||
from mindspore.nn.layer import normalization
|
||||
|
|
@ -3993,7 +3994,14 @@ test_case_sparse_ops = [
|
|||
'desc_inputs': [Tensor(np.array([[0, 0], [1, 1]]), mstype.int64),
|
||||
Tensor(np.array([1, 1]), mstype.int64),
|
||||
Tensor(np.array([[1, 2], [3, 4]]), mstype.int64)],
|
||||
'skip': ['backward']})
|
||||
'skip': ['backward']}),
|
||||
('SparseTensorDenseAdd', {
|
||||
'block': SparseTensorDenseAdd(),
|
||||
'desc_inputs': [Tensor([[0]], mstype.int32),
|
||||
Tensor([1], mstype.float32),
|
||||
Tensor([1], mstype.int32),
|
||||
Tensor([1], mstype.float32)],
|
||||
'desc_bprop': [Tensor([1], mstype.float32)]}),
|
||||
]
|
||||
|
||||
test_case_lists = [test_case_nn_ops, test_case_math_ops, test_case_array_ops,
|
||||
|
|
|
|||
Loading…
Reference in New Issue