def _transform_1x1_conv_gemm_rcr(sorted_graph: List[Tensor]) -> List[Tensor]:
    """
    Replace a 1x1 conv2d with an equivalent gemm_rcr call.
    """

    conv_to_gemm = {
        "conv2d": ops.gemm_rcr,
        "conv2d_bias": ops.gemm_rcr_bias,
        "conv2d_bias_relu": ops.gemm_rcr_bias_relu,
        "conv2d_bias_sigmoid": ops.gemm_rcr_bias_sigmoid,
        "conv2d_bias_hardswish": ops.gemm_rcr_bias_hardswish,
        "conv2d_bias_add_relu": ops.gemm_rcr_bias_add_relu,
    }

    def match_func(tensor: Tensor) -> bool:
        src_ops = tensor._attrs["src_ops"]
        if src_ops is None or len(src_ops) != 1:
            return False

        src_op = list(src_ops)[0]
        if src_op._attrs["op"] not in conv_to_gemm:
            return False

        valid_pad = src_op._attrs["pad"] == 0 or src_op._attrs["pad"] == (0, 0)
        valid_dilate = src_op._attrs["dilate"] == 1 or src_op._attrs["dilate"] == (1, 1)
        valid_stride = src_op._attrs["stride"] == 1 or src_op._attrs["stride"] == (1, 1)
        valid_group = src_op._attrs["group"] == 1
        if not valid_pad or not valid_dilate or not valid_stride or not valid_group:
            return False

        # Check that the filter is 1x1
        w_shape = src_op._attrs["inputs"][1]._attrs["shape"]
        if not is_singleton_dimension(w_shape[1]) or not is_singleton_dimension(
            w_shape[2]
        ):
            return False

        return True

    def replace_func(old_tensor: Tensor) -> List[Tensor]:
        src_op = list(old_tensor._attrs["src_ops"])[0]
        inputs = src_op._attrs["inputs"]
        X = inputs[0]
        W = inputs[1]

        # Get rid of the singleton dimensions
        W_squeeze_0 = ops.squeeze(1)(W)
        W_squeeze_1 = ops.squeeze(1)(W_squeeze_0)

        batch, HH, WW, CI = ops.size()(X)
        CO = ops.size()(W, dim=0)

        X_reshape = ops.reshape()(X, (-1, CI))
        new_op = conv_to_gemm[src_op._attrs["op"]]
        if len(inputs) == 2:
            new_tensor = new_op()(X_reshape, W_squeeze_1)
        elif len(inputs) == 3:
            new_tensor = new_op()(X_reshape, W_squeeze_1, inputs[2])
        elif len(inputs) == 4:
            new_tensor = new_op()(X_reshape, W_squeeze_1, inputs[2], inputs[3])
        else:
            raise NotImplementedError(f"Unsupported number of inputs: {len(inputs)}")

        new_tensor_reshape = ops.reshape()(new_tensor, (batch, HH, WW, CO))

        copy_tensor_attributes(new_tensor_reshape, old_tensor)
        copy_src_op_attributes(new_tensor_reshape, old_tensor)

        remove_dst_op_from_tensor(inputs, src_op)
        replace_tensor(old_tensor, new_tensor_reshape)

        return [
            W_squeeze_0,
            W_squeeze_1,
            batch,
            HH,
            WW,
            CI,
            CO,
            X_reshape,
            new_tensor,
            new_tensor_reshape,
        ]

    return _simple_transform_with_constraint(sorted_graph, match_func, replace_func)
Replace a 1x1 conv2d with an equivalent gemm_rcr call.
python/aitemplate/compiler/transform/transform_special_ops.py