Function _fuse_quantized_cat

exir/passes/quant_fusion_pass.py:21–94

fuse "dequantize -> cat -> quantize" pattern to cat operator, only happens if the quantization parameters for dequantize for all the inputs matches, and it also matches the quantization parameters for the quantize node after cat

(model: GraphModule)

def _fuse_quantized_cat(model: GraphModule) -> None:
    """fuse "dequantize -> cat -> quantize" pattern to cat operator, only happens if the quantization
    parameters for dequantize for all the inputs matches, and it also matches the quantization
    parameters for the quantize node after cat
    """

    # get quantization parameters for the node, either for quantize or dequantize node
    def _get_qparams(node):
        assert node.target in (
            exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
            exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
        )
        args = list(node.args)
        # skip input
        qparams = args[1:]
        return qparams

    for n in model.graph.nodes:
        if (
            n.op != "call_function"
            or n.target
            != exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
        ):

            continue
        qnode = n
        maybe_cat = qnode.args[0]
        if (
            maybe_cat.op != "call_function"
            or maybe_cat.target != exir_ops.edge.aten.cat.default
        ):

            continue
        tensor_args = maybe_cat.args[0]
        if not isinstance(tensor_args, (tuple, list)):
            continue

        matched_quantized_cat = True
        output_qparams = _get_qparams(qnode)
        for tensor_arg in tensor_args:
            if (
                tensor_arg.op != "call_function"
                or tensor_arg.target
                != exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default
            ):

                matched_quantized_cat = False
                break

            # make sure the input qparams for each input tensor in the concat list
            # matches the output qparams
            current_input_qparams = _get_qparams(tensor_arg)
            if not current_input_qparams == output_qparams:
                matched_quantized_cat = False
                break

        if not matched_quantized_cat:
            continue