wehub-resource-sync
86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
"""A compiler pass that fuses dequantize + matmul + elementwise."""
|
|
|
|
import tvm
|
|
from tvm import IRModule, relax
|
|
from tvm.relax.dpl.pattern import GlobalVarPattern, TuplePattern, is_op, wildcard
|
|
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeMatmulEwise")
|
|
class FuseDequantizeMatmulEwise:
|
|
"""A compiler pass that fuses dequantize + matmul + elementwise."""
|
|
|
|
def transform_module(
|
|
self,
|
|
mod: IRModule,
|
|
_ctx: tvm.transform.PassContext,
|
|
) -> IRModule:
|
|
"""IRModule-level transformation"""
|
|
seq = []
|
|
for n_aux_tensor in [0, 1, 2, 3, 4]:
|
|
for match_ewise in [0, 1, 2, 3, 6]:
|
|
if match_ewise == 6 and n_aux_tensor != 4:
|
|
continue
|
|
seq.append(
|
|
relax.transform.FuseOpsByPattern(
|
|
[
|
|
(
|
|
"dequantize_matmul",
|
|
*_pattern(match_ewise, n_aux_tensor),
|
|
)
|
|
]
|
|
)
|
|
)
|
|
seq.append(relax.transform.FuseTIR())
|
|
return tvm.transform.Sequential(seq)(mod)
|
|
|
|
|
|
def _pattern(match_ewise: int, n_aux_tensor: int):
|
|
w_scaled = wildcard()
|
|
x = wildcard()
|
|
w = is_op("relax.call_tir")(
|
|
GlobalVarPattern(),
|
|
TuplePattern([w_scaled] + [wildcard() for _ in range(n_aux_tensor)]),
|
|
add_constraint=False,
|
|
)
|
|
matmul = is_op("relax.call_tir")(
|
|
GlobalVarPattern(),
|
|
TuplePattern([x, w] + [wildcard() for _ in range(match_ewise)]),
|
|
add_constraint=False,
|
|
)
|
|
annotations = {
|
|
"w_scaled": w_scaled,
|
|
"x": x,
|
|
"w": w,
|
|
"matmul": matmul,
|
|
}
|
|
|
|
def _check_decoding(ctx: relax.transform.PatternCheckContext) -> bool:
|
|
call = ctx.annotated_expr["w"]
|
|
if not isinstance(call, relax.Call):
|
|
return False
|
|
g_var = call.args[0]
|
|
if not isinstance(g_var, relax.GlobalVar):
|
|
return False
|
|
return g_var.name_hint.startswith("dequantize") or g_var.name_hint.startswith(
|
|
"fused_dequantize"
|
|
)
|
|
|
|
def _check_matmul(ctx: relax.transform.PatternCheckContext) -> bool:
|
|
call = ctx.annotated_expr["matmul"]
|
|
if not isinstance(call, relax.Call):
|
|
return False
|
|
g_var = call.args[0]
|
|
if not isinstance(g_var, relax.GlobalVar):
|
|
return False
|
|
return (
|
|
g_var.name_hint.startswith("matmul")
|
|
or g_var.name_hint.startswith("fused_matmul")
|
|
or g_var.name_hint.startswith("NT_matmul")
|
|
or g_var.name_hint.startswith("fused_NT_matmul")
|
|
)
|
|
|
|
def _check(ctx: relax.transform.PatternCheckContext) -> bool:
|
|
return _check_decoding(ctx) and _check_matmul(ctx)
|
|
|
|
return matmul, annotations, _check
|