[IO] Add TCP backend and benchmark/test coverage

include/mori/io/backend.hpp

@@ -21,6 +21,8 @@
 // SOFTWARE.
 #pragma once
 
+#include <ostream>
+
 #include "mori/io/common.hpp"
 #include "mori/io/enum.hpp"
 
@@ -79,6 +81,52 @@ inline std::ostream& operator<<(std::ostream& os, const XgmiBackendConfig& c) {
   return os << "numStreams[" << c.numStreams << "] numEvents[" << c.numEvents << "]";
 }
 
+struct TcpBackendConfig : public BackendConfig {
+  TcpBackendConfig() : BackendConfig(BackendType::TCP) {}
+  TcpBackendConfig(int sockSndbufBytes_, int sockRcvbufBytes_, int opTimeoutMs_,
+                   bool enableKeepalive_, int keepaliveIdleSec_, int keepaliveIntvlSec_,
+                   int keepaliveCnt_, bool enableCtrlNodelay_, int numDataConns_,
+                   int stripingThresholdBytes_)
+      : BackendConfig(BackendType::TCP),
+        sockSndbufBytes(sockSndbufBytes_),
+        sockRcvbufBytes(sockRcvbufBytes_),
+        opTimeoutMs(opTimeoutMs_),
+        enableKeepalive(enableKeepalive_),
+        keepaliveIdleSec(keepaliveIdleSec_),
+        keepaliveIntvlSec(keepaliveIntvlSec_),
+        keepaliveCnt(keepaliveCnt_),
+        enableCtrlNodelay(enableCtrlNodelay_),
+        numDataConns(numDataConns_),
+        stripingThresholdBytes(stripingThresholdBytes_) {}
+
+  int sockSndbufBytes{32 * 1024 * 1024};
+  int sockRcvbufBytes{32 * 1024 * 1024};
+
+  int opTimeoutMs{30 * 1000};
+
+  bool enableKeepalive{true};
+  int keepaliveIdleSec{30};
+  int keepaliveIntvlSec{10};
+  int keepaliveCnt{3};
+
+  bool enableCtrlNodelay{true};
+
+  // Number of parallel DATA TCP connections per peer (iperf-like multi-stream striping).
+  // Effective only when peer has >= numDataConns established and transfer is contiguous.
+  int numDataConns{8};
+  // Stripe large contiguous transfers; keep small transfers on a single stream for latency.
+  int stripingThresholdBytes{64 * 1024};
+};
+
+inline std::ostream& operator<<(std::ostream& os, const TcpBackendConfig& c) {
+  return os << "sockSndbufBytes[" << c.sockSndbufBytes << "] sockRcvbufBytes[" << c.sockRcvbufBytes
+            << "] opTimeoutMs[" << c.opTimeoutMs << "] enableKeepalive[" << c.enableKeepalive
+            << "] keepaliveIdleSec[" << c.keepaliveIdleSec << "] keepaliveIntvlSec["
+            << c.keepaliveIntvlSec << "] keepaliveCnt[" << c.keepaliveCnt << "] enableCtrlNodelay["
+            << c.enableCtrlNodelay << "] numDataConns[" << c.numDataConns
+            << "] stripingThresholdBytes[" << c.stripingThresholdBytes << "]";
+}
+
 /* ---------------------------------------------------------------------------------------------- */
 /*                                         BackendSession                                         */
 /* ---------------------------------------------------------------------------------------------- */

python/mori/io/__init__.py

@@ -29,6 +29,7 @@
     MemoryLocationType,
     PollCqMode,
     RdmaBackendConfig,
+    TcpBackendConfig,
     XgmiBackendConfig,
     set_log_level,
 )

python/mori/io/engine.py

@@ -91,6 +91,8 @@ def create_backend(self, type: mori_cpp.BackendType, config=None):
                 config = mori_cpp.RdmaBackendConfig()
             elif type is mori_cpp.BackendType.XGMI:
                 config = mori_cpp.XgmiBackendConfig()
+            elif type is mori_cpp.BackendType.TCP:
+                config = mori_cpp.TcpBackendConfig()
             else:
                 raise NotImplementedError("backend not implemented yet")
         return self._engine.CreateBackend(type, config)

src/io/CMakeLists.txt

