| 211 | |
| 212 | |
| 213 | class FuseQuantizedOpsTransform(ExportPass): |
| 214 | def __init__(self) -> None: |
| 215 | super().__init__() |
| 216 | self._exported_program: Optional[ExportedProgram] = None |
| 217 | |
| 218 | def call(self, graph_module: torch.fx.GraphModule) -> PassResult: |
| 219 | assert self._exported_program is not None |
| 220 | |
| 221 | for node in graph_module.graph.nodes: |
| 222 | # Check for linear_qcnw pattern (weight-only quantization) |
| 223 | qcnw_details = matches_linear_qcnw_pattern(self._exported_program, node) |
| 224 | if qcnw_details is not None: |
| 225 | qcnw_method, qcnw_nbits = qcnw_details |
| 226 | fuse_into_linear_qcnw_node( |
| 227 | self._exported_program, graph_module, node, qcnw_method, qcnw_nbits |
| 228 | ) |
| 229 | continue |
| 230 | |
| 231 | graph_module.recompile() |
| 232 | dead_code_elimination_pass(graph_module) |
| 233 | |
| 234 | # Re-trace the graph since new nodes were (potentially) inserted |
| 235 | graph_module = super().call(graph_module).graph_module |
| 236 | return PassResult(graph_module, True) |