# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
"""Configure pytest"""

import numpy as np
import pytest

pytest.importorskip("scipy")  # tvm.topi.testing imports scipy

import tvm
import tvm.testing
import tvm.topi.testing
from tvm import te
from tvm.contrib import cblas, dnnl, mkl


def verify_matmul_add(
    matrix_m, matrix_l, matrix_n, lib, transa=False, transb=False, dtype="float32"
):
    """Tests matmul+add op"""
    bias = te.var("bias", dtype=dtype)
    ashape = (matrix_l, matrix_n) if transa else (matrix_n, matrix_l)
    bshape = (matrix_m, matrix_l) if transb else (matrix_l, matrix_m)
    input1_data = te.placeholder(ashape, name="input1_data", dtype=dtype)
    input2_data = te.placeholder(bshape, name="input2_data", dtype=dtype)
    matmul_result = lib.matmul(input1_data, input2_data, transa, transb)
    final_result = te.compute(
        matmul_result.shape, lambda i, j: matmul_result[i, j] + bias, name="final_result"
    )

    def get_numpy(a, b, matrix_bias, transa, transb):
        if transa:
            a = a.transpose()
        if transb:
            b = b.transpose()
        return np.dot(a, b) + matrix_bias

    def compiling(f, name="test_matmul_add", ext=".so"):
        path = name + ext
        f.export_library(path)
        mod = tvm.runtime.load_module(path)
        f = mod[name]
        return f

    def verify(target="llvm"):
        if not tvm.testing.device_enabled(target):
            print(f"skip because {target} is not enabled...")
            return
        if not tvm.get_global_func(lib.__name__ + ".matmul", True):
            print("skip because extern function is not available")
            return
        dev = tvm.cpu(0)
        name = "test_matmul_add"
        f = tvm.compile(
            te.create_prim_func([input1_data, input2_data, final_result, bias]).with_attr(
                "global_symbol", name
            ),
            target=target,
        )
        if target == "c":
            f = compiling(f, name)
        matrix_input1 = tvm.runtime.tensor(
            np.random.uniform(size=ashape).astype(input1_data.dtype), dev
        )
        matrix_input2 = tvm.runtime.tensor(
            np.random.uniform(size=bshape).astype(input2_data.dtype), dev
        )
        matrix_result = tvm.runtime.tensor(
            np.zeros((matrix_n, matrix_m), dtype=final_result.dtype), dev
        )
        matrix_bias = 10.0
        f(matrix_input1, matrix_input2, matrix_result, matrix_bias)
        tvm.testing.assert_allclose(
            matrix_result.numpy(),
            get_numpy(matrix_input1.numpy(), matrix_input2.numpy(), matrix_bias, transa, transb),
            rtol=1e-5,
        )

    verify("llvm")
    verify("c")


def test_matmul_add():
    """Tests of matmul+add op"""
    verify_matmul_add(235, 128, 1024, cblas)
    verify_matmul_add(235, 128, 1024, cblas, True, False)
    verify_matmul_add(235, 128, 1024, cblas, False, True)
    verify_matmul_add(235, 128, 1024, cblas, True, True)
    verify_matmul_add(235, 128, 1024, mkl)
    verify_matmul_add(235, 128, 1024, mkl, True, False)
    verify_matmul_add(235, 128, 1024, mkl, False, True)
    verify_matmul_add(235, 128, 1024, mkl, True, True)
    verify_matmul_add(235, 128, 1024, dnnl)
    verify_matmul_add(235, 128, 1024, dnnl, True, False)
    verify_matmul_add(235, 128, 1024, dnnl, False, True)
    verify_matmul_add(235, 128, 1024, dnnl, True, True)
    verify_matmul_add(1, 16, 4, cblas)
    verify_matmul_add(1, 16, 3, cblas, True, False)
    verify_matmul_add(1, 16, 3, cblas, False, False)
    verify_matmul_add(1, 16, 3, cblas, True, True)
    verify_matmul_add(1, 16, 4, mkl)
    verify_matmul_add(1, 16, 3, mkl, True, False)
    verify_matmul_add(1, 16, 3, mkl, False, False)
    verify_matmul_add(1, 16, 3, mkl, True, True)
    verify_matmul_add(1, 16, 4, dnnl)
    verify_matmul_add(1, 16, 3, dnnl, True, False)
    verify_matmul_add(1, 16, 3, dnnl, False, False)
    verify_matmul_add(1, 16, 3, dnnl, True, True)


