From 3d6e7317a7fcbf7d29784898a888010eaff4554f Mon Sep 17 00:00:00 2001 From: Xiaoyu <85524621+xiaoyu-work@users.noreply.github.com> Date: Wed, 15 Jan 2025 15:31:25 -0800 Subject: [PATCH] Set onnxoptimizer.optimize as default for peepholeoptimizer (#1550) ## Describe your changes Set onnxoptimizer.optimize as default for peepholeoptimizer ## 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. - [ ] Is this PR including examples changes? If yes, please remember to update [example documentation](https://github.com/microsoft/Olive/blob/main/docs/source/examples.md) in a follow-up PR. ## (Optional) Issue link --- .../model-opt-and-transform/onnx.md | 4 ++ olive/passes/onnx/peephole_optimizer.py | 45 ++++++------------- .../passes/onnx/test_peephole_optimizer.py | 31 +++++++------ 3 files changed, 33 insertions(+), 47 deletions(-) diff --git a/docs/source/how-to/configure-workflows/model-opt-and-transform/onnx.md b/docs/source/how-to/configure-workflows/model-opt-and-transform/onnx.md index 538eb05b4..20788cb7b 100644 --- a/docs/source/how-to/configure-workflows/model-opt-and-transform/onnx.md +++ b/docs/source/how-to/configure-workflows/model-opt-and-transform/onnx.md @@ -55,6 +55,10 @@ The `OnnxPeepholeOptimizer` leverages `onnxscript` (https://onnxscript.ai/tutori | **Fuse Consecutive Slices** | Fuses consecutive Slice operations. | | **Eliminate Unused Initializer** | Eliminates unused initializers. | | **Eliminate Duplicate Initializer** | Eliminates duplicate initializers. | +| **Broadcast to MatMul** | Converts broadcast patterns into MatMul operations for better efficiency. | +| **Cast Constant of Shape** | Simplifies constant casting for shape operations. | +| **GEMM to MatMul+Add** | Converts GEMM operations into MatMul and Add for improved compatibility. | +| **No-Op Removal** | Removes redundant or no-op operations in the computation graph. | Please refer to [OnnxPeepholeOptimizer](../../../reference/pass.rst#onnx_peephole_optimizer) for more details about the pass and its config parameters. diff --git a/olive/passes/onnx/peephole_optimizer.py b/olive/passes/onnx/peephole_optimizer.py index d1fa3ca9c..4a6af4241 100644 --- a/olive/passes/onnx/peephole_optimizer.py +++ b/olive/passes/onnx/peephole_optimizer.py @@ -4,7 +4,7 @@ # -------------------------------------------------------------------------- import logging from pathlib import Path -from typing import Any, Dict, List +from typing import Any, Dict import numpy as np import onnx @@ -236,35 +236,27 @@ def onnxscript_optimize(self): try: import onnxscript except ImportError: - logger.warning("Please install onnxscript to use the ONNX optimizer feature. Skip onnxscript optimization.") + logger.warning("Please install `onnxscript` to apply more optimization.") return onnxscript.optimizer.optimize(self.model) + def onnxoptimizer_optimize(self): + try: + from onnxoptimizer import optimize + except ImportError: + logger.warning("Please install `onnxoptimizer` to apply more optimization.") + return + + self.model = optimize(self.model) + class OnnxPeepholeOptimizer(Pass): """Optimize ONNX model by fusing nodes.""" @classmethod def _default_config(cls, accelerator_spec: AcceleratorSpec) -> Dict[str, PassConfigParam]: - return { - "onnxoptimizer": PassConfigParam( - type_=bool, - default_value=False, - description="Whether to run the ONNX optimizer (https://github.com/onnx/optimizer/). Default is False.", - ), - "passes": PassConfigParam( - type_=List[str], - default_value=None, - description="List of passes of ONNX optimizer to run. Default is None.", - ), - "fixed_point": PassConfigParam( - type_=bool, - default_value=False, - description="Whether to run the fixed-point optimization of ONNX optimizer. Default is False.", - ), - **get_external_data_config(), - } + return get_external_data_config() def _run_for_config( self, model: ONNXModelHandler, config: Dict[str, Any], output_model_path: str @@ -273,20 +265,11 @@ def _run_for_config( # optimize model peephole_optimizer = ModelOptimizer(model.model_path) + peephole_optimizer.onnxscript_optimize() + peephole_optimizer.onnxoptimizer_optimize() peephole_optimizer.fuse_transpose_qat() peephole_optimizer.patch_unsupported_argmax_operator() peephole_optimizer.fuse_reshape_operations() - peephole_optimizer.onnxscript_optimize() - - if config["onnxoptimizer"]: - try: - from onnxoptimizer import optimize - - peephole_optimizer.model = optimize(peephole_optimizer.model, config["passes"], config["fixed_point"]) - except ImportError: - logger.warning( - "Please install onnxoptimizer to use the ONNX optimizer feature. Skip onnxoptimizer optimization." - ) # save the model to the output path and return the model return model_proto_to_olive_model(peephole_optimizer.model, output_model_path, config) diff --git a/test/unit_test/passes/onnx/test_peephole_optimizer.py b/test/unit_test/passes/onnx/test_peephole_optimizer.py index 985dd2a89..e27232e50 100644 --- a/test/unit_test/passes/onnx/test_peephole_optimizer.py +++ b/test/unit_test/passes/onnx/test_peephole_optimizer.py @@ -65,9 +65,10 @@ def _make_model_for_patch_unsupported_argmax_operator( return model_proto_to_olive_model(model, filepath, config) +@patch("onnxoptimizer.optimize") @patch("onnxscript.optimizer.optimize") def test_onnx_peephole_optimizer_pass_patch_unsupported_argmax_operator_modified( - mock_optimize, tmp_path, external_data_config + mock_onnxscript, mock_onnxoptimizer, tmp_path, external_data_config ): m = _make_model_for_patch_unsupported_argmax_operator( TensorProto.INT64, str(tmp_path / "input.onnx"), external_data_config @@ -75,6 +76,8 @@ def test_onnx_peephole_optimizer_pass_patch_unsupported_argmax_operator_modified p = create_pass_from_dict( OnnxPeepholeOptimizer, external_data_config, disable_search=True, accelerator_spec=DEFAULT_GPU_CUDA_ACCELERATOR ) + mock_onnxscript.return_value = m.load_model() + mock_onnxoptimizer.return_value = m.load_model() actual_model = p.run(m, str(tmp_path / "onnx")) assert Path(actual_model.model_path).exists() @@ -193,40 +196,36 @@ def test_onnx_peephole_optimizer_pass_fuse_reshape_operations(tmp_path, external @patch("olive.passes.onnx.peephole_optimizer.model_proto_to_olive_model") +@patch("onnxoptimizer.optimize") @patch("onnxscript.optimizer.optimize") -def test_onnxscript(mock_optimize, mock_model_proto_to_olive_model, tmp_path): +def test_onnxscript(mock_onnxscript, mock_onnxoptimizer, mock_model_proto_to_olive_model, tmp_path): # setup input_model = get_onnx_model() p = create_pass_from_dict(OnnxPeepholeOptimizer, {}, disable_search=True) + mock_onnxscript.return_value = input_model.load_model() + mock_onnxoptimizer.return_value = input_model.load_model() output_folder = str(tmp_path / "onnx") # execute p.run(input_model, output_folder) # assert - mock_optimize.assert_called_once_with(input_model.load_model()) + mock_onnxscript.assert_called_once_with(input_model.load_model()) @patch("olive.passes.onnx.peephole_optimizer.model_proto_to_olive_model") @patch("onnxoptimizer.optimize") -def test_onnxoptimizer(mock_optimize, mock_model_proto_to_olive_model, tmp_path): +@patch("onnxscript.optimizer.optimize") +def test_onnxoptimizer(mock_onnxscript, mock_onnxoptimizer, mock_model_proto_to_olive_model, tmp_path): # setup input_model = get_onnx_model() - passes = ["pass"] - fixed_point = True - p = create_pass_from_dict( - OnnxPeepholeOptimizer, - { - "onnxoptimizer": True, - "passes": passes, - "fixed_point": fixed_point, - }, - disable_search=True, - ) + p = create_pass_from_dict(OnnxPeepholeOptimizer, {}, disable_search=True) + mock_onnxscript.return_value = input_model.load_model() + mock_onnxoptimizer.return_value = input_model.load_model() output_folder = str(tmp_path / "onnx") # execute p.run(input_model, output_folder) # assert - mock_optimize.assert_called_once_with(input_model.load_model(), passes, fixed_point) + mock_onnxoptimizer.assert_called_once()