From 7df6bfe7dd87e4b970a138446ea64b74eb4ab365 Mon Sep 17 00:00:00 2001 From: xuguoyang Date: Mon, 22 Mar 2021 16:14:30 +0800 Subject: [PATCH] support sparse tensor dense matmul fot CPU --- .../sparse_tensor_dense_matmul_cpu_kernel.cc | 75 ++++++ .../sparse_tensor_dense_matmul_cpu_kernel.h | 243 ++++++++++++++++++ mindspore/nn/sparse/__init__.py | 3 +- mindspore/nn/sparse/sparse.py | 49 +++- mindspore/ops/operations/__init__.py | 3 +- mindspore/ops/operations/sparse_ops.py | 71 ++++- .../cpu/test_sparse_tensor_dense_matmul_op.py | 53 ++++ 7 files changed, 493 insertions(+), 4 deletions(-) create mode 100644 mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.cc create mode 100644 mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.h create mode 100644 tests/st/ops/cpu/test_sparse_tensor_dense_matmul_op.py diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.cc new file mode 100644 index 0000000000..8860d5d18b --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.cc @@ -0,0 +1,75 @@ +/** + * Copyright 2020-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 "backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.h" + +namespace mindspore { +namespace kernel { +template +void SparseTensorDenseMatmulCPUKernel::InitKernel(const CNodePtr &kernel_node) { + output_size_ = 1; + auto output_shape = AnfAlgo::GetOutputInferShape(kernel_node, 0); + for (auto &dim : output_shape) { + output_size_ *= dim; + } + + aValues_size_ = 1; + auto aValues_shape = AnfAlgo::GetInputDeviceShape(kernel_node, 1); + for (auto &dim : aValues_shape) { + aValues_size_ *= dim; + } + + b_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, 3); + output_shape_ = AnfAlgo::GetOutputDeviceShape(kernel_node, 0); +} + +template +bool SparseTensorDenseMatmulCPUKernel::Launch(const std::vector &inputs, + const std::vector & /*workspace*/, + const std::vector &outputs) { + auto a_indices = reinterpret_cast(inputs[0]->addr); + auto a_values = reinterpret_cast(inputs[1]->addr); + auto b = reinterpret_cast(inputs[3]->addr); + auto out = reinterpret_cast(outputs[0]->addr); + + memset(out, 0, output_size_); + + const size_t nnz = aValues_size_; + const size_t rhs_right = b_shape_[1]; + const size_t lhs_right = b_shape_[0]; + + for (size_t i = 0; i < nnz; ++i) { + const size_t m = a_indices[i * 2]; + const size_t k = a_indices[i * 2 + 1]; + + if (k > lhs_right) { + MS_LOG(ERROR) << "Invalid value: k: " << k << ", lhs_right: " << lhs_right; + return false; + } + if (m > output_shape_[0]) { + MS_LOG(ERROR) << "Invalid value: m: " << m << ", output_shape: " << output_shape_[0]; + return false; + } + + for (size_t n = 0; n < rhs_right; ++n) { + const float b_value = b[k * lhs_right + n]; + out[m * output_shape_[0] + n] += a_values[i] * b_value; + } + } + return true; +} +} // namespace kernel +} // namespace mindspore diff --git a/mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.h new file mode 100644 index 0000000000..bb08051e2c --- /dev/null +++ b/mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.h @@ -0,0 +1,243 @@ +/** + * Copyright 2020-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_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSE_TENSOR_DENSE_MATMUL_CPU_KERNEL_H_ +#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSE_TENSOR_DENSE_MATMUL_CPU_KERNEL_H_ + +#include +#include "backend/kernel_compiler/cpu/cpu_kernel.h" +#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h" + +namespace mindspore { +namespace kernel { +template +class SparseTensorDenseMatmulCPUKernel : public CPUKernel { + public: + SparseTensorDenseMatmulCPUKernel() = default; + ~SparseTensorDenseMatmulCPUKernel() override = default; + + void InitKernel(const CNodePtr &kernel_node) override; + + bool Launch(const std::vector &inputs, const std::vector &workspace, + const std::vector &outputs) override; + + private: + std::vector output_shape_; + std::vector b_shape_; + size_t output_size_; + size_t aValues_size_; +}; +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeBool) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeBool) + .AddOutputAttr(kNumberTypeBool), + SparseTensorDenseMatmulCPUKernel, int32_t, bool); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt8) + .AddOutputAttr(kNumberTypeUInt8), + SparseTensorDenseMatmulCPUKernel, int32_t, uint8_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt16) + .AddOutputAttr(kNumberTypeUInt16), + SparseTensorDenseMatmulCPUKernel, int32_t, uint16_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddOutputAttr(kNumberTypeUInt32), + SparseTensorDenseMatmulCPUKernel, int32_t, uint32_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt64) + .AddOutputAttr(kNumberTypeUInt64), + SparseTensorDenseMatmulCPUKernel, int32_t, uint64_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt8) + .AddOutputAttr(kNumberTypeInt8), + SparseTensorDenseMatmulCPUKernel, int32_t, int8_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt16) + .AddOutputAttr(kNumberTypeInt16), + SparseTensorDenseMatmulCPUKernel, int32_t, int16_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeInt32), + SparseTensorDenseMatmulCPUKernel, int32_t, int32_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeInt64), + SparseTensorDenseMatmulCPUKernel, int32_t, int64_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32), + SparseTensorDenseMatmulCPUKernel, int32_t, float); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat64) + .AddOutputAttr(kNumberTypeFloat64), + SparseTensorDenseMatmulCPUKernel, int32_t, double); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeBool) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeBool) + .AddOutputAttr(kNumberTypeBool), + SparseTensorDenseMatmulCPUKernel, int64_t, bool); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt8) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt8) + .AddOutputAttr(kNumberTypeUInt8), + SparseTensorDenseMatmulCPUKernel, int64_t, uint8_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt16) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt16) + .AddOutputAttr(kNumberTypeUInt16), + SparseTensorDenseMatmulCPUKernel, int64_t, uint16_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt32) + .AddOutputAttr(kNumberTypeUInt32), + SparseTensorDenseMatmulCPUKernel, int64_t, uint32_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeUInt64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeUInt64) + .AddOutputAttr(kNumberTypeUInt64), + SparseTensorDenseMatmulCPUKernel, int64_t, uint64_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt8) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt8) + .AddOutputAttr(kNumberTypeInt8), + SparseTensorDenseMatmulCPUKernel, int64_t, int8_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt16) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt16) + .AddOutputAttr(kNumberTypeInt16), + SparseTensorDenseMatmulCPUKernel, int64_t, int16_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt32) + .AddOutputAttr(kNumberTypeInt32), + SparseTensorDenseMatmulCPUKernel, int64_t, int32_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeInt64) + .AddOutputAttr(kNumberTypeInt64), + SparseTensorDenseMatmulCPUKernel, int64_t, int64_t); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeFloat32) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat32) + .AddOutputAttr(kNumberTypeFloat32), + SparseTensorDenseMatmulCPUKernel, int64_t, float); + +MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul, + KernelAttr() + .AddInputAttr(kNumberTypeInt64) + .AddInputAttr(kNumberTypeFloat64) + .AddInputAttr(kNumberTypeInt32) + .AddInputAttr(kNumberTypeFloat64) + .AddOutputAttr(kNumberTypeFloat64), + SparseTensorDenseMatmulCPUKernel, int64_t, double); + +} // namespace kernel +} // namespace mindspore +#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_RMSPROP_CPU_KERNEL_H_ diff --git a/mindspore/nn/sparse/__init__.py b/mindspore/nn/sparse/__init__.py index 1a4c350a0d..938064cfab 100644 --- a/mindspore/nn/sparse/__init__.py +++ b/mindspore/nn/sparse/__init__.py @@ -15,8 +15,9 @@ """ Sparse related transformation. """ -from .sparse import SparseToDense +from .sparse import (SparseToDense, SparseTensorDenseMatmul) __all__ = [ "SparseToDense", + "SparseTensorDenseMatmul", ] diff --git a/mindspore/nn/sparse/sparse.py b/mindspore/nn/sparse/sparse.py index ce50905cf8..8a1bd85678 100644 --- a/mindspore/nn/sparse/sparse.py +++ b/mindspore/nn/sparse/sparse.py @@ -1,4 +1,4 @@ -# Copyright 2020 Huawei Technologies Co., Ltd +# Copyright 2020-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. @@ -52,3 +52,50 @@ class SparseToDense(Cell): return self.sparse_to_dense(sparse_tensor.indices, sparse_tensor.values, sparse_tensor.dense_shape) + +class SparseTensorDenseMatmul(Cell): + """ + Multiply SparseTensor(of rank 2) "A" by dense tensor. + The shape of sparse tensor is :math:`(N, C)`, and the shape of dense tensor is :math:`(C, M)`, then the shape of + output tensor is :math:`(N, M)`.The output data type is the same as "values". + + Args: + - *adjoint_st** (Bool) - If true, SparseTensor is transposed before multiplication. Default: False. + - *adjoint_dt** (Bool) - If true, DenseTensor is transposed before multiplication. Default: False. + + Inputs: + - **indices** (Tensor) - The indices of sparse representation, support int32/int64. + - **values** (Tensor) - Values corresponding to each row of indices. + - **dense_shape** (tuple) - An int tuple which specifies the shape of dense tensor. The dense_shape is : + math:`(N, C)`. If `adjoint_st` is True, its shape must be :math:`(N, C)` after transpose. + - **dense** (Tensor) - Dense Matrix. The shape of the tensor is :math:`(C, M)`. If + `adjoint_dt` is True, its shape must be :math:`(C, M)` after transpose. + + Returns: + Tensor, the shape of tensor is :math:`(N, M)`.The output data type is the same as "values". + + Examples: + >>> class NetSparseDenseMatmul(nn.Cell): + ... def __init__(self): + ... super(NetSparseDenseMatmul, self).__init__() + ... self.matmul = nn.SparseTensorDenseMatmul() + ... + ... def construct(self, indices, values, dens_shape, dt): + ... return self.matmul(indices, values, dens_shape, dt) + ... + >>> indices = Tensor([[0, 1], [1, 2]], dtype=ms.int32) + >>> values = Tensor([1, 2], dtype=ms.float32) + >>> dense_shape = (3, 4) + >>> dsMatrix = Tensor([[1, 1], [2, 2], [3, 3], [4, 4]], dtype=ms.float32) + >>> test_SparseDenseMatmul = NetSparseDenseMatmul() + >>> out = test_SparseDenseMatmul(indices, values, dens_shape, dsMatrix) + """ + def __init__(self, adjoint_st=False, adjoint_dt=False): + """Initialize SparseTensorDenseMatmul""" + super(SparseTensorDenseMatmul, self).__init__() + self.adjst = adjoint_st + self.adjdt = adjoint_dt + self.matmul = P.SparseTensorDenseMatmul(adjoint_st=self.adjst, adjoint_dt=self.adjdt) + + def construct(self, indices, values, dense_shape, dense): + return self.matmul(indices, values, dense_shape, dense) diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py index 92470c7e80..c6be73ec5d 100644 --- a/mindspore/ops/operations/__init__.py +++ b/mindspore/ops/operations/__init__.py @@ -94,7 +94,7 @@ from ._thor_ops import (CusBatchMatMul, CusCholeskyTrsm, CusFusedAbsMax1, CusImg CusMatMulCubeDenseRight, CusMatMulCubeFraczLeftCast, Im2Col, UpdateThorGradient, Cholesky, CholeskyTrsm, DetTriangle, ProdForceSeA) -from .sparse_ops import SparseToDense +from .sparse_ops import (SparseToDense, SparseTensorDenseMatmul) from ._embedding_cache_ops import (CacheSwapHashmap, SearchCacheIdx, CacheSwapTable, UpdateCache, MapCacheIdx, SubAndFilter, MapUniform, DynamicAssign, PadAndShift) @@ -428,6 +428,7 @@ __all__ = [ "Pull", "ReLUV2", "SparseToDense", + "SparseTensorDenseMatmul", "MatrixInverse", "Range", "IndexAdd", diff --git a/mindspore/ops/operations/sparse_ops.py b/mindspore/ops/operations/sparse_ops.py index b1b4bd4e7e..d3f187a30b 100644 --- a/mindspore/ops/operations/sparse_ops.py +++ b/mindspore/ops/operations/sparse_ops.py @@ -1,6 +1,6 @@ # coding: utf-8 -# Copyright 2020 Huawei Technologies Co., Ltd +# Copyright 2020-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. @@ -53,3 +53,72 @@ class SparseToDense(PrimitiveWithInfer): 'dtype': values['dtype'], 'value': None} return out + +class SparseTensorDenseMatmul(PrimitiveWithInfer): + """ + Multiply SparseTensor(of rank 2) "A" by dense tensor. + The shape of sparse tensor is :math:`(N, C)`, and the shape of dense tensor is :math:`(C, M)`, then the shape of + output tensor is :math:`(N, M)`.The output data type is the same as "values". + tensors. + + Args: + - *adjoint_st** (Bool) - If true, SparseTensor is transposed before multiplication. Default: False. + - *adjoint_dt** (Bool) - If true, DenseTensor is transposed before multiplication. Default: False. + + Inputs: + - **indices** (Tensor) - The indices of sparse representation, support int32/int64. + - **values** (Tensor) - Values corresponding to each row of indices. + - **dense_shape** (tuple) - An int tuple which specifies the shape of dense tensor. The dense_shape is : + math:`(N, C)`. If `adjoint_st` is True, its shape must be :math:`(N, C)` after transpose. + - **dense** (Tensor) - Dense Matrix. The shape of the tensor is :math:`(C, M)`. If + `adjoint_dt` is True, its shape must be :math:`(C, M)` after transpose. + + Outputs: + Tensor, the shape of tensor is :math:`(N, M)`. The output data type is the same as "values". + + Raises: + TypeError: If `indices` is neither int32 nor int64. + TypeError: If 'values' is not boot, uint8-64, int8-64, float16-64. + TypeError: If 'dense' is not boot, uint8-64, int8-64, float16-64. + ValueError: If length of shape of `SparseTensor` or `DenseTensor` is not equal to 2 + + Supported Platforms: + ``CPU`` + + Examples: + >>> indices = Tensor([[0, 1], [1, 2]], dtype=ms.int32) + >>> values = Tensor([1, 2], dtype=ms.float32) + >>> dense_shape = (3, 4) + >>> dsMatrix = Tensor([[1,1], [2,2], [3,3 ], [4, 4]], dtype=ms.float32) + >>> out = ops.SparseTensorDenseMatmul(indices, values, dense_shape, dsMatrix) + """ + @prim_attr_register + def __init__(self, adjoint_st=False, adjoint_dt=False): + """Initialize SparseTensorDenseMatmul""" + self.adjoint_st = adjoint_st + self.adjoint_dt = adjoint_dt + self.init_prim_io_names(inputs=['indices', 'values', 'dense_shape', 'dense'], + outputs=['output']) + self.add_prim_attr('adjoint_st', self.adjoint_st) + self.add_prim_attr('adjoint_dt', self.adjoint_dt) + validator.check_value_type("adjoint_st", adjoint_st, [bool], self.name) + validator.check_value_type("adjoint_dt", adjoint_dt, [bool], self.name) + + def __infer__(self, indices, values, dense_shape, dense): + validator.check_tensor_dtype_valid('indices', indices['dtype'], [mstype.int32, mstype.int64], self.name) + valid_types = mstype.number_type + (mstype.bool_,) + args = {'values': values['dtype'], 'dense': dense['dtype']} + validator.check_tensors_dtypes_same_and_valid(args, valid_types, self.name) + a_shape = dense_shape['value'] + b_shape = dense['shape'] + if len(a_shape) != 2 or len(b_shape) != 2: + raise ValueError('SparseTensorDenseMatmul SparseTensor, DenseTensor should have the same dimension size ' + + f'and equal to 2, while SparseTensor size is ({len(a_shape)}) and DenseTensor size is ' + + f'({len(b_shape)}).') + out_shape = [] + out_shape.append(a_shape[0]) + out_shape.append(b_shape[1]) + out = {'shape': tuple(out_shape), + 'dtype': values['dtype'], + 'value': None} + return out diff --git a/tests/st/ops/cpu/test_sparse_tensor_dense_matmul_op.py b/tests/st/ops/cpu/test_sparse_tensor_dense_matmul_op.py new file mode 100644 index 0000000000..26cbdfd593 --- /dev/null +++ b/tests/st/ops/cpu/test_sparse_tensor_dense_matmul_op.py @@ -0,0 +1,53 @@ +# Copyright 2020-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. +# ============================================================================ + +import numpy as np +import mindspore as ms +import mindspore.context as context +import mindspore.nn as nn +from mindspore import Tensor +from mindspore import SparseTensor + +context.set_context(mode=context.GRAPH_MODE, device_target="CPU") + +class NetSparseDenseMatmul(nn.Cell): + def __init__(self): + super(NetSparseDenseMatmul, self).__init__() + self.matmul = nn.SparseTensorDenseMatmul() + + def construct(self, indices, values, dens_shape, dt): + return self.matmul(indices, values, dens_shape, dt) + +class NetSparseTensor(nn.Cell): + def __init__(self, dense_shape): + super(NetSparseTensor, self).__init__() + self.dense_shape = dense_shape + def construct(self, indices, values): + x = SparseTensor(indices, values, self.dense_shape) + return x.values, x.indices, x.dense_shape + +def test_sparse_tensor_dense_matmul(): + indices = Tensor([[0, 1], [1, 1]]) + values = Tensor([5, 5], dtype=ms.float32) + dens_shape = (3, 3) + spMatrix = np.array([[5, 0, 0], [0, 5, 0], [0, 0, 5]], dtype=np.float32) + dsMatrix = np.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]], dtype=np.float32) + test_SparseDenseMatmul = NetSparseDenseMatmul() + + out_ms = test_SparseDenseMatmul(indices, values, dens_shape, Tensor(dsMatrix)) + out_np = np.matmul(spMatrix, dsMatrix) + error = np.ones(shape=dsMatrix.shape) * 10e-6 + diff = out_ms.asnumpy() - out_np + assert np.all(diff < error)