def verify_quantized_matmul_add(matrix_m, matrix_l, matrix_n, transa=False, transb=False):
    """Tests quantized matmul+add op"""
    if not tvm.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True):
        pytest.skip("Quantized dense is supported only for MKL. TVM GPU CI uses openblas")
    data_dtype = "uint8"
    kernel_dtype = "int8"
    out_dtype = "int32"
    bias = te.var("bias", dtype=out_dtype)
    ashape = (matrix_l, matrix_n) if transa else (matrix_n, matrix_l)
    bshape = (matrix_m, matrix_l) if transb else (matrix_l, matrix_m)
    input1_data = te.placeholder(ashape, name="input1_data", dtype=data_dtype)
    input2_data = te.placeholder(bshape, name="input2_data", dtype=kernel_dtype)
    matmul_result = mkl.matmul_u8s8s32(input1_data, input2_data, transa, transb, dtype=out_dtype)
    final_result = te.compute(
        matmul_result.shape, lambda i, j: matmul_result[i, j] + bias, name="final_result"
    )

    def get_numpy(a, b, matrix_bias, transa, transb):
        if transa:
            a = a.transpose()
        if transb:
            b = b.transpose()
        return np.dot(a, b) + matrix_bias

    def verify(target="llvm"):
        if not tvm.testing.device_enabled(target):
            print(f"skip because {target} is not enabled...")
            return
        if not tvm.get_global_func("tvm.contrib.mkl.matmul_u8s8s32", True):
            print("skip because extern function is not available")
            return
        dev = tvm.cpu(0)
        f = tvm.compile(
            te.create_prim_func([input1_data, input2_data, final_result, bias]), target=target
        )
        matrix_input1 = tvm.runtime.tensor(
            np.random.randint(low=0, high=50, size=ashape).astype(input1_data.dtype), dev
        )
        matrix_input2 = tvm.runtime.tensor(
            np.random.randint(low=0, high=50, size=bshape).astype(input2_data.dtype), dev
        )
        matrix_result = tvm.runtime.tensor(
            np.zeros((matrix_n, matrix_m), dtype=final_result.dtype), dev
        )
        matrix_bias = 10
        f(matrix_input1, matrix_input2, matrix_result, matrix_bias)
        tvm.testing.assert_allclose(
            matrix_result.numpy(),
            get_numpy(
                matrix_input1.numpy().astype("int32"),
                matrix_input2.numpy().astype("int32"),
                matrix_bias,
                transa,
                transb,
            ),
            rtol=1e-5,
        )

    verify()


def test_quantized_matmul_add():
    """Tests of quantized matmul+add op"""
    verify_quantized_matmul_add(235, 128, 1024)
    verify_quantized_matmul_add(235, 128, 1024, True, False)
    verify_quantized_matmul_add(235, 128, 1024, False, True)
    verify_quantized_matmul_add(235, 128, 1024, True, True)
    verify_quantized_matmul_add(1, 16, 4)
    verify_quantized_matmul_add(1, 16, 3, True, False)
    verify_quantized_matmul_add(1, 16, 3, False, True)
    verify_quantized_matmul_add(1, 16, 3, True, True)


