Returns the path to the program schema file after copying it and its deps into out_dir. May patch the schema contents depending on the parameters to this function.
(
out_dir: str,
constant_tensor_alignment: Optional[int] = None,
delegate_alignment: Optional[int] = None,
)
| 172 | |
| 173 | |
| 174 | def _prepare_schema( |
| 175 | out_dir: str, |
| 176 | constant_tensor_alignment: Optional[int] = None, |
| 177 | delegate_alignment: Optional[int] = None, |
| 178 | ) -> _SchemaInfo: |
| 179 | """Returns the path to the program schema file after copying it and its deps |
| 180 | into out_dir. May patch the schema contents depending on the parameters to |
| 181 | this function. |
| 182 | """ |
| 183 | program_schema = "program.fbs" |
| 184 | # Included by the root program schema; must also be present. |
| 185 | deps = ["scalar_type.fbs"] |
| 186 | |
| 187 | schemas = _ResourceFiles([program_schema] + deps) |
| 188 | |
| 189 | # Update annotated alignments in the schema files. |
| 190 | schemas.patch_files( |
| 191 | lambda data: _patch_schema_alignment( |
| 192 | schema=data, |
| 193 | constant_tensor_alignment=constant_tensor_alignment, |
| 194 | delegate_alignment=delegate_alignment, |
| 195 | ), |
| 196 | ) |
| 197 | # Find the largest alignment used in the patched schema files. |
| 198 | get_alignments = _SchemaMaxAlignmentGetter() |
| 199 | schemas.patch_files(get_alignments) |
| 200 | get_file_identifier = _SchemaFileIdentifierGetter() |
| 201 | schemas.patch_files(get_file_identifier) |
| 202 | if get_file_identifier.file_identifier is None: |
| 203 | raise ValueError("Missing file_identifier in schema files.") |
| 204 | |
| 205 | def extract_alignment(schema: bytes, marker: bytes) -> int: |
| 206 | for line in schema.splitlines(): |
| 207 | if marker in line: |
| 208 | match = re.search(rb"force_align\s*:\s*(\d+)", line) |
| 209 | if match: |
| 210 | return int(match.group(1)) |
| 211 | raise RuntimeError(f"Failed to find marker {marker!r} in program.fbs") |
| 212 | |
| 213 | program_schema_data = schemas.get(program_schema) |
| 214 | tensor_alignment = extract_alignment( |
| 215 | program_schema_data, b"@executorch-tensor-alignment" |
| 216 | ) |
| 217 | effective_delegate_alignment = extract_alignment( |
| 218 | program_schema_data, b"@executorch-delegate-alignment" |
| 219 | ) |
| 220 | |
| 221 | # Write the patched schema files to the filesystem. |
| 222 | schemas.write_to(out_dir) |
| 223 | |
| 224 | return _SchemaInfo( |
| 225 | root_path=os.path.join(out_dir, program_schema), |
| 226 | max_alignment=get_alignments.max_alignment, |
| 227 | file_identifier=get_file_identifier.file_identifier, |
| 228 | tensor_alignment=tensor_alignment, |
| 229 | delegate_alignment=effective_delegate_alignment, |
| 230 | ) |
| 231 |
no test coverage detected