# 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.

"""
Codegen tests for VLA extensions
"""

import re

import pytest

import tvm
import tvm.testing
from tvm.script import tirx as T
from tvm.target.codegen import llvm_version_major


@pytest.mark.skipif(
    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
)
@pytest.mark.parametrize(
    "target",
    [
        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
        {
            "kind": "llvm",
            "device": "riscv_cpu",
            "mtriple": "riscv64-linux-gnu",
            "mcpu": "generic-rv64",
            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
        },
    ],
)
def test_codegen_vscale(target):
    if not tvm.testing.device_enabled(target):
        pytest.skip(f"{target} not enabled")
    vscale = tvm.tirx.vscale()

    @T.prim_func(s_tir=True)
    def main(A: T.Buffer((5,), "int32")):
        for i in range(5):
            A[i] = 2 * vscale

    with tvm.target.Target(target):
        build_mod = tvm.tirx.build(main)

    llvm = build_mod.inspect_source()
    assert re.findall(r"llvm.vscale.i32", llvm), "No vscale in generated LLVM."


@pytest.mark.skipif(
    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
)
@pytest.mark.parametrize(
    "target",
    [
        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
        {
            "kind": "llvm",
            "device": "riscv_cpu",
            "mtriple": "riscv64-linux-gnu",
            "mcpu": "generic-rv64",
            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
        },
    ],
)
def test_scalable_buffer_load_store(target):
    if not tvm.testing.device_enabled(target):
        pytest.skip(f"{target} not enabled")

    @T.prim_func(s_tir=True)
    def my_func(a: T.handle, b: T.handle):
        A = T.match_buffer(a, (128,), "float32")
        B = T.match_buffer(b, (128,), "float32")
        T.func_attr({"global_symbol": "my_module", "tirx.noalias": True})
        B[T.ramp(0, 1, 4 * T.vscale())] = A[T.ramp(0, 1, 4 * T.vscale())]

    with tvm.target.Target(target):
        mod = tvm.tirx.build(my_func)

    llvm = mod.inspect_source("ll")
    assert re.findall(r"load <vscale x 4 x float>", llvm), "No scalable load in generated LLVM."
    assert re.findall(r" store <vscale x 4 x float>", llvm), "No scalable store in generated LLVM."


@pytest.mark.skipif(
    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
)
@pytest.mark.parametrize(
    "target",
    [
        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
        {
            "kind": "llvm",
            "device": "riscv_cpu",
            "mtriple": "riscv64-linux-gnu",
            "mcpu": "generic-rv64",
            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
        },
    ],
)
def test_scalable_broadcast(target):
    if not tvm.testing.device_enabled(target):
        pytest.skip(f"{target} not enabled")

    @T.prim_func(s_tir=True)
    def my_func(a: T.handle):
        A = T.match_buffer(a, (128,), "float32")
        T.func_attr({"global_symbol": "my_module", "tirx.noalias": True})
        A[T.ramp(0, 1, 4 * T.vscale())] = T.broadcast(1, 4 * T.vscale())

    with tvm.target.Target(target):
        mod = tvm.tirx.build(my_func)

    llvm = mod.inspect_source("ll")
    # Older LLVM versions print the broadcast as a shufflevector of an insertelement,
    # newer ones print it as a splat constant.
    assert (
        "shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float>" in llvm
        or "store <vscale x 4 x float> splat (float 1.000000e+00)" in llvm
    ), "No scalable broadcast in generated LLVM."
    assert re.findall(r" store <vscale x 4 x float>", llvm), "No scalable store in generated LLVM."


@pytest.mark.skip(
    reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
)
@pytest.mark.parametrize(
    "target",
    [
        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
        {
            "kind": "llvm",
            "device": "riscv_cpu",
            "mtriple": "riscv64-linux-gnu",
            "mcpu": "generic-rv64",
            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
        },
    ],
)
def test_get_active_lane_mask(target):
    if not tvm.testing.device_enabled(target):
        pytest.skip(f"{target} not enabled")

    @T.prim_func(s_tir=True)
    def before(a: T.handle):
        A = T.match_buffer(a, (30,), "int1")
        for i in range(T.ceildiv(30, T.vscale() * 4)):
            A[i : i + T.vscale() * 4] = T.get_active_lane_mask("uint1xvscalex4", i, 30)

    with tvm.target.Target(target):
        out = tvm.tirx.build(before)

    ll = out.inspect_source("ll")
    assert "get.active.lane.mask" in ll


@pytest.mark.skip(
    reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
)
@pytest.mark.parametrize(
    "target",
    [
        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
        {
            "kind": "llvm",
            "device": "riscv_cpu",
            "mtriple": "riscv64-linux-gnu",
            "mcpu": "generic-rv64",
            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
        },
    ],
)
def test_predicated_scalable_buffer(target):
    if not tvm.testing.device_enabled(target):
        pytest.skip(f"{target} not enabled")

    @T.prim_func(s_tir=True)
    def before(a: T.handle, b: T.handle):
        A = T.match_buffer(a, (16,), "float32")
        B = T.match_buffer(b, (16,), "float32")
        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
        for i_0 in T.serial(T.ceildiv(16, 4 * T.vscale())):
            for i_1 in T.vectorized(4 * T.vscale()):
                if i_0 * 4 * T.vscale() + i_1 < 14:
                    B[i_0 * 4 * T.vscale() + i_1] = A[i_0 * 4 * T.vscale() + i_1] + 1.0

    with tvm.target.Target(target):
        out = tvm.tirx.build(before)

    ll = out.inspect_source("ll")
    assert "get.active.lane.mask" in ll
    assert "llvm.masked.load" in ll
    assert "llvm.masked.store" in ll


if __name__ == "__main__":
    tvm.testing.main()