From f694d52ee988810d1928669e09d9498b4cb76c1a Mon Sep 17 00:00:00 2001
From: Stephen Baione <stbaione@amd.com>
Date: Fri, 17 Jan 2025 15:37:55 +0000
Subject: [PATCH] Default `mmap` to `false`, Add private `LoadMmap` function
 that uses legacy loading method to still allow mmaping param files

---
 shortfin/python/lib_ext.cc             |  6 ++---
 shortfin/src/shortfin/local/program.cc | 37 ++++++++++++++++++++++++++
 shortfin/src/shortfin/local/program.h  |  6 +++--
 3 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/shortfin/python/lib_ext.cc b/shortfin/python/lib_ext.cc
index c668e6a8b..af7aba767 100644
--- a/shortfin/python/lib_ext.cc
+++ b/shortfin/python/lib_ext.cc
@@ -672,7 +672,7 @@ void BindLocal(py::module_ &m) {
       // Methods not on System but on child objects, taking System as an arg.
       // Emitted here for convenience.
       .def("load_module", &local::ProgramModule::Load, py::arg("path"),
-           py::arg("mmap") = true);
+           py::arg("mmap") = false);
 
   // Support classes.
   py::class_<local::Node>(m, "Node")
@@ -731,7 +731,7 @@ void BindLocal(py::module_ &m) {
       .def_prop_ro("exports", &local::ProgramModule::exports)
       .def("__repr__", &local::ProgramModule::to_s)
       .def_static("load", &local::ProgramModule::Load, py::arg("system"),
-                  py::arg("path"), py::arg("mmap") = true)
+                  py::arg("path"), py::arg("mmap") = false)
       .def_static(
           "parameter_provider",
           [](local::System &system, py::args params) {
@@ -817,7 +817,7 @@ void BindLocal(py::module_ &m) {
           },
           py::arg("file_path"), py::arg("format") = std::string_view(),
           py::arg("readable") = true, py::arg("writable") = false,
-          py::arg("mmap") = true);
+          py::arg("mmap") = false);
 
   struct DevicesSet {
     DevicesSet(py::object fiber_obj, std::optional<size_t> index = {})
diff --git a/shortfin/src/shortfin/local/program.cc b/shortfin/src/shortfin/local/program.cc
index 55bd6ba67..9744c1691 100644
--- a/shortfin/src/shortfin/local/program.cc
+++ b/shortfin/src/shortfin/local/program.cc
@@ -633,6 +633,11 @@ void StaticProgramParameters::Load(std::filesystem::path file_path,
     options.format = file_path.extension().string();
   }
 
+  if (options.mmap) {
+    this->LoadMmap(file_path, options);
+    return;
+  }
+
   auto file_path_string = file_path.string();
   const iree_string_view_t path =
       iree_make_cstring_view(file_path_string.c_str());
@@ -646,6 +651,38 @@ void StaticProgramParameters::Load(std::filesystem::path file_path,
                                index_.get(), host_allocator_));
 }
 
+void StaticProgramParameters::LoadMmap(std::filesystem::path file_path,
+                                       LoadOptions options) {
+  SHORTFIN_TRACE_SCOPE_NAMED("StaticProgramParameters::LoadMmap");
+
+  iree_file_read_flags_t read_flags = IREE_FILE_READ_FLAG_MMAP;
+  iree_file_contents_t *file_contents = nullptr;
+  SHORTFIN_THROW_IF_ERROR(iree_file_read_contents(
+      file_path.string().c_str(), read_flags, host_allocator_, &file_contents));
+  iree_io_file_handle_release_callback_t release_callback = {
+      +[](void *user_data, iree_io_file_handle_primitive_t handle_primitive) {
+        iree_file_contents_t *file_contents = (iree_file_contents_t *)user_data;
+        iree_file_contents_free(file_contents);
+      },
+      file_contents,
+  };
+
+  // Wrap contents.
+  iree::io_file_handle_ptr file_handle;
+  iree_status_t status = iree_io_file_handle_wrap_host_allocation(
+      IREE_IO_FILE_ACCESS_READ, file_contents->buffer, release_callback,
+      host_allocator_, file_handle.for_output());
+  if (!iree_status_is_ok(status)) {
+    iree_file_contents_free(file_contents);
+    SHORTFIN_THROW_IF_ERROR(status);
+  }
+
+  // Parse.
+  SHORTFIN_THROW_IF_ERROR(iree_io_parse_file_index(
+      to_iree_string_view(options.format), file_handle.get(), index_.get(),
+      host_allocator_));
+}
+
 // -------------------------------------------------------------------------- //
 // ProgramIsolate
 // -------------------------------------------------------------------------- //
diff --git a/shortfin/src/shortfin/local/program.h b/shortfin/src/shortfin/local/program.h
index 450b29736..f83b83cbd 100644
--- a/shortfin/src/shortfin/local/program.h
+++ b/shortfin/src/shortfin/local/program.h
@@ -265,7 +265,7 @@ class SHORTFIN_API ProgramModule {
 
   // Loads a dynamic bytecode module (VMFB) from a path on the file system.
   static ProgramModule Load(System &system, const std::filesystem::path &path,
-                            bool mmap = true);
+                            bool mmap = false);
 
   // Creates a ProgramModule that will provide the given list of parameters
   // to modules loaded after it. In IREE parlance, this produces an
@@ -386,7 +386,7 @@ class SHORTFIN_API StaticProgramParameters : public BaseProgramParameters {
     // Whether the backing file can be written.
     bool writable = false;
     // Whether to mmap the file.
-    bool mmap = true;
+    bool mmap = false;
   };
   // Load parameters from a supported file format, applying no name
   // transformation.
@@ -396,6 +396,8 @@ class SHORTFIN_API StaticProgramParameters : public BaseProgramParameters {
  private:
   iree_allocator_t host_allocator_;
   iree::io_parameter_index_ptr index_;
+
+  void LoadMmap(std::filesystem::path file_path, LoadOptions options);
 };
 
 namespace detail {