Skip to content

完成作业，笔记尚未整理完成 #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
scbz4learning wants to merge 1 commit into
base: main
Choose a base branch
from
+623 −8
Open

完成作业，笔记尚未整理完成 #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions helper/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
This dir is for understanding the source files by greating some test or helper scripts.
12 changes: 12 additions & 0 deletions helper/TensorObj.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
#include "core/tensor.h"
#include <iostream>

using namespace infini;

TensorObj t;

int main() {
std::cout << t << std::endl;
return 0;
}

71 changes: 71 additions & 0 deletions include/core/allocator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,77 @@ namespace infini {
// HINT: 可以使用一个 map 来存储 free block，key 为 block 的起始/结尾地址，value 为 block 的大小
// =================================== 作业 ===================================

// 尝试了 first-fit，一个 map 的 best-fit
// 最后还是参考了 InfiniTensor 的实现。
// 笔记见 https://scbz4learning.github.io/Infinitensor/Stage_1/5_TinyInfiniTensor.md

// Finding a suitable block need O(n), as only linear search is available.
// std::map<size_t, size_t> free_blocks;

// However, referring to the original project
// Using an ordered set to sort the blocksize from smallest to largest
// can reduce the time in O(log n)
// And the blockAddr can use the hash table with doubly linked list
// to avoid maintain multiple datastructures

// 这样也不行
// 一个是必须要维护尾指针块信息，不然就没法知道扩容要不要merge
// 一个是unorder_map 存储addrInfo是不能找到临近块的
// map能，但是O(n)了
// struct BlockInfo {
// size_t addr; // 当前块的起始地址
// size_t size; // 当前块的大小
// size_t prevAddr; // 前一个块的起始地址，如果没有则设为 SIZE_MAX
// size_t nextAddr; // 后一个块的起始地址，如果没有则设为 SIZE_MAX

// // unordered_map needs default constructer
// BlockInfo()
// : addr(0), size(0), prevAddr(SIZE_MAX), nextAddr(SIZE_MAX) {}

// BlockInfo(size_t a, size_t s, size_t prev = SIZE_MAX, size_t next = SIZE_MAX)
// : addr(a), size(s), prevAddr(prev), nextAddr(next) {}

// // overload <
// bool operator<(const BlockInfo &other) const {
// if (size != other.size) return size < other.size;
// return addr < other.addr;
// }
// };

// // balanced tree ordered by size in descending order
// std::set<BlockInfo> freeBlocks;
// // hash map to find block info by addr
// // key addr
// // value blockInfo
// std::unordered_map<size_t, BlockInfo> addrInfo;

// // add a tailblock pointer for merging when increasing peak
// BlockInfo* tailBlock = nullptr;

struct BlockInfo {
size_t addr_;
size_t size_;

BlockInfo() : addr_(0), size_(0){}
BlockInfo(size_t s) : addr_(0), size_(s){}
BlockInfo(size_t a, size_t s): addr_(a), size_(s){}

bool operator<(const BlockInfo &other) const {
return (size_ != other.size_) ? (size_ < other.size_) : (addr_ < other.addr_);
}
};

// Balanced tree for free memory blocks, sorted by size and address
std::set<BlockInfo> freeBlocks;

// Key: Starting address of the free memory block
// Value: Size of the block
std::unordered_map<size_t, size_t> blockStartToSize;

// Key: Ending address of the free memory block
// Value: Size of the block
std::unordered_map<size_t, size_t> blockEndToSize;

public:
Allocator(Runtime runtime);

Expand Down
173 changes: 172 additions & 1 deletion src/core/allocator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,88 @@ namespace infini
}
}

// size_t Allocator::alloc(size_t size)
// {
// IT_ASSERT(this->ptr == nullptr);
// // pad the size to the multiple of alignment
// size = this->getAlignedSize(size);

// // =================================== 作业 ===================================
// // TODO: 设计一个算法来分配内存，返回起始地址偏移量
// // =================================== 作业 ===================================
// size_t offset = 0;

// // iterate through free blocks (first-fit / best-fit)
// auto it = freeBlocks.lower_bound(BlockInfo(0, size)); // find first block >= size
// if (it != freeBlocks.end()) {
// BlockInfo block = *it;
// offset = block.addr;

// // remove old block
// freeBlocks.erase(it);
// addrInfo.erase(block.addr);

