Merge pull request #37252 from vespa-engine/arnej/configurable-onnx-optimize

arnej27959 · web-flow · commit 1cb54bb32631 · 2026-06-19T11:33:16.000+02:00
Make ONNX graph optimization configurable per model (default off)
diff --git a/config-model/src/main/java/com/yahoo/schema/OnnxModel.java b/config-model/src/main/java/com/yahoo/schema/OnnxModel.java
@@ -30,6 +30,7 @@ public class OnnxModel extends DistributableResource implements Cloneable {
 
     // Runtime options
     private OnnxModelOptions onnxModelOptions = OnnxModelOptions.empty();
+    private boolean optimizeModel = false;
 
     public OnnxModel(String name) {
         super(name);
@@ -174,4 +175,12 @@ public Optional<OnnxModelOptions.GpuDevice> getGpuDevice() {
 
     public OnnxModelOptions onnxModelOptions() { return onnxModelOptions; }
 
+    public void setOptimizeModel(boolean optimizeModel) {
+        this.optimizeModel = optimizeModel;
+    }
+
+    public boolean getOptimizeModel() {
+        return optimizeModel;
+    }
+
 }
diff --git a/config-model/src/main/java/com/yahoo/schema/derived/FileDistributedOnnxModels.java b/config-model/src/main/java/com/yahoo/schema/derived/FileDistributedOnnxModels.java
@@ -49,6 +49,7 @@ public FileDistributedOnnxModels clone() {
     private static OnnxModelsConfig.Model.Builder toConfig(OnnxModel model) {
         OnnxModelsConfig.Model.Builder builder = new OnnxModelsConfig.Model.Builder();
         builder.dry_run_on_setup(true);
+        builder.optimize_model(model.getOptimizeModel());
         builder.name(model.getName());
         builder.fileref(model.getFileReference());
         model.getInputMap().forEach((name, source) -> builder.input(new OnnxModelsConfig.Model.Input.Builder().name(name).source(source)));
diff --git a/config-model/src/main/javacc/SchemaParser.jj b/config-model/src/main/javacc/SchemaParser.jj
@@ -201,6 +201,7 @@ TOKEN :
 | < INTEROP_THREADS: "interop-threads">
 | < GPU_DEVICE: "gpu-device">
 | < EXECUTION_MODE: "execution-mode">
+| < OPTIMIZE_MODEL: "optimize-model">
 | < PARALLEL: "parallel">
 | < SEQUENTIAL: "sequential">
 | < MODEL: "model" >
@@ -1735,6 +1736,7 @@ void onnxModelItem(OnnxModel onnxModel) :
 {
     String path = null;
     int num;
+    boolean enable;
 }
 {
     (
@@ -1745,6 +1747,7 @@ void onnxModelItem(OnnxModel onnxModel) :
         <INTEROP_THREADS> <COLON> num = integer() { onnxModel.setStatelessInterOpThreads(num); } |
         <EXECUTION_MODE> <COLON> ( <PARALLEL> { onnxModel.setStatelessExecutionMode("parallel"); }
                                   | <SEQUENTIAL> { onnxModel.setStatelessExecutionMode("sequential"); } ) |
+        <OPTIMIZE_MODEL> <COLON> enable = bool() { onnxModel.setOptimizeModel(enable); } |
         (<ONNX_INPUT_SL>) {
             String name = token.image.substring(5, token.image.lastIndexOf(":")).trim();
             if (name.startsWith("\"")) { name = name.substring(1, name.length() - 1); }
diff --git a/config-model/src/test/derived/globalphase_onnx_inside/onnx-models.cfg b/config-model/src/test/derived/globalphase_onnx_inside/onnx-models.cfg
@@ -10,6 +10,7 @@ model[].input[].source "constant(xx)"
 model[].output[].name "vector_Y"
 model[].output[].as "out"
 model[].dry_run_on_setup true
+model[].optimize_model false
 model[].stateless_execution_mode ""
 model[].stateless_interop_threads -1
 model[].stateless_intraop_threads -1
@@ -26,6 +27,7 @@ model[].input[].source "rankingExpression(indirect_x)"
 model[].output[].name "vector_Y"
 model[].output[].as "foobar"
 model[].dry_run_on_setup true
+model[].optimize_model false
 model[].stateless_execution_mode "parallel"
 model[].stateless_interop_threads 5
 model[].stateless_intraop_threads 3
@@ -42,6 +44,7 @@ model[].input[].source "rankingExpression(indirect_x)"
 model[].output[].name "vector_Y"
 model[].output[].as "foobar"
 model[].dry_run_on_setup true
+model[].optimize_model false
 model[].stateless_execution_mode ""
 model[].stateless_interop_threads -1
 model[].stateless_intraop_threads -1
@@ -58,6 +61,7 @@ model[].input[].source "rankingExpression(indirect_x)"
 model[].output[].name "vector_Y"
 model[].output[].as "foobar"
 model[].dry_run_on_setup true
+model[].optimize_model false
 model[].stateless_execution_mode ""
 model[].stateless_interop_threads -1
 model[].stateless_intraop_threads -1
diff --git a/config-model/src/test/derived/globalphase_token_functions/onnx-models.cfg b/config-model/src/test/derived/globalphase_token_functions/onnx-models.cfg
@@ -10,6 +10,7 @@ model[].input[].source "rankingExpression(token_type_ids)"
 model[].output[].name "score"
 model[].output[].as "score"
 model[].dry_run_on_setup true
+model[].optimize_model false
 model[].stateless_execution_mode ""
 model[].stateless_interop_threads -1
 model[].stateless_intraop_threads -1
diff --git a/config-model/src/test/derived/vector_constant/onnx-models.cfg b/config-model/src/test/derived/vector_constant/onnx-models.cfg
@@ -10,6 +10,7 @@ model[].input[].source "constant(xx)"
 model[].output[].name "vector_Y"
 model[].output[].as "foobar"
 model[].dry_run_on_setup true
+model[].optimize_model false
 model[].stateless_execution_mode ""
 model[].stateless_interop_threads -1
 model[].stateless_intraop_threads -1
diff --git a/config-model/src/test/java/com/yahoo/schema/SchemaTestCase.java b/config-model/src/test/java/com/yahoo/schema/SchemaTestCase.java
@@ -125,6 +125,33 @@ function vespa_input() {
         assertEquals("vespa_input", input);
     }
 
+    @Test
+    void testOnnxModelOptimizeModel() throws Exception {
+        String schema =
+                """
+                schema test {
+                    document test {
+                        field id type string {
+                            indexing: summary | attribute
+                        }
+                    }
+                    onnx-model default_model {
+                        file: files/default.onnx
+                    }
+                    onnx-model optimized_model {
+                        file: files/optimized.onnx
+                        optimize-model: true
+                    }
+                }""";
+        ApplicationBuilder builder = new ApplicationBuilder(new DeployLoggerStub());
+        builder.processorsToSkip().add(OnnxModelTypeResolver.class); // Avoid discovering the Onnx model referenced does not exist
+        builder.addSchema(schema);
+        var application = builder.build(true);
+        var models = application.schemas().get("test").onnxModels();
+        assertFalse(models.get("default_model").getOptimizeModel());
+        assertTrue(models.get("optimized_model").getOptimizeModel());
+    }
+
     @Test
     void testSchemaInheritance() throws ParseException {
         String parentLines = joinLines(
diff --git a/config-model/src/test/java/com/yahoo/schema/processing/RankingExpressionWithOnnxModelTestCase.java b/config-model/src/test/java/com/yahoo/schema/processing/RankingExpressionWithOnnxModelTestCase.java
@@ -18,6 +18,7 @@
 import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class RankingExpressionWithOnnxModelTestCase {
@@ -75,6 +76,7 @@ private void assertGeneratedConfig(VespaModel vespaModel) {
         assertEquals(6, config.model().size());
         for (OnnxModelsConfig.Model model : config.model()) {
             assertTrue(model.dry_run_on_setup());
+            assertFalse(model.optimize_model());
         }
 
         OnnxModelsConfig.Model model = config.model(0);
diff --git a/configdefinitions/src/vespa/onnx-models.def b/configdefinitions/src/vespa/onnx-models.def
@@ -11,6 +11,7 @@ model[].input[].source              string
 model[].output[].name               string
 model[].output[].as                 string
 model[].dry_run_on_setup            bool default=false
+model[].optimize_model              bool default=false
 model[].stateless_execution_mode    string default=""
 model[].stateless_interop_threads   int default=-1
 model[].stateless_intraop_threads   int default=-1
diff --git a/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp b/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp
@@ -501,6 +501,22 @@ TEST(OnnxModelCacheTest, share_and_evict_onnx_models) {
     EXPECT_EQ(OnnxModelCache::count_refs(), 0);
 }
 
+TEST(OnnxModelCacheTest, optimize_setting_is_part_of_cache_key) {
+    {
+        auto enabled1 = OnnxModelCache::load(simple_model, Onnx::Optimize::ENABLE);
+        auto enabled2 = OnnxModelCache::load(simple_model, Onnx::Optimize::ENABLE);
+        auto disabled = OnnxModelCache::load(simple_model, Onnx::Optimize::DISABLE);
+        // same file + same optimize setting shares the loaded model
+        EXPECT_EQ(&(enabled1->get()), &(enabled2->get()));
+        // same file but different optimize setting is a distinct cache entry
+        EXPECT_NE(&(enabled1->get()), &(disabled->get()));
+        EXPECT_EQ(OnnxModelCache::num_cached(), 2);
+        EXPECT_EQ(OnnxModelCache::count_refs(), 3);
+    }
+    EXPECT_EQ(OnnxModelCache::num_cached(), 0);
+    EXPECT_EQ(OnnxModelCache::count_refs(), 0);
+}
+
 TensorSpec val(const std::string& expr) {
     auto result = TensorSpec::from_expr(expr);
     EXPECT_FALSE(ValueType::from_spec(result.type()).is_error());
diff --git a/eval/src/vespa/eval/onnx/onnx_model_cache.cpp b/eval/src/vespa/eval/onnx/onnx_model_cache.cpp
@@ -14,12 +14,13 @@ void OnnxModelCache::release(Map::iterator entry) {
     }
 }
 
-OnnxModelCache::Token::UP OnnxModelCache::load(const std::string& model_file) {
+OnnxModelCache::Token::UP OnnxModelCache::load(const std::string& model_file, Onnx::Optimize optimize) {
     std::lock_guard<std::mutex> guard(_lock);
-    auto                        pos = _cached.find(model_file);
+    Key                         key{model_file, optimize};
+    auto                        pos = _cached.find(key);
     if (pos == _cached.end()) {
-        auto model = std::make_unique<Onnx>(model_file, Onnx::Optimize::ENABLE);
-        auto res = _cached.emplace(model_file, std::move(model));
+        auto model = std::make_unique<Onnx>(model_file, optimize);
+        auto res = _cached.emplace(std::move(key), std::move(model));
         assert(res.second);
         pos = res.first;
     }
diff --git a/eval/src/vespa/eval/onnx/onnx_model_cache.h b/eval/src/vespa/eval/onnx/onnx_model_cache.h
@@ -19,7 +19,7 @@ namespace vespalib::eval {
 class OnnxModelCache {
 private:
     struct ctor_tag {};
-    using Key = std::string;
+    using Key = std::pair<std::string, Onnx::Optimize>;
     struct Value {
         size_t                num_refs;
         std::unique_ptr<Onnx> model;
@@ -48,7 +48,7 @@ class OnnxModelCache {
         ~Token() { OnnxModelCache::release(_entry); }
     };
 
-    static Token::UP load(const std::string& model_file);
+    static Token::UP load(const std::string& model_file, Onnx::Optimize optimize = Onnx::Optimize::DISABLE);
     static size_t num_cached();
     static size_t count_refs();
 };
diff --git a/integration/schema-language-server/language-server/src/main/ccc/SchemaParser.ccc b/integration/schema-language-server/language-server/src/main/ccc/SchemaParser.ccc
@@ -426,6 +426,7 @@ TOKEN :
 | < INTEROP_THREADS: "interop-threads">
 | < GPU_DEVICE: "gpu-device">
 | < EXECUTION_MODE: "execution-mode">
+| < OPTIMIZE_MODEL: "optimize-model">
 | < PARALLEL: "parallel">
 | < SEQUENTIAL: "sequential">
 | < MODEL: "model" >
@@ -2094,6 +2095,7 @@ void onnxModelItem(OnnxModel onnxModel) :
 {
     String path = null;
     int num;
+    boolean enable;
     Node inputNode = null;
 }
 
@@ -2147,6 +2149,14 @@ void onnxModelItem(OnnxModel onnxModel) :
                 }
             }
         )
+        | <OPTIMIZE_MODEL> <COLON> enable = bool()
+            {
+                try {
+                    onnxModel.setOptimizeModel(enable);
+                } catch (IllegalArgumentException ex) {
+                    lastConsumedToken.addIllegalArgumentException(ex);
+                }
+            }
         | onnxModelInput(onnxModel)
         | onnxModelOutput(onnxModel)
     )
@@ -3485,6 +3495,7 @@ String identifierWithDashStr :
     | <ONNX_MODEL>
     | <ON_SECOND_PHASE>
     | <ON_SUMMARY>
+    | <OPTIMIZE_MODEL>
     | <POST_FILTER_THRESHOLD>
     | <PRE_POST_FILTER_TIPPING_POINT>
     | <QUERY_COMMAND>
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
@@ -211,6 +211,7 @@ private OnnxModel readOnnxModelConfig(OnnxModelsConfig.Model onnxModelConfig) {
                     .setInterOpThreads(onnxModelConfig.stateless_interop_threads())
                     .setIntraOpThreads(onnxModelConfig.stateless_intraop_threads())
                     .setGpuDevice(onnxModelConfig.gpu_device(), onnxModelConfig.gpu_device_required())
+                    .setOptimizeModel(onnxModelConfig.optimize_model())
                     .build();
             var m =  new OnnxModel(name, file, options, onnx);
             for (var spec : onnxModelConfig.input()) {
diff --git a/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/EmbeddedOnnxRuntime.java b/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/EmbeddedOnnxRuntime.java
@@ -96,7 +96,10 @@ static boolean isRuntimeAvailable(String modelPath) {
 
     private static OrtSession.SessionOptions createSessionOptions(OnnxEvaluatorOptions vespaOpts, boolean loadCuda) throws OrtException {
         var sessionOpts = new OrtSession.SessionOptions();
-        sessionOpts.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT);
+        var optLevel = vespaOpts.optimizeModel()
+                ? OrtSession.SessionOptions.OptLevel.ALL_OPT
+                : OrtSession.SessionOptions.OptLevel.NO_OPT;
+        sessionOpts.setOptimizationLevel(optLevel);
         var execMode = vespaOpts.executionMode() == OnnxEvaluatorOptions.ExecutionMode.PARALLEL ? PARALLEL : SEQUENTIAL;
         sessionOpts.setExecutionMode(execMode);
         sessionOpts.setInterOpNumThreads(execMode == PARALLEL ? vespaOpts.interOpThreads() : 1);
diff --git a/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluatorOptions.java b/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluatorOptions.java
@@ -23,6 +23,7 @@ public record OnnxEvaluatorOptions(
         int intraOpThreads,
         int gpuDeviceNumber,
         boolean gpuDeviceRequired,
+        boolean optimizeModel,
         int batchingMaxSize,
         Optional<Duration> batchingMaxDelay,
         int numModelInstances,
@@ -98,6 +99,7 @@ public static class Builder {
         private int intraOpThreads;
         private int gpuDeviceNumber;
         private boolean gpuDeviceRequired;
+        private boolean optimizeModel;
         private int batchingMaxSize;
         private Optional<Duration> batchingMaxDelay;
         private int numModelInstances;
@@ -117,6 +119,7 @@ public Builder(int availableProcessors) {
             intraOpThreads = calculateThreads(-4, availableProcessors);
             gpuDeviceNumber = -1;
             gpuDeviceRequired = false;
+            optimizeModel = false;
             batchingMaxSize = 1;
             batchingMaxDelay = Optional.empty();
             numModelInstances = 1;
@@ -131,6 +134,7 @@ public Builder(OnnxEvaluatorOptions options) {
             this.intraOpThreads = options.intraOpThreads();
             this.gpuDeviceNumber = options.gpuDeviceNumber();
             this.gpuDeviceRequired = options.gpuDeviceRequired();
+            this.optimizeModel = options.optimizeModel();
             this.batchingMaxSize = options.batchingMaxSize();
             this.batchingMaxDelay = options.batchingMaxDelay;
             this.numModelInstances = options.numModelInstances();
@@ -186,6 +190,11 @@ public Builder setGpuDevice(int deviceNumber) {
             return this;
         }
 
+        public Builder setOptimizeModel(boolean optimizeModel) {
+            this.optimizeModel = optimizeModel;
+            return this;
+        }
+
         public Builder setBatchingMaxSize(int maxSize) {
             this.batchingMaxSize = maxSize;
             return this;
@@ -223,6 +232,7 @@ public OnnxEvaluatorOptions build() {
                     intraOpThreads,
                     gpuDeviceNumber,
                     gpuDeviceRequired,
+                    optimizeModel,
                     batchingMaxSize,
                     batchingMaxDelay,
                     numModelInstances,
diff --git a/searchlib/src/vespa/searchlib/features/onnx_feature.cpp b/searchlib/src/vespa/searchlib/features/onnx_feature.cpp
@@ -119,7 +119,8 @@ bool OnnxBlueprint::setup(const IIndexEnvironment& env, const ParameterList& par
             _debug_model = std::make_unique<Onnx>(model_cfg->file_path(), Optimize::DISABLE);
             _model = _debug_model.get();
         } else {
-            _cache_token = OnnxModelCache::load(model_cfg->file_path());
+            auto optimize = model_cfg->optimize_model() ? Optimize::ENABLE : Optimize::DISABLE;
+            _cache_token = OnnxModelCache::load(model_cfg->file_path(), optimize);
             _model = &(_cache_token->get());
         }
     } catch (const Ort::Exception& ex) {
diff --git a/searchlib/src/vespa/searchlib/fef/onnx_model.cpp b/searchlib/src/vespa/searchlib/fef/onnx_model.cpp
@@ -7,7 +7,12 @@
 namespace search::fef {
 
 OnnxModel::OnnxModel(const std::string& name_in, const std::string& file_path_in)
-    : _name(name_in), _file_path(file_path_in), _input_features(), _output_names(), _dry_run_on_setup(false) {
+    : _name(name_in),
+      _file_path(file_path_in),
+      _input_features(),
+      _output_names(),
+      _dry_run_on_setup(false),
+      _optimize_model(false) {
 }
 
 OnnxModel::OnnxModel(OnnxModel&&) noexcept = default;
@@ -29,6 +34,11 @@ OnnxModel& OnnxModel::dry_run_on_setup(bool value) {
     return *this;
 }
 
+OnnxModel& OnnxModel::optimize_model(bool value) {
+    _optimize_model = value;
+    return *this;
+}
+
 std::optional<std::string> OnnxModel::input_feature(const std::string& model_input_name) const {
     auto pos = _input_features.find(model_input_name);
     if (pos != _input_features.end()) {
@@ -46,8 +56,9 @@ std::optional<std::string> OnnxModel::output_name(const std::string& model_outpu
 }
 
 bool OnnxModel::operator==(const OnnxModel& rhs) const {
-    return (std::tie(_name, _file_path, _input_features, _output_names, _dry_run_on_setup) ==
-            std::tie(rhs._name, rhs._file_path, rhs._input_features, rhs._output_names, rhs._dry_run_on_setup));
+    return (std::tie(_name, _file_path, _input_features, _output_names, _dry_run_on_setup, _optimize_model) ==
+            std::tie(rhs._name, rhs._file_path, rhs._input_features, rhs._output_names, rhs._dry_run_on_setup,
+                     rhs._optimize_model));
 }
 
 } // namespace search::fef
diff --git a/searchlib/src/vespa/searchlib/fef/onnx_model.h b/searchlib/src/vespa/searchlib/fef/onnx_model.h
@@ -20,6 +20,7 @@ class OnnxModel {
     std::map<std::string, std::string> _input_features;
     std::map<std::string, std::string> _output_names;
     bool                               _dry_run_on_setup;
+    bool                               _optimize_model;
 
 public:
     OnnxModel(const std::string& name_in, const std::string& file_path_in);
@@ -34,9 +35,11 @@ class OnnxModel {
     OnnxModel& input_feature(const std::string& model_input_name, const std::string& input_feature);
     OnnxModel& output_name(const std::string& model_output_name, const std::string& output_name);
     OnnxModel& dry_run_on_setup(bool value);
+    OnnxModel& optimize_model(bool value);
     std::optional<std::string> input_feature(const std::string& model_input_name) const;
     std::optional<std::string> output_name(const std::string& model_output_name) const;
     bool dry_run_on_setup() const { return _dry_run_on_setup; }
+    bool optimize_model() const { return _optimize_model; }
     bool operator==(const OnnxModel& rhs) const;
     const std::map<std::string, std::string>& inspect_input_features() const { return _input_features; }
     const std::map<std::string, std::string>& inspect_output_names() const { return _output_names; }
diff --git a/searchlib/src/vespa/searchlib/fef/onnx_models.cpp b/searchlib/src/vespa/searchlib/fef/onnx_models.cpp

Original file line number	Diff line number	Diff line change
`@@ -201,6 +201,7 @@ TOKEN :`
`201`	`201`	`\| < INTEROP_THREADS: "interop-threads">`
`202`	`202`	`\| < GPU_DEVICE: "gpu-device">`
`203`	`203`	`\| < EXECUTION_MODE: "execution-mode">`
	`204`	`+\| < OPTIMIZE_MODEL: "optimize-model">`
`204`	`205`	`\| < PARALLEL: "parallel">`
`205`	`206`	`\| < SEQUENTIAL: "sequential">`
`206`	`207`	`\| < MODEL: "model" >`
`@@ -1735,6 +1736,7 @@ void onnxModelItem(OnnxModel onnxModel) :`
`1735`	`1736`	`{`
`1736`	`1737`	`String path = null;`
`1737`	`1738`	`int num;`
	`1739`	`+ boolean enable;`
`1738`	`1740`	`}`
`1739`	`1741`	`{`
`1740`	`1742`	`(`
`@@ -1745,6 +1747,7 @@ void onnxModelItem(OnnxModel onnxModel) :`
`1745`	`1747`	`<INTEROP_THREADS> <COLON> num = integer() { onnxModel.setStatelessInterOpThreads(num); } \|`
`1746`	`1748`	`<EXECUTION_MODE> <COLON> ( <PARALLEL> { onnxModel.setStatelessExecutionMode("parallel"); }`
`1747`	`1749`	`\| <SEQUENTIAL> { onnxModel.setStatelessExecutionMode("sequential"); } ) \|`
	`1750`	`+ <OPTIMIZE_MODEL> <COLON> enable = bool() { onnxModel.setOptimizeModel(enable); } \|`
`1748`	`1751`	`(<ONNX_INPUT_SL>) {`
`1749`	`1752`	`String name = token.image.substring(5, token.image.lastIndexOf(":")).trim();`
`1750`	`1753`	`if (name.startsWith("\"")) { name = name.substring(1, name.length() - 1); }`