Migrate bnb4, io_datatype_converter, dynamic_to_fixed_shape, and static_llm passes to onnx_ir

[justinchuby]

Describe your changes

Continues the ongoing migration of ONNX passes from onnx.ModelProto APIs to onnx_ir (following the pattern established by the weight-only quantizers rtn/hqq/kquant). Migrates four self-contained graph-transform passes:

• OnnxBnb4Quantization (bnb_quantization.py): Traverses the graph with onnx_ir and replaces MatMul nodes with MatMulBnb4 nodes via ir.convenience.replace_nodes_and_values, instead of going through the ORT MatMulBnb4Quantizer proto wrapper. The native quantize_matmul_bnb4 kernel is retained for the FP4/NF4 bit-packing (that packing is defined by the kernel). Graph output names are preserved, consistent with the other weight-only quant passes.
• OnnxIODataTypeConverter (io_datatype_converter.py): Reimplements the input/output Cast insertion and consumer rewiring on onnx_ir.
• DynamicToFixedShape (dynamic_to_fixed_shape.py): Reimplements ORT's make_dim_param_fixed / make_input_shape_fixed / remove_invalid_dim_values natively over ir.Model (traversing subgraphs via ir.Model.graphs()). Output shape inference uses the onnx-shape-inference package (infer_symbolic_shapes), which runs directly on the ir.Model with no proto round-trip and refines shapes in place, replacing the earlier ORT SymbolicShapeInference (via ir.to_proto).
• StaticLLM (static_llm.py): Now operates entirely on ir.Model. All onnx.ModelProto / onnx.load / onnx.save usage is removed in favor of ir.load / ONNXModelHandler.load_ir_model() and ir.save. fix_shape runs on the IR model and reuses the IR-based fix_dim_params from dynamic_to_fixed_shape, removing the indirect dependency on onnxruntime's make_dim_param_fixed / onnx_model_utils (output shape inference uses the onnx-shape-inference package, as with the other migrated passes). The external-data sharing between the context and iterator component models is preserved by lazily loading the intermediate model so both fixed-shape outputs reference the same transformer*.onnx.data file.

Also adds OpType.MatMulBnb4 and OpType.Cast to olive/constants.py, adds an IR-based add_version_metadata_to_ir_model helper to common.py, and removes the now-unused proto-based fix_dim_params / fix_input_shapes / _fix_output_shapes helpers from common.py (they were the only remaining callers of the onnxruntime onnx_model_utils shape helpers, and static_llm.py no longer depends on them). Adds onnx-shape-inference>=0.2.0 to requirements.txt.

Other remaining proto-based ONNX passes (e.g. mixed_precision, mixed_precision_overrides, add_metadata, and external-tool wrappers like aimet/inc/nvmo/transformer_optimization/optimum) are intentionally left as-is: they are thin wrappers over third-party proto APIs where an IR migration would only force an IR→proto→IR round-trip with no benefit.

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all tests can pass.
• Update documents if necessary.
• Lint and apply fixes to your code by running lintrunner -a
• Is this a user-facing change? If yes, give a description of this change to be included in the release notes.

Existing pass tests (test_bnb_quantization.py, test_io_datatype_converter.py, test_dynamic_to_fixed_shape.py) all pass (27 total). Verified end-to-end with ONNX Runtime inference, exact parity against the ORT reference for shape fixing (including subgraph recursion), and onnx.checker validity. For static_llm.py, verified shape-fixing parity on a symbolic-dim model and that the context/iterator external-data sharing is preserved (both components load, pass onnx.checker, and carry olive_version metadata). No new public config/behavior, so not user-facing.

(Optional) Issue link

[justinchuby]

@copilot Update olive/passes/onnx/static_llm.py to use onnx-ir completely. Remove any import/usage of OnnxModel from onnxruntime, and update / recreate the logic using onnx-ir / onnxscript rewriter instead

[Copilot]

Pull request overview

This PR continues Olive’s ONNX pass migration from onnx.ModelProto-based graph edits to onnx_ir, aligning these passes with the newer IR-based weight-only quantization pattern already used elsewhere in the codebase.

Changes:

