test_sparse_matmul_op.py

# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# 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 paddle
from paddle.fluid.framework import _test_eager_guard

import numpy as np
import scipy
import scipy.sparse as sp
import unittest
import os
import re

np.random.seed(2022)


def get_cuda_version():
    result = os.popen("nvcc --version").read()
    regex = r'release (\S+),'
    match = re.search(regex, result)
    if match:
        num = str(match.group(1))
        integer, decimal = num.split('.')
        return int(integer) * 1000 + int(float(decimal) * 10)
    else:
        return -1


@unittest.skipIf(
    not paddle.is_compiled_with_cuda() or get_cuda_version() < 11000,
    "paddle is not compiled with CUDA and cuda version need to >= 11.0")
class TestCsrDenseMatmul2D(unittest.TestCase):
    # x: csr, y: dense, out: dense
    def test_matmul(self):
        with _test_eager_guard():
            mask = np.random.rand(10, 12) < 0.2
            np_x = np.random.rand(10, 12) * mask

            np_csr = sp.csr_matrix(np_x)
            np_dense = np.random.rand(12, 6)
            np_out = np_csr @ np_dense

            np_out_grad = np.ones([10, 6])

            # dx(csr) = dout(dense) * y'(dense) * mask
            np_csr_grad = sp.csr_matrix(
                np.matmul(np_out_grad, np_dense.transpose(1, 0)) * mask)
            # dy(dense) = x'(csr) * dout(dense)
            np_dense_grad = np_csr.transpose() @ np_out_grad

            csr = paddle.to_tensor(np_x, stop_gradient=False).to_sparse_csr()
            dense = paddle.to_tensor(np_dense, stop_gradient=False)
            out = paddle.incubate.sparse.matmul(csr, dense)

            self.assertTrue(np.allclose(np_out, out.numpy()))

            if get_cuda_version() >= 11030:
                out.backward()
                self.assertTrue(
                    np.allclose(np_csr_grad.indptr,
                                csr.grad.crows().numpy()))
                self.assertTrue(
                    np.allclose(np_csr_grad.indices,
                                csr.grad.cols().numpy()))
                self.assertTrue(
                    np.allclose(np_csr_grad.data,
                                csr.grad.values().numpy()))

                self.assertTrue(np.allclose(np_dense_grad, dense.grad.numpy()))


@unittest.skipIf(
    not paddle.is_compiled_with_cuda() or get_cuda_version() < 11030,
    "paddle is not compiled with CUDA and cuda version need to >= 11.3")
class TestCsrMaskedMatmul2D(unittest.TestCase):
    # x: dense, y: dense, out: csr
    def test_matmul(self):
        with _test_eager_guard():
            np_mask = np.random.rand(10, 6) < 0.2

            np_x = np.random.rand(10, 12)
            np_y = np.random.rand(12, 6)
            np_out = sp.csr_matrix(np.matmul(np_x, np_y) * np_mask)

            np_out_grad = sp.csr_matrix(np.ones([10, 6]) * np_mask)
            # dx(dense) = dout(csr) * y'(dense)
            np_x_grad = np_out_grad @ np_y.transpose(1, 0)
            # dy(dense) = x'(dense) * dout(csr) -> dy'(dense) = dout'(csr) * x(dense)
            np_y_grad = (np_out_grad.transpose() @ np_x).transpose(1, 0)

            x = paddle.to_tensor(np_x, stop_gradient=False)
            y = paddle.to_tensor(np_y, stop_gradient=False)
            mask = paddle.to_tensor(np.ones([10, 6]) * np_mask).to_sparse_csr()
            out = paddle.incubate.sparse.masked_matmul(x, y, mask)

            self.assertTrue(np.allclose(np_out.indptr, out.crows().numpy()))
            self.assertTrue(np.allclose(np_out.indices, out.cols().numpy()))
            self.assertTrue(np.allclose(np_out.data, out.values().numpy()))

            out.backward()
            self.assertTrue(np.allclose(out.is_sparse_csr(), True))
            self.assertTrue(np.allclose(np_x_grad, x.grad.numpy()))
            self.assertTrue(np.allclose(np_y_grad, y.grad.numpy()))


@unittest.skipIf(
    not paddle.is_compiled_with_cuda() or get_cuda_version() < 11070,
    "paddle is not compiled with CUDA and cuda version need to >= 11.7")
class TestCsrDenseMatmul3D(unittest.TestCase):
    # x: csr, y: dense, out: dense
    def test_matmul(self):
        with _test_eager_guard():
            paddle.set_default_dtype('float32')
            origin_x = paddle.rand([16, 16, 12])
            mask = paddle.randint(0, 2, [16, 12])
            origin_x = origin_x * mask
            origin_y = paddle.rand([16, 12, 10])

            dense_x = origin_x.detach()
            dense_x.stop_gradient = False
            dense_y = origin_y.detach()
            dense_y.stop_gradient = False
            dense_out = paddle.matmul(dense_x, dense_y)
            dense_out.backward()

            sp_x = origin_x.detach().to_sparse_csr()
            sp_x.stop_gradient = False
            sp_y = origin_y.detach()
            sp_y.stop_gradient = False
            sp_out = paddle.incubate.sparse.matmul(sp_x, sp_y)
            sp_out.backward()

            self.assertTrue(np.allclose(sp_out.numpy(), dense_out.numpy()))
            self.assertTrue(
                np.allclose(sp_x.grad.to_dense().numpy(),
                            (dense_x.grad * mask).numpy()))
            self.assertTrue(np.allclose(sp_y.grad.numpy(),
                                        dense_y.grad.numpy()))


@unittest.skipIf(
    not paddle.is_compiled_with_cuda() or get_cuda_version() < 11070,
    "paddle is not compiled with CUDA and cuda version need to >= 11.7")
class TestCsrMaskedMatmul3D(unittest.TestCase):
    # x: dense, y: dense, out: csr
    def test_matmul(self):
        with _test_eager_guard():
            paddle.set_default_dtype('float64')
            origin_x = paddle.rand([16, 16, 12])
            origin_y = paddle.rand([16, 12, 10])

            mask = paddle.randint(0, 2, [16, 10])

            dense_x = origin_x.detach()
            dense_x.stop_gradient = False
            dense_y = origin_y.detach()
            dense_y.stop_gradient = False
            dense_out = paddle.matmul(dense_x, dense_y)
            dense_out = dense_out * mask
            dense_out.backward()

            sp_x = origin_x.detach()
            sp_x.stop_gradient = False
            sp_y = origin_y.detach()
            sp_y.stop_gradient = False
            sp_out = paddle.incubate.sparse.masked_matmul(
                sp_x, sp_y, dense_out.to_sparse_csr())
            sp_out.backward()

            self.assertTrue(
                np.allclose(sp_out.to_dense().numpy(), dense_out.numpy()))
            self.assertTrue(np.allclose(sp_x.grad.numpy(),
                                        dense_x.grad.numpy()))
            self.assertTrue(np.allclose(sp_y.grad.numpy(),
                                        dense_y.grad.numpy()))


if __name__ == "__main__":
    unittest.main()