From 69f7cc6494cb88f2b51acee70ff7b1f0c3b0505e Mon Sep 17 00:00:00 2001 From: petermcaughan Date: Wed, 7 Sep 2022 11:15:00 -0700 Subject: [PATCH] Add pybind support for all memory config options in OrtArenaCfg (#12658) * Add support for initial_growth_chunk_size_bytes setting in OrtArenaCfg pybind * Add overloaded constructor for KVP, UT still in progress * Fix class member access in pybind, fix unit test * Resolve linter warnings * Improve formatting * Simplify UT * Fix linter formatting Co-authored-by: Peter Mcaughan --- .../python/onnxruntime_pybind_state.cc | 29 ++++++++++++- .../test/python/onnxruntime_test_python.py | 41 ++++++++++++++++++- 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 4c530b89fe..77d06ab014 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -1083,6 +1083,8 @@ void addObjectMethods(py::module& m, Environment& env, ExecutionProviderRegistra .def_static("default_memory", []() { return OrtDevice::MemType::DEFAULT; }); py::class_ ort_arena_cfg_binding(m, "OrtArenaCfg"); + // Note: Doesn't expose initial_growth_chunk_sizes_bytes option. This constructor kept for + // backwards compatibility, key-value pair constructor overload exposes all options // There is a global var: arena_extend_strategy, which means we can't use that var name here // See docs/C_API.md for details on what the following parameters mean and how to choose these values ort_arena_cfg_binding.def(py::init([](size_t max_mem, int arena_extend_strategy_local, @@ -1093,7 +1095,32 @@ void addObjectMethods(py::module& m, Environment& env, ExecutionProviderRegistra ort_arena_cfg->initial_chunk_size_bytes = initial_chunk_size_bytes; ort_arena_cfg->max_dead_bytes_per_chunk = max_dead_bytes_per_chunk; return ort_arena_cfg; - })); + })) + .def(py::init([](const py::dict& feeds) { + auto ort_arena_cfg = std::make_unique(); + for (const auto kvp : feeds) { + std::string key = kvp.first.cast(); + if (key == "max_mem") { + ort_arena_cfg->max_mem = kvp.second.cast(); + } else if (key == "arena_extend_strategy") { + ort_arena_cfg->arena_extend_strategy = kvp.second.cast(); + } else if (key == "initial_chunk_size_bytes") { + ort_arena_cfg->initial_chunk_size_bytes = kvp.second.cast(); + } else if (key == "max_dead_bytes_per_chunk") { + ort_arena_cfg->max_dead_bytes_per_chunk = kvp.second.cast(); + } else if (key == "initial_growth_chunk_size_bytes") { + ort_arena_cfg->initial_growth_chunk_size_bytes = kvp.second.cast(); + } else { + ORT_THROW("Invalid OrtArenaCfg option: ", key); + } + } + return ort_arena_cfg; + })) + .def_readwrite("max_mem", &OrtArenaCfg::max_mem) + .def_readwrite("arena_extend_strategy", &OrtArenaCfg::arena_extend_strategy) + .def_readwrite("initial_chunk_size_bytes", &OrtArenaCfg::initial_chunk_size_bytes) + .def_readwrite("max_dead_bytes_per_chunk", &OrtArenaCfg::max_dead_bytes_per_chunk) + .def_readwrite("initial_growth_chunk_size_bytes", &OrtArenaCfg::initial_growth_chunk_size_bytes); py::class_ ort_memory_info_binding(m, "OrtMemoryInfo"); ort_memory_info_binding.def(py::init([](const char* name, OrtAllocatorType type, int id, OrtMemType mem_type) { diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py index 2abdc12f19..cff7d2b337 100644 --- a/onnxruntime/test/python/onnxruntime_test_python.py +++ b/onnxruntime/test/python/onnxruntime_test_python.py @@ -1154,7 +1154,9 @@ class TestInferenceSession(unittest.TestCase): def testSharedAllocatorUsingCreateAndRegisterAllocator(self): # Create and register an arena based allocator - # ort_arena_cfg = onnxrt.OrtArenaCfg(0, -1, -1, -1) (create an OrtArenaCfg like this template if you want to use non-default parameters) + # To create an OrtArenaCfg using non-default parameters, use one of below templates: + # ort_arena_cfg = onnxrt.OrtArenaCfg(0, -1, -1, -1) - Note: doesn't expose initial_growth_chunk_size_bytes option + # ort_arena_cfg = onnxrt.OrtArenaCfg({"max_mem": -1, ""arena_extend_strategy": 1, etc..}) ort_memory_info = onnxrt.OrtMemoryInfo( "Cpu", onnxrt.OrtAllocatorType.ORT_ARENA_ALLOCATOR, @@ -1308,6 +1310,43 @@ class TestInferenceSession(unittest.TestCase): ) print("Create session with customize execution provider successfully!") + def testCreateAllocator(self): + def verify_allocator(allocator, expected_config): + for key, val in expected_config.items(): + if key == "max_mem": + self.assertEqual(allocator.max_mem, val) + elif key == "arena_extend_strategy": + self.assertEqual(allocator.arena_extend_strategy, val) + elif key == "initial_chunk_size_bytes": + self.assertEqual(allocator.initial_chunk_size_bytes, val) + elif key == "max_dead_bytes_per_chunk": + self.assertEqual(allocator.max_dead_bytes_per_chunk, val) + elif key == "initial_growth_chunk_size_bytes": + self.assertEqual(allocator.initial_growth_chunk_size_bytes, val) + else: + raise ValueError("Invalid OrtArenaCfg option: " + key) + + # Verify ordered parameter initialization + ort_arena_cfg = onnxrt.OrtArenaCfg(8, 0, 4, 2) + expected_allocator = { + "max_mem": 8, + "arena_extend_strategy": 0, + "initial_chunk_size_bytes": 4, + "max_dead_bytes_per_chunk": 2, + } + verify_allocator(ort_arena_cfg, expected_allocator) + + # Verify key-value pair initialization + expected_kvp_allocator = { + "max_mem": 16, + "arena_extend_strategy": 1, + "initial_chunk_size_bytes": 8, + "max_dead_bytes_per_chunk": 4, + "initial_growth_chunk_size_bytes": 2, + } + ort_arena_cfg_kvp = onnxrt.OrtArenaCfg(expected_kvp_allocator) + verify_allocator(ort_arena_cfg_kvp, expected_kvp_allocator) + if __name__ == "__main__": unittest.main(verbosity=1)