buffer_manager.cpp

/**
 * MIT License
 *
 * Copyright (c) 2026 Huawei Technologies Co., Ltd. All rights reserved.
 *
 * 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 "buffer_manager.h"
#include <acl/acl.h>
#include <cstdlib>
#include <cstring>
#include "trans/ascend/ascend_buffer.h"

namespace UC::ASU {
namespace {

struct BufferRegion {
    std::shared_ptr<void> owner;
    void* localAddr{nullptr};
    void* deviceAddr{nullptr};
    TransProvider::MemType providerMemType{TransProvider::MemType::MEM_HOST};
};

class BufferRegionCreator : public Trans::AscendBuffer {
public:
    Status MakeRegion(MemoryType type, std::size_t size, BufferRegion& region)
    {
        switch (type) {
            case MemoryType::HOST: {
                auto owner = MakeHostBuffer(size);
                if (!owner) { return AllocationFailed("host"); }
                region = {owner, owner.get(), owner.get(), TransProvider::MemType::MEM_HOST};
                return Status::OK();
            }
            case MemoryType::HOST_PINNED: return MakeHostPinnedBuffer(size, region);
            case MemoryType::ASCEND_DEVICE: {
                auto owner = MakeDeviceBuffer(size);
                if (!owner) { return AllocationFailed("device"); }
                region = {owner, owner.get(), owner.get(), TransProvider::MemType::MEM_DEVICE};
                return Status::OK();
            }
            default:
                return Status::Error(StatusCode::INVALID_ARGUMENT, "unsupported memory type");
        }
    }

private:
    static Status AllocationFailed(const char* type)
    {
        return Status::Error(StatusCode::INTERNAL_ERROR,
                             std::string("failed to allocate ") + type + " memory");
    }

    Status MakeHostPinnedBuffer(std::size_t size, BufferRegion& region)
    {
        void* hostAddr = nullptr;
        auto ret = aclrtMallocHost(&hostAddr, size);
        if (ret != ACL_SUCCESS) { return AllocationFailed("host-pinned"); }

        ret = aclrtHostRegisterV2(hostAddr, size, ACL_HOST_REG_MAPPED | ACL_HOST_REG_PINNED);
        if (ret != ACL_SUCCESS) {
            aclrtFreeHost(hostAddr);
            return Status::Error(StatusCode::INTERNAL_ERROR,
                                 "failed to register host-pinned memory with ACL");
        }

        void* deviceAddr = nullptr;
        ret = aclrtHostGetDevicePointer(hostAddr, &deviceAddr, 0);
        if (ret != ACL_SUCCESS) {
            aclrtHostUnregister(hostAddr);
            aclrtFreeHost(hostAddr);
            return Status::Error(StatusCode::INTERNAL_ERROR,
                                 "failed to get host-pinned device address");
        }

        // The owner keeps the ACL registration alive until after HCOMM has
        // unregistered the region in BufferManager's destructor.
        auto owner = std::shared_ptr<void>(hostAddr, [](void* addr) {
            aclrtHostUnregister(addr);
            aclrtFreeHost(addr);
        });
        region = {owner, hostAddr, deviceAddr, TransProvider::MemType::MEM_DEVICE};
        return Status::OK();
    }
};

}  // namespace

BufferManager::~BufferManager()
{
    if (provider_ && memHandle_) {
        std::vector<TransProvider::UnregisterMemoryDesc> descs{
            {nullptr, memHandle_}
        };
        provider_->UnregisterMemory(descs);
    }
    memory_.reset();
    device_memory_ = nullptr;
    slot_size_ = 0;
    slot_num_ = 0;
}

Status BufferManager::Init(std::string name, MemoryType type, std::size_t slot_size,
                           std::size_t slot_num, TransProvider* provider)
{
    if (memory_) {
        return Status::Error(StatusCode::INVALID_ARGUMENT, name + " already initialized");
    }
    if (slot_size == 0 || slot_num == 0) {
        return Status::Error(StatusCode::INVALID_ARGUMENT,
                             name + ": slot_size and slot_num must be non-zero");
    }

    name_ = std::move(name);
    memory_type_ = type;
    slot_size_ = slot_size;
    slot_num_ = slot_num;

    std::size_t total = slot_size * slot_num;

    BufferRegionCreator regionCreator;
    BufferRegion region;
    auto allocStatus = regionCreator.MakeRegion(memory_type_, total, region);
    if (!allocStatus.ok()) { return allocStatus; }
    memory_ = std::move(region.owner);
    device_memory_ = region.deviceAddr;
    provider_mem_type_ = region.providerMemType;

    if (memory_type_ == MemoryType::ASCEND_DEVICE) {
        if (aclrtMemset(memory_.get(), total, 0, total) != ACL_SUCCESS) {
            memory_.reset();
            device_memory_ = nullptr;
            return Status::Error(StatusCode::INTERNAL_ERROR,
                                 name_ + ": failed to zero device memory");
        }
    } else {
        std::memset(memory_.get(), 0, total);
    }

    index_pool_.Setup(static_cast<IndexPool::Index>(slot_num));

    if (provider) {
        provider_ = provider;
        auto regStatus = RegisterMemory();
        if (!regStatus.ok()) {
            provider_ = nullptr;
            memory_.reset();
            device_memory_ = nullptr;
            return regStatus;
        }
    }

    return Status::OK();
}

Status BufferManager::RegisterMemory()
{
    std::size_t total = slot_size_ * slot_num_;
    std::vector<TransProvider::RegisterMemoryDesc> descs{
        {provider_mem_type_, reinterpret_cast<uintptr_t>(device_memory_), total}
    };
    std::vector<TransProvider::MemHandle> memHandles;
    auto regStatus = provider_->RegisterMemory(nullptr, descs, memHandles);
    if (!regStatus.ok() || memHandles.empty()) {
        return Status::Error(StatusCode::INTERNAL_ERROR,
                             name_ + ": failed to register memory: " + regStatus.message);
    }

    auto tokenStatus = provider_->GetMemTokenId(memHandles[0], tokenId_);
    if (!tokenStatus.ok()) {
        std::vector<TransProvider::UnregisterMemoryDesc> unregDescs{
            {nullptr, memHandles[0]}
        };
        provider_->UnregisterMemory(unregDescs);
        return Status::Error(StatusCode::INTERNAL_ERROR,
                             name_ + ": failed to get token id: " + tokenStatus.message);
    }

    memHandle_ = memHandles[0];
    return Status::OK();
}

Status BufferManager::Allocate(std::size_t size, ScatterGatherEntry& sge)
{
    if (!memory_) { return Status::Error(StatusCode::NOT_INITIALIZED, name_ + " not initialized"); }
    if (size == 0) {
        return Status::Error(StatusCode::INVALID_ARGUMENT, name_ + ": size must be non-zero");
    }
    if (size > slot_size_) {
        return Status::Error(StatusCode::INVALID_ARGUMENT, name_ + ": size exceeds slot_size");
    }

    auto idx = index_pool_.Acquire();
    if (idx == IndexPool::npos) {
        return Status::Error(StatusCode::RESOURCE_BUSY, name_ + ": no free slots");
    }
    const auto offset = idx * slot_size_;
    sge.addr = reinterpret_cast<std::uint64_t>(static_cast<char*>(memory_.get()) + offset);
    sge.device_addr =
        reinterpret_cast<std::uint64_t>(static_cast<char*>(device_memory_) + offset);
    sge.length = static_cast<std::uint32_t>(size);
    sge.tokenId = tokenId_;
    sge.slot_index = idx;
    return Status::OK();
}

Status BufferManager::Free(std::uint32_t slot_index)
{
    if (!memory_) { return Status::Error(StatusCode::NOT_INITIALIZED, name_ + " not initialized"); }
    if (slot_index >= slot_num_) {
        return Status::Error(StatusCode::INVALID_ARGUMENT, name_ + ": slot_index out of range");
    }
    auto* p = static_cast<char*>(memory_.get()) + slot_index * slot_size_;
    if (memory_type_ == MemoryType::ASCEND_DEVICE) {
        if (aclrtMemset(p, slot_size_, 0, slot_size_) != ACL_SUCCESS) {
            return Status::Error(StatusCode::INTERNAL_ERROR,
                                 name_ + ": failed to zero device memory");
        }
    } else {
        std::memset(p, 0, slot_size_);
    }
    index_pool_.Release(static_cast<IndexPool::Index>(slot_index));
    return Status::OK();
}

bool BufferManager::IsValidPointer(const void* ptr) const
{
    if (!ptr || !memory_) { return false; }
    auto* base = static_cast<const char*>(memory_.get());
    auto* p = static_cast<const char*>(ptr);
    if (p < base || p >= base + slot_size_ * slot_num_) { return false; }
    auto offset = static_cast<std::size_t>(p - base);
    return (offset % slot_size_) == 0;
}

}  // namespace UC::ASU