@@ -5,6 +5,8 @@ add_library(
   rdma/backend_impl.cpp
   rdma/executor.cpp
   rdma/common.cpp
+  tcp/backend_impl.cpp
+  tcp/transport.cpp
   xgmi/backend_impl.cpp
   xgmi/hip_resource_pool.cpp)
 

src/io/engine.cpp

@@ -27,6 +27,7 @@
 
 #include "mori/io/logging.hpp"
 #include "src/io/rdma/backend_impl.hpp"
+#include "src/io/tcp/backend_impl.hpp"
 #include "src/io/xgmi/backend_impl.hpp"
 
 namespace mori {
@@ -132,6 +133,27 @@ void IOEngine::CreateBackend(BackendType type, const BackendConfig& beConfig) {
                                                  static_cast<const XgmiBackendConfig&>(beConfig));
     backends.insert({type, std::move(backend)});
     InvalidateRouteCache();
+  } else if (type == BackendType::TCP) {
+    assert(backends.find(type) == backends.end());
+    auto backend = std::make_unique<TcpBackend>(desc.key, config,
+                                                static_cast<const TcpBackendConfig&>(beConfig));
+
+    if (config.port == 0) {
+      auto bound_port_opt = backend->GetListenPort();
+      if (!bound_port_opt.has_value() || bound_port_opt.value() == 0) {
+        MORI_IO_ERROR("IOEngine key {} failed to retrieve bound port after TCP backend init",
+                      desc.key);
+        assert(false && "Failed to retrieve bound port after TCP backend init");
+      } else {
+        uint16_t bound_port = bound_port_opt.value();
+        desc.port = bound_port;
+        this->config.port = bound_port;
+        MORI_IO_INFO("IOEngine key {} bound ephemeral port {}", desc.key, bound_port);
+      }
+    }
+
+    backends.insert({type, std::move(backend)});
+    InvalidateRouteCache();
   } else {
     assert(false && "not implemented");
   }

src/io/tcp/backend_impl.cpp

@@ -0,0 +1,93 @@
+// Copyright © Advanced Micro Devices, Inc. All rights reserved.
+//
+// MIT License
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+#include "src/io/tcp/backend_impl.hpp"
+
+#include "src/io/tcp/transport.hpp"
+
+namespace mori {
+namespace io {
+
+TcpBackendSession::TcpBackendSession(const TcpBackendConfig& cfg, const MemoryDesc& l,
+                                     const MemoryDesc& r, TcpTransport* t)
+    : config(cfg), local(l), remote(r), transport(t) {}
+
+void TcpBackendSession::ReadWrite(size_t localOffset, size_t remoteOffset, size_t size,
+                                  TransferStatus* status, TransferUniqueId id, bool isRead) {
+  MORI_IO_FUNCTION_TIMER;
+  transport->SubmitReadWrite(local, localOffset, remote, remoteOffset, size, status, id, isRead);
+}
+
+void TcpBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets,
+                                       const SizeVec& sizes, TransferStatus* status,
+                                       TransferUniqueId id, bool isRead) {
+  MORI_IO_FUNCTION_TIMER;
+  transport->SubmitBatchReadWrite(local, localOffsets, remote, remoteOffsets, sizes, status, id,
+                                  isRead);
+}
+
+bool TcpBackendSession::Alive() const { return true; }
+
+TcpBackend::TcpBackend(EngineKey k, const IOEngineConfig& engCfg, const TcpBackendConfig& cfg)
+    : myEngKey(std::move(k)), config(cfg) {
+  transport = std::make_unique<TcpTransport>(myEngKey, engCfg, cfg);
+  transport->Start();
+  MORI_IO_INFO("TcpBackend created key={}", myEngKey);
+}
+
+TcpBackend::~TcpBackend() { transport->Shutdown(); }
+
+std::optional<uint16_t> TcpBackend::GetListenPort() const { return transport->GetListenPort(); }
+
+void TcpBackend::RegisterRemoteEngine(const EngineDesc& desc) {
+  transport->RegisterRemoteEngine(desc);
+}
+void TcpBackend::DeregisterRemoteEngine(const EngineDesc& desc) {
+  transport->DeregisterRemoteEngine(desc);
+}
+void TcpBackend::RegisterMemory(MemoryDesc& desc) { transport->RegisterMemory(desc); }
+void TcpBackend::DeregisterMemory(const MemoryDesc& desc) { transport->DeregisterMemory(desc); }
+
+void TcpBackend::ReadWrite(const MemoryDesc& ld, size_t lo, const MemoryDesc& rs, size_t ro,
+                           size_t sz, TransferStatus* st, TransferUniqueId id, bool isRead) {
+  MORI_IO_FUNCTION_TIMER;
+  transport->SubmitReadWrite(ld, lo, rs, ro, sz, st, id, isRead);
+}
+
+void TcpBackend::BatchReadWrite(const MemoryDesc& ld, const SizeVec& lo, const MemoryDesc& rs,
+                                const SizeVec& ro, const SizeVec& sz, TransferStatus* st,
+                                TransferUniqueId id, bool isRead) {
+  MORI_IO_FUNCTION_TIMER;
+  transport->SubmitBatchReadWrite(ld, lo, rs, ro, sz, st, id, isRead);
+}
+
+BackendSession* TcpBackend::CreateSession(const MemoryDesc& local, const MemoryDesc& remote) {
+  return new TcpBackendSession(config, local, remote, transport.get());
+}
+
+bool TcpBackend::PopInboundTransferStatus(EngineKey remote, TransferUniqueId id,
+                                          TransferStatus* status) {
+  return transport->PopInboundTransferStatus(remote, id, status);
+}
+
+}  // namespace io
+}  // namespace mori

