[Feat] Asu: use host-pinned memory with dual CPU/device addresses for transport buffers

## Purpose
Switch ASU send/flag buffers from plain host memory to host-pinned
memory so that CPU code packs SQEs through the local mapping while
HCOMM/RDMA uses the device-visible mapping of the same allocation.

## Modifications
1. BufferManager: allocate host-pinned memory via aclrtMallocHost +
aclrtHostRegisterV2, obtain device pointer via
aclrtHostGetDevicePointer. ScatterGatherEntry gains device_addr.
RegisterMemory uses device address for host-pinned regions.
2. AsuTransportImpl: send/flag buffers use HOST_PINNED instead of HOST.
3. asu_submit_flow: pass device_addr to SendIoBatch.
4. sqe_request: use flagBuffer.device_addr for response_buffer_addr.
5. Move IsTransportBufferReady from asu_submit_flow to buffer_manager.
6. Tests: added host-pinned dual-address and device_addr assertions.

## Test
- buffer_manager_test: HostPinnedRegistersDeviceAddress.
- asu_submit_flow_test: BuildSubBatchSendBuffersUsesHostPinnedDeviceAddresses.
- sqe_request_test: packed response address matches device_addr.

Author: yumingyue624

ucm/shared/trans/ascend/ascend_buffer.cc

