No public description PiperOrigin-RevId: 797844866 Change-Id: If7867865eca03b12866ae45055a78cbe8a947b12

diff --git a/mpact/sim/util/program_loader/BUILD b/mpact/sim/util/program_loader/BUILD
index 421ffdf..3a3f113 100644
--- a/mpact/sim/util/program_loader/BUILD
+++ b/mpact/sim/util/program_loader/BUILD

@@ -47,6 +47,7 @@
         "//mpact/sim/util/memory",
         "@com_github_serge1_elfio//:elfio",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

diff --git a/mpact/sim/util/program_loader/elf_program_loader.cc b/mpact/sim/util/program_loader/elf_program_loader.cc
index 4c7e5a6..ec4beea 100644
--- a/mpact/sim/util/program_loader/elf_program_loader.cc
+++ b/mpact/sim/util/program_loader/elf_program_loader.cc

@@ -20,6 +20,7 @@
 #include <cstring>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -41,6 +42,10 @@
                                    util::MemoryInterface* data_memory)
     : code_memory_(code_memory), data_memory_(data_memory) {}
 
+ElfProgramLoader::ElfProgramLoader(
+    const std::vector<MemoryDescriptor>& memories)
+    : memories_(&memories) {}
+
 ElfProgramLoader::ElfProgramLoader(util::MemoryInterface* memory)
     : code_memory_(memory), data_memory_(memory) {}
 
@@ -95,7 +100,7 @@
       if (type == STT_FUNC) {
         fcn_symbol_map_.emplace(value, name);
         function_range_map_.insert(
-            std::make_pair(AddressRange(value, size), name));
+            std::make_pair(AddressRange(value, size / text_size_scale_), name));
       }
     }
   }
@@ -123,14 +128,30 @@
     if (segment->get_file_size() == 0) continue;
     // Read the data from the elf file.
     if (dbg_if_ == nullptr) {  // Use memory interfaces.
-      auto* db = db_factory.Allocate(segment->get_file_size());
-      std::memcpy(db->raw_ptr(), segment->get_data(), segment->get_file_size());
-
-      if (segment->get_flags() &
-          PF_X) {  // Executable, so write to code memory.
-        code_memory_->Store(segment->get_virtual_address(), db);
-      } else {  // Write to data memory.
-        data_memory_->Store(segment->get_virtual_address(), db);
+      auto size = segment->get_file_size();
+      auto* db = db_factory.Allocate(size);
+      std::memcpy(db->raw_ptr(), segment->get_data(), size);
+      if (memories_ == nullptr) {
+        if (segment->get_flags() &
+            PF_X) {  // Executable, so write to code memory.
+          code_memory_->Store(segment->get_virtual_address(), db);
+        } else {  // Write to data memory.
+          data_memory_->Store(segment->get_virtual_address(), db);
+        }
+      } else {
+        int i = 0;
+        for (auto& memory : *memories_) {
+          if (memory.predicate_fcn(*segment)) {
+            if (memory.address_fcn) {
+              memory.memory->Store(
+                  memory.address_fcn(segment->get_virtual_address()), db);
+            } else {
+              memory.memory->Store(segment->get_virtual_address(), db);
+            }
+            break;
+          }
+          i++;
+        }
       }
       db->DecRef();
       continue;

diff --git a/mpact/sim/util/program_loader/elf_program_loader.h b/mpact/sim/util/program_loader/elf_program_loader.h
index 760f11b..a1ae625 100644
--- a/mpact/sim/util/program_loader/elf_program_loader.h
+++ b/mpact/sim/util/program_loader/elf_program_loader.h

@@ -22,8 +22,10 @@
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
+#include "absl/functional/any_invocable.h"
 #include "absl/status/statusor.h"
 #include "elfio/elfio.hpp"
+#include "elfio/elfio_segment.hpp"
 #include "elfio/elfio_symbols.hpp"
 #include "mpact/sim/generic/core_debug_interface.h"
 #include "mpact/sim/util/memory/memory_interface.h"
@@ -51,6 +53,20 @@
   }
 };
 
+// This struct is used to describe a memory to be used by the elf program loader
+// when there is more than one or two memories that have to be loaded.
+struct MemoryDescriptor {
+  // Pointer to the memory interface.
+  util::MemoryInterface* memory;
+  // Predicate function that takes a segment and return true if it should be
+  // loaded into this memory.
+  absl::AnyInvocable<bool(const ELFIO::segment&) const> predicate_fcn;
+  // Function that takes a segment load address and returns the address it
+  // should be loaded at in the memory. If not specified, the load address is
+  // used unmodified.
+  absl::AnyInvocable<uint64_t(uint64_t) const> address_fcn;
+};
+
 // This class wraps the elfio class to provide an easy interface to load the
 // segments of an elf executable file into memory. If both code and data
 // memories are given, then executable segments are loaded into code memory and
@@ -63,6 +79,7 @@
  public:
   ElfProgramLoader(util::MemoryInterface* code_memory,
                    util::MemoryInterface* data_memory);
+  ElfProgramLoader(const std::vector<MemoryDescriptor>& memories);
   explicit ElfProgramLoader(util::MemoryInterface* memory);
   explicit ElfProgramLoader(generic::CoreDebugInterface* dbg_if);
   ElfProgramLoader() = delete;
@@ -82,9 +99,13 @@
   // If the GNU stack size program header exists, return the memory size.
   absl::StatusOr<uint64_t> GetStackSize() const;
 
+  void set_text_size_scale(uint64_t scale) { text_size_scale_ = scale; }
+  void set_data_size_scale(uint64_t scale) { data_size_scale_ = scale; }
+
   const ELFIO::elfio* elf_reader() const { return &elf_reader_; }
 
  private:
+  const std::vector<MemoryDescriptor>* memories_ = nullptr;
   bool loaded_ = false;
   ELFIO::elfio elf_reader_;
   util::MemoryInterface* code_memory_ = nullptr;
@@ -95,6 +116,8 @@
   std::map<AddressRange, std::string, AddressRangeComp> function_range_map_;
   uint64_t has_stack_size_ = false;
   uint64_t stack_size_ = 0;
+  uint64_t text_size_scale_ = 1;
+  uint64_t data_size_scale_ = 1;
 };
 
 }  // namespace util