[ESI runtime] Host memory service (#7367)

This service is responsible for managing host memory which is accessible from the accelerator.

integration_test/Dialect/ESI/runtime/loopback.mlir.py

@@ -7,6 +7,13 @@
 platform = sys.argv[1]
 acc = esiaccel.AcceleratorConnection(platform, sys.argv[2])
 
+hostmem = acc.get_host_memory()
+if hostmem is not None:
+  mem1 = hostmem.allocate(1024)
+  assert mem1.size == 1024
+  print(f"mem1: {mem1.ptr} size {mem1.size}")
+  mem1 = None
+
 assert acc.sysinfo().esi_version() == 0
 m = acc.manifest()
 assert m.api_version == 0

lib/Dialect/ESI/runtime/cpp/include/esi/Services.h

@@ -113,6 +113,43 @@ class MMIOSysInfo final : public SysInfo {
   const MMIO *mmio;
 };
 
+class HostMem : public Service {
+public:
+  virtual ~HostMem() = default;
+  virtual std::string getServiceSymbol() const override;
+
+  /// RAII memory region for host memory. Automatically frees the memory when
+  /// deconstructed.
+  struct HostMemRegion {
+    virtual ~HostMemRegion() = default;
+    virtual void *getPtr() const = 0;
+    operator void *() const { return getPtr(); }
+    virtual std::size_t getSize() const = 0;
+  };
+
+  /// Options for allocating host memory.
+  struct Options {
+    bool writeable = false;
+    bool useLargePages = false;
+  };
+
+  /// Allocate a region of host memory in accelerator accessible address space.
+  virtual std::unique_ptr<HostMemRegion> allocate(std::size_t size,
+                                                  Options opts) const = 0;
+
+  /// Try to make a region of host memory accessible to the accelerator. Returns
+  /// 'false' on failure. It is optional for an accelerator backend to implement
+  /// this, so client code needs to have a fallback for when this returns
+  /// 'false'. On success, it is the client's responsibility to ensure that the
+  /// memory eventually gets unmapped.
+  virtual bool mapMemory(void *ptr, std::size_t size, Options opts) const {
+    return false;
+  }
+  /// Unmap memory which was previously mapped with 'mapMemory'. Undefined
+  /// behavior when called with a pointer which was not previously mapped.
+  virtual void unmapMemory(void *ptr) const {}
+};
+
 /// Service for calling functions.
 class FuncService : public Service {
 public:

lib/Dialect/ESI/runtime/cpp/lib/Services.cpp

@@ -72,6 +72,8 @@ std::vector<uint8_t> MMIOSysInfo::getCompressedManifest() const {
   return manifest;
 }
 
+std::string HostMem::getServiceSymbol() const { return "__builtin_HostMem"; }
+
 CustomService::CustomService(AppIDPath idPath,
                              const ServiceImplDetails &details,
                              const HWClientDetails &clients)

lib/Dialect/ESI/runtime/cpp/lib/backends/Cosim.cpp

@@ -394,6 +394,35 @@ class CosimMMIO : public MMIO {
   std::unique_ptr<FuncService::Function> readMMIO;
 };
 
+class CosimHostMem : public HostMem {
+public:
+  CosimHostMem() {}
+
+  struct CosimHostMemRegion : public HostMemRegion {
+    CosimHostMemRegion(std::size_t size) {
+      ptr = malloc(size);
+      this->size = size;
+    }
+    virtual ~CosimHostMemRegion() { free(ptr); }
+    virtual void *getPtr() const { return ptr; }
+    virtual std::size_t getSize() const { return size; }
+
+  private:
+    void *ptr;
+    std::size_t size;
+  };
+
+  virtual std::unique_ptr<HostMemRegion> allocate(std::size_t size,
+                                                  HostMem::Options opts) const {
+    return std::unique_ptr<HostMemRegion>(new CosimHostMemRegion(size));
+  }
+  virtual bool mapMemory(void *ptr, std::size_t size,
+                         HostMem::Options opts) const {
+    return true;
+  }
+  virtual void unmapMemory(void *ptr) const {}
+};
+
 } // namespace
 
 Service *CosimAccelerator::createService(Service::Type svcType,
@@ -420,6 +449,8 @@ Service *CosimAccelerator::createService(Service::Type svcType,
 
   if (svcType == typeid(services::MMIO)) {
     return new CosimMMIO(getCtxt(), rpcClient);
+  } else if (svcType == typeid(services::HostMem)) {
+    return new CosimHostMem();
   } else if (svcType == typeid(SysInfo)) {
     switch (manifestMethod) {
     case ManifestMethod::Cosim:

lib/Dialect/ESI/runtime/cpp/lib/backends/Trace.cpp

@@ -75,6 +75,10 @@ struct esi::backends::trace::TraceAccelerator::Impl {
 
   void write(const AppIDPath &id, const std::string &portName, const void *data,
              size_t size);
+  std::ostream &write(std::string service) {
+    *traceWrite << "[" << service << "] ";
+    return *traceWrite;
+  }
 
 private:
   std::ofstream *traceWrite;
@@ -271,12 +275,63 @@ TraceAccelerator::requestChannelsFor(AppIDPath idPath,
   return impl->requestChannelsFor(idPath, bundleType);
 }
 
+class TraceHostMem : public HostMem {
+public:
+  TraceHostMem(TraceAccelerator::Impl &impl) : impl(impl) {}
+
+  struct TraceHostMemRegion : public HostMemRegion {
+    TraceHostMemRegion(std::size_t size, TraceAccelerator::Impl &impl)
+        : impl(impl) {
+      ptr = malloc(size);
+      this->size = size;
+    }
+    virtual ~TraceHostMemRegion() {
+      impl.write("HostMem") << "free " << ptr << std::endl;
+      free(ptr);
+    }
+    virtual void *getPtr() const { return ptr; }
+    virtual std::size_t getSize() const { return size; }
+
+  private:
+    void *ptr;
+    std::size_t size;
+    TraceAccelerator::Impl &impl;
+  };
+
+  virtual std::unique_ptr<HostMemRegion> allocate(std::size_t size,
+                                                  HostMem::Options opts) const {
+    auto ret =
+        std::unique_ptr<HostMemRegion>(new TraceHostMemRegion(size, impl));
+    impl.write("HostMem 0x")
+        << ret->getPtr() << " allocate " << size
+        << " bytes. Writeable: " << opts.writeable
+        << ", useLargePages: " << opts.useLargePages << std::endl;
+    return ret;
+  }
+  virtual bool mapMemory(void *ptr, std::size_t size,
+                         HostMem::Options opts) const {
+    impl.write("HostMem") << "map 0x" << ptr << " size " << size
+                          << " bytes. Writeable: " << opts.writeable
+                          << ", useLargePages: " << opts.useLargePages
+                          << std::endl;
+    return true;
+  }
+  virtual void unmapMemory(void *ptr) const {
+    impl.write("HostMem") << "unmap 0x" << ptr << std::endl;
+  }
+
+private:
+  TraceAccelerator::Impl &impl;
+};
+
 Service *
 TraceAccelerator::Impl::createService(Service::Type svcType, AppIDPath idPath,
                                       const ServiceImplDetails &details,
                                       const HWClientDetails &clients) {
   if (svcType == typeid(SysInfo))
     return new TraceSysInfo(manifestJson);
+  if (svcType == typeid(HostMem))
+    return new TraceHostMem(*this);
   if (svcType == typeid(CustomService))
     return new TraceCustomService(*this, idPath, details, clients);
   return nullptr;

lib/Dialect/ESI/runtime/python/esiaccel/accelerator.py

@@ -38,6 +38,12 @@ def build_accelerator(self) -> "Accelerator":
   def get_service_mmio(self) -> cpp.MMIO:
     return self.cpp_accel.get_service_mmio()
 
+  def get_service_hostmem(self) -> cpp.HostMem:
+    return self.cpp_accel.get_service_hostmem()
+
+
+from .esiCppAccel import HostMemOptions
+
 
 class HWModule:
   """Represents either the top level or an instance of a hardware module."""

lib/Dialect/ESI/runtime/python/esiaccel/esiCppAccel.cpp

@@ -104,6 +104,25 @@ PYBIND11_MODULE(esiCppAccel, m) {
       .def("read", &services::MMIO::read)
       .def("write", &services::MMIO::write);
 
+  py::class_<services::HostMem::HostMemRegion>(m, "HostMemRegion")
+      .def_property_readonly("ptr", &services::HostMem::HostMemRegion::getPtr)
+      .def_property_readonly("size",
+                             &services::HostMem::HostMemRegion::getSize);
+
+  py::class_<services::HostMem::Options>(m, "HostMemOptions")
+      .def(py::init<>())
+      .def_readwrite("writeable", &services::HostMem::Options::writeable)
+      .def_readwrite("use_large_pages",
+                     &services::HostMem::Options::useLargePages);
+
+  py::class_<services::HostMem>(m, "HostMem")
+      .def("allocate", &services::HostMem::allocate, py::arg("size"),
+           py::arg("options") = services::HostMem::Options(),
+           py::return_value_policy::take_ownership)
+      .def("map_memory", &services::HostMem::mapMemory, py::arg("ptr"),
+           py::arg("size"), py::arg("options") = services::HostMem::Options())
+      .def("unmap_memory", &services::HostMem::unmapMemory, py::arg("ptr"));
+
   py::class_<AppID>(m, "AppID")
       .def(py::init<std::string, std::optional<uint32_t>>(), py::arg("name"),
            py::arg("idx") = std::nullopt)
@@ -216,6 +235,12 @@ PYBIND11_MODULE(esiCppAccel, m) {
                          [](AcceleratorConnection &acc) {
                            return acc.getService<services::MMIO>({});
                          },
+                         py::return_value_policy::reference)
+                     .def(
+                         "get_service_hostmem",
+                         [](AcceleratorConnection &acc) {
+                           return acc.getService<services::HostMem>({});
+                         },
                          py::return_value_policy::reference);
 
   py::class_<Manifest>(m, "Manifest")