diff --git a/include/core/allocator.h b/include/core/allocator.h index 002601d2..2a601acf 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 keyed by their starting offset. Keeping the map ordered makes + // adjacent-block coalescing deterministic. + std::map freeBlocks; public: Allocator(Runtime runtime); diff --git a/src/core/allocator.cc b/src/core/allocator.cc index ff593aef..49b36992 100644 --- a/src/core/allocator.cc +++ b/src/core/allocator.cc @@ -26,24 +26,91 @@ namespace infini size_t Allocator::alloc(size_t size) { IT_ASSERT(this->ptr == nullptr); + IT_ASSERT(size > 0); // pad the size to the multiple of alignment size = this->getAlignedSize(size); - // =================================== 作业 =================================== - // TODO: 设计一个算法来分配内存,返回起始地址偏移量 - // =================================== 作业 =================================== + size_t addr = 0; + auto fit = freeBlocks.end(); + for (auto it = freeBlocks.begin(); it != freeBlocks.end(); ++it) + { + if (it->second >= size) + { + fit = it; + break; + } + } - return 0; + if (fit != freeBlocks.end()) + { + addr = fit->first; + const size_t blockSize = fit->second; + freeBlocks.erase(fit); + if (blockSize > size) + freeBlocks.emplace(addr + size, blockSize - size); + } + else if (!freeBlocks.empty()) + { + auto tail = std::prev(freeBlocks.end()); + if (tail->first + tail->second == peak) + { + addr = tail->first; + const size_t available = tail->second; + freeBlocks.erase(tail); + peak += size - available; + } + else + { + addr = peak; + peak += size; + } + } + else + { + addr = peak; + peak += size; + } + + used += size; + return addr; } void Allocator::free(size_t addr, size_t size) { IT_ASSERT(this->ptr == nullptr); + IT_ASSERT(size > 0); size = getAlignedSize(size); + const size_t releasedSize = size; + IT_ASSERT(addr % alignment == 0); + IT_ASSERT(addr <= peak && size <= peak - addr); + IT_ASSERT(used >= size); + + auto next = freeBlocks.lower_bound(addr); + if (next != freeBlocks.end()) + IT_ASSERT(addr + size <= next->first, + "Freed block overlaps an existing free block"); - // =================================== 作业 =================================== - // TODO: 设计一个算法来回收内存 - // =================================== 作业 =================================== + if (next != freeBlocks.begin()) + { + auto previous = std::prev(next); + IT_ASSERT(previous->first + previous->second <= addr, + "Freed block overlaps an existing free block"); + if (previous->first + previous->second == addr) + { + addr = previous->first; + size += previous->second; + freeBlocks.erase(previous); + } + } + + next = freeBlocks.lower_bound(addr); + if (next != freeBlocks.end() && addr + size == next->first) + { + size += next->second; + freeBlocks.erase(next); + } + freeBlocks.emplace(addr, size); + used -= releasedSize; } void *Allocator::getPtr() diff --git a/src/core/graph.cc b/src/core/graph.cc index 3a906370..deec5732 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,141 @@ namespace infini void GraphObj::optimize() { - // =================================== 作业 =================================== - // TODO: 设计一个算法来实现指定的图优化规则 - // 图优化规则如下: - // 1. 去除冗余的算子(例如,两个相邻的算子都是 transpose 算子,且做的是相反的操作,可以将其全部删除) - // 2. 合并算子(例如,矩阵乘算子中含有属性transA、transB,如果其输入存在transpose,且对最后两个维度做交换,就可以将transpose融入到矩阵乘算子的属性中去) - // =================================== 作业 =================================== + IT_ASSERT(topo_sort()); + + std::unordered_set removedOps; + std::unordered_set removedTensors; + + auto permutationsCancel = [](const vector &first, + const vector &second) + { + if (first.size() != second.size()) + return false; + for (size_t axis = 0; axis < first.size(); ++axis) + { + const int intermediate = second[axis]; + if (intermediate < 0 || + static_cast(intermediate) >= first.size() || + first[intermediate] != static_cast(axis)) + return false; + } + return true; + }; + + // Remove adjacent transpose pairs whose composed permutation is the + // identity. Restrict the rewrite to internal tensors so graph outputs keep + // their object identity. + for (const auto &op : ops) + { + auto first = as(op); + if (!first || removedOps.count(first.get())) + continue; + + const Tensor firstOutput = first->getOutput(); + const OpVec firstTargets = firstOutput->getTargets(); + if (firstTargets.size() != 1) + continue; + + auto second = as(firstTargets[0]); + if (!second || removedOps.count(second.get()) || + second->getInputs(0) != firstOutput) + continue; + + const Tensor secondOutput = second->getOutput(); + const OpVec consumers = secondOutput->getTargets(); + if (consumers.empty() || + !permutationsCancel(first->getPermute(), second->getPermute())) + continue; + + const Tensor replacement = first->getInputs(0); + for (const auto &consumer : consumers) + consumer->replaceInput(secondOutput, replacement); + + removedOps.insert(first.get()); + removedOps.insert(second.get()); + removedTensors.insert(firstOutput.get()); + removedTensors.insert(secondOutput.get()); + } + + auto swapsLastTwoAxes = [](const vector &permute) + { + const size_t rank = permute.size(); + if (rank < 2) + return false; + for (size_t axis = 0; axis + 2 < rank; ++axis) + if (permute[axis] != static_cast(axis)) + return false; + return permute[rank - 2] == static_cast(rank - 1) && + permute[rank - 1] == static_cast(rank - 2); + }; + + // Fold a single-use last-two-axis transpose into each MatMul input. + for (const auto &op : ops) + { + auto matmul = as(op); + if (!matmul || removedOps.count(matmul.get())) + continue; + + for (size_t inputIndex = 0; inputIndex < 2; ++inputIndex) + { + const Tensor input = matmul->getInputs(inputIndex); + auto transpose = as(input->getSource()); + if (!transpose || removedOps.count(transpose.get()) || + !swapsLastTwoAxes(transpose->getPermute())) + continue; + + const OpVec consumers = input->getTargets(); + if (consumers.size() != 1 || consumers[0] != matmul) + continue; + + matmul->replaceInput(input, transpose->getInputs(0)); + if (inputIndex == 0) + matmul->setTransA(!matmul->getTransA()); + else + matmul->setTransB(!matmul->getTransB()); + + removedOps.insert(transpose.get()); + removedTensors.insert(input.get()); + } + } + + ops.erase(std::remove_if(ops.begin(), ops.end(), [&](const Operator &op) + { return removedOps.count(op.get()) != 0; }), + ops.end()); + tensors.erase( + std::remove_if(tensors.begin(), tensors.end(), [&](const Tensor &tensor) + { return removedTensors.count(tensor.get()) != 0; }), + tensors.end()); + + // Rebuild all bidirectional relationships from the retained operators. + for (const auto &tensor : tensors) + { + tensor->targets.clear(); + tensor->source.reset(); + } + for (const auto &op : ops) + { + op->predecessors.clear(); + op->successors.clear(); + } + for (const auto &op : ops) + { + for (const auto &input : op->inputs) + { + input->targets.emplace_back(op); + if (auto predecessor = input->source.lock()) + { + predecessor->successors.emplace_back(op); + op->predecessors.emplace_back(predecessor); + } + } + for (const auto &output : op->outputs) + output->source = op; + } + + sorted = false; + IT_ASSERT(topo_sort()); + IT_ASSERT(checkValid()); } Tensor GraphObj::getTensor(int fuid) const @@ -148,10 +279,41 @@ 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.emplace(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())); + }; + + for (const auto &op : ops) + { + for (const auto &input : op->getInputs()) + allocateTensor(input); + for (const auto &output : op->getOutputs()) + allocateTensor(output); + + for (const auto &input : op->getInputs()) + { + auto &uses = remainingUses.at(input.get()); + IT_ASSERT(uses > 0); + if (--uses == 0) + allocator.free(offsets.at(input.get()), input->getBytes()); + } + } + + auto *base = static_cast(allocator.getPtr()); + for (const auto &tensor : tensors) + { + auto it = offsets.find(tensor.get()); + IT_ASSERT(it != offsets.end()); + tensor->setDataBlob( + make_ref(runtime, static_cast(base + it->second))); + } allocator.info(); } @@ -227,4 +389,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..d1f4043e 100644 --- a/src/operators/concat.cc +++ b/src/operators/concat.cc @@ -10,13 +10,27 @@ ConcatObj::ConcatObj(GraphObj *graph, TensorVec inputs, Tensor output, int _dim) } optional> ConcatObj::inferShape(const TensorVec &inputs) { + if (inputs.empty()) + return std::nullopt; + Shape dims = inputs[0]->getDims(); auto rank = inputs[0]->getRank(); - // =================================== 作业 =================================== - // TODO:修改 dims,返回正确的 concat 后的 shape - // REF: https://onnx.ai/onnx/operators/onnx__Concat.html#concat-13 - // =================================== 作业 =================================== + if (dim < 0 || static_cast(dim) >= rank) + return std::nullopt; + + dims[dim] = 0; + for (const auto &input : inputs) { + const auto inputDims = input->getDims(); + if (inputDims.size() != rank) + return std::nullopt; + + for (size_t axis = 0; axis < rank; ++axis) { + if (static_cast(axis) != dim && inputDims[axis] != dims[axis]) + return std::nullopt; + } + dims[dim] += inputDims[dim]; + } return {{dims}}; } diff --git a/src/operators/matmul.cc b/src/operators/matmul.cc index 7a16ca27..8c9ea16c 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,33 @@ 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; + if (inputs.size() != 2) + return std::nullopt; + + const Shape aDims = inputs[0]->getDims(); + const Shape bDims = inputs[1]->getDims(); + if (aDims.size() < 2 || bDims.size() < 2) + return std::nullopt; + + const size_t aRank = aDims.size(); + const size_t bRank = bDims.size(); + const int aRows = transA ? aDims[aRank - 1] : aDims[aRank - 2]; + const int aCols = transA ? aDims[aRank - 2] : aDims[aRank - 1]; + const int bRows = transB ? bDims[bRank - 1] : bDims[bRank - 2]; + const int bCols = transB ? bDims[bRank - 2] : bDims[bRank - 1]; + if (aCols != bRows) + return std::nullopt; + + Shape aBatch(aDims.begin(), aDims.end() - 2); + Shape bBatch(bDims.begin(), bDims.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..08e526d8 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 - i - 1); } } else @@ -29,12 +30,20 @@ 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 - // =================================== 作业 =================================== + if (static_cast(transposePermute.size()) != rank) + return std::nullopt; - return std::nullopt; + vector seen(rank, false); + for (int i = 0; i < rank; ++i) + { + const int axis = transposePermute[i]; + if (axis < 0 || axis >= rank || seen[axis]) + return std::nullopt; + 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..5138a9d8 100644 --- a/src/operators/unary.cc +++ b/src/operators/unary.cc @@ -35,11 +35,9 @@ 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; + if (inputs.size() != 1) + return std::nullopt; + return {{inputs[0]->getDims()}}; } std::string ClipObj::toString() const @@ -61,21 +59,16 @@ 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 {}; + if (inputs.size() != 1) + return {}; + 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; + if (inputs.size() != 1) + 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..d33974ac 100644 --- a/src/utils/operator_utils.cc +++ b/src/utils/operator_utils.cc @@ -4,13 +4,24 @@ 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); + const size_t aOffset = rank - A.size(); + const size_t bOffset = rank - B.size(); - // =================================== 作业 =================================== - // TODO:对 A 和 B 进行双向广播,返回广播后的形状。 - // REF: https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md - // =================================== 作业 =================================== - - return {}; + for (size_t axis = 0; axis < rank; ++axis) { + const ShapeElem a = axis < aOffset ? 1 : A[axis - aOffset]; + const ShapeElem b = axis < bOffset ? 1 : B[axis - bOffset]; + + IT_ASSERT(a == b || a == 1 || b == 1, + "Shapes cannot be broadcast together"); + if (a == b) + result[axis] = a; + else + result[axis] = a == 1 ? b : a; + } + + return result; } int get_real_axis(const int &axis, const int &rank) {