Skip to content

Conversation

@stomfaig
Copy link
Contributor

@stomfaig stomfaig commented Dec 24, 2025

This PR is implementing a new pass for quantising graphs. We repeatedly evaluate each node, check if we can decide if it can be quantised, try to quantise it, and repeat the process. For more details see the code, it is extensively documented.

Note: the code also contains the code from #641, so that it is runnable locally. This poc only supports simple models like below, and is intended as a demonstration of the approach. The changes strictly associated with this pr are only in frontend/Python/graph/transform/quantise.py

class SimpleNet(nn.Module):
    def __init__(self):
      super().__init__()
      self.linear = torch.nn.Linear(28, 28)

    def forward(self, x):
        return self.linear(x)

Test driver to see the changes:

def main():

    model = SimpleNet()

    dynamo_compiler = DynamoCompiler(
        primary_registry=tosa.ops_registry,
        aot_autograd_decomposition=inductor_decomp,
        func_name="simple",
    )

    graphs = dynamo_compiler.importer(
        model,
        x=torch.rand((1, 1, 28, 28)),
    )

    assert len(graphs) == 1

    for graph in graphs:
        graph.init_op_group()

    graphs[0].perform(
        [quantise_graph]
    )

    g_body = graphs[0]._body
    for op in g_body:
        print("\n---")
        print(f"{op._name} : {type(op)}")
        print(f"children: {op._children}")
        print(f"parants: {op._parents}")
        print(f"tensor_meta: {op._tensor_meta}")

Depends on #641, #644

@stomfaig stomfaig marked this pull request as draft December 24, 2025 16:35
@stomfaig
Copy link
Contributor Author

cc: @R-Tars

@stomfaig stomfaig changed the title Add quantisation pass [Frontend] Add quantisation pass Dec 26, 2025
@zhanghb97
Copy link
Member

Looks like progress is going well. Just a reminder: we upgraded to LLVM 22. Don't forget to rebase to avoid potential compatibility issues.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants