[IR] Implement save/load functions in IR and handle external data properly

[justinchuby]

Implement efficient save/load and handle loading external data properly in the IR.

Before this change, when a ModelProto containing external data is converted to IR, the external tensor objects will load the data from a path relative to the working directory, not the ONNX file. This is because we do not store the onnx file path and thus have no way to look for the external data file.

With the change, a base_dir property is added to ExternalTensor that we can set, in a separate pass when the directory is available, so the object has full information to find the data file on disk. The base_dir is not serialized to the proto to maintain a relative path in the "location" field in TensorProto.

#1701, #1792

Example:

>>> m.graph.initializers["model.model.decoder.layers.2.encoder_attn.v_proj.weight"].const_value.display()
ExternalTensor<FLOAT,[512,512]>(path='model.onnx.data', 
name='model.model.decoder.layers.2.encoder_attn.v_proj.weight', offset=245864448, length=1048576, 
base_dir='/home/justinchu/dev/ONNXConverter/docker/dump_bash_bench/BlenderbotSmallForConditionalGeneration-torch
-onnx-detailed-cpu-')

Min: -0.08586505800485611, Max: 0.09103105217218399, NaN count: 0, Inf count: 0
Sparsity (abs<1e-06): 0.00
Histogram:
   11504 ┼
   10226 ┤                                  ╭───────╮
    8948 ┤                                ╭─╯       ╰─╮
    7670 ┤                              ╭─╯           ╰─╮
    6392 ┤                            ╭─╯               ╰─╮
    5113 ┤                          ╭─╯                   ╰─╮
    3835 ┤                        ╭─╯                       ╰─╮
    2557 ┤                     ╭──╯                           ╰─╮
    1279 ┤                ╭────╯                                ╰────╮
       1 ┼────────────────╯                                          ╰───────────────────
    -0.0859  -0.0682  -0.0505  -0.0306  -0.0129  0.0070  0.0225  0.0402  0.0557  0.0733  0.0910

[titaiwangms]

Is it only the path change that makes it more efficient?

[justinchuby]

Is it only the path change that makes it more efficient?

I added PR description

[codecov]

Codecov Report

Attention: Patch coverage is 77.38095% with 19 lines in your changes missing coverage. Please review.

Project coverage is 75.03%. Comparing base (87aee66) to head (fb65fa4).

Files	Patch %	Lines
onnxscript/ir/_external_data.py	60.86%	4 Missing and 5 partials ⚠️
onnxscript/ir/_io.py	50.00%	7 Missing ⚠️
onnxscript/ir/_external_data_test.py	90.47%	1 Missing and 1 partial ⚠️
onnxscript/ir/_core.py	92.85%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1801      +/-   ##
==========================================
- Coverage   75.07%   75.03%   -0.04%     
==========================================
  Files         245      248       +3     
  Lines       26555    26633      +78     
  Branches     4873     4887      +14     
==========================================
+ Hits        19936    19984      +48     
- Misses       5685     5707      +22     
- Partials      934      942       +8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[github-actions]

Test Results

     24 files  ±     0       24 suites  ±0   3h 32m 17s ⏱️ + 7m 30s
 13 506 tests  -  2 220   11 936 ✅  - 1 783    1 543 💤  -    434   27 ❌  - 3 
482 068 runs  +36 066  100 122 ✅ +5 187  381 718 💤 +30 882  228 ❌  - 3 

For more details on these failures, see this check.

Results for commit 7887b9e. ± Comparison against base commit 87aee66.

This pull request removes 2220 tests.

onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_basic_defs
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_doc_string
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_if_defs
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_if_loop_defs
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_loop_defs
onnxscript._internal.analysis_test.TestExposedUses ‑ test_basic
onnxscript._internal.analysis_test.TestExposedUses ‑ test_called_function
onnxscript._internal.analysis_test.TestExposedUses ‑ test_doc_string
onnxscript._internal.analysis_test.TestExposedUses ‑ test_for_loop
onnxscript._internal.analysis_test.TestExposedUses ‑ test_if
…