Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #include "l2_mem_pool.h"
      12              : #include "caching_mem_allocator.h"
      13              : #include "common/checker.h"
      14              : #include "core/executor/multi_thread_topological/executor/schedule/scheduler/task_scheduler.h"
      15              : #include "rts_caching_mem_allocator.h"
      16              : #include "utils/rt2_utils.h"
      17              : 
      18              : namespace gert {
      19              : namespace memory {
      20              : L2MemPool::L2MemPool(ge::Allocator *allocator, aclrtStream stream, TypedContinuousVector<L2MemPool *> *all_l2_mem_pool)
      21              :     : first_level_pool_(allocator, stream, all_l2_mem_pool),
      22              :       memory_pool_(new ScalableAllocator(span_allocator_, first_level_pool_, ScalableConfig())),
      23              :       stream_(stream) {
      24              :   GELOGI("create l2 allocator:%s", memory_pool_->GetId().c_str());
      25              : }
      26              : 
      27              : L2MemPool::~L2MemPool() {
      28              :   (void)Finalize(true);
      29              : }
      30              : 
      31              : ge::MemBlock *L2MemPool::Malloc(size_t size) {
      32              :   auto addr = memory_pool_->Alloc(*this, size);
      33              :   if (addr != nullptr) {
      34              :     return addr;
      35              :   }
      36              :   GELOGE(ge::MEMALLOC_FAILED,
      37              :          "stream %p 's L2 allocator failed to apply for memory. We will try to free memory from memory pool, the above "
      38              :          "error log can be ignored. Try to free cached memory...",
      39              :          stream_);
      40              :   GE_ASSERT_SUCCESS(Synchronize());
      41              :   Recycle();
      42              :   addr = memory_pool_->Alloc(*this, size);
      43              :   return addr;
      44              : }
      45              : 
      46              : void L2MemPool::Free(ge::MemBlock *block) {
      47              :   memory_pool_->Free(reinterpret_cast<PageSpan *>(block));
      48              : }
      49              : 
      50              : void L2MemPool::Recycle() {
      51              :   memory_pool_->Recycle();
      52              : }
      53              : 
      54              : ge::Status L2MemPool::Synchronize() const {
      55              :   const auto wait_status = WaitForLaunchSubmissions();
      56              :   if (wait_status != ge::SUCCESS) {
      57              :     return wait_status;
      58              :   }
      59              :   GE_ASSERT_SUCCESS(DoRtStreamSyncWithTimeout(stream_));
      60              :   return ge::SUCCESS;
      61              : }
      62              : 
      63              : ge::Status L2MemPool::WaitForLaunchSubmissions() const {
      64              :   auto scheduler = TaskScheduler::GetCurrentScheduler();
      65              :   if (scheduler == nullptr) {
      66              :     return ge::SUCCESS;
      67              :   }
      68              :   return scheduler->WaitForLaunchSubmissions();
      69              : }
      70              : 
      71              : ge::Status L2MemPool::Finalize(bool no_log) {
      72              :   return memory_pool_->Finalize(no_log);
      73              : }
      74              : 
      75            4 : aclrtStream L2MemPool::GetStream() const {
      76              :   return stream_;
      77              : }
      78              : 
      79              : void L2MemPool::SetStream(aclrtStream stream) {
      80              :   stream_ = stream;
      81              :   first_level_pool_.SetStream(stream);
      82              : }
      83              : 
      84              : ge::MemBlock *L2MemPool::MoveL2ToL1(ge::MemBlock *block) {
      85              :   GE_ASSERT_NOTNULL(block);
      86              :   auto l1_block = memory_pool_->ConvertToRootBlock(block);
      87              :   if (l1_block == nullptr) {
      88              :     auto size = block->GetSize();
      89              :     l1_block = first_level_pool_.Alloc(size);
      90              :     GE_ASSERT_NOTNULL(l1_block);
      91              :     GE_ASSERT_RT_OK(aclrtMemcpyAsync(l1_block->GetAddr(), size, block->GetAddr(), block->GetSize(),
      92              :                                      ACL_MEMCPY_DEVICE_TO_DEVICE, stream_));
      93              :     GELOGI("l2 block %p addr %p is split, it has been moved to L1 block %p addr %p", block, block->GetAddr(), l1_block,
      94              :            l1_block->GetAddr());
      95              :     block->Free();
      96              :   }
      97              :   return l1_block;
      98              : }
      99              : 
     100              : BlockAddr MultiStreamL1Allocator::Alloc(const MemSize size) {
     101              :   auto block = l1_allocator_->Malloc(size);
     102              :   if ((block == nullptr) && (all_l2_mem_pool_ != nullptr) && !is_rt2_multi_thread_) {
     103              :     GELOGI("malloc memory not success, try to free l2 mem pool and malloc again");
     104              :     for (size_t i = 0U; i < all_l2_mem_pool_->GetSize(); ++i) {
     105              :       auto l2_mem_pool = all_l2_mem_pool_->MutableData()[i];
     106              :       GE_ASSERT_NOTNULL(l2_mem_pool);
     107              :       GE_ASSERT_SUCCESS(l2_mem_pool->Synchronize());
     108              :       l2_mem_pool->Recycle();
     109              :       block = l1_allocator_->Malloc(size);
     110              :       if (block == nullptr) {
     111              :         continue;
     112              :       }
     113              :       break;
     114              :     }
     115              :   }
     116              :   GE_ASSERT_NOTNULL(block,
     117              :                     "Failed to expand memory for l2 allocator, stream %p, size %zu, is enable rt2 multi thread: %zu",
     118              :                     l2_stream_, size, static_cast<size_t>(is_rt2_multi_thread_));
     119              :   GE_ASSERT_NOTNULL(block->GetAddr());
     120              :   GELOGI("[MEM]Expand memory pool at stream %p, address %p, size %zu, block %p. allocator addr %p", l2_stream_,
     121              :          block->GetAddr(), size, block, l1_allocator_);
     122              :   return block;
     123              : }
     124              : bool MultiStreamL1Allocator::Free(ge::MemBlock *const block) {
     125              :   if (block != nullptr) {
     126              :     GELOGI("[MEM]Shrink memory pool at stream %p, address %p, block %p", l2_stream_, block->GetAddr(), block);
     127              :     block->Free();
     128              :   }
     129              :   return true;
     130              : }
     131              : DeviceId MultiStreamL1Allocator::GetDeviceId() const {
     132              :   return -1;
     133              : }
     134              : }  // namespace memory
     135              : }  // namespace gert