# 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.
import tvm
from tvm import te
from tvm.contrib import nvcc
import numpy as np
import time
import tvm.testing


@tvm.testing.requires_gpu
def test_exp():
    # graph
    n = tvm.runtime.convert(1024)
    A = te.placeholder((n,), name="A")
    B = te.compute(A.shape, lambda *i: te.exp(A(*i)), name="B")
    s = te.create_schedule(B.op)
    # create iter var and assign them tags.
    num_thread = 8
    bx, tx = s[B].split(B.op.axis[0], factor=num_thread)
    s[B].bind(bx, te.thread_axis("blockIdx.x"))
    s[B].bind(tx, te.thread_axis("threadIdx.x"))

    # one line to build the function.
    def check_device(device, host="stackvm"):
        if not tvm.testing.device_enabled(host):
            return
        ctx = tvm.context(device, 0)
        fexp = tvm.build(s, [A, B], device, host, name="myexp")
        ctx = tvm.context(device, 0)
        # launch the kernel.
        n = 1024
        a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
        b = tvm.nd.array(np.zeros(n, dtype=B.dtype), ctx)
        fexp(a, b)
        tvm.testing.assert_allclose(b.asnumpy(), np.exp(a.asnumpy()), rtol=1e-5)

    check_device("opencl -device=intel_graphics")
    check_device("cuda", "llvm")
    check_device("vulkan")


@tvm.testing.requires_gpu
def test_fmod():
    # graph
    def run(dtype):
        n = te.size_var("n")
        A = te.placeholder((n,), name="A", dtype=dtype)
        B = te.placeholder((n,), name="B", dtype=dtype)
        C = te.compute(A.shape, lambda *i: te.fmod(A(*i), B(*i)), name="C")
        s = te.create_schedule(C.op)
        # create iter var and assign them tags.
        num_thread = 8
        bx, tx = s[C].split(C.op.axis[0], factor=num_thread)

        def check_device(device):
            ctx = tvm.context(device, 0)
            if not tvm.testing.device_enabled(device):
                print("skip because %s is not enabled.." % device)
                return
            target = tvm.target.Target(device)
            if "cpu" not in target.keys:
                s[C].bind(bx, te.thread_axis("blockIdx.x"))
                s[C].bind(tx, te.thread_axis("threadIdx.x"))
            fmod = tvm.build(s, [A, B, C], device, name="myfmod")

            # launch the kernel.
            n = 1024
            a_np = (np.random.uniform(size=n) * 256).astype(A.dtype)
            b_np = (np.random.uniform(size=n) * 256).astype(B.dtype)

            # "fix" the values in a and b to avoid the result being too small
            b_np += (b_np < 2.0) * 2
            a_np[np.abs(np.fmod(a_np, b_np)) < 1] += 1

            a = tvm.nd.array(a_np, ctx)
            b = tvm.nd.array(b_np, ctx)
            c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
            ftimer = fmod.time_evaluator(fmod.entry_name, ctx, number=1)
            tcost = ftimer(a, b, c).mean
            # fmod(a, b, c)
            np.testing.assert_allclose(c.asnumpy(), np.mod(a.asnumpy(), b.asnumpy()), rtol=1e-5)

        check_device("cuda")
        check_device("opencl -device=intel_graphics")
        check_device("metal")

    run("float32")


@tvm.testing.requires_gpu
def test_multiple_cache_write():
    # graph
    n = tvm.runtime.convert(1024)
    A0 = te.placeholder((n,), name="A0", dtype="float32")
    A1 = te.placeholder((n,), name="A1", dtype="float32")
    B0, B1 = te.compute((n,), lambda *i: (A0(*i) + A1(*i), A0(*i) * A1(*i)), name="B")
    C = te.compute((n,), lambda *i: B0(*i) + B1(*i), name="C")
    s = te.create_schedule(C.op)
    # create iter var and assign them tags.
    num_thread = 8
    B0_cache, B1_cache = s.cache_write([B0, B1], "local")
    bx, tx = s[C].split(C.op.axis[0], factor=num_thread)
    s[B0].compute_at(s[C], bx)
    s[B0_cache].compute_at(s[C], bx)
    s[C].bind(bx, te.thread_axis("blockIdx.x"))
    s[C].bind(tx, te.thread_axis("threadIdx.x"))
    # one line to build the function.
    def check_device(device, host="stackvm"):
        if not tvm.testing.device_enabled(host):
            return
        ctx = tvm.context(device, 0)
        if not tvm.testing.device_enabled(device):
            return
        func = tvm.build(s, [A0, A1, C], device, host, name="multiple_cache_write")
        ctx = tvm.context(device, 0)
        # launch the kernel.
        n = 1024
        a0 = tvm.nd.array(np.random.uniform(size=n).astype(A0.dtype), ctx)
        a1 = tvm.nd.array(np.random.uniform(size=n).astype(A1.dtype), ctx)
        c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
        func(a0, a1, c)
        tvm.testing.assert_allclose(
            c.asnumpy(), a0.asnumpy() + a1.asnumpy() + (a0.asnumpy() * a1.asnumpy()), rtol=1e-5
        )

    check_device("cuda", "llvm")
    check_device("vulkan")
    check_device("opencl")


def test_log_pow_llvm():
    # graph
    n = te.size_var("n")
    A = te.placeholder((n,), name="A")
    B = te.compute(A.shape, lambda *i: te.power(te.log(A(*i)), 2.0), name="B")
    s = te.create_schedule(B.op)
    # create iter var and assign them tags.
    bx, tx = s[B].split(B.op.axis[0], factor=32)
    # one line to build the function.
    if not tvm.testing.device_enabled("llvm"):
        return

    flog = tvm.build(s, [A, B], "llvm", name="mylog")
    ctx = tvm.cpu(0)
    # launch the kernel.
    n = 1028
    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx)
    b = tvm.nd.array(np.zeros(n, dtype=B.dtype), ctx)
    repeat = 10
    ftimer = flog.time_evaluator(flog.entry_name, ctx, number=1, repeat=repeat)
    res = ftimer(a, b)
    assert len(res.results) == repeat
    tvm.testing.assert_allclose(b.asnumpy(), np.power(np.log(a.asnumpy()), 2.0), rtol=1e-5)


