mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-06-25 02:50:42 +00:00
454f77cd94
# Motivation Currently, ORT minimal builds use kernel def hashes to map from nodes to kernels to execute when loading the model. As the kernel def hashes must be known ahead of time, this works for statically registered kernels. This works well for the CPU EP. For this approach to work, the kernel def hashes must also be known at ORT format model conversion time, which means the EP with statically registered kernels must also be enabled then. This is not an issue for the always-available CPU EP. However, we do not want to require that any EP which statically registers kernels is always available too. Consequently, we explore another approach to match nodes to kernels that does not rely on kernel def hashes. An added benefit of this is the possibility of moving away from kernel def hashes completely, which would eliminate the maintenance burden of keeping the hashes stable. # Approach In a full build, ORT uses some information from the ONNX op schema to match a node to a kernel. We want to avoid including the ONNX op schema in a minimal build to reduce binary size. Essentially, we take the necessary information from the ONNX op schema and make it available in a minimal build. We decouple the ONNX op schema from the kernel matching logic. The kernel matching logic instead relies on per-op information which can either be obtained from the ONNX op schema or another source. This per-op information must be available in a minimal build when there are no ONNX op schemas. We put it in the ORT format model. Existing uses of kernel def hashes to look up kernels are replaced with the updated kernel matching logic. We no longer store kernel def hashes in the ORT format model’s session state and runtime optimization representations. We no longer keep the logic to generate and ensure stability of kernel def hashes.
92 lines
3.9 KiB
Python
92 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
|
|
import argparse
|
|
import pathlib
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
|
|
SCRIPT_DIR = pathlib.Path(__file__).parent.resolve()
|
|
|
|
|
|
def update_namespace(schema_path: pathlib.Path, updated_schema_path: pathlib.Path):
|
|
# create a copy of the schema so we can replace the namespace so that the generated module name doesn't clash
|
|
# with the 'onnxruntime' package.
|
|
with open(schema_path, 'r') as input, open(updated_schema_path, 'w') as output:
|
|
for line in input:
|
|
# convert any line with the namespace to use ort_flatbuffers_py instead of onnxruntime as the top level
|
|
# namespace. this doesn't change how anything works - it just avoids a naming clash with the 'real'
|
|
# onnxruntime python package
|
|
output.write(line.replace('onnxruntime.fbs', 'ort_flatbuffers_py.fbs'))
|
|
|
|
|
|
def generate_python(flatc: pathlib.Path, schema_path: pathlib.Path, output_dir: pathlib.Path):
|
|
# run flatc to generate Python code
|
|
cmd = [str(flatc), '--python', str(schema_path)]
|
|
subprocess.run(cmd, check=True, cwd=output_dir)
|
|
|
|
|
|
def create_init_py(output_dir: pathlib.Path):
|
|
# create an __init__.py that imports all the py files so we can just 'import ort_flatbuffers_py.fbs'
|
|
# in a script that wants to process an ORT format model
|
|
init_py_path = output_dir / 'ort_flatbuffers_py/fbs/__init__.py'
|
|
with open(init_py_path, 'w') as init_py:
|
|
init_py.write('''from os.path import dirname, basename, isfile, join, splitext
|
|
import glob
|
|
modules = glob.glob(join(dirname(__file__), "*.py"))
|
|
__all__ = [splitext(basename(f))[0] for f in modules if isfile(f) and not f.endswith('__init__.py')]
|
|
|
|
from . import * # noqa
|
|
''')
|
|
|
|
|
|
def generate_cpp(flatc: pathlib.Path, schema_path: pathlib.Path):
|
|
# run flatc to generate C++ code
|
|
cmd = [str(flatc), '--cpp', '--scoped-enums', '--filename-suffix', '.fbs', str(schema_path)]
|
|
subprocess.run(cmd, check=True, cwd=SCRIPT_DIR)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Generate language bindings for the ORT flatbuffers schema.',
|
|
usage='Provide the path to the flatbuffers flatc executable. '
|
|
'Script can be executed from anywhere but must be located in its original '
|
|
'directory in the ONNX Runtime enlistment.')
|
|
|
|
parser.add_argument('-f', '--flatc', required=True, type=pathlib.Path,
|
|
help='Path to flatbuffers flatc executable. '
|
|
'Can be found in the build directory under external/flatbuffers/<config>/')
|
|
|
|
all_languages = ['python', 'cpp']
|
|
parser.add_argument('-l', '--language', action='append', dest='languages', choices=all_languages,
|
|
help='Specify which language bindings to generate.')
|
|
|
|
args = parser.parse_args()
|
|
languages = args.languages if args.languages is not None else all_languages
|
|
flatc = args.flatc.resolve(strict=True)
|
|
schema_path = SCRIPT_DIR / 'ort.fbs'
|
|
|
|
if 'python' in languages:
|
|
with tempfile.TemporaryDirectory() as temp_dir_name:
|
|
temp_dir = pathlib.Path(temp_dir_name).resolve()
|
|
updated_schema_path = temp_dir / 'ort.py.fbs'
|
|
update_namespace(schema_path, updated_schema_path)
|
|
|
|
output_dir = temp_dir / 'out'
|
|
output_dir.mkdir()
|
|
generate_python(flatc, updated_schema_path, output_dir)
|
|
create_init_py(output_dir)
|
|
|
|
# replace generated files in repo
|
|
target_dir = SCRIPT_DIR.parent / 'ort_flatbuffers_py'
|
|
if target_dir.is_dir():
|
|
shutil.rmtree(target_dir)
|
|
shutil.move(str(output_dir / 'ort_flatbuffers_py'), str(target_dir))
|
|
|
|
if 'cpp' in languages:
|
|
generate_cpp(flatc, schema_path)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|