"""A compiler pass that fuses transpose + dequantize."""

import tvm
from tvm import relax, s_tir, tirx
from tvm.ir.module import IRModule
from tvm.relax.analysis import remove_all_unused
from tvm.relax.expr_functor import PyExprMutator, mutator


@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeTranspose")
class FuseDequantizeTranspose:
    """A compiler pass that fuses transpose + dequantize."""

    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
        """IRModule-level transformation"""
        return _DequantizeTransposeFuser(mod).transform()


@mutator
class _DequantizeTransposeFuser(PyExprMutator):
    def __init__(
        self,
        mod: IRModule,
    ):
        super().__init__(mod)
        self.mod = mod

    def transform(self) -> IRModule:
        """Entry point"""
        for g_var, func in self.mod.functions_items():
            if isinstance(func, relax.Function):
                updated_func = self.visit_expr(func)
                updated_func = remove_all_unused(updated_func)
                self.builder_.update_func(g_var, updated_func)
        return self.builder_.get()

    def visit_call_(
        self,
        call: relax.Call,
    ) -> relax.Expr:
        call = self.visit_expr_post_order(call)
        if call.op != tvm.ir.Op.get("relax.matmul"):
            return call
        # Do not fuse dequantize-transpose for GeMM
        if (
            call.args[0].ty.ndim < 2
            or not isinstance(call.args[0].ty.shape[-2], tirx.IntImm)
            or call.args[0].ty.shape[-2].value != 1
        ):
            return call

        matmul_rhs = self.lookup_binding(call.args[1])
        if (
            not isinstance(matmul_rhs, relax.Call)
            or matmul_rhs.op != tvm.ir.Op.get("relax.permute_dims")
            or matmul_rhs.args[0].ty.ndim != 2
            or matmul_rhs.attrs.axes is not None
        ):
            return call

        transpose_input = self.lookup_binding(matmul_rhs.args[0])
        if (
            not isinstance(transpose_input, relax.Call)
            or transpose_input.op != tvm.ir.Op.get("relax.call_tir")
            or not transpose_input.args[0].name_hint.startswith("dequantize")
            or not isinstance(transpose_input.ty, relax.TensorType)
        ):
            return call

        dequantize_tir_func = self.mod[transpose_input.args[0]]
        assert isinstance(dequantize_tir_func, tirx.PrimFunc)
        if (
            len(dequantize_tir_func.body.block.alloc_buffers) != 1
            or not isinstance(dequantize_tir_func.body.block.body, tirx.SeqStmt)
            or len(dequantize_tir_func.body.block.body) != 2
            or not isinstance(dequantize_tir_func.body.block.body[1], tirx.For)
            or not isinstance(dequantize_tir_func.body.block.body[1].body.body, tirx.SBlockRealize)
            or dequantize_tir_func.body.block.body[1].body.body.block.name_hint != "T_transpose"
        ):
            return call

        new_func_buffers = [
            dequantize_tir_func.buffer_map[var] for var in dequantize_tir_func.params
        ]
        new_func_buffers[-1] = dequantize_tir_func.body.block.alloc_buffers[0]
        new_func = tirx.PrimFunc(
            params=new_func_buffers,
            body=tirx.SBlockRealize(
                iter_values=[],
                predicate=True,
                block=tirx.SBlock(
                    iter_vars=[],
                    reads=[],
                    writes=[],
                    name_hint="root",
                    body=dequantize_tir_func.body.block.body[0],
                ),
            ),
        )
        # Call `renew_defs` for deep-copy to avoid IR node duplication in
        # different PrimFuncs of an IRModule.
        new_func = s_tir.renew_defs(new_func)
        g_var = self.builder_.add_func(new_func, func_name="dequantize")
        dequantize_matmul_rhs = self.builder_.emit(
            relax.call_tir(g_var, transpose_input.args[1], out_ty=matmul_rhs.ty)
        )
        return relax.op.matmul(call.args[0], dequantize_matmul_rhs, out_dtype=call.attrs.out_dtype)