diff --git a/src/generators.cpp b/src/generators.cpp index 30d4e9844e..53ab91cc92 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -290,7 +290,8 @@ GeneratorParams::GeneratorParams(const Config& config) } GeneratorParams::GeneratorParams(const Model& model) - : config{*model.config_.get()}, + : model_{model.shared_from_this()}, + config{*model_->config_.get()}, use_graph_capture{IsGraphCaptureEnabled(model.config_->model.decoder.session_options)}, use_multi_profile{IsMultiProfileEnabled(model.config_->model.decoder.session_options)}, p_device{model.p_device_scoring_} { diff --git a/src/generators.h b/src/generators.h index f70d599b9f..2bc6bd0292 100644 --- a/src/generators.h +++ b/src/generators.h @@ -73,7 +73,10 @@ struct GeneratorParams : std::enable_shared_from_this, LeakChec GeneratorParams(const Config& config); // This constructor is only used for internal generator benchmarks GeneratorParams(const Model& model); - const Config& config; // The model outlives the GeneratorParams + // Co-owns the model so the aliased Config below cannot be freed while this + // params object is alive. Null for the benchmark-only Config constructor. + std::shared_ptr model_; + const Config& config; // Aliases model-owned Config; kept alive by model_ Config::Search search{config.search}; // Copy of the search parameters from the config // Query the params to get the value set for a param diff --git a/test/c_api_tests.cpp b/test/c_api_tests.cpp index 90056e79ed..167077e6a1 100644 --- a/test/c_api_tests.cpp +++ b/test/c_api_tests.cpp @@ -1182,6 +1182,34 @@ TEST(CAPITests, TopKTopPExceedsVocabSizeThrows) { } } +// Regression test: a GeneratorParams created from a model must keep that model +// alive, so destroying the model handle before creating the generator does not +// cause a use-after-free (GeneratorParams aliases the model-owned Config, and +// Generator::Generator calls model.shared_from_this()). +TEST(CAPITests, CreateGeneratorAfterDestroyModel) { + OgaModel* model = nullptr; + ASSERT_EQ(OgaCreateModel(PHI2_PATH, &model), nullptr); + ASSERT_NE(model, nullptr); + + OgaGeneratorParams* params = nullptr; + ASSERT_EQ(OgaCreateGeneratorParams(model, ¶ms), nullptr); + ASSERT_NE(params, nullptr); + + // Drop the external reference to the model by destroying its handle. Because + // params co-owns the underlying Model (and its Config) via shared ownership, + // the object itself stays alive, so dereferencing the raw model pointer below + // remains valid. This does NOT imply the handle is generally usable after + // OgaDestroyModel; it is valid here only because another owner keeps it alive. + OgaDestroyModel(model); + + OgaGenerator* generator = nullptr; + ASSERT_EQ(OgaCreateGenerator(model, params, &generator), nullptr); + ASSERT_NE(generator, nullptr); + + OgaDestroyGenerator(generator); + OgaDestroyGeneratorParams(params); +} + TEST(CAPITests, AdaptersTest) { #ifdef USE_CUDA using OutputType = Ort::Float16_t;