From 4297ede17ac1d48a1aec14e6a9939b42c6e16ca6 Mon Sep 17 00:00:00 2001 From: wjia8591-dou Date: Wed, 12 Aug 2026 23:27:29 +0800 Subject: [PATCH] Complete tinyinfinitensor assignments --- include/core/allocator.h | 7 +- src/core/allocator.cc | 68 ++++++++++++++++--- src/core/graph.cc | 129 +++++++++++++++++++++++++++++++++--- src/operators/concat.cc | 11 ++- src/operators/matmul.cc | 25 ++++++- src/operators/transpose.cc | 14 +++- src/operators/unary.cc | 9 ++- src/utils/operator_utils.cc | 13 +++- 8 files changed, 239 insertions(+), 37 deletions(-) diff --git a/include/core/allocator.h b/include/core/allocator.h index 002601d2..01c6d697 100644 --- a/include/core/allocator.h +++ b/include/core/allocator.h @@ -23,10 +23,9 @@ namespace infini { // pointer to the memory actually allocated void *ptr; - // =================================== 作业 =================================== - // TODO：可能需要设计一个数据结构来存储free block，以便于管理和合并 - // HINT: 可以使用一个 map 来存储 free block，key 为 block 的起始/结尾地址，value 为 block 的大小 - // =================================== 作业 =================================== + // Free blocks indexed by their starting offset. Adjacent blocks are + // coalesced on free so that long-running graphs do not become fragmented. + std::map freeBlocks; public: Allocator(Runtime runtime); diff --git a/src/core/allocator.cc b/src/core/allocator.cc index ff593aef..229ded4c 100644 --- a/src/core/allocator.cc +++ b/src/core/allocator.cc @@ -29,11 +29,41 @@ namespace infini // pad the size to the multiple of alignment size = this->getAlignedSize(size); - // =================================== 作业 =================================== - // TODO: 设计一个算法来分配内存，返回起始地址偏移量 - // =================================== 作业 =================================== + // First fit is deterministic and reuses the lowest suitable address. + for (auto it = freeBlocks.begin(); it != freeBlocks.end(); ++it) + { + if (it->second < size) + continue; + const size_t addr = it->first; + const size_t remainder = it->second - size; + freeBlocks.erase(it); + if (remainder) + freeBlocks.emplace(addr + size, remainder); + used += size; + return addr; + } + + // A free block at the current high-water mark can grow in place. This + // preserves its address while extending the backing allocation only by + // the missing bytes. + if (!freeBlocks.empty()) + { + auto last = std::prev(freeBlocks.end()); + if (last->first + last->second == peak) + { + const size_t addr = last->first; + const size_t available = last->second; + freeBlocks.erase(last); + peak += size - available; + used += size; + return addr; + } + } - return 0; + const size_t addr = peak; + peak += size; + used += size; + return addr; } void Allocator::free(size_t addr, size_t size) @@ -41,9 +71,31 @@ namespace infini IT_ASSERT(this->ptr == nullptr); size = getAlignedSize(size); - // =================================== 作业 =================================== - // TODO: 设计一个算法来回收内存 - // =================================== 作业 =================================== + IT_ASSERT(size <= used); + used -= size; + + auto next = freeBlocks.lower_bound(addr); + if (next != freeBlocks.end()) + IT_ASSERT(addr + size <= next->first, "Overlapping free block"); + if (next != freeBlocks.begin()) + { + auto prev = std::prev(next); + IT_ASSERT(prev->first + prev->second <= addr, + "Overlapping free block"); + if (prev->first + prev->second == addr) + { + addr = prev->first; + size += prev->second; + freeBlocks.erase(prev); + } + } + next = freeBlocks.lower_bound(addr); + if (next != freeBlocks.end() && addr + size == next->first) + { + size += next->second; + freeBlocks.erase(next); + } + freeBlocks.emplace(addr, size); } void *Allocator::getPtr() @@ -58,7 +110,7 @@ namespace infini size_t Allocator::getAlignedSize(size_t size) { - return ((size - 1) / this->alignment + 1) * this->alignment; + return size == 0 ? 0 : ((size - 1) / this->alignment + 1) * this->alignment; } void Allocator::info() diff --git a/src/core/graph.cc b/src/core/graph.cc index 3a906370..91160a40 100644 --- a/src/core/graph.cc +++ b/src/core/graph.cc @@ -1,4 +1,6 @@ #include "core/graph.h" +#include "operators/matmul.h" +#include "operators/transpose.h" #include #include #include @@ -100,12 +102,96 @@ namespace infini void GraphObj::optimize() { - // =================================== 作业 =================================== - // TODO: 设计一个算法来实现指定的图优化规则 - // 图优化规则如下： - // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除） - // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去） - // =================================== 作业 =================================== + auto inversePermutations = [](const vector &a, + const vector &b) { + if (a.size() != b.size()) + return false; + for (size_t i = 0; i < a.size(); ++i) + if (a[b[i]] != static_cast(i)) + return false; + return true; + }; + auto swapsLastTwo = [](const vector &perm) { + if (perm.size() < 2) + return false; + for (size_t i = 0; i + 2 < perm.size(); ++i) + if (perm[i] != static_cast(i)) + return false; + return perm[perm.size() - 2] == static_cast(perm.size() - 1) && + perm[perm.size() - 1] == static_cast(perm.size() - 2); + }; + + bool changed; + do { + changed = false; + // Cancel adjacent inverse transposes when the intermediate value has + // no other consumer. + for (auto second : ops) { + auto t2 = as(second); + if (!t2) + continue; + auto middle = t2->getInputs(0); + auto first = as(middle->getSource()); + if (!first || middle->getTargets().size() != 1 || + !inversePermutations(first->getPermute(), t2->getPermute())) + continue; + auto original = first->getInputs(0); + auto result = t2->getOutput(); + for (auto consumer : result->getTargets()) + consumer->replaceInput(result, original); + removeOperator(first); + removeOperator(second); + removeTensor(middle); + removeTensor(result); + changed = true; + break; + } + } while (changed); + + // Fold a final-two-axis transpose into either matmul input. Only fold a + // single-consumer tensor, so the rewrite cannot affect another branch. + for (auto op : ops) { + auto matmul = as(op); + if (!matmul) + continue; + for (size_t inputIndex = 0; inputIndex < 2; ++inputIndex) { + auto transposed = matmul->getInputs(inputIndex); + auto transpose = as(transposed->getSource()); + if (!transpose || transposed->getTargets().size() != 1 || + !swapsLastTwo(transpose->getPermute())) + continue; + matmul->replaceInput(transposed, transpose->getInputs(0)); + if (inputIndex == 0) + matmul->setTransA(!matmul->getTransA()); + else + matmul->setTransB(!matmul->getTransB()); + removeOperator(transpose); + removeTensor(transposed); + } + } + + // Rebuild every reverse edge after the rewrites. + for (auto &tensor : tensors) { + tensor->targets.clear(); + tensor->source.reset(); + } + for (auto &op : ops) { + op->predecessors.clear(); + op->successors.clear(); + } + for (auto &op : ops) { + for (auto &input : op->inputs) { + input->addTarget(op); + if (auto pred = input->getSource()) { + pred->addSuccessors(op); + op->addPredecessors(pred); + } + } + for (auto &output : op->outputs) + output->setSource(op); + } + sorted = false; + IT_ASSERT(checkValid()); } Tensor GraphObj::getTensor(int fuid) const @@ -148,10 +234,31 @@ namespace infini // topological sorting first IT_ASSERT(topo_sort() == true); - // =================================== 作业 =================================== - // TODO：利用 allocator 给计算图分配内存 - // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存 - // =================================== 作业 =================================== + std::unordered_map offsets; + std::unordered_map remainingUses; + for (auto &tensor : tensors) + remainingUses[tensor.get()] = tensor->getTargets().size(); + + // Graph inputs are live before the first operator executes. + for (auto &tensor : tensors) + if (!tensor->getSource()) + offsets[tensor.get()] = allocator.alloc(tensor->getBytes()); + + for (auto &op : ops) { + for (auto &output : op->getOutputs()) + offsets[output.get()] = allocator.alloc(output->getBytes()); + for (auto &input : op->getInputs()) { + auto &uses = remainingUses[input.get()]; + IT_ASSERT(uses > 0); + if (--uses == 0 && !input->getTargets().empty()) + allocator.free(offsets.at(input.get()), input->getBytes()); + } + } + + auto *base = static_cast(allocator.getPtr()); + for (auto &tensor : tensors) + tensor->setDataBlob( + make_ref(runtime, base + offsets.at(tensor.get()))); allocator.info(); } @@ -227,4 +334,4 @@ namespace infini return true; } -} // namespace infini \ No newline at end of file +} // namespace infini diff --git a/src/operators/concat.cc b/src/operators/concat.cc index d1963308..0f686ee7 100644 --- a/src/operators/concat.cc +++ b/src/operators/concat.cc @@ -16,8 +16,15 @@ optional> ConcatObj::inferShape(const TensorVec &inputs) { // =================================== 作业 =================================== // TODO：修改 dims，返回正确的 concat 后的 shape // REF: https://onnx.ai/onnx/operators/onnx__Concat.html#concat-13 - // =================================== 作业 =================================== - + IT_ASSERT(!inputs.empty()); + for (size_t i = 1; i < inputs.size(); ++i) { + const auto current = inputs[i]->getDims(); + IT_ASSERT(inputs[i]->getRank() == rank); + for (size_t axis = 0; axis < rank; ++axis) + if (static_cast(axis) != dim) + IT_ASSERT(current[axis] == dims[axis]); + dims[dim] += current[dim]; + } return {{dims}}; } diff --git a/src/operators/matmul.cc b/src/operators/matmul.cc index 7a16ca27..943f8fbc 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 { @@ -26,8 +27,26 @@ namespace infini // =================================== 作业 =================================== // TODO：返回经过 matmul 操作后的 shape // REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm - // =================================== 作业 =================================== - return std::nullopt; + IT_ASSERT(inputs.size() == 2); + auto a = inputs[0]->getDims(); + auto b = inputs[1]->getDims(); + IT_ASSERT(a.size() >= 2 && b.size() >= 2); + + const int aRows = transA ? a[a.size() - 1] : a[a.size() - 2]; + const int aCols = transA ? a[a.size() - 2] : a[a.size() - 1]; + const int bRows = transB ? b[b.size() - 1] : b[b.size() - 2]; + const int bCols = transB ? b[b.size() - 2] : b[b.size() - 1]; + IT_ASSERT(aCols == bRows, "Matmul inner dimensions mismatch"); + + Shape aBatch(a.begin(), a.end() - 2); + Shape bBatch(b.begin(), b.end() - 2); + Shape output = infer_broadcast(aBatch, bBatch); + output.push_back(aRows); + output.push_back(bCols); + m = aRows; + n = bCols; + k = aCols; + return {{output}}; } -} // namespace infini \ No newline at end of file +} // namespace infini diff --git a/src/operators/transpose.cc b/src/operators/transpose.cc index faab2b69..77a29ef3 100644 --- a/src/operators/transpose.cc +++ b/src/operators/transpose.cc @@ -9,6 +9,7 @@ namespace infini auto rank = input->getRank(); if (permute.empty()) { + transposePermute.resize(rank); for (size_t i = 0; i < rank; ++i) { transposePermute[i] = i; @@ -32,9 +33,16 @@ namespace infini // =================================== 作业 =================================== // TODO：修改 output_dim，返回正确的 transpose 后的 shape // REF: https://onnx.ai/onnx/operators/onnx__Transpose.html#transpose-21 - // =================================== 作业 =================================== - - return std::nullopt; + IT_ASSERT(static_cast(transposePermute.size()) == rank); + vector seen(rank, false); + for (int i = 0; i < rank; ++i) + { + const int axis = transposePermute[i]; + IT_ASSERT(axis >= 0 && axis < rank && !seen[axis]); + seen[axis] = true; + output_dim[i] = input_dim[axis]; + } + return {{output_dim}}; } std::string TransposeObj::toString() const diff --git a/src/operators/unary.cc b/src/operators/unary.cc index 3daad361..def3187d 100644 --- a/src/operators/unary.cc +++ b/src/operators/unary.cc @@ -39,7 +39,8 @@ namespace infini // TODO：返回经过 clip 操作后的 shape // REF: https://onnx.ai/onnx/operators/onnx__Clip.html#clip-13 // =================================== 作业 =================================== - return std::nullopt; + IT_ASSERT(inputs.size() == 1); + return {{inputs[0]->getDims()}}; } std::string ClipObj::toString() const @@ -66,7 +67,8 @@ namespace infini // REF_FILE: src/core/operator.cc // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 // =================================== 作业 =================================== - return {}; + IT_ASSERT(inputs.size() == 1); + return {getOutputDataType()}; } optional> CastObj::inferShape(const TensorVec &inputs) @@ -75,7 +77,8 @@ namespace infini // TODO：返回经过 cast 操作后的 shape // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 // =================================== 作业 =================================== - return std::nullopt; + IT_ASSERT(inputs.size() == 1); + return {{inputs[0]->getDims()}}; } std::string CastObj::toString() const diff --git a/src/utils/operator_utils.cc b/src/utils/operator_utils.cc index edbd2c82..4c7b0f6b 100644 --- a/src/utils/operator_utils.cc +++ b/src/utils/operator_utils.cc @@ -8,9 +8,16 @@ Shape infer_broadcast(const Shape &A, const Shape &B) { // =================================== 作业 =================================== // TODO：对 A 和 B 进行双向广播，返回广播后的形状。 // REF: https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md - // =================================== 作业 =================================== - - return {}; + const size_t rank = std::max(A.size(), B.size()); + Shape result(rank, 1); + for (size_t i = 0; i < rank; ++i) { + const int a = i < A.size() ? A[A.size() - 1 - i] : 1; + const int b = i < B.size() ? B[B.size() - 1 - i] : 1; + IT_ASSERT(a == b || a == 1 || b == 1, + "Shapes cannot be bidirectionally broadcast"); + result[rank - 1 - i] = std::max(a, b); + } + return result; } int get_real_axis(const int &axis, const int &rank) {