From 3e29c620c512030edeabff105e61b8df10d5f6a2 Mon Sep 17 00:00:00 2001 From: ule0p Date: Mon, 10 Aug 2026 10:49:31 +0800 Subject: [PATCH] first complete assignments --- include/core/allocator.h | 8 +- src/core/allocator.cc | 69 ++++++++++++-- src/core/graph.cc | 177 +++++++++++++++++++++++++++++++++--- src/operators/concat.cc | 19 +++- src/operators/matmul.cc | 45 +++++++-- src/operators/transpose.cc | 19 ++-- src/operators/unary.cc | 20 +--- src/utils/operator_utils.cc | 21 +++-- 8 files changed, 317 insertions(+), 61 deletions(-) diff --git a/include/core/allocator.h b/include/core/allocator.h index 002601d2..4a4a1cca 100644 --- a/include/core/allocator.h +++ b/include/core/allocator.h @@ -23,10 +23,10 @@ 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 offsets. + std::map freeBlocks; + // The current end of the simulated heap. `peak` is its high-water mark. + size_t cursor; public: Allocator(Runtime runtime); diff --git a/src/core/allocator.cc b/src/core/allocator.cc index ff593aef..a2fe4999 100644 --- a/src/core/allocator.cc +++ b/src/core/allocator.cc @@ -1,4 +1,5 @@ #include "core/allocator.h" +#include #include namespace infini @@ -8,6 +9,7 @@ namespace infini used = 0; peak = 0; ptr = nullptr; + cursor = 0; // 'alignment' defaults to sizeof(uint64_t), because it is the length of // the longest data type currently supported by the DataType field of @@ -29,11 +31,28 @@ namespace infini // pad the size to the multiple of alignment size = this->getAlignedSize(size); - // =================================== 作业 =================================== - // TODO: 设计一个算法来分配内存，返回起始地址偏移量 - // =================================== 作业 =================================== + if (size == 0) + return cursor; - return 0; + for (auto it = freeBlocks.begin(); it != freeBlocks.end(); ++it) + { + if (it->second < size) + continue; + + const size_t addr = it->first; + const size_t remaining = it->second - size; + freeBlocks.erase(it); + if (remaining > 0) + freeBlocks.emplace(addr + size, remaining); + used += size; + return addr; + } + + const size_t addr = cursor; + cursor += size; + peak = std::max(peak, cursor); + used += size; + return addr; } void Allocator::free(size_t addr, size_t size) @@ -41,9 +60,43 @@ namespace infini IT_ASSERT(this->ptr == nullptr); size = getAlignedSize(size); - // =================================== 作业 =================================== - // TODO: 设计一个算法来回收内存 - // =================================== 作业 =================================== + if (size == 0) + return; + IT_ASSERT(addr % alignment == 0); + IT_ASSERT(addr + size <= cursor); + IT_ASSERT(size <= used); + + size_t blockStart = addr; + size_t blockSize = size; + auto next = freeBlocks.lower_bound(addr); + if (next != freeBlocks.end()) + IT_ASSERT(addr + size <= next->first, "Free block overlaps"); + + if (next != freeBlocks.begin()) + { + auto prev = std::prev(next); + IT_ASSERT(prev->first + prev->second <= addr, + "Free block overlaps"); + if (prev->first + prev->second == addr) + { + blockStart = prev->first; + blockSize += prev->second; + freeBlocks.erase(prev); + } + } + + next = freeBlocks.lower_bound(blockStart); + if (next != freeBlocks.end() && blockStart + blockSize == next->first) + { + blockSize += next->second; + freeBlocks.erase(next); + } + + used -= size; + if (blockStart + blockSize == cursor) + cursor = blockStart; + else + freeBlocks.emplace(blockStart, blockSize); } void *Allocator::getPtr() @@ -58,6 +111,8 @@ namespace infini size_t Allocator::getAlignedSize(size_t size) { + if (size == 0) + return 0; return ((size - 1) / this->alignment + 1) * this->alignment; } diff --git a/src/core/graph.cc b/src/core/graph.cc index 3a906370..ba733b4b 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,131 @@ namespace infini void GraphObj::optimize() { - // =================================== 作业 =================================== - // TODO: 设计一个算法来实现指定的图优化规则 - // 图优化规则如下： - // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除） - // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去） - // =================================== 作业 =================================== + auto rebuildConnections = [this]() { + 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); + for (auto &output : op->outputs) + output->setSource(op); + } + for (auto &op : ops) { + for (auto &input : op->inputs) { + if (auto predecessor = input->getSource()) { + predecessor->addSuccessors(op); + op->addPredecessors(predecessor); + } + } + } + }; + + auto eraseOperator = [this](const Operator &target) { + ops.erase(std::remove(ops.begin(), ops.end(), target), ops.end()); + }; + auto eraseTensor = [this](const Tensor &target) { + tensors.erase(std::remove(tensors.begin(), tensors.end(), target), + tensors.end()); + }; + + // Cancel adjacent transpose pairs whose composed permutation is identity. + bool changed = true; + while (changed) { + changed = false; + rebuildConnections(); + for (const auto &candidate : ops) { + auto second = as(candidate); + if (!second) + continue; + const Tensor intermediate = second->getInputs(0); + auto first = as(intermediate->getSource()); + if (!first || intermediate->getTargets().size() != 1) + continue; + + const auto firstPerm = first->getPermute(); + const auto secondPerm = second->getPermute(); + if (firstPerm.size() != secondPerm.size()) + continue; + bool identity = true; + for (size_t i = 0; i < firstPerm.size(); ++i) { + if (firstPerm[secondPerm[i]] != static_cast(i)) { + identity = false; + break; + } + } + if (!identity) + continue; + + const Tensor original = first->getInputs(0); + const Tensor secondOutput = second->getOutput(); + for (auto &op : ops) { + if (op != first && op != second) + op->replaceInput(secondOutput, original); + } + eraseOperator(first); + eraseOperator(second); + eraseTensor(intermediate); + eraseTensor(secondOutput); + changed = true; + break; + } + } + + // Fold a last-two-axes transpose used only by MatMul into transA/transB. + changed = true; + while (changed) { + changed = false; + rebuildConnections(); + for (const auto &candidate : ops) { + auto matmul = as(candidate); + if (!matmul) + continue; + for (size_t inputId = 0; inputId < 2; ++inputId) { + const Tensor transposed = matmul->getInputs(inputId); + auto transpose = as(transposed->getSource()); + if (!transpose || transposed->getTargets().size() != 1) + continue; + + const auto perm = transpose->getPermute(); + const size_t rank = perm.size(); + if (rank < 2) + continue; + bool swapsLastTwo = true; + for (size_t i = 0; i + 2 < rank; ++i) + swapsLastTwo = swapsLastTwo && + perm[i] == static_cast(i); + swapsLastTwo = swapsLastTwo && + perm[rank - 2] == static_cast(rank - 1) && + perm[rank - 1] == static_cast(rank - 2); + if (!swapsLastTwo) + continue; + + if (inputId == 0) + matmul->setTransA(!matmul->getTransA()); + else + matmul->setTransB(!matmul->getTransB()); + matmul->replaceInput(transposed, transpose->getInputs(0)); + eraseOperator(transpose); + eraseTensor(transposed); + changed = true; + break; + } + if (changed) + break; + } + } + + rebuildConnections(); + sorted = false; + IT_ASSERT(topo_sort()); + shape_infer(); + IT_ASSERT(checkValid()); } Tensor GraphObj::getTensor(int fuid) const @@ -148,10 +269,44 @@ 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 (const auto &tensor : tensors) + remainingUses[tensor.get()] = tensor->getTargets().size(); + + auto allocateTensor = [&](const Tensor &tensor) { + if (offsets.find(tensor.get()) == offsets.end()) + offsets.emplace(tensor.get(), allocator.alloc(tensor->getBytes())); + }; + + // Graph inputs must exist before the first operator is executed. + for (const auto &tensor : tensors) { + if (!tensor->getSource()) + allocateTensor(tensor); + } + + for (const auto &op : ops) { + // Do not alias an output with one of the current operator's inputs. + for (const auto &output : op->getOutputs()) + allocateTensor(output); + + for (const auto &input : op->getInputs()) { + auto &uses = remainingUses.at(input.get()); + IT_ASSERT(uses > 0); + --uses; + if (uses == 0 && !input->getTargets().empty()) { + const auto offset = offsets.at(input.get()); + allocator.free(offset, input->getBytes()); + } + } + } + + unsigned char *base = static_cast(allocator.getPtr()); + for (const auto &tensor : tensors) { + const auto it = offsets.find(tensor.get()); + IT_ASSERT(it != offsets.end(), "Tensor was not assigned memory"); + tensor->setDataBlob(make_ref(runtime, base + it->second)); + } allocator.info(); } @@ -227,4 +382,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..681170ce 100644 --- a/src/operators/concat.cc +++ b/src/operators/concat.cc @@ -10,13 +10,22 @@ ConcatObj::ConcatObj(GraphObj *graph, TensorVec inputs, Tensor output, int _dim) } optional> ConcatObj::inferShape(const TensorVec &inputs) { + IT_ASSERT(!inputs.empty()); Shape dims = inputs[0]->getDims(); - auto rank = inputs[0]->getRank(); + const auto rank = inputs[0]->getRank(); + IT_ASSERT(dim >= 0 && static_cast(dim) < rank); - // =================================== 作业 =================================== - // TODO：修改 dims，返回正确的 concat 后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Concat.html#concat-13 - // =================================== 作业 =================================== + for (size_t inputId = 1; inputId < inputs.size(); ++inputId) { + const Shape current = inputs[inputId]->getDims(); + IT_ASSERT(current.size() == rank); + for (size_t axis = 0; axis < rank; ++axis) { + if (axis == static_cast(dim)) + dims[axis] += current[axis]; + else + IT_ASSERT(dims[axis] == current[axis], + "Concat dimensions do not match"); + } + } return {{dims}}; } diff --git a/src/operators/matmul.cc b/src/operators/matmul.cc index 7a16ca27..9bd92b6c 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 { @@ -23,11 +24,43 @@ namespace infini optional> MatmulObj::inferShape(const TensorVec &inputs) { - // =================================== 作业 =================================== - // TODO：返回经过 matmul 操作后的 shape - // REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm - // =================================== 作业 =================================== - return std::nullopt; + IT_ASSERT(inputs.size() == 2); + const Shape a = inputs[0]->getDims(); + const Shape b = inputs[1]->getDims(); + IT_ASSERT(!a.empty() && !b.empty()); + + const bool aVector = a.size() == 1; + const bool bVector = b.size() == 1; + IT_ASSERT(!(aVector && transA), "Cannot transpose a rank-1 MatMul input"); + IT_ASSERT(!(bVector && transB), "Cannot transpose a rank-1 MatMul input"); + + const int aM = aVector ? 1 + : (transA ? a.back() : a[a.size() - 2]); + const int aK = aVector ? a.back() + : (transA ? a[a.size() - 2] : a.back()); + const int bK = bVector ? b.front() + : (transB ? b.back() : b[b.size() - 2]); + const int bN = bVector ? 1 + : (transB ? b[b.size() - 2] : b.back()); + IT_ASSERT(aK == bK, "MatMul reduction dimensions do not match"); + + Shape aBatch; + Shape bBatch; + if (!aVector) + aBatch.assign(a.begin(), a.end() - 2); + if (!bVector) + bBatch.assign(b.begin(), b.end() - 2); + + Shape output = infer_broadcast(aBatch, bBatch); + if (!aVector) + output.push_back(aM); + if (!bVector) + output.push_back(bN); + + m = aM; + n = bN; + k = aK; + 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..344ff9a8 100644 --- a/src/operators/transpose.cc +++ b/src/operators/transpose.cc @@ -9,9 +9,10 @@ namespace infini auto rank = input->getRank(); if (permute.empty()) { + transposePermute.resize(rank); for (size_t i = 0; i < rank; ++i) { - transposePermute[i] = i; + transposePermute[i] = static_cast(rank - 1 - i); } } else @@ -19,6 +20,14 @@ namespace infini IT_ASSERT(rank == permute.size()); transposePermute = std::move(permute); } + + vector seen(rank, false); + for (int axis : transposePermute) + { + IT_ASSERT(axis >= 0 && static_cast(axis) < rank); + IT_ASSERT(!seen[axis], "Transpose permutation contains duplicates"); + seen[axis] = true; + } IT_ASSERT(checkValid(graph)); } @@ -29,12 +38,10 @@ namespace infini auto output_dim = input_dim; int rank = A->getRank(); - // =================================== 作业 =================================== - // TODO：修改 output_dim，返回正确的 transpose 后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Transpose.html#transpose-21 - // =================================== 作业 =================================== + 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..21130a86 100644 --- a/src/operators/unary.cc +++ b/src/operators/unary.cc @@ -35,11 +35,7 @@ namespace infini optional> ClipObj::inferShape(const TensorVec &inputs) { - // =================================== 作业 =================================== - // TODO：返回经过 clip 操作后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Clip.html#clip-13 - // =================================== 作业 =================================== - return std::nullopt; + return {{inputs[0]->getDims()}}; } std::string ClipObj::toString() const @@ -61,21 +57,13 @@ 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 - // =================================== 作业 =================================== - return {}; + IT_ASSERT(inputs.size() == 1); + return {getOutputDataType()}; } optional> CastObj::inferShape(const TensorVec &inputs) { - // =================================== 作业 =================================== - // TODO：返回经过 cast 操作后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Cast.html#cast-21 - // =================================== 作业 =================================== - return std::nullopt; + 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..326bb708 100644 --- a/src/utils/operator_utils.cc +++ b/src/utils/operator_utils.cc @@ -4,13 +4,22 @@ namespace infini { Shape infer_broadcast(const Shape &A, const Shape &B) { + const size_t rank = std::max(A.size(), B.size()); + Shape result(rank, 1); - // =================================== 作业 =================================== - // TODO：对 A 和 B 进行双向广播，返回广播后的形状。 - // REF: https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md - // =================================== 作业 =================================== - - return {}; + for (size_t offset = 0; offset < rank; ++offset) { + const int a = offset < A.size() ? A[A.size() - 1 - offset] : 1; + const int b = offset < B.size() ? B[B.size() - 1 - offset] : 1; + IT_ASSERT(a == b || a == 1 || b == 1, + "Shapes are not broadcast-compatible"); + + if (a == 1) + result[rank - 1 - offset] = b; + else + result[rank - 1 - offset] = a; + } + + return result; } int get_real_axis(const int &axis, const int &rank) {