Method call

backends/arm/_passes/insert_table_ops.py:251–322

(self, graph_module: GraphModule)

                )

    def call(self, graph_module: GraphModule) -> PassResult:
        modified = False
        for node in graph_module.graph.nodes:
            if node.op != "call_function" or node not in self.table_ops:
                continue
            input_qparams = node.meta.get("input_qparams", {})
            output_qparams = node.meta.get("output_qparams", {})
            if len(input_qparams) == 0 or len(output_qparams) == 0:
                # We only want to replace the node if it's quantized
                continue
            # Create table node
            insert_pos = list(node.graph.nodes)[0]
            with graph_module.graph.inserting_before(insert_pos):
                # Expect exactly one quantization parameter for input and output
                if len(input_qparams) != 1:
                    raise ValueError(
                        f"InsertTableOpsPass expected exactly one input quantization parameter, "
                        f"got {len(input_qparams)} for node {node.name}"
                    )
                if len(output_qparams) != 1:
                    raise ValueError(
                        f"InsertTableOpsPass expected exactly one output quantization parameter, "
                        f"got {len(output_qparams)} for node {node.name}"
                    )

                # Generate table buffer and how much to lshift the table output.
                buffer, lshift = self.generate_table_values(
                    torch_op=self.table_ops[node],
                    in_quantargs=input_qparams[0],
                    out_quantargs=output_qparams[0],
                )
                # Register buffer in self.exported_program.state_dict
                # b_ prefix is important to be recognized as a constant in RemovePermutesAroundElementwiseOps
                const_table_node = create_constant_placeholder(
                    exp_program=self.exported_program,
                    graph=node.graph,
                    kind=InputKind.BUFFER,
                    name="b_" + node.name + "_table_constant",
                    data=buffer,
                    persistent_buffer=True,
                )

            # Create table node
            with graph_module.graph.inserting_before(node):
                table_op_node = create_node(
                    graph=graph_module.graph,
                    op_target=exir_ops.backend.tosa.TABLE.default,
                    args=(node.args[0], const_table_node),
                )
                output_node = table_op_node

                if lshift != 0:
                    scale = 2.0**lshift
                    rescale_node = create_node(
                        graph=graph_module.graph,
                        op_target=exir_ops.backend.tosa.RESCALE.default,
                        args=(table_op_node, output_qparams[0].dtype, [scale], 0, 0),
                    )