compile_fx

torch_tvarant.compiler.compile_fx(gm, example_inputs=None) → GraphModule

Optimize a captured FX graph for Tvarant.

Runs GEMM-epilogue fusion then pointwise fusion.

Parameters:

Name	Type	Description	Default
gm	GraphModule	FX graph module to optimize in place / recompile.	required
example_inputs	Any	Unused; accepted for Dynamo backend compatibility.	None

Returns:

Type	Description
GraphModule	torch.fx.GraphModule: Fused graph ready to run on device="tvarant".

Source code in torch_tvarant/compiler.py

def compile_fx(gm: torch.fx.GraphModule, example_inputs: Any = None) -> torch.fx.GraphModule:
    """Optimize a captured FX graph for Tvarant.

    Runs GEMM-epilogue fusion then pointwise fusion.

    Args:
        gm: FX graph module to optimize in place / recompile.
        example_inputs: Unused; accepted for Dynamo backend compatibility.

    Returns:
        torch.fx.GraphModule: Fused graph ready to run on ``device="tvarant"``.
    """
    _ = example_inputs
    gm = fuse_gemm_epilogue(gm)
    gm = fuse_pointwise(gm)
    return gm