From 16f83597d9f191b84c3067ae95d114a240d930f6 Mon Sep 17 00:00:00 2001 From: duwei Date: Tue, 11 Aug 2026 13:40:49 +0800 Subject: [PATCH] success --- include/core/allocator.h | 6 +- src/core/allocator.cc | 52 +++++++++- src/core/graph.cc | 187 ++++++++++++++++++++++++++++++++++-- src/operators/concat.cc | 8 +- src/operators/matmul.cc | 27 +++++- src/operators/transpose.cc | 10 +- src/operators/unary.cc | 18 ++-- src/utils/operator_utils.cc | 21 +++- 8 files changed, 294 insertions(+), 35 deletions(-) diff --git a/include/core/allocator.h b/include/core/allocator.h index 002601d2..bbe66370 100644 --- a/include/core/allocator.h +++ b/include/core/allocator.h @@ -7,7 +7,6 @@ #include #include #include - namespace infini { class Allocator { @@ -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 free_blocks; public: Allocator(Runtime runtime); diff --git a/src/core/allocator.cc b/src/core/allocator.cc index ff593aef..6847928a 100644 --- a/src/core/allocator.cc +++ b/src/core/allocator.cc @@ -1,4 +1,5 @@ #include "core/allocator.h" +#include #include namespace infini @@ -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) @@ -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); // =================================== 作业 =================================== } diff --git a/src/core/graph.cc b/src/core/graph.cc index 3a906370..066e48e4 100644 --- a/src/core/graph.cc +++ b/src/core/graph.cc @@ -1,7 +1,11 @@ #include "core/graph.h" +#include "operators/matmul.h" +#include "operators/transpose.h" #include #include #include +#include +#include namespace infini { @@ -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 &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 &P, + const std::vector &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 deadOps; + std::unordered_set deadTensors; + + // ---- 规则 2：把 matmul 输入前的 transpose 融合进 transA / transB ---- + for (auto &op : ops) + { + if (op->getOpType() != OpType::MatMul) + continue; + auto matmul = as(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(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(op); + auto predTrans = as(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 removeOps; + for (auto &op : ops) + if (deadOps.count(op.get())) + removeOps.emplace_back(op); + for (auto &op : removeOps) + removeOp(op); + + std::vector 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 @@ -149,8 +285,47 @@ namespace infini IT_ASSERT(topo_sort() == true); // =================================== 作业 =================================== - // TODO：利用 allocator 给计算图分配内存 - // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存 + // 记录每个 tensor 剩余的消费者个数：减到 0 说明之后不会再被使用，内存可回收 + std::unordered_map refCount; + for (auto &t : tensors) + refCount[t->getFuid()] = t->getTargets().size(); + + // 每个 tensor 在 allocator 中分到的偏移 + std::unordered_map 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(runtime, (char *)base + it->second)); + } // =================================== 作业 =================================== allocator.info(); diff --git a/src/operators/concat.cc b/src/operators/concat.cc index d1963308..db499b25 100644 --- a/src/operators/concat.cc +++ b/src/operators/concat.cc @@ -14,9 +14,13 @@ optional> 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}}; } diff --git a/src/operators/matmul.cc b/src/operators/matmul.cc index 7a16ca27..4cd7f5c1 100644 --- a/src/operators/matmul.cc +++ b/src/operators/matmul.cc @@ -1,4 +1,5 @@ #include "operators/matmul.h" +#include "utils/operator_utils.h" namespace infini { @@ -24,10 +25,30 @@ namespace infini optional> 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 \ No newline at end of file diff --git a/src/operators/transpose.cc b/src/operators/transpose.cc index faab2b69..f3cdf8c3 100644 --- a/src/operators/transpose.cc +++ b/src/operators/transpose.cc @@ -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 diff --git a/src/operators/unary.cc b/src/operators/unary.cc index 3daad361..26e570aa 100644 --- a/src/operators/unary.cc +++ b/src/operators/unary.cc @@ -36,10 +36,10 @@ namespace infini optional> ClipObj::inferShape(const TensorVec &inputs) { // =================================== 作业 =================================== - // TODO：返回经过 clip 操作后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Clip.html#clip-13 + const auto A = inputs[0]; + return {{A->getDims()}}; // =================================== 作业 =================================== - return std::nullopt; + } std::string ClipObj::toString() const @@ -62,20 +62,18 @@ namespace infini vector CastObj::inferDataType(const TensorVec &inputs) const { // =================================== 作业 =================================== - // TODO：返回经过 cast 操作后, 输出 tensor 的数目和数据类型 - // REF_FILE: src/core/operator.cc - // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 + // cast 按 castType 转换数据类型；输出只有一个 tensor，类型由 getOutputDataType 给出 // =================================== 作业 =================================== - return {}; + return {getOutputDataType()}; } optional> CastObj::inferShape(const TensorVec &inputs) { // =================================== 作业 =================================== - // TODO：返回经过 cast 操作后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 + // cast 是逐元素操作，输出的 shape 与输入一致 // =================================== 作业 =================================== - return std::nullopt; + const auto A = inputs[0]; + return {{A->getDims()}}; } std::string CastObj::toString() const diff --git a/src/utils/operator_utils.cc b/src/utils/operator_utils.cc index edbd2c82..ff12b4d6 100644 --- a/src/utils/operator_utils.cc +++ b/src/utils/operator_utils.cc @@ -1,16 +1,29 @@ #include "utils/operator_utils.h" #include "core/runtime.h" +#include namespace infini { Shape infer_broadcast(const Shape &A, const Shape &B) { // =================================== 作业 =================================== - // TODO：对 A 和 B 进行双向广播，返回广播后的形状。 - // REF: https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md + // ONNX 双向广播（numpy 规则）： + // - 从右往左逐维对齐 + // - 每维要么相等，要么有一个是 1（1 会广播成另一边） + // - 输出 shape = 逐维取 max + // 例如 {1,3,2,4} 与 {4} 广播 → {1,3,2,4} // =================================== 作业 =================================== - - return {}; + int rank = std::max(A.size(), B.size()); + Shape result(rank, 1); + for (int i = 1; i <= rank; ++i) + { + int a = i <= (int)A.size() ? A[A.size() - i] : 1; + int b = i <= (int)B.size() ? B[B.size() - i] : 1; + IT_ASSERT(a == b || a == 1 || b == 1, + "Incompatible shapes for broadcast"); + result[rank - i] = std::max(a, b); + } + return result; } int get_real_axis(const int &axis, const int &rank) {