def verify_batch_matmul(
    batch_a,
    batch_b,
    matrix_m,
    matrix_l,
    matrix_n,
    lib,
    transa=False,
    transb=False,
    dtype="float32",
):
    """Tests matmul op where matrices are in batch"""
    batch = max(batch_a, batch_b)
    ashape = (batch_a, matrix_l, matrix_n) if transa else (batch_a, matrix_n, matrix_l)
    bshape = (batch_b, matrix_m, matrix_l) if transb else (batch_b, matrix_l, matrix_m)
    input1_data = te.placeholder(ashape, name="input1_data", dtype=dtype)
    input2_data = te.placeholder(bshape, name="input2_data", dtype=dtype)
    matmul_result = lib.batch_matmul(input1_data, input2_data, transa, transb)
    final_result = te.compute(
        matmul_result.shape, lambda k, i, j: matmul_result[k, i, j], name="final_result"
    )

    def get_numpy(a, b, transa, transb):
        if transa:
            a = a.transpose(0, 2, 1)
        if not transb:
            b = b.transpose(0, 2, 1)
        return tvm.topi.testing.batch_matmul(a, b)

    def compiling(f, name="test_batch_matmul", ext=".so"):
        path = name + ext
        f.export_library(path)
        mod = tvm.runtime.load_module(path)
        f = mod[name]
        return f

    def verify(target="llvm"):
        if not tvm.testing.device_enabled(target):
            print(f"skip because {target} is not enabled...")
            return
        if not tvm.get_global_func(lib.__name__ + ".matmul", True):
            print("skip because extern function is not available")
            return
        dev = tvm.cpu(0)
        name = "test_batch_matmul"
        f = tvm.compile(
            te.create_prim_func([input1_data, input2_data, final_result]), target=target
        )
        if target == "c":
            f = compiling(f, name)
        matrix_input1 = tvm.runtime.tensor(
            np.random.uniform(size=ashape).astype(input1_data.dtype), dev
        )
        matrix_input2 = tvm.runtime.tensor(
            np.random.uniform(size=bshape).astype(input2_data.dtype), dev
        )
        matrix_result = tvm.runtime.tensor(
            np.zeros((batch, matrix_n, matrix_m), dtype=final_result.dtype), dev
        )
        f(matrix_input1, matrix_input2, matrix_result)
        tvm.testing.assert_allclose(
            matrix_result.numpy(),
            get_numpy(matrix_input1.numpy(), matrix_input2.numpy(), transa, transb),
            rtol=1e-5,
        )

    verify("llvm")
    verify("c")


def test_batch_matmul():
    """Tests of matmul op where matrices are in batch"""
    verify_batch_matmul(16, 16, 235, 128, 1024, cblas)
    verify_batch_matmul(16, 16, 235, 128, 1024, cblas, True, False)
    verify_batch_matmul(16, 16, 235, 128, 1024, cblas, False, True)
    verify_batch_matmul(16, 16, 235, 128, 1024, cblas, True, True)
    verify_batch_matmul(16, 16, 235, 128, 1024, mkl)
    verify_batch_matmul(16, 16, 235, 128, 1024, mkl, True, False)
    verify_batch_matmul(16, 16, 235, 128, 1024, mkl, False, True)
    verify_batch_matmul(16, 16, 235, 128, 1024, mkl, True, True)
    verify_batch_matmul(16, 1, 235, 128, 1024, cblas)
    verify_batch_matmul(1, 16, 235, 128, 1024, cblas)
    verify_batch_matmul(16, 1, 235, 128, 1024, cblas)
    verify_batch_matmul(1, 16, 235, 128, 1024, cblas)
    verify_batch_matmul(16, 1, 235, 128, 1024, mkl)
    verify_batch_matmul(1, 16, 235, 128, 1024, mkl)
    verify_batch_matmul(16, 1, 235, 128, 1024, mkl)
    verify_batch_matmul(1, 16, 235, 128, 1024, mkl)
    verify_batch_matmul(1, 1, 1, 16, 3, cblas)
    verify_batch_matmul(1, 1, 1, 16, 3, cblas, True, False)
    verify_batch_matmul(1, 1, 1, 16, 3, cblas, False, False)
    verify_batch_matmul(1, 1, 1, 16, 3, cblas, True, True)
    verify_batch_matmul(1, 1, 1, 16, 3, cblas)
    verify_batch_matmul(1, 1, 1, 16, 3, mkl)
    verify_batch_matmul(1, 1, 1, 16, 3, mkl, True, False)
    verify_batch_matmul(1, 1, 1, 16, 3, mkl, False, False)
    verify_batch_matmul(1, 1, 1, 16, 3, mkl, True, True)
    verify_batch_matmul(1, 1, 1, 16, 3, mkl)


if __name__ == "__main__":
    test_matmul_add()
    test_quantized_matmul_add()
    test_batch_matmul()