// // split if block is larger than needed
// if (block.size > size) {
// BlockInfo newBlock(block.addr + size, block.size - size, offset, block.nextAddr);
// if (block.nextAddr != SIZE_MAX) {
// addrInfo[block.nextAddr].prevAddr = newBlock.addr;
// } else {
// tailBlock = &newBlock;
// }
// freeBlocks.insert(newBlock);
// addrInfo[newBlock.addr] = newBlock;
// } else {
// // update prev/next of neighbors if the block is removed
// if (block.prevAddr != SIZE_MAX) {
// addrInfo[block.prevAddr].nextAddr = block.nextAddr;
// }
// if (block.nextAddr != SIZE_MAX) {
// addrInfo[block.nextAddr].prevAddr = block.prevAddr;
// } else {
// tailBlock = &addrInfo[block.prevAddr];
// }
// }
// } else {
// // no suitable free block found in freeBlocks
// if (tailBlock && tailBlock->addr + tailBlock->size == peak) {
// // Merge new allocation with the current tail block
// freeBlocks.erase(*tailBlock);
// tailBlock->size += size;
// freeBlocks.insert(*tailBlock);
// offset = tailBlock->addr;
// // peak doesn't need to be increased since it's already included
// } else {
// // Extend peak: create a new block at the end
// offset = peak;
// BlockInfo newBlock(peak, size, tailBlock ? tailBlock->addr : SIZE_MAX, SIZE_MAX);
// if (tailBlock) {
// tailBlock->nextAddr = newBlock.addr; // link previous tail
// }
// addrInfo[newBlock.addr] = newBlock;
// tailBlock = &addrInfo[newBlock.addr]; // update tailBlock pointer
// peak += size; // move peak to include new allocation
// }
// }

// used += size;
// return offset;
// }

// void Allocator::free(size_t addr, size_t size)
// {
// IT_ASSERT(this->ptr == nullptr);
// size = getAlignedSize(size);

// // =================================== 作业 ===================================
// // TODO: 设计一个算法来回收内存
// // =================================== 作业 ===================================
// BlockInfo newBlock(addr, size, SIZE_MAX, SIZE_MAX);

// auto it = freeBlocks.lower_bound(addr);

// }

size_t Allocator::alloc(size_t size)
{
IT_ASSERT(this->ptr == nullptr);
Expand All @@ -33,7 +115,65 @@ namespace infini
// TODO: 设计一个算法来分配内存，返回起始地址偏移量
// =================================== 作业 ===================================

return 0;
size_t retAddr;
auto it = freeBlocks.lower_bound(BlockInfo(size));

// 如果找到符合条件的空闲块
if (it != freeBlocks.end()) {
BlockInfo bestFitBlock = *it;
retAddr = bestFitBlock.addr_;

// 删除旧块
freeBlocks.erase(it);
blockStartToSize.erase(bestFitBlock.addr_);
blockEndToSize.erase(bestFitBlock.addr_ + bestFitBlock.size_);

// 如果剩下碎片
if (bestFitBlock.size_ > size){
// 更新空闲块列表
BlockInfo remainBlock = {bestFitBlock.addr_ + size, bestFitBlock.size_ - size};
freeBlocks.insert(remainBlock);
blockStartToSize[remainBlock.addr_] = remainBlock.size_;
blockEndToSize[remainBlock.addr_ + remainBlock.size_] = remainBlock.size_;
}

// 更新已使用内存
used += size;

// 返回分配的起始地址
return retAddr;
}
else {
// 没有找到大小够的空闲块，尝试扩展内存池
// 如果尾块存在，要合并
auto itEnd = blockEndToSize.find(peak);
if (itEnd != blockEndToSize.end()) {
size_t endBlockSize = itEnd->second;

// 更新freeBlocks
BlockInfo endBlock(peak - endBlockSize, endBlockSize);
freeBlocks.erase(endBlock);

// 更新 hash map
blockStartToSize.erase(endBlock.addr_);
blockEndToSize.erase(peak);

// Update used & peak, just margin
used += size - endBlockSize;
peak += size - endBlockSize;

// retAddr
retAddr = endBlock.addr_;
} else {
// retAddr
retAddr = peak;

// Update used & peak
used += size;
peak += size;
}
}
return retAddr;
}

void Allocator::free(size_t addr, size_t size)
Expand All @@ -44,8 +184,39 @@ namespace infini
// =================================== 作业 ===================================
// TODO: 设计一个算法来回收内存
// =================================== 作业 ===================================

BlockInfo newBlock(addr, size);

// 如果有邻接前块，合并并删除前块
auto itAdjBlock = blockEndToSize.find(addr);
if (itAdjBlock != blockEndToSize.end()) {
// 拓展 newBlock 空间
newBlock.addr_ -= itAdjBlock->second;
newBlock.size_ += itAdjBlock->second;

// 删除 prev block
freeBlocks.erase(BlockInfo(itAdjBlock->first - itAdjBlock->second,
itAdjBlock->second));
blockStartToSize.erase(itAdjBlock->first - itAdjBlock->second);
blockEndToSize.erase(itAdjBlock->first);
}

// 如果有邻接后块，合并并删除后块
itAdjBlock = blockStartToSize.find(newBlock.addr_ + newBlock.size_);
if (itAdjBlock != blockStartToSize.end()) {
newBlock.size_ += itAdjBlock->second;

freeBlocks.erase(BlockInfo(itAdjBlock->first, itAdjBlock->second));
blockStartToSize.erase(itAdjBlock->first);
blockEndToSize.erase(itAdjBlock->first + itAdjBlock->second);
}

freeBlocks.emplace(newBlock);
blockStartToSize.emplace(newBlock.addr_, newBlock.size_);
blockEndToSize.emplace(newBlock.addr_ + newBlock.size_, newBlock.size_);
}


void *Allocator::getPtr()
{
if (this->ptr == nullptr)
Expand Down
Loading