@@ -28,6 +28,27 @@
 
 namespace UC::Trans {
 
+namespace {
+
+constexpr std::uintptr_t HOST_REGISTER_PAGE_SIZE = 4096;
+
+void FreeHostMemory(void* host)
+{
+    auto ret = aclrtFreeHost(host);
+    if (ret != ACL_SUCCESS) { UC_ERROR("Failed to free host memory addr={} ret={}", host, ret); }
+}
+
+void ReleaseHostPinnedMemory(void* host)
+{
+    auto ret = aclrtHostUnregister(host);
+    if (ret != ACL_SUCCESS) {
+        UC_ERROR("Failed to unregister host-pinned memory addr={} ret={}", host, ret);
+    }
+    FreeHostMemory(host);
+}
+
+}  // namespace
+
 class HostHugePages : public std::enable_shared_from_this<HostHugePages> {
     struct ConstructorKey {};
     static constexpr auto HUGE_PAGE_SIZE = 2UL << 20;
@@ -124,6 +145,33 @@ std::shared_ptr<void> Trans::AscendBuffer::MakeHostBuffer(size_t size)
     return nullptr;
 }
 
+std::shared_ptr<void> Trans::AscendBuffer::MakeHostPinnedBuffer(size_t size, void** pDevice)
+{
+    if (pDevice) { *pDevice = nullptr; }
+
+    void* host = nullptr;
+    auto ret = aclrtMallocHost(&host, size);
+    if (ret != ACL_SUCCESS) { return nullptr; }
+
+    if (reinterpret_cast<std::uintptr_t>(host) % HOST_REGISTER_PAGE_SIZE != 0) {
+        UC_ERROR("Host-pinned memory is not 4K page aligned addr={} size={}", host, size);
+        FreeHostMemory(host);
+        return nullptr;
+    }
+
+    void* device = nullptr;
+    auto status = Buffer::RegisterHostBuffer(host, size, &device);
+    if (status.Failure()) {
+        UC_ERROR("Failed to register host-pinned memory addr={} size={} status={}", host, size,
+                 status);
+        FreeHostMemory(host);
+        return nullptr;
+    }
+
+    if (pDevice) { *pDevice = device; }
+    return std::shared_ptr<void>(host, ReleaseHostPinnedMemory);
+}
+
 std::shared_ptr<void> Trans::AscendBuffer::MakeHostBuffer4DirectIo(size_t size)
 {
     try {

ucm/shared/trans/ascend/ascend_buffer.h

@@ -32,6 +32,7 @@ class AscendBuffer : public ReservedBuffer {
 public:
     std::shared_ptr<void> MakeDeviceBuffer(size_t size) override;
     std::shared_ptr<void> MakeHostBuffer(size_t size) override;
+    std::shared_ptr<void> MakeHostPinnedBuffer(size_t size, void** pDevice = nullptr);
     std::shared_ptr<void> MakeHostBuffer4DirectIo(size_t size) override;
 };
 

ucm/transport/kv/asu/test/case/buffer_manager_test.cc

@@ -71,6 +71,20 @@ TEST_F(BufferManagerTest, InitWithZeroSlotNum)
     ASSERT_EQ(status.code, StatusCode::INVALID_ARGUMENT);
 }
 
+TEST_F(BufferManagerTest, InitHostWithUnalignedSlotCapacity)
+{
+    BufferManager mgr;
+    auto status = mgr.Init("test_buffer", MemoryType::HOST, 1000, 10);
+    ASSERT_TRUE(status.ok()) << status.message;
+}
+
+TEST_F(BufferManagerTest, InitDeviceWithUnalignedSlotCapacity)
+{
+    BufferManager mgr;
+    auto status = mgr.Init("test_buffer", MemoryType::ASCEND_DEVICE, 1000, 10);
+    ASSERT_TRUE(status.ok()) << status.message;
+}
+
 TEST_F(BufferManagerTest, DoubleInit)
 {
     BufferManager mgr;
@@ -123,12 +137,13 @@ TEST_F(BufferManagerTest, SingleAllocateAndFree)
     ScatterGatherEntry sge;
     status = mgr.Allocate(64, sge);
     ASSERT_TRUE(status.ok()) << status.message;
-    ASSERT_NE(sge.addr, 0);
+    ASSERT_NE(sge.local_addr, 0);
     ASSERT_EQ(sge.length, 64);
     ASSERT_EQ(sge.tokenId, 0);
     ASSERT_NE(sge.slot_index, UINT32_MAX);
+    ASSERT_EQ(sge.memory_type, MemoryType::HOST);
 
-    auto* ptr = reinterpret_cast<void*>(sge.addr);
+    auto* ptr = reinterpret_cast<void*>(sge.local_addr);
     std::memset(ptr, 0xAB, 64);
 
     status = mgr.Free(sge.slot_index);
@@ -147,12 +162,12 @@ TEST_F(BufferManagerTest, MultipleAllocatesAndFrees)
     for (int i = 0; i < kCount; ++i) {
         status = mgr.Allocate(128, sges[i]);
         ASSERT_TRUE(status.ok()) << "Failed at i=" << i << ": " << status.message;
-        ASSERT_NE(sges[i].addr, 0);
-        std::memset(reinterpret_cast<void*>(sges[i].addr), i, 128);
+        ASSERT_NE(sges[i].local_addr, 0);
+        std::memset(reinterpret_cast<void*>(sges[i].local_addr), i, 128);
     }
 
     for (int i = 0; i < kCount; ++i) {
-        auto* data = reinterpret_cast<unsigned char*>(sges[i].addr);
+        auto* data = reinterpret_cast<unsigned char*>(sges[i].local_addr);
         for (int j = 0; j < 128; ++j) { ASSERT_EQ(data[j], static_cast<unsigned char>(i)); }
     }
 
@@ -191,10 +206,61 @@ TEST_F(BufferManagerTest, AllocateFullSlotSize)
     status = mgr.Allocate(1024, sge);
     ASSERT_TRUE(status.ok()) << status.message;
     ASSERT_EQ(sge.length, 1024);
+}
+
+TEST_F(BufferManagerTest, AllocateFull4160ByteSlotCapacity)
+{
+    BufferManager mgr;
+    auto status = mgr.Init("test_buffer", MemoryType::HOST, 4160, 10);
+    ASSERT_TRUE(status.ok());
 
-    std::memset(reinterpret_cast<void*>(sge.addr), 0xFF, 1024);
+    ScatterGatherEntry sge;
+    status = mgr.Allocate(4160, sge);
+    ASSERT_TRUE(status.ok()) << status.message;
+    ASSERT_EQ(sge.length, 4160);
+}
 
-    mgr.Free(sge.slot_index);
+TEST_F(BufferManagerTest, AllocateExceeds4160ByteSlotCapacity)
+{
+    BufferManager mgr;
+    auto status = mgr.Init("test_buffer", MemoryType::HOST, 4160, 10);
+    ASSERT_TRUE(status.ok());
+
+    ScatterGatherEntry sge;
+    status = mgr.Allocate(4161, sge);
+    ASSERT_FALSE(status.ok());
+    ASSERT_EQ(status.code, StatusCode::INVALID_ARGUMENT);
+}
+
+TEST_F(BufferManagerTest, AllMemoryTypesUseAlignedSlotStride)
+{
+    for (const auto type : {MemoryType::HOST, MemoryType::HOST_PINNED, MemoryType::ASCEND_DEVICE}) {
+        BufferManager mgr;
+        auto status = mgr.Init("test_buffer", type, 4160, 2);
+        ASSERT_TRUE(status.ok()) << status.message;
+
+        ScatterGatherEntry first;
+        ScatterGatherEntry second;
+        ASSERT_TRUE(mgr.Allocate(4160, first).ok());
+        ASSERT_TRUE(mgr.Allocate(4160, second).ok());
+        ASSERT_EQ(second.local_addr - first.local_addr, 4160);
+        ASSERT_EQ(second.device_addr - first.device_addr, 4160);
+    }
+}
+
+TEST_F(BufferManagerTest, FlagBufferCapacity71Uses128ByteStride)
+{
+    BufferManager mgr;
+    auto status = mgr.Init("flag_buffer", MemoryType::HOST_PINNED, 71, 2);
+    ASSERT_TRUE(status.ok()) << status.message;
+
+    ScatterGatherEntry first;
+    ScatterGatherEntry second;
+    ASSERT_TRUE(mgr.Allocate(71, first).ok());
+    ASSERT_TRUE(mgr.Allocate(71, second).ok());
+    ASSERT_EQ(first.length, 71);
+    ASSERT_EQ(second.local_addr - first.local_addr, 128);
+    ASSERT_EQ(second.device_addr - first.device_addr, 128);
 }
 
 TEST_F(BufferManagerTest, ReuseAfterFree)
@@ -212,7 +278,7 @@ TEST_F(BufferManagerTest, ReuseAfterFree)
     ScatterGatherEntry sge2;
     status = mgr.Allocate(64, sge2);
     ASSERT_TRUE(status.ok());
-    ASSERT_EQ(sge2.addr, sge1.addr);
+    ASSERT_EQ(sge2.local_addr, sge1.local_addr);
     ASSERT_EQ(sge2.slot_index, sge1.slot_index);
 
     mgr.Free(sge2.slot_index);
@@ -233,7 +299,7 @@ TEST_F(BufferManagerTest, ConcurrentAllocateAndFree)
             auto s = mgr.Allocate(64, sge);
             ASSERT_TRUE(s.ok()) << "Thread " << thread_id << " op " << i << ": " << s.message;
 
-            std::memset(reinterpret_cast<void*>(sge.addr), thread_id, 64);
+            std::memset(reinterpret_cast<void*>(sge.local_addr), thread_id, 64);
 
             s = mgr.Free(sge.slot_index);
             ASSERT_TRUE(s.ok()) << s.message;
@@ -260,10 +326,10 @@ TEST_F(BufferManagerTest, ConcurrentStressTest)
             auto s = mgr.Allocate(128, sge);
             ASSERT_TRUE(s.ok());
 
-            std::memset(reinterpret_cast<void*>(sge.addr), thread_id, 128);
+            std::memset(reinterpret_cast<void*>(sge.local_addr), thread_id, 128);
 
             for (int j = 0; j < 128; ++j) {
-                ASSERT_EQ(reinterpret_cast<unsigned char*>(sge.addr)[j], thread_id);
+                ASSERT_EQ(reinterpret_cast<unsigned char*>(sge.local_addr)[j], thread_id);
             }
 
             s = mgr.Free(sge.slot_index);
@@ -286,7 +352,7 @@ TEST_F(BufferManagerTest, FreeZeroesMemory)
     status = mgr.Allocate(64, sge1);
     ASSERT_TRUE(status.ok());
 
-    auto* ptr = reinterpret_cast<uint8_t*>(sge1.addr);
+    auto* ptr = reinterpret_cast<uint8_t*>(sge1.local_addr);
     std::memset(ptr, 0xAB, 1024);
 
     status = mgr.Free(sge1.slot_index);
@@ -295,10 +361,10 @@ TEST_F(BufferManagerTest, FreeZeroesMemory)
     ScatterGatherEntry sge2;
     status = mgr.Allocate(64, sge2);
     ASSERT_TRUE(status.ok());
-    ASSERT_EQ(sge2.addr, sge1.addr);
+    ASSERT_EQ(sge2.local_addr, sge1.local_addr);
     ASSERT_EQ(sge2.slot_index, sge1.slot_index);
 
-    auto* ptr2 = reinterpret_cast<uint8_t*>(sge2.addr);
+    auto* ptr2 = reinterpret_cast<uint8_t*>(sge2.local_addr);
     for (size_t i = 0; i < 1024; ++i) {
         ASSERT_EQ(ptr2[i], 0) << "byte " << i << " not zeroed after free";
     }
@@ -395,6 +461,33 @@ TEST_F(BufferManagerTest, InitWithProviderRegistersMemory)
     ASSERT_NE(provider.lastAddr, 0);
     ASSERT_EQ(provider.lastSize, 1024 * 10);
     ASSERT_EQ(mgr.GetTokenId(), 42);
+
+    ScatterGatherEntry sge;
+    ASSERT_TRUE(mgr.Allocate(64, sge).ok());
+    ASSERT_EQ(sge.local_addr, sge.device_addr);
+}
+
+TEST_F(BufferManagerTest, HostPinnedRegistersDeviceAddress)
+{
+    StubTransProvider provider;
+
+    BufferManager mgr;
+    auto status = mgr.Init("test_rdma_pinned", MemoryType::HOST_PINNED, 4096, 1, &provider);
+    ASSERT_TRUE(status.ok()) << status.message;
+    ASSERT_EQ(provider.registerCount, 1);
+    ASSERT_EQ(provider.lastMemType, TransProvider::MemType::MEM_DEVICE);
+
+    ScatterGatherEntry sge;
+    ASSERT_TRUE(mgr.Allocate(64, sge).ok());
+    ASSERT_NE(sge.local_addr, 0);
+    ASSERT_NE(sge.device_addr, 0);
+    ASSERT_NE(sge.local_addr, sge.device_addr);
+    ASSERT_EQ(sge.local_addr % 4096, 0);
+    ASSERT_EQ(provider.lastAddr, sge.device_addr);
+
+    // The CPU writes through addr while HCOMM and remote RDMA use device_addr.
+    std::memset(reinterpret_cast<void*>(sge.local_addr), 0x5A, sge.length);
+    ASSERT_EQ(*reinterpret_cast<unsigned char*>(sge.local_addr), 0x5A);
 }
 
 TEST_F(BufferManagerTest, InitWithProviderAllocateReturnsTokenId)

ucm/transport/kv/asu/trans/include/asu_transport/asu_transport.h

@@ -72,9 +72,12 @@ struct TransportConfig {
     bool preconnect{true};
     bool bindCqPoller{true};
 
+    // Slot sizes are caller-visible capacities; BufferManager computes the
+    // aligned physical stride used for allocation and memory registration.
     std::size_t sendBufferSlotSize{4160};
     std::size_t sendBufferSlotNum{128};
-    std::size_t flagBufferSlotSize{128};
+    // Maximum memory required by a batch store/retrieve response flag buffer.
+    std::size_t flagBufferSlotSize{71};
     std::size_t flagBufferSlotNum{4096};
     std::size_t asuBatchLoadIoNum{110};
     std::size_t asuBatchStoreIoNum{110};

ucm/transport/kv/asu/trans/src/asu_submit_flow.cpp

@@ -130,24 +130,31 @@ Status AsuTransportImpl::BuildSubBatchSendBuffers(
             continue;
         }
 
-        if (subBatchContext.flagBuffer.addr == 0 || subBatchContext.flagBuffer.length == 0) {
-            const auto subBatchStatus =
-                Status::Error(StatusCode::NOT_INITIALIZED, "sub-batch flag buffer is not ready");
+        if (subBatchContext.channel == nullptr ||
+            !IsTransportBufferReady(subBatchContext.sendSge) ||
+            !IsTransportBufferReady(subBatchContext.flagBuffer)) {
+            const auto subBatchStatus = Status::Error(StatusCode::NOT_INITIALIZED,
+                                                      "sub-batch transport buffers are not ready");
             UC_ERROR(
-                "Sub-batch flag buffer is not ready index={} cid={} flag_addr={} flag_length={}",
-                index, subBatchContext.cid, subBatchContext.flagBuffer.addr,
-                subBatchContext.flagBuffer.length);
+                "Sub-batch transport buffers are not ready index={} cid={} channel={} "
+                "send_local_addr={} send_device_addr={} send_length={} send_slot={} "
+                "flag_local_addr={} flag_device_addr={} flag_length={} flag_slot={}",
+                index, subBatchContext.cid, subBatchContext.channel != nullptr,
+                subBatchContext.sendSge.local_addr, subBatchContext.sendSge.device_addr,
+                subBatchContext.sendSge.length, subBatchContext.sendSge.slot_index,
+                subBatchContext.flagBuffer.local_addr, subBatchContext.flagBuffer.device_addr,
+                subBatchContext.flagBuffer.length, subBatchContext.flagBuffer.slot_index);
             SetSubBatchSendFailed(subBatchContext, subBatchStatus);
             if (status.ok()) { status = subBatchStatus; }
             ReleaseSubBatchResources(subBatchContext);
             continue;
         }
 
-        ioBatches.push_back(
-            TransProvider::SendIoBatch{subBatchContext.channel->GetConnection(),
-                                       reinterpret_cast<void*>(subBatchContext.sendSge.addr),
-                                       reinterpret_cast<void*>(subBatchContext.flagBuffer.addr),
-                                       subBatchContext.sendSge.length});
+        ioBatches.push_back(TransProvider::SendIoBatch{
+            subBatchContext.channel->GetConnection(),
+            reinterpret_cast<void*>(subBatchContext.sendSge.device_addr),
+            reinterpret_cast<void*>(subBatchContext.flagBuffer.device_addr),
+            subBatchContext.sendSge.length});
         subBatchIndexes.emplace_back(index);
     }