# 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.
# pylint: disable=missing-docstring, unused-variable
# ruff: noqa: F841

# The test attempts to eliminate redundant pad branch and overcompute the value for elementwise ops.
# This helps to expose more opportunities to vectorize the code.

import tvm_ffi

import tvm
import tvm.script
import tvm.testing
from tvm.script import relax as R
from tvm.script import tirx as T


@tvm.script.ir_module
class AddBefore:
    @T.prim_func(private=True, s_tir=True)
    def add(
        a: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        b: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        compute: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
    ):
        T.func_attr(
            {
                "op_attrs": {"lhs_axis": 0, "op_name": "qnn.add", "rhs_axis": 0},
                "op_pattern": 0,
                "operator_name": "add",
                "tirx.noalias": True,
            }
        )
        # with T.sblock("root"):
        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_A_assumptions"):
                v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    not (
                        (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
                        or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
                    )
                    or a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_B_assumptions"):
                v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    not (
                        (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
                        or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
                    )
                    or b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("compute"):
                v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(
                    a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                    b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                )
                T.writes(compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6] = (
                    T.if_then_else(
                        (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
                        or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5),
                        T.uint8(0),
                        a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                        + b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                    )
                )

    @R.function
    def main(
        a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
        b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
    ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
        out = R.call_tir(
            AddBefore.add,
            (a, b),
            out_ty=R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8"),
        )
        return out


@tvm.script.ir_module
class AddExpected:
    @T.prim_func(private=True, s_tir=True)
    def add(
        a: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        b: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        compute: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
    ):
        T.func_attr(
            {
                "op_attrs": {"lhs_axis": 0, "op_name": "qnn.add", "rhs_axis": 0},
                "op_pattern": 0,
                "operator_name": "add",
                "tirx.noalias": True,
            }
        )
        # with T.sblock("root"):
        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_A_assumptions"):
                v_axis0 = T.axis.spatial(T.int64(1), T.int64(0))
                v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    (
                        (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
                        and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
                    )
                    or a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_B_assumptions"):
                v_axis0 = T.axis.spatial(T.int64(1), T.int64(0))
                v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    (
                        (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
                        and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
                    )
                    or b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5_0 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(2)
        ):
            for axis5_1_axis6_fused in T.vectorized(T.int64(128)):
                with T.sblock("compute"):
                    v_axis0 = T.axis.spatial(T.int64(1), T.int64(0))
                    v_axis1, v_axis2, v_axis3, v_axis4 = T.axis.remap(
                        "SSSS", [axis1, axis2, axis3, axis4]
                    )
                    v_axis5 = T.axis.spatial(
                        T.int64(8), axis5_0 * T.int64(4) + axis5_1_axis6_fused // T.int64(32)
                    )
                    v_axis6 = T.axis.spatial(T.int64(32), axis5_1_axis6_fused % T.int64(32))
                    T.reads(
                        a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                        b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                    )
                    T.writes(
                        compute[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    )
                    compute[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6] = (
                        a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                        + b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    )

    @R.function
    def main(
        a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
        b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
    ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
        out = R.call_tir(
            AddExpected.add,
            (a, b),
            out_ty=R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8"),
        )
        return out


@tvm.script.ir_module
class SubBefore:
    @T.prim_func(private=True, s_tir=True)
    def sub(
        a: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        b: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        compute: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
    ):
        T.func_attr(
            {
                "op_attrs": {"lhs_axis": 0, "op_name": "qnn.subtract", "rhs_axis": 0},
                "op_pattern": 0,
                "operator_name": "sub",
                "tirx.noalias": True,
            }
        )
        # with T.sblock("root"):
        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_A_assumptions"):
                v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    not (
                        (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
                        or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
                    )
                    or a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_B_assumptions"):
                v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    not (
                        (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
                        or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
                    )
                    or b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("compute"):
                v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(
                    a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                    b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                )
                T.writes(compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6] = (
                    T.if_then_else(
                        (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
                        or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5),
                        T.uint8(0),
                        a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                        - b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                    )
                )

    @R.function
    def main(
        a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
        b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
    ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
        out = R.call_tir(
            SubBefore.sub,
            (a, b),
            out_ty=R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8"),
        )
        return out


@tvm.script.ir_module
class SubExpected:
    @T.prim_func(private=True, s_tir=True)
    def sub(
        a: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        b: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        compute: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
    ):
        T.func_attr(
            {
                "op_attrs": {"lhs_axis": 0, "op_name": "qnn.subtract", "rhs_axis": 0},
                "op_pattern": 0,
                "operator_name": "sub",
                "tirx.noalias": True,
            }
        )
        # with T.sblock("root"):
        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_A_assumptions"):
                v_axis0 = T.axis.spatial(T.int64(1), T.int64(0))
                v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    (
                        (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
                        and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
                    )
                    or a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_B_assumptions"):
                v_axis0 = T.axis.spatial(T.int64(1), T.int64(0))
                v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    (
                        (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
                        and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
                    )
                    or b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5_0 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(2)
        ):
            for axis5_1_axis6_fused in T.vectorized(T.int64(128)):
                with T.sblock("compute"):
                    v_axis0 = T.axis.spatial(T.int64(1), T.int64(0))
                    v_axis1, v_axis2, v_axis3, v_axis4 = T.axis.remap(
                        "SSSS", [axis1, axis2, axis3, axis4]
                    )
                    v_axis5 = T.axis.spatial(
                        T.int64(8), axis5_0 * T.int64(4) + axis5_1_axis6_fused // T.int64(32)
                    )
                    v_axis6 = T.axis.spatial(T.int64(32), axis5_1_axis6_fused % T.int64(32))
                    T.reads(
                        a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                        b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                    )
                    T.writes(
                        compute[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    )
                    compute[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6] = (
                        a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                        - b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    )

    @R.function
    def main(
        a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
        b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
    ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
        out = R.call_tir(
            SubExpected.sub,
            (a, b),
            out_ty=R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8"),
        )
        return out


@tvm.script.ir_module
class MulBefore:
    @T.prim_func(private=True, s_tir=True)
    def mul(
        a: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        b: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        compute: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
    ):
        T.func_attr(
            {
                "op_attrs": {"lhs_axis": 0, "op_name": "qnn.mul", "rhs_axis": 0},
                "op_pattern": 0,
                "operator_name": "mul",
                "tirx.noalias": True,
            }
        )
        # with T.sblock("root"):
        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_A_assumptions"):
                v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    not (
                        (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
                        or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
                    )
                    or a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_B_assumptions"):
                v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    not (
                        (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
                        or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5)
                    )
                    or b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("compute"):
                v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSSS", [axis0, axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(
                    a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                    b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                )
                T.writes(compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                compute[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6] = (
                    T.if_then_else(
                        (v_axis1 == T.int64(3) and T.int64(4) <= v_axis4)
                        or (v_axis2 == T.int64(3) and T.int64(4) <= v_axis5),
                        T.uint8(0),
                        a[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                        * b[v_axis0, v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                    )
                )

    @R.function
    def main(
        a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
        b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
    ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
        out = R.call_tir(
            MulBefore.mul,
            (a, b),
            out_ty=R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8"),
        )
        return out


@tvm.script.ir_module
class MulExpected:
    @T.prim_func(private=True, s_tir=True)
    def mul(
        a: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        b: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
        compute: T.Buffer(
            (T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)),
            "uint8",
        ),
    ):
        T.func_attr(
            {
                "op_attrs": {"lhs_axis": 0, "op_name": "qnn.mul", "rhs_axis": 0},
                "op_pattern": 0,
                "operator_name": "mul",
                "tirx.noalias": True,
            }
        )
        # with T.sblock("root"):
        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_A_assumptions"):
                v_axis0 = T.axis.spatial(T.int64(1), T.int64(0))
                v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    (
                        (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
                        and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
                    )
                    or a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5, axis6 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(8), T.int64(32)
        ):
            with T.sblock("buffer_B_assumptions"):
                v_axis0 = T.axis.spatial(T.int64(1), T.int64(0))
                v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6 = T.axis.remap(
                    "SSSSSS", [axis1, axis2, axis3, axis4, axis5, axis6]
                )
                T.reads(b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6])
                T.writes()
                T.assume(
                    (
                        (v_axis1 < T.int64(3) or v_axis4 < T.int64(4))
                        and (v_axis2 < T.int64(3) or v_axis5 < T.int64(4))
                    )
                    or b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    == T.uint8(0)
                )

        for axis0, axis1, axis2, axis3, axis4, axis5_0 in T.grid(
            T.int64(1), T.int64(4), T.int64(4), T.int64(16), T.int64(8), T.int64(2)
        ):
            for axis5_1_axis6_fused in T.vectorized(T.int64(128)):
                with T.sblock("compute"):
                    v_axis0 = T.axis.spatial(T.int64(1), T.int64(0))
                    v_axis1, v_axis2, v_axis3, v_axis4 = T.axis.remap(
                        "SSSS", [axis1, axis2, axis3, axis4]
                    )
                    v_axis5 = T.axis.spatial(
                        T.int64(8), axis5_0 * T.int64(4) + axis5_1_axis6_fused // T.int64(32)
                    )
                    v_axis6 = T.axis.spatial(T.int64(32), axis5_1_axis6_fused % T.int64(32))
                    T.reads(
                        a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                        b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6],
                    )
                    T.writes(
                        compute[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    )
                    compute[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6] = (
                        a[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                        * b[T.int64(0), v_axis1, v_axis2, v_axis3, v_axis4, v_axis5, v_axis6]
                    )

    @R.function
    def main(
        a: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
        b: R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"),
    ) -> R.Tensor((1, 4, 4, 16, 8, 8, 32), "uint8"):
        out = R.call_tir(
            MulExpected.mul,
            (a, b),
            out_ty=R.Tensor((1, 4, 4, 16, 8, 8, 32), dtype="uint8"),
        )
        return out


def test_add_primfunc_overcompute():
    add_after = tvm.s_tir.transform.UseAssumeToReduceBranches()(AddBefore)
    tvm_ffi.structural_equal(add_after["add"], AddExpected["add"], map_free_vars=True)


def test_sub_primfunc_overcompute():
    sub_after = tvm.s_tir.transform.UseAssumeToReduceBranches()(SubBefore)
    tvm_ffi.structural_equal(sub_after["sub"], SubExpected["sub"], map_free_vars=True)


def test_mul_primfunc_overcompute():
    mul_after = tvm.s_tir.transform.UseAssumeToReduceBranches()(MulBefore)
    tvm_ffi.structural_equal(mul_after["mul"], MulExpected["mul"], map_free_vars=True)


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