Skip to content

success #19

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
maybeavailable wants to merge 1 commit into
base: main
Choose a base branch
from
Open

success #19

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
6 changes: 3 additions & 3 deletions include/core/allocator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
#include <cstddef>
#include <map>
#include <unordered_set>

namespace infini {
class Allocator
{
Expand All @@ -24,9 +23,10 @@ namespace infini {
void *ptr;

// =================================== 作业 ===================================
// TODO：可能需要设计一个数据结构来存储free block，以便于管理和合并
// HINT: 可以使用一个 map 来存储 free block，key 为 block 的起始/结尾地址，value 为 block 的大小
// 空闲块表：key 为 block 的起始地址，value 为 block 的大小。
// 利用 map 的有序性，可以快速找到相邻的空闲块用于合并
// =================================== 作业 ===================================
std::map<size_t, size_t> free_blocks;

public:
Allocator(Runtime runtime);
Expand Down
52 changes: 48 additions & 4 deletions src/core/allocator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#include "core/allocator.h"
#include <algorithm>
#include <utility>

namespace infini
Expand Down Expand Up @@ -30,10 +31,26 @@ namespace infini
size = this->getAlignedSize(size);

// =================================== 作业 ===================================
// TODO: 设计一个算法来分配内存，返回起始地址偏移量
// =================================== 作业 ===================================
// 策略：first-fit，从头到尾找到第一个大小足够的空闲块并复用
for (auto it = free_blocks.begin(); it != free_blocks.end(); ++it)
{
if (it->second >= size)
{
size_t offset = it->first;
// 若空闲块比所需更大，把剩余部分作为一个新的空闲块放回表里
if (it->second > size)
free_blocks.emplace(offset + size, it->second - size);
free_blocks.erase(it);
return offset;
}
}

return 0;
// 没有可复用的空闲块，就在内存尾部（bump pointer）扩展分配
size_t offset = used;
used += size;
peak = std::max(peak, used);
return offset;
// =================================== 作业 ===================================
}

void Allocator::free(size_t addr, size_t size)
Expand All @@ -42,7 +59,34 @@ namespace infini
size = getAlignedSize(size);

// =================================== 作业 ===================================
// TODO: 设计一个算法来回收内存
// 先与右侧相邻的空闲块合并：addr+size 正好是某个空闲块的起点
auto rightIt = free_blocks.find(addr + size);
if (rightIt != free_blocks.end())
{
size += rightIt->second;
free_blocks.erase(rightIt);
}

// 再与左侧相邻的空闲块合并：找起点小于 addr 的最大空闲块，
// 若它的终点正好是 addr，则两者相邻
auto it = free_blocks.lower_bound(addr);
if (it != free_blocks.begin())
{
--it;
if (it->first + it->second == addr)
{
addr = it->first;
size += it->second;
free_blocks.erase(it);
}
}

// 若合并后的空闲块正好位于内存尾部，则可以直接回退 bump pointer，
// 下次 alloc 从 used 处扩展时自然能复用这段空间，无需记入空闲表
if (addr + size == used)
used = addr;
else
free_blocks.emplace(addr, size);
// =================================== 作业 ===================================
}

Expand Down
187 changes: 181 additions & 6 deletions src/core/graph.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
#include "core/graph.h"
#include "operators/matmul.h"
#include "operators/transpose.h"
#include <algorithm>
#include <numeric>
#include <queue>
#include <unordered_map>
#include <unordered_set>

namespace infini
{
Expand Down Expand Up @@ -101,11 +105,143 @@ namespace infini
void GraphObj::optimize()
{
// =================================== 作业 ===================================
// TODO: 设计一个算法来实现指定的图优化规则
// 图优化规则如下：
// 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除）
// 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去）
// 图优化规则：
// 1. 删除相邻且互逆的 transpose 算子（两个 transpose 复合为单位置换时互相抵消）
// 2. 把 matmul 输入前的、仅交换最后两维的 transpose 融合进 transA / transB
// =================================== 作业 ===================================
// perm 是否恰好是"交换最后两维"的置换
auto isSwapLastTwo = [](const std::vector<int> &perm, int rank) -> bool {
if ((int)perm.size() != rank)
return false;
for (int d = 0; d < rank; ++d)
{
int expected =
(d == rank - 2) ? rank - 1 : (d == rank - 1) ? rank - 2 : d;
if (perm[d] != expected)
return false;
}
return true;
};
// 两个置换 P、Q 复合后是否为单位置换：transpose(Q, transpose(P, x)) == x
auto composeIsIdentity = [](const std::vector<int> &P,
const std::vector<int> &Q) -> bool {
if (P.size() != Q.size())
return false;
for (size_t j = 0; j < Q.size(); ++j)
if (P[Q[j]] != (int)j)
return false;
return true;
};
// 删除 op 时同步清理所有反向连接，避免残留悬空的 weak_ptr
auto removeOp = [this](const Operator &op) {
for (auto &in : op->getInputs())
if (in)
in->removeTarget(op);
for (auto &out : op->getOutputs())
if (out)
out->setSource(nullptr);
for (auto &pred : op->getPredecessors())
if (pred)
pred->removeSuccessors(op);
for (auto &succ : op->getSuccessors())
if (succ)
succ->removePredecessors(op);
removeOperator(op);
};

std::unordered_set<OperatorObj *> deadOps;
std::unordered_set<UidBaseType> deadTensors;

// ---- 规则 2：把 matmul 输入前的 transpose 融合进 transA / transB ----
for (auto &op : ops)
{
if (op->getOpType() != OpType::MatMul)
continue;
auto matmul = as<MatmulObj>(op);

// 沿输入 0（-> transA）或输入 1（-> transB）一路跳过可融合的 transpose
auto fold = [&](int inputIdx, bool &flipped) {
auto in = matmul->getInputs(inputIdx);
while (in && in->getSource() &&
in->getSource()->getOpType() == OpType::Transpose)
{
auto trans = as<TransposeObj>(in->getSource());
if (!isSwapLastTwo(trans->getPermute(), in->getRank()))
break;

flipped = !flipped;
auto next = trans->getInputs(0);
matmul->replaceInput(in, next); // 输入改为跳过这个 transpose
next->addTarget(matmul); // matmul 现在消费 next
in->removeTarget(matmul); // matmul 不再消费 in
deadOps.insert(trans.get());
deadTensors.insert(in->getFuid());
in = next;
}
};
bool flipA = false, flipB = false;
fold(0, flipA);
fold(1, flipB);
if (flipA)
matmul->setTransA(!matmul->getTransA());
if (flipB)
matmul->setTransB(!matmul->getTransB());
}

// ---- 规则 1：删除相邻且互逆的 transpose 对 ----
// 用 while 循环，因为消掉一对后可能暴露出新的可消除对
bool changed = true;
while (changed)
{
changed = false;
for (auto &op : ops)
{
if (deadOps.count(op.get()) || op->getOpType() != OpType::Transpose)
continue;
auto out = op->getOutput();
auto in = op->getInputs(0);
if (!in || !in->getSource() || deadOps.count(in->getSource().get()))
continue;
auto pred = in->getSource();
if (pred->getOpType() != OpType::Transpose)
continue;

auto opTrans = as<TransposeObj>(op);
auto predTrans = as<TransposeObj>(pred);
if (!composeIsIdentity(predTrans->getPermute(), opTrans->getPermute()))
continue;

// pred: pIn -> in ；op: in -> out ，两者复合为单位置换，
// 所以所有消费 out 的算子可以直接改用 pIn
auto pIn = predTrans->getInputs(0);
for (auto &consumer : out->getTargets())
{
consumer->replaceInput(out, pIn);
out->removeTarget(consumer);
pIn->addTarget(consumer);
}
deadOps.insert(op.get());
deadOps.insert(pred.get());
deadTensors.insert(in->getFuid());
deadTensors.insert(out->getFuid());
changed = true;
}
}

// ---- 清理：删除死 op 与死 tensor ----
std::vector<Operator> removeOps;
for (auto &op : ops)
if (deadOps.count(op.get()))
removeOps.emplace_back(op);
for (auto &op : removeOps)
removeOp(op);

std::vector<Tensor> removeTensors;
for (auto &t : tensors)
if (deadTensors.count(t->getFuid()))
removeTensors.emplace_back(t);
for (auto &t : removeTensors)
removeTensor(t);
}

Tensor GraphObj::getTensor(int fuid) const
Expand Down Expand Up @@ -149,8 +285,47 @@ namespace infini
IT_ASSERT(topo_sort() == true);

// =================================== 作业 ===================================
// TODO：利用 allocator 给计算图分配内存
// HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存
// 记录每个 tensor 剩余的消费者个数：减到 0 说明之后不会再被使用，内存可回收
std::unordered_map<UidBaseType, int> refCount;
for (auto &t : tensors)
refCount[t->getFuid()] = t->getTargets().size();

// 每个 tensor 在 allocator 中分到的偏移
std::unordered_map<UidBaseType, size_t> offsets;

// ① 图的输入 tensor（没有 source）全程存活，最先分配、永不回收
for (auto &input : getInputs())
offsets[input->getFuid()] = allocator.alloc(input->getBytes());

// ② 按拓扑序遍历 op：先回收本 op 消费完的输入，再给输出分配
for (auto &op : ops)
{
// 本 op 消费一次输入，引用计数减到 0 后内存可交给 allocator 回收
for (auto &input : op->getInputs())
{
if (input && input->getSource()) // 图输入无 source，不回收
{
auto fuid = input->getFuid();
if (--refCount[fuid] == 0)
allocator.free(offsets[fuid], input->getBytes());
}
}

// 给输出 tensor 分配内存
for (auto &output : op->getOutputs())
{
if (output)
offsets[output->getFuid()] = allocator.alloc(output->getBytes());
}
}

// ③ 记账已完成，peak 已定；真正 malloc 一次（getPtr），再给所有 tensor 绑定
void *base = allocator.getPtr();
for (auto &t : tensors)
{
if (auto it = offsets.find(t->getFuid()); it != offsets.end())
t->setDataBlob(make_ref<BlobObj>(runtime, (char *)base + it->second));
}
// =================================== 作业 ===================================

allocator.info();
Expand Down
8 changes: 6 additions & 2 deletions src/operators/concat.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,13 @@ optional<vector<Shape>> ConcatObj::inferShape(const TensorVec &inputs) {
auto rank = inputs[0]->getRank();

// =================================== 作业 ===================================
// TODO：修改 dims，返回正确的 concat 后的 shape
// REF: https://onnx.ai/onnx/operators/onnx__Concat.html#concat-13
// concat：除 dim 维度外，其余维度必须一致；输出在 dim 维上是所有输入的求和
// 例如 {1,3,2,4} 与 {1,3,2,5} 在 dim=3 上 concat → {1,3,2,9}
// =================================== 作业 ===================================
int sum = 0;
for (auto &input : inputs)
sum += input->getDims()[dim];
dims[dim] = sum;

return {{dims}};
}
Expand Down
27 changes: 24 additions & 3 deletions src/operators/matmul.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#include "operators/matmul.h"
#include "utils/operator_utils.h"

namespace infini
{
Expand All @@ -24,10 +25,30 @@ namespace infini
optional<vector<Shape>> MatmulObj::inferShape(const TensorVec &inputs)
{
// =================================== 作业 ===================================
// TODO：返回经过 matmul 操作后的 shape
// REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm
// matmul：最后两维是矩阵，前面的维度是 batch，batch 之间做广播。
// transA/transB 只交换矩阵维，不改变 batch 维。
// 输出 = broadcast(batchA, batchB) + {m, n}
// =================================== 作业 ===================================
return std::nullopt;
const auto A = inputs[0], B = inputs[1];
auto dimsA = A->getDims(), dimsB = B->getDims();
auto rankA = A->getRank(), rankB = B->getRank();

// 有效矩阵维度
m = transA ? dimsA[rankA - 1] : dimsA[rankA - 2];
int kA = transA ? dimsA[rankA - 2] : dimsA[rankA - 1];
int kB = transB ? dimsB[rankB - 1] : dimsB[rankB - 2];
n = transB ? dimsB[rankB - 2] : dimsB[rankB - 1];
IT_ASSERT(kA == kB, "Matmul: inner dims mismatch");
k = kA;

// batch 维度广播
Shape batchA(dimsA.begin(), dimsA.end() - 2);
Shape batchB(dimsB.begin(), dimsB.end() - 2);
auto batch = infer_broadcast(batchA, batchB);

batch.emplace_back(m);
batch.emplace_back(n);
return {{batch}};
}

} // namespace infini
10 changes: 7 additions & 3 deletions src/operators/transpose.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,17 @@ namespace infini
auto input_dim = A->getDims();
auto output_dim = input_dim;
int rank = A->getRank();
if (rank <= 0)
return std::nullopt;

// =================================== 作业 ===================================
// TODO：修改 output_dim，返回正确的 transpose 后的 shape
// REF: https://onnx.ai/onnx/operators/onnx__Transpose.html#transpose-21
// transpose：把 input 的维度按 perm 重排，output_dim[i] = input_dim[perm[i]]
// 例如 perm=[0,2,1,3]、input=[1,2,3,4] → output=[1,3,2,4]
// =================================== 作业 ===================================
for (int i = 0; i < rank; ++i)
output_dim[i] = input_dim[transposePermute[i]];

return std::nullopt;
return {{output_dim}};
}

std::string TransposeObj::toString() const
Expand Down
Loading