• Migrates OnnxBnb4Quantization to traverse/modify graphs via onnx_ir and replace MatMul with MatMulBnb4.
• Reimplements OnnxIODataTypeConverter input/output Cast insertion + rewiring using onnx_ir.
• Reimplements DynamicToFixedShape dimension fixing over ir.Model (including subgraph traversal), while keeping ORT SymbolicShapeInference for output shape inference.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
test/passes/onnx/test_bnb_quantization.py	Updates the test to use IR loading (load_ir_model) for MatMul discovery.
olive/passes/onnx/io_datatype_converter.py	Moves I/O Cast insertion and consumer/provider rewiring to onnx_ir.
olive/passes/onnx/dynamic_to_fixed_shape.py	Rebuilds dim-param and input-shape fixing logic on ir.Model, with ORT-based inference preserved.
olive/passes/onnx/bnb_quantization.py	Migrates BNB4 MatMul quantization to IR and performs node/value replacement in the IR graph.
olive/constants.py	Adds OpType.MatMulBnb4 and OpType.Cast to support the migrated passes.

[Copilot]

@copilot Update olive/passes/onnx/static_llm.py to use onnx-ir completely. Remove any import/usage of OnnxModel from onnxruntime, and update / recreate the logic using onnx-ir / onnxsc...

Done. static_llm.py now works entirely on onnx_ir:

• Removed all onnx.ModelProto / onnx.load / onnx.save usage; models are loaded with ir.load / ONNXModelHandler.load_ir_model() and saved with ir.save.
• fix_shape now operates on ir.Model and reuses the IR-based fix_dim_params from dynamic_to_fixed_shape, so the indirect dependency on onnxruntime's make_dim_param_fixed / onnx_model_utils is gone. (Output shape inference still uses ORT SymbolicShapeInference, as it does for the other IR passes.)
• The context/iterator external-data sharing is preserved by lazily loading the intermediate model so both fixed-shape models reference the same transformer*.onnx.data file.
• Added an IR-based add_version_metadata_to_ir_model helper and removed the now-unused proto fix_dim_params / fix_input_shapes / _fix_output_shapes from common.py.

Verified the shape-fixing parity and external-data sharing (both components load, pass onnx.checker, and carry olive_version metadata).

[justinchuby]

@copilot is ir.external_data.load_to_model needed? I don't think so?

[Copilot]

You're right — removed it. This pass only edits shape metadata, not tensor data, so lazily-loaded external tensors don't need to be materialized; ir.save reads them from the source file on save. Verified external data (and weight values) are preserved end-to-end and the test_dynamic_to_fixed_shape.py tests still pass.

[justinchuby]

@copilot error:

test/passes/onnx/test_compose.py:90: in test_compose_onnx_models_llm_pipeline
    llm_model = create_pass_from_dict(StaticLLM, {"batch_size": 1, "context_length": 64}, disable_search=True).run(
olive/passes/olive_pass.py:243: in run
    output_model = self._run_for_config(model, self.config, output_model_path)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
olive/passes/onnx/static_llm.py:90: in _run_for_config
    return self._run_generic(model, config, output_model_path)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
olive/passes/onnx/static_llm.py:190: in _run_generic
    return process_llm_pipeline(
olive/passes/onnx/common.py:629: in process_llm_pipeline
    new_groups = process_func(component_models, llm_pipeline, output_dir)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
olive/passes/onnx/static_llm.py:150: in process_context_iterator
    self.fix_shape(component_ir, param_mapping)
olive/passes/onnx/static_llm.py:267: in fix_shape
    for old_dim, new_dim in zip(original_shape, new_shape):
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E   TypeError: 'NoneType' object is not iterable

[Copilot]

Pull request overview

Copilot reviewed 7 out of 8 changed files in this pull request and generated 4 comments.

[Copilot]

@copilot Update olive/passes/onnx/static_llm.py to use onnx-ir completely. Remove any import/usage of OnnxModel from onnxruntime, and update / recreate the logic using onnx-ir / onnxsc...

static_llm.py is already fully migrated to onnx-ir — there is no remaining import or use of OnnxModel/onnxruntime or onnx.ModelProto/onnx.load/onnx.save. Models are loaded via ir.load/ONNXModelHandler.load_ir_model() and saved with ir.save; shape fixing reuses the IR-based fix_dim_params (which now uses the onnx-shape-inference package for output shape inference). Let me know if there's a specific remaining usage you'd like changed.

[Copilot]

@copilot error:

...

Root cause: in StaticLLM.fix_shape, a graph output with no shape metadata makes _ir_io_shape(output) return None, and the loop then did zip(original_shape, new_shape) on None. Fixed by skipping outputs whose original or new shape is None (there are no symbolic dims to map for them). Reproduced the crash synthetically and confirmed the guard resolves it; test_dynamic_to_fixed_shape.py still passes.

Note: I have the fix committed locally, but the push is being rejected by a repository ruleset on justinchu/ir-script (GH013 ... Cannot update this protected ref), so it hasn't landed on the PR yet. Could you check the branch protection rules / re-run so the commit can be pushed?