@tvm.testing.uses_gpu
def test_popcount():
    def run(dtype):
        # graph
        n = tvm.runtime.convert(1024)
        A = te.placeholder((n,), name="A", dtype=dtype)
        B = te.compute(A.shape, lambda *i: tvm.tir.popcount(A(*i)), name="B")
        s = te.create_schedule(B.op)
        # simple schedule
        num_thread = 8
        bx, tx = s[B].split(B.op.axis[0], factor=num_thread)

        def check_device(device):
            ctx = tvm.context(device, 0)
            if not tvm.testing.device_enabled(device):
                print("skip because %s is not enabled.." % device)
                return
            target = tvm.target.Target(device)
            if "cpu" not in target.keys:
                s[B].bind(bx, te.thread_axis("blockIdx.x"))
                s[B].bind(tx, te.thread_axis("threadIdx.x"))
            func = tvm.build(s, [A, B], device)
            # launch the kernel.
            n = 1024
            a = tvm.nd.array(np.random.randint(low=0, high=1000, size=n, dtype=A.dtype), ctx)
            b = tvm.nd.array(np.zeros(shape=n, dtype=B.dtype), ctx)
            func(a, b)
            tvm.testing.assert_allclose(
                b.asnumpy(), list(map(lambda x: bin(x).count("1"), a.asnumpy())), rtol=1e-5
            )

        check_device("llvm")
        check_device("cuda")
        check_device("opencl")
        if dtype == "uint32":
            check_device("metal")
            check_device("vulkan")

    run("uint32")
    run("uint64")


@tvm.testing.requires_gpu
def test_add():
    def run(dtype):
        # graph
        n = te.size_var("n")
        A = te.placeholder((n,), name="A", dtype=dtype)
        B = te.placeholder((n,), name="B", dtype=dtype)
        bias = te.var("bias", dtype=dtype)
        scale = te.var("scale", dtype=dtype)
        C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name="C")
        # schedule
        s = te.create_schedule(C.op)
        # create iter var and assign them tags.
        num_thread = 16
        bx, x = s[C].split(C.op.axis[0], factor=num_thread * 4)
        tx, x = s[C].split(x, nparts=num_thread)
        _, x = s[C].split(x, factor=4)
        s[C].bind(bx, te.thread_axis("blockIdx.x"))
        s[C].bind(tx, te.thread_axis("threadIdx.x"))
        s[C].vectorize(x)

        # one line to build the function.
        def check_device(device):
            ctx = tvm.context(device, 0)
            if not tvm.testing.device_enabled(device):
                print("skip because %s is not enabled.." % device)
                return
            fadd = tvm.build(s, [A, B, C], device, name="myadd")

            # launch the kernel.
            n = 1024
            a = tvm.nd.array((np.random.uniform(size=n) * 256).astype(A.dtype), ctx)
            b = tvm.nd.array((np.random.uniform(size=n) * 256).astype(B.dtype), ctx)
            c = tvm.nd.array(np.zeros(n, dtype=C.dtype), ctx)
            ftimer = fadd.time_evaluator(fadd.entry_name, ctx, number=1)
            tcost = ftimer(a, b, c).mean
            tvm.testing.assert_allclose(c.asnumpy(), a.asnumpy() + b.asnumpy(), rtol=1e-6)

        check_device("opencl")
        check_device("cuda")
        if dtype == "float32":
            check_device("metal")
            check_device("vulkan")

    run("float32")
    run("int32")
    run("int64")
    run("uint64")


@tvm.testing.requires_gpu
def try_warp_memory():
    """skip this in default test because it require higher arch"""
    m = 128
    A = te.placeholder((m,), name="A")
    B = te.compute((m,), lambda i: A[i] + 3, name="B")
    warp_size = 32
    s = te.create_schedule(B.op)
    AA = s.cache_read(A, "warp", [B])
    xo, xi = s[B].split(B.op.axis[0], warp_size * 2)
    xi0, xi1 = s[B].split(xi, factor=warp_size)
    tx = te.thread_axis("threadIdx.x")
    s[B].bind(xi1, tx)
    s[B].bind(xo, te.thread_axis("blockIdx.x"))
    s[AA].compute_at(s[B], xo)
    xo, xi = s[AA].split(s[AA].op.axis[0], warp_size)
    s[AA].bind(xi, tx)

    @tvm.register_func
    def tvm_callback_cuda_compile(code):
        ptx = nvcc.compile_cuda(code, target="ptx")
        return ptx

    # one line to build the function.
    def check_device(device):
        ctx = tvm.context(device, 0)
        if not tvm.testing.device_enabled(device):
            print("skip because %s is not enabled.." % device)
            return
        f = tvm.build(s, [A, B], device)
        a = tvm.nd.array((np.random.uniform(size=m) * 256).astype(A.dtype), ctx)
        b = tvm.nd.array(np.zeros(m, dtype=B.dtype), ctx)
        f(a, b)
        tvm.testing.assert_allclose(b.asnumpy(), a.asnumpy() + 3, rtol=1e-6)

    check_device("cuda")


if __name__ == "__main__":
    test_exp()
    try_warp_memory()
    test_multiple_cache_write()
    test_add()
    test_log_pow_llvm()
    test_popcount()
    test_fmod()
