Function _patch_schema_alignment

exir/_serialize/_flatbuffer.py:36–79

Modifies annotated "force_align" values in a flatbuffer schema. Args: schema: The flatbuffer schema to modify. constant_tensor_alignment: If provided, the alignment to use for lines annotated with "@executorch-tensor-alignment". If not provided, does not patch

(
    schema: bytes,
    constant_tensor_alignment: Optional[int],
    delegate_alignment: Optional[int],
)

def _patch_schema_alignment(
    schema: bytes,
    constant_tensor_alignment: Optional[int],
    delegate_alignment: Optional[int],
) -> bytes:
    """Modifies annotated "force_align" values in a flatbuffer schema.

    Args:
        schema: The flatbuffer schema to modify.
        constant_tensor_alignment: If provided, the alignment to use for lines annotated
            with "@executorch-tensor-alignment". If not provided, does not patch
            tensor alignment.
        delegate_alignment: If provided, the alignment to use for lines
            annotated with "@executorch-delegate-alignment". If not provided,
            does not patch delegate alignment.

    Returns:
        The possibly-modified flatbuffer schema.
    """

    def assert_valid_alignment(alignment: Optional[int], name: str) -> None:
        if not (alignment is None or _is_valid_alignment(alignment)):
            raise ValueError(f"Bad {name} {alignment}")

    assert_valid_alignment(constant_tensor_alignment, "constant_tensor_alignment")
    assert_valid_alignment(delegate_alignment, "delegate_alignment")

    def patch_alignment(line: bytes, alignment: int) -> bytes:
        """Replaces an existing alignment with a new alignment."""
        return re.sub(
            rb"\(\s*force_align\s*:\s*\d+\s*\)",
            f"(force_align: {alignment})".encode("utf-8"),
            line,
        )

    lines = []
    for line in schema.splitlines():
        if constant_tensor_alignment and b"@executorch-tensor-alignment" in line:
            lines.append(patch_alignment(line, constant_tensor_alignment))
        elif delegate_alignment and b"@executorch-delegate-alignment" in line:
            lines.append(patch_alignment(line, delegate_alignment))
        else:
            lines.append(line)
    return b"\n".join(lines)


class _SchemaMaxAlignmentGetter: