"""A compiler pass that dispatch low-batch-gemm to gemv schedule."""

import tvm
import tvm_ffi
from tvm import tirx
from tvm.ir.module import IRModule
from tvm.s_tir import dlight as dl


@tvm.transform.module_pass(opt_level=0, name="LowBatchGemvSpecialize")
class LowBatchGemvSpecialize:
    """A compiler pass that dispatch low-batch-gemm to gemv schedule."""

    def transform_module(
        self,
        mod: IRModule,
        _ctx: tvm.transform.PassContext,
    ) -> IRModule:
        """IRModule-level transformation"""
        for g_var, func in mod.functions_items():
            if isinstance(func, tirx.PrimFunc):
                low_batch_range = [2, 8]
                buckets = [2, 4]
                low_batch_funcs = []
                for bucket in buckets:
                    low_batch_mod = IRModule({})
                    low_batch_mod["main"] = func
                    low_batch_mod = dl.ApplyDefaultSchedule(
                        dl.gpu.LowBatchGEMV(bucket),
                    )(low_batch_mod)
                    low_batch_funcs.append(low_batch_mod["main"])
                if any(
                    tvm_ffi.structural_equal(low_batch_func, func)
                    for low_batch_func in low_batch_funcs
                ):
                    continue
                buffers = func.buffer_map.values()
                shapes = [buffer.shape for buffer in buffers]
                symbolic_vars = set(
                    expr for shape in shapes for expr in shape if isinstance(expr, tirx.Var)
                )
                if len(symbolic_vars) != 1:
                    continue
                gemm_mod = IRModule({})
                gemm_mod["main"] = func
                gemm_mod = dl.ApplyDefaultSchedule(
                    dl.gpu.Matmul(),
                )(gemm_mod)
                gemm_func = gemm_mod["main"]
                sym_var = next(iter(symbolic_vars))
                body = gemm_func.body
                for i, range_limit in reversed(list(enumerate(low_batch_range))):
                    body = tirx.IfThenElse(
                        tirx.op.tvm_thread_invariant(sym_var <= range_limit),
                        low_batch_funcs[i].body,
                        body,
                    )
                body = tirx.SBlock([], [], [], "root", body)
                body = tirx.SBlockRealize([], True, body)
                new_func = func.with_body(body)
                new_func = new_func.with_attr("tirx.is_scheduled", 1)
                new_func = new_func.with_attr("tirx.HoistIfThenElseExprWithBlock", 1)
                mod.update_func(g_var, new_func)
        return mod