From 0b4f3f0ff723e4745d80d45e1393ab402e462e71 Mon Sep 17 00:00:00 2001
From: Sai Kiran Polisetty <spolisetty@nvidia.com>
Date: Wed, 29 Jan 2025 10:21:03 +0530
Subject: [PATCH] Add support for
 `session.use_device_allocator_for_initializers` in onnxruntime_backend (#294)

* Add support for ArenaCfg configuration options
---
 README.md          |  3 ++-
 src/onnxruntime.cc | 27 ++++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index e61c31f..ef3efd1 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 <!--
-# Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2020-2025, NVIDIA CORPORATION. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -283,6 +283,7 @@ for more information.
 * `memory.enable_memory_arena_shrinkage`:
 See [this](https://github.com/microsoft/onnxruntime/blob/master/include/onnxruntime/core/session/onnxruntime_run_options_config_keys.h)
 for more information.
+* `session.use_device_allocator_for_initializers`: Use "1" to enable using device allocator for allocating initialized tensor memory and "0" to disable. The default is "0". See [this](https://onnxruntime.ai/docs/get-started/with-c.html) for more information.
 
 ### Command line options
 
diff --git a/src/onnxruntime.cc b/src/onnxruntime.cc
index acdb81e..6eae3c4 100644
--- a/src/onnxruntime.cc
+++ b/src/onnxruntime.cc
@@ -1,4 +1,4 @@
-// Copyright 2019-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// Copyright 2019-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 //
 // Redistribution and use in source and binary forms, with or without
 // modification, are permitted provided that the following conditions
@@ -302,6 +302,31 @@ ModelState::ModelState(TRITONBACKEND_Model* triton_model)
     }
   }
 
+  // Enable/disable use_device_allocator_for_initializers
+  {
+    triton::common::TritonJson::Value params;
+    if (ModelConfig().Find("parameters", &params)) {
+      triton::common::TritonJson::Value json_value;
+      const char* use_device_allocator_for_initializers_key =
+          "session.use_device_allocator_for_initializers";
+      if (params.Find(use_device_allocator_for_initializers_key, &json_value)) {
+        std::string string_value;
+        THROW_IF_BACKEND_MODEL_ERROR(
+            json_value.MemberAsString("string_value", &string_value));
+
+        LOG_MESSAGE(
+            TRITONSERVER_LOG_VERBOSE,
+            (std::string("Configuring '") +
+             use_device_allocator_for_initializers_key + "' to '" +
+             string_value + "' for '" + Name() + "'")
+                .c_str());
+        THROW_IF_BACKEND_MODEL_ORT_ERROR(ort_api->AddSessionConfigEntry(
+            soptions, use_device_allocator_for_initializers_key,
+            string_value.c_str()));
+      }
+    }
+  }
+
   // memory configs
   // enable/disable mem arena
   {