src/io/tcp/backend_impl.hpp

@@ -0,0 +1,90 @@
+// Copyright © Advanced Micro Devices, Inc. All rights reserved.
+//
+// MIT License
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+#pragma once
+
+#include <memory>
+#include <optional>
+#include <string>
+
+#include "mori/io/backend.hpp"
+#include "mori/io/common.hpp"
+#include "mori/io/engine.hpp"
+
+namespace mori {
+namespace io {
+
+class TcpTransport;
+
+class TcpBackendSession : public BackendSession {
+ public:
+  TcpBackendSession() = default;
+  TcpBackendSession(const TcpBackendConfig& config, const MemoryDesc& local,
+                    const MemoryDesc& remote, TcpTransport* transport);
+  ~TcpBackendSession() override = default;
+
+  void ReadWrite(size_t localOffset, size_t remoteOffset, size_t size, TransferStatus* status,
+                 TransferUniqueId id, bool isRead) override;
+  void BatchReadWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets,
+                      const SizeVec& sizes, TransferStatus* status, TransferUniqueId id,
+                      bool isRead) override;
+  bool Alive() const override;
+
+ private:
+  TcpBackendConfig config{};
+  MemoryDesc local{};
+  MemoryDesc remote{};
+  TcpTransport* transport{nullptr};
+};
+
+class TcpBackend : public Backend {
+ public:
+  TcpBackend(EngineKey, const IOEngineConfig&, const TcpBackendConfig&);
+  ~TcpBackend() override;
+
+  std::optional<uint16_t> GetListenPort() const;
+
+  void RegisterRemoteEngine(const EngineDesc&) override;
+  void DeregisterRemoteEngine(const EngineDesc&) override;
+  void RegisterMemory(MemoryDesc& desc) override;
+  void DeregisterMemory(const MemoryDesc& desc) override;
+
+  void ReadWrite(const MemoryDesc& localDest, size_t localOffset, const MemoryDesc& remoteSrc,
+                 size_t remoteOffset, size_t size, TransferStatus* status, TransferUniqueId id,
+                 bool isRead) override;
+  void BatchReadWrite(const MemoryDesc& localDest, const SizeVec& localOffsets,
+                      const MemoryDesc& remoteSrc, const SizeVec& remoteOffsets,
+                      const SizeVec& sizes, TransferStatus* status, TransferUniqueId id,
+                      bool isRead) override;
+
+  BackendSession* CreateSession(const MemoryDesc& local, const MemoryDesc& remote) override;
+  bool PopInboundTransferStatus(EngineKey remote, TransferUniqueId id,
+                                TransferStatus* status) override;
+
+ private:
+  EngineKey myEngKey;
+  TcpBackendConfig config{};
+  std::unique_ptr<TcpTransport> transport{nullptr};
+};
+
+}  // namespace io
+}  // namespace mori