Skip to content

Commit 933a8fa

Browse files
yumingyue624yumingyue624
committed
[Refactor] Move host-pinned allocation to AscendBuffer, Support HBM transport staging buffers
1 parent 2ec300b commit 933a8fa

7 files changed

Lines changed: 171 additions & 89 deletions

File tree

‎ucm/shared/trans/ascend/ascend_buffer.cc‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,27 @@
2828

2929
namespace UC::Trans {
3030

31+
namespace {
32+
33+
constexpr std::uintptr_t HOST_REGISTER_PAGE_SIZE = 4096;
34+
35+
void FreeHostMemory(void* host)
36+
{
37+
auto ret = aclrtFreeHost(host);
38+
if (ret != ACL_SUCCESS) { UC_ERROR("Failed to free host memory addr={} ret={}", host, ret); }
39+
}
40+
41+
void ReleaseHostPinnedMemory(void* host)
42+
{
43+
auto ret = aclrtHostUnregister(host);
44+
if (ret != ACL_SUCCESS) {
45+
UC_ERROR("Failed to unregister host-pinned memory addr={} ret={}", host, ret);
46+
}
47+
FreeHostMemory(host);
48+
}
49+
50+
} // namespace
51+
3152
class HostHugePages : public std::enable_shared_from_this<HostHugePages> {
3253
struct ConstructorKey {};
3354
static constexpr auto HUGE_PAGE_SIZE = 2UL << 20;
@@ -124,6 +145,33 @@ std::shared_ptr<void> Trans::AscendBuffer::MakeHostBuffer(size_t size)
124145
return nullptr;
125146
}
126147

148+
std::shared_ptr<void> Trans::AscendBuffer::MakeHostPinnedBuffer(size_t size, void** pDevice)
149+
{
150+
if (pDevice) { *pDevice = nullptr; }
151+
152+
void* host = nullptr;
153+
auto ret = aclrtMallocHost(&host, size);
154+
if (ret != ACL_SUCCESS) { return nullptr; }
155+
156+
if (reinterpret_cast<std::uintptr_t>(host) % HOST_REGISTER_PAGE_SIZE != 0) {
157+
UC_ERROR("Host-pinned memory is not 4K page aligned addr={} size={}", host, size);
158+
FreeHostMemory(host);
159+
return nullptr;
160+
}
161+
162+
void* device = nullptr;
163+
auto status = Buffer::RegisterHostBuffer(host, size, &device);
164+
if (status.Failure()) {
165+
UC_ERROR("Failed to register host-pinned memory addr={} size={} status={}", host, size,
166+
status);
167+
FreeHostMemory(host);
168+
return nullptr;
169+
}
170+
171+
if (pDevice) { *pDevice = device; }
172+
return std::shared_ptr<void>(host, ReleaseHostPinnedMemory);
173+
}
174+
127175
std::shared_ptr<void> Trans::AscendBuffer::MakeHostBuffer4DirectIo(size_t size)
128176
{
129177
try {

‎ucm/shared/trans/ascend/ascend_buffer.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ class AscendBuffer : public ReservedBuffer {
3232
public:
3333
std::shared_ptr<void> MakeDeviceBuffer(size_t size) override;
3434
std::shared_ptr<void> MakeHostBuffer(size_t size) override;
35+
std::shared_ptr<void> MakeHostPinnedBuffer(size_t size, void** pDevice = nullptr);
3536
std::shared_ptr<void> MakeHostBuffer4DirectIo(size_t size) override;
3637
};
3738

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,7 @@ TEST_F(BufferManagerTest, SingleAllocateAndFree)
141141
ASSERT_EQ(sge.length, 64);
142142
ASSERT_EQ(sge.tokenId, 0);
143143
ASSERT_NE(sge.slot_index, UINT32_MAX);
144+
ASSERT_EQ(sge.memory_type, MemoryType::HOST);
144145

145146
auto* ptr = reinterpret_cast<void*>(sge.local_addr);
146147
std::memset(ptr, 0xAB, 64);
@@ -242,8 +243,8 @@ TEST_F(BufferManagerTest, AllMemoryTypesUseAlignedSlotStride)
242243
ScatterGatherEntry second;
243244
ASSERT_TRUE(mgr.Allocate(4160, first).ok());
244245
ASSERT_TRUE(mgr.Allocate(4160, second).ok());
245-
ASSERT_EQ(second.local_addr - first.local_addr, 4224);
246-
ASSERT_EQ(second.device_addr - first.device_addr, 4224);
246+
ASSERT_EQ(second.local_addr - first.local_addr, 4160);
247+
ASSERT_EQ(second.device_addr - first.device_addr, 4160);
247248
}
248249
}
249250

@@ -458,7 +459,7 @@ TEST_F(BufferManagerTest, InitWithProviderRegistersMemory)
458459
ASSERT_EQ(provider.registerCount, 1);
459460
ASSERT_EQ(provider.lastMemType, TransProvider::MemType::MEM_HOST);
460461
ASSERT_NE(provider.lastAddr, 0);
461-
ASSERT_EQ(provider.lastSize, 1088 * 10);
462+
ASSERT_EQ(provider.lastSize, 1024 * 10);
462463
ASSERT_EQ(mgr.GetTokenId(), 42);
463464

464465
ScatterGatherEntry sge;

‎ucm/transport/kv/asu/trans/src/asu_transport_impl.cpp‎

Lines changed: 65 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,13 @@
2222
* SOFTWARE.
2323
* */
2424
#include "asu_transport_impl.h"
25+
#include <acl/acl.h>
2526
#include <algorithm>
27+
#include <array>
2628
#include <chrono>
29+
#include <cstdint>
2730
#include <memory>
31+
#include <string>
2832
#include <thread>
2933
#include <utility>
3034
#include "aicpu_trans_provider.h"
@@ -38,6 +42,29 @@
3842

3943
namespace UC::ASU {
4044

45+
namespace {
46+
47+
constexpr std::size_t kFlagBufferHeaderCopySize = kCqeDwordCount * sizeof(std::uint32_t);
48+
49+
Status CopyDeviceToHost(const ScatterGatherEntry& sge, void* host, std::size_t size,
50+
const char* name)
51+
{
52+
if (size > sge.length) {
53+
return Status::Error(StatusCode::INVALID_ARGUMENT,
54+
std::string(name) + ": copy size exceeds buffer length");
55+
}
56+
const auto ret = aclrtMemcpy(host, size, reinterpret_cast<void*>(sge.device_addr), size,
57+
ACL_MEMCPY_DEVICE_TO_HOST);
58+
if (ret != ACL_SUCCESS) {
59+
return Status::Error(
60+
StatusCode::INTERNAL_ERROR,
61+
std::string(name) + ": copy device memory to host failed ret=" + std::to_string(ret));
62+
}
63+
return Status::OK();
64+
}
65+
66+
} // namespace
67+
4168
AsuTransportImpl::~AsuTransportImpl() { Shutdown(); }
4269

4370
Status AsuTransportImpl::Init(const std::string& configPath)
@@ -436,23 +463,54 @@ void AsuTransportImpl::PollTaskCompletions(const TransportTaskContextPtr& ctx)
436463
for (auto& subBatchContext : ctx->subBatchContexts) {
437464
if (subBatchContext.state != TransportSubBatchState::PENDING) { continue; }
438465

466+
auto completeWithError = [this, &ctx, &subBatchContext](const Status& status) {
467+
std::fill(subBatchContext.entryStatus.begin(), subBatchContext.entryStatus.end(),
468+
status);
469+
CompleteSubBatch(*ctx, subBatchContext, status);
470+
};
471+
439472
std::uint16_t completedCid = 0;
440-
if (const auto status = protocolManager_->PollResponseCid(
441-
reinterpret_cast<void*>(subBatchContext.flagBuffer.local_addr), completedCid);
473+
const void* responseData = nullptr;
474+
std::array<std::uint8_t, kFlagBufferHeaderCopySize> flagHeader{};
475+
std::vector<std::uint8_t> flagBuffer;
476+
if (IsCpuAccessible(subBatchContext.flagBuffer.memory_type)) {
477+
responseData = reinterpret_cast<void*>(subBatchContext.flagBuffer.local_addr);
478+
} else {
479+
auto status = CopyDeviceToHost(subBatchContext.flagBuffer, flagHeader.data(),
480+
flagHeader.size(), "flag buffer header");
481+
if (!status.ok()) {
482+
// Without a readable header, this sub-batch cannot be polled or unpacked.
483+
completeWithError(status);
484+
continue;
485+
}
486+
responseData = flagHeader.data();
487+
}
488+
489+
if (const auto status = protocolManager_->PollResponseCid(responseData, completedCid);
442490
!status.ok()) {
443491
continue;
444492
}
445493
if (completedCid == 0 || completedCid != subBatchContext.cid) { continue; }
446494

495+
if (!IsCpuAccessible(subBatchContext.flagBuffer.memory_type)) {
496+
// The header matched; copy the full CQE before unpacking entry status.
497+
flagBuffer.resize(subBatchContext.flagBuffer.length);
498+
auto status = CopyDeviceToHost(subBatchContext.flagBuffer, flagBuffer.data(),
499+
flagBuffer.size(), "flag buffer");
500+
if (!status.ok()) {
501+
// The matched CQE cannot be decoded without the complete flag buffer.
502+
completeWithError(status);
503+
continue;
504+
}
505+
responseData = flagBuffer.data();
506+
}
507+
447508
KvResponse response;
448509
const auto batchNumber = static_cast<std::uint16_t>(subBatchContext.entryStatus.size());
449510
if (const auto status = protocolManager_->UnpackResponse(
450-
reinterpret_cast<void*>(subBatchContext.flagBuffer.local_addr),
451-
ToKvOpcode(subBatchContext.opType), batchNumber, response);
511+
responseData, ToKvOpcode(subBatchContext.opType), batchNumber, response);
452512
!status.ok()) {
453-
std::fill(subBatchContext.entryStatus.begin(), subBatchContext.entryStatus.end(),
454-
status);
455-
CompleteSubBatch(*ctx, subBatchContext, status);
513+
completeWithError(status);
456514
continue;
457515
}
458516

‎ucm/transport/kv/asu/trans/src/buffer_manager.cpp‎

Lines changed: 32 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -26,48 +26,26 @@
2626
#include <cstdlib>
2727
#include <cstring>
2828
#include <limits>
29-
#include "logger.h"
3029
#include "trans/ascend/ascend_buffer.h"
3130

3231
namespace UC::ASU {
3332

34-
constexpr std::uintptr_t kHostRegisterPageSize = 4096;
35-
constexpr std::size_t kSlotSizeAlignment = 32;
36-
constexpr std::size_t kSlotPadding = 32;
3733
constexpr std::size_t kSlotAddressAlignment = 64;
3834

39-
void FreeHostMemory(void* addr)
40-
{
41-
auto ret = aclrtFreeHost(addr);
42-
if (ret != ACL_SUCCESS) { UC_ERROR("Failed to free host memory addr={} ret={}", addr, ret); }
43-
}
44-
45-
void ReleaseHostPinnedMemory(void* addr)
46-
{
47-
auto ret = aclrtHostUnregister(addr);
48-
if (ret != ACL_SUCCESS) {
49-
UC_ERROR("Failed to unregister host-pinned memory addr={} ret={}", addr, ret);
50-
return;
51-
}
52-
FreeHostMemory(addr);
53-
}
54-
5535
bool GetSlotStride(std::size_t capacity, std::size_t& stride)
5636
{
57-
// Keep one layout for every memory type: reserve ALIGN_UP(capacity, 32) + 32
58-
// bytes, then align the next slot address to a 64-byte boundary.
37+
// NOTE: Ascend ACL documents an aclrtMallocHost large-block suballocation
38+
// layout of ALIGN_UP(len, 32) + 32 bytes with 64-byte-aligned segment
39+
// starts. Current HCOMM/RDMA validation did not reproduce failures without
40+
// the extra 32-byte tail room, so ASU keeps only the 64-byte slot-start
41+
// alignment for now.
42+
// Keep one layout for every memory type by aligning each slot start to a
43+
// 64-byte boundary.
5944
constexpr auto kMaxSize = std::numeric_limits<std::size_t>::max();
60-
if (capacity > kMaxSize - (kSlotSizeAlignment - 1)) { return false; }
61-
62-
const auto alignedSize =
63-
(capacity + kSlotSizeAlignment - 1) / kSlotSizeAlignment * kSlotSizeAlignment;
64-
if (alignedSize > kMaxSize - kSlotPadding) { return false; }
65-
66-
const auto paddedSize = alignedSize + kSlotPadding;
67-
if (paddedSize > kMaxSize - (kSlotAddressAlignment - 1)) { return false; }
45+
if (capacity > kMaxSize - (kSlotAddressAlignment - 1)) { return false; }
6846

6947
stride =
70-
(paddedSize + kSlotAddressAlignment - 1) / kSlotAddressAlignment * kSlotAddressAlignment;
48+
(capacity + kSlotAddressAlignment - 1) / kSlotAddressAlignment * kSlotAddressAlignment;
7149
return true;
7250
}
7351

@@ -77,16 +55,30 @@ Status BufferManager::BufferRegion::Create(MemoryType type, std::size_t size, Bu
7755
switch (type) {
7856
case MemoryType::HOST: {
7957
auto owner = ascendBuffer.MakeHostBuffer(size);
80-
if (!owner) { return AllocationFailed("host"); }
58+
if (!owner) {
59+
return Status::Error(StatusCode::INTERNAL_ERROR, "failed to allocate host memory");
60+
}
8161
// HOST has one CPU-visible address, which is also passed to the
8262
// provider when it registers the region as MEM_HOST.
8363
region = {owner, owner.get(), owner.get(), TransProvider::MemType::MEM_HOST};
8464
return Status::OK();
8565
}
86-
case MemoryType::HOST_PINNED: return MakeHostPinned(size, region);
66+
case MemoryType::HOST_PINNED: {
67+
void* deviceAddr = nullptr;
68+
auto owner = ascendBuffer.MakeHostPinnedBuffer(size, &deviceAddr);
69+
if (!owner) {
70+
return Status::Error(StatusCode::INTERNAL_ERROR,
71+
"failed to allocate host-pinned memory");
72+
}
73+
region = {owner, owner.get(), deviceAddr, TransProvider::MemType::MEM_DEVICE};
74+
return Status::OK();
75+
}
8776
case MemoryType::ASCEND_DEVICE: {
8877
auto owner = ascendBuffer.MakeDeviceBuffer(size);
89-
if (!owner) { return AllocationFailed("device"); }
78+
if (!owner) {
79+
return Status::Error(StatusCode::INTERNAL_ERROR,
80+
"failed to allocate device memory");
81+
}
9082
region = {owner, owner.get(), owner.get(), TransProvider::MemType::MEM_DEVICE};
9183
return Status::OK();
9284
}
@@ -102,51 +94,17 @@ void BufferManager::BufferRegion::Reset()
10294
providerMemType = TransProvider::MemType::MEM_HOST;
10395
}
10496

105-
Status BufferManager::BufferRegion::AllocationFailed(const char* type)
106-
{
107-
return Status::Error(StatusCode::INTERNAL_ERROR,
108-
std::string("failed to allocate ") + type + " memory");
109-
}
110-
111-
Status BufferManager::BufferRegion::MakeHostPinned(std::size_t size, BufferRegion& region)
112-
{
113-
void* hostAddr = nullptr;
114-
auto ret = aclrtMallocHost(&hostAddr, size);
115-
if (ret != ACL_SUCCESS) { return AllocationFailed("host-pinned"); }
116-
if (reinterpret_cast<std::uintptr_t>(hostAddr) % kHostRegisterPageSize != 0) {
117-
FreeHostMemory(hostAddr);
118-
return Status::Error(StatusCode::INTERNAL_ERROR,
119-
"host-pinned memory is not 4K page aligned");
120-
}
121-
122-
ret = aclrtHostRegisterV2(hostAddr, size, ACL_HOST_REG_MAPPED | ACL_HOST_REG_PINNED);
123-
if (ret != ACL_SUCCESS) {
124-
FreeHostMemory(hostAddr);
125-
return Status::Error(StatusCode::INTERNAL_ERROR,
126-
"failed to register host-pinned memory with ACL");
127-
}
128-
129-
void* deviceAddr = nullptr;
130-
ret = aclrtHostGetDevicePointer(hostAddr, &deviceAddr, 0);
131-
if (ret != ACL_SUCCESS) {
132-
ReleaseHostPinnedMemory(hostAddr);
133-
return Status::Error(StatusCode::INTERNAL_ERROR,
134-
"failed to get host-pinned device address");
135-
}
136-
137-
// The owner keeps the ACL registration alive until after HCOMM has
138-
// unregistered the region in BufferManager's destructor.
139-
auto owner = std::shared_ptr<void>(hostAddr, ReleaseHostPinnedMemory);
140-
region = {owner, hostAddr, deviceAddr, TransProvider::MemType::MEM_DEVICE};
141-
return Status::OK();
142-
}
143-
14497
bool IsTransportBufferReady(const ScatterGatherEntry& sge)
14598
{
14699
return sge.local_addr != 0 && sge.device_addr != 0 && sge.length != 0 &&
147100
sge.slot_index != UINT32_MAX;
148101
}
149102

103+
bool IsCpuAccessible(MemoryType type)
104+
{
105+
return type == MemoryType::HOST || type == MemoryType::HOST_PINNED;
106+
}
107+
150108
BufferManager::~BufferManager()
151109
{
152110
if (provider_ && memHandle_) {
@@ -262,6 +220,7 @@ Status BufferManager::Allocate(std::size_t size, ScatterGatherEntry& sge)
262220
sge.length = static_cast<std::uint32_t>(size);
263221
sge.tokenId = tokenId_;
264222
sge.slot_index = idx;
223+
sge.memory_type = memory_type_;
265224
return Status::OK();
266225
}
267226

‎ucm/transport/kv/asu/trans/src/buffer_manager.h‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,11 @@ struct ScatterGatherEntry {
4242
std::uint32_t length{0};
4343
std::uint32_t tokenId{0};
4444
std::uint32_t slot_index{UINT32_MAX};
45+
MemoryType memory_type{MemoryType::HOST};
4546
};
4647

4748
bool IsTransportBufferReady(const ScatterGatherEntry& sge);
49+
bool IsCpuAccessible(MemoryType type);
4850

4951
class BufferManager {
5052
public:
@@ -75,10 +77,6 @@ class BufferManager {
7577
void* localAddr{nullptr};
7678
void* deviceAddr{nullptr};
7779
TransProvider::MemType providerMemType{TransProvider::MemType::MEM_HOST};
78-
79-
private:
80-
static Status AllocationFailed(const char* type);
81-
static Status MakeHostPinned(std::size_t size, BufferRegion& region);
8280
};
8381

8482
Status RegisterMemory();

0 commit comments

Comments
 (0)