Skip to content

[onnx_importer] Fix bug with path-based shape inference when --temp-dir is set and --data-dir is not #4375

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tkocmathla wants to merge 4 commits into
base: main
Choose a base branch
from
+37 −2
Open

[onnx_importer] Fix bug with path-based shape inference when --temp-dir is set and --data-dir is not #4375

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4027976
Add unit test covering issue.
tkocmathla Nov 12, 2025
5009096
Format
tkocmathla Nov 12, 2025
42ac53f
Check data_dir, not temp_dir
tkocmathla Nov 12, 2025
5f2d9f6
Fix test
tkocmathla Dec 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion python/torch_mlir/tools/import_onnx/__main__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ def load_onnx_model(args: argparse.Namespace) -> onnx.ModelProto:

# Load the temp file and the external data.
inferred_model = onnx.load(temp_inferred_file, load_external_data=False)
data_dir = Path(input_dir if args.temp_dir is None else args.data_dir)
data_dir = Path(input_dir if args.data_dir is None else args.data_dir)
Copy link
Collaborator

@zjgarvey zjgarvey Dec 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this needs to be moved up above the condition on line 127 (if raw_model_modified), since the external data may end up being organized differently and this gets saved to the temp_dir.

E.g.,

data_dir = Path(args.data_dir or input_dir)
if raw_model_modified:
    ...
    data_dir = Path(temp_dir)
...

onnx.load_external_data_for_model(inferred_model, str(data_dir))

# Remove the inferred shape file unless asked to keep it
Expand Down
37 changes: 36 additions & 1 deletion test/python/onnx_importer/command_line_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,26 @@ def linear_model() -> onnx.ModelProto:
return onnx_model


ALL_MODELS = [const_model, linear_model]
def path_based_shape_inference_model() -> onnx.ModelProto:
# Create a model with a serialized form that's large enough to require
# path-based shape inference.
dtype = numpy.float32
byte_size = numpy.dtype(dtype).itemsize
tensor_size = onnx.checker.MAXIMUM_PROTOBUF // byte_size + 1
large_tensor = numpy.random.rand(tensor_size).astype(dtype)
assert large_tensor.nbytes > onnx.checker.MAXIMUM_PROTOBUF
Comment on lines +94 to +97
Copy link
Collaborator

@zjgarvey zjgarvey Dec 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It does seem rather unfortunate that we need to make the CI create a 2GB tensor just to test this.

How long does this test take to run? If it takes more than a minute, I think we should refactor the load_onnx_model code so we can test the file-based shape inference directly.

Copy link
Collaborator

@zjgarvey zjgarvey Dec 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe pulling out the file based shape inference into a utility is just a good idea in general. https://github.com/llvm/torch-mlir/pull/4375/files#r2608194510

Copy link
Collaborator

@zjgarvey zjgarvey Dec 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alternatively, we could try to mirror what is currently being done in IREE:

https://github.com/iree-org/iree/blob/d0dd8893758dcf40558a57916b869ba0babc0a95/compiler/bindings/python/iree/compiler/tools/import_onnx/__main__.py#L86

Ignore the params (this is primarily the reason the importer code has diverged in IREE, but there are some improvements which haven't found their way back to torch-mlir yet).

node1 = make_node(
"Constant",
[],
["large_const"],
value=numpy_helper.from_array(large_tensor, name="large_const"),
)
X = make_tensor_value_info("large_const", TensorProto.FLOAT, [tensor_size])
graph = make_graph([node1], "large_const_graph", [], [X])
return make_model(graph)


ALL_MODELS = [const_model, linear_model, path_based_shape_inference_model]


class CommandLineTest(unittest.TestCase):
Expand Down Expand Up @@ -141,6 +160,20 @@ def run_model_extern(self, onnx_model: onnx.ModelProto, model_name: str):
)
__main__.main(args)

def run_model_explicit_temp_implicit_data(
self, onnx_model: onnx.ModelProto, model_name: str
):
run_path = self.get_run_path(model_name)
model_file = run_path / f"{model_name}-explicit_temp_implicit_data.onnx"
mlir_file = run_path / f"{model_name}-explicit_temp_implicit_data.torch.mlir"
onnx.save(onnx_model, model_file)
temp_dir = run_path / "temp"
temp_dir.mkdir(exist_ok=True)
Copy link
Collaborator

@zjgarvey zjgarvey Dec 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

May be good to use tempfile.TemporaryDirectory in case there is an exception thrown before the temp files are cleaned up by the importer.

args = __main__.parse_arguments(
[str(model_file), "-o", str(mlir_file), "--temp-dir", str(temp_dir)]
)
__main__.main(args)

def test_all(self):
for model_func in ALL_MODELS:
model_name = model_func.__name__
Expand All @@ -150,6 +183,8 @@ def test_all(self):
self.run_model_intern(model, model_name)
with self.subTest("External data"):
self.run_model_extern(model, model_name)
with self.subTest("Explicit temp dir, implicit data dir"):
self.run_model_explicit_temp_implicit_data(model, model_name)


if __name__ == "__main__":
Expand Down
Loading