def fuse_gemm_epilogue(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
"""Fuse mm/addmm/linear + optional bias + relu/silu into linear_act."""
graph = gm.graph
fused = 0
for node in list(graph.nodes):
act = _act_name(node)
if act is None or not node.args:
continue
prod = node.args[0]
if not isinstance(prod, torch.fx.Node) or len(prod.users) != 1:
continue
def emit(x, w, bias, trans_b: bool) -> None:
nonlocal fused
with graph.inserting_before(node):
out = graph.call_function(
torch.ops.tvarant.linear_act,
args=(x, w, bias, act, trans_b),
)
node.replace_all_uses_with(out)
graph.erase_node(node)
fused += 1
if _is_op(prod, "linear"):
x, w = prod.args[0], prod.args[1]
bias = prod.args[2] if len(prod.args) > 2 else prod.kwargs.get("bias")
emit(x, w, bias, True)
graph.erase_node(prod)
continue
if _is_op(prod, "addmm"):
alpha = prod.kwargs.get("alpha", 1)
beta = prod.kwargs.get("beta", 1)
try:
if float(alpha) != 1.0 or float(beta) != 1.0:
continue
except (TypeError, ValueError):
continue
bias, x, w = prod.args[0], prod.args[1], prod.args[2]
emit(x, w, bias, False)
graph.erase_node(prod)
continue
if _is_op(prod, "mm"):
x, w = prod.args[0], prod.args[1]
emit(x, w, None, False)
graph.erase_node(prod)
continue
if _is_op(prod, "add") and len(prod.args) >= 2:
left, right = prod.args[0], prod.args[1]
gemm = bias = None
for cand, other in ((left, right), (right, left)):
if _is_op(cand, "mm") and len(cand.users) == 1:
gemm, bias = cand, other
break
if gemm is None:
continue
x, w = gemm.args[0], gemm.args[1]
emit(x, w, bias, False)
graph.erase_node(prod)
graph.erase_node(gemm)
graph.lint()
gm.recompile()
last_log["gemm_epilogue"] = fused
return gm