Implement shape inference and optimize graph memory planning

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<size_t, size_t> freeBlocks;
 
   public:
     Allocator(Runtime runtime);

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()

src/core/graph.cc

@@ -1,4 +1,6 @@
 #include "core/graph.h"
+#include "operators/matmul.h"
+#include "operators/transpose.h"
 #include <algorithm>
 #include <numeric>
 #include <queue>
@@ -100,12 +102,141 @@ namespace infini
 
     void GraphObj::optimize()
     {
-        // =================================== 作业 ===================================
-        // TODO: 设计一个算法来实现指定的图优化规则
-        // 图优化规则如下：
-        // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除）
-        // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去）
-        // =================================== 作业 ===================================
+        IT_ASSERT(topo_sort());
+
+        std::unordered_set<OperatorObj *> removedOps;
+        std::unordered_set<TensorObj *> removedTensors;
+
+        auto permutationsCancel = [](const vector<int> &first,
+                                     const vector<int> &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<size_t>(intermediate) >= first.size() ||
+                    first[intermediate] != static_cast<int>(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<TransposeObj>(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<TransposeObj>(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<int> &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<int>(axis))
+                    return false;
+            return permute[rank - 2] == static_cast<int>(rank - 1) &&
+                   permute[rank - 1] == static_cast<int>(rank - 2);
+        };
+
+        // Fold a single-use last-two-axis transpose into each MatMul input.
+        for (const auto &op : ops)
+        {
+            auto matmul = as<MatmulObj>(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<TransposeObj>(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<TensorObj *, size_t> offsets;
+        std::unordered_map<TensorObj *, size_t> 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<uint8_t *>(allocator.getPtr());
+        for (const auto &tensor : tensors)
+        {
+            auto it = offsets.find(tensor.get());
+            IT_ASSERT(it != offsets.end());
+            tensor->setDataBlob(
+                make_ref<BlobObj>(runtime, static_cast<void *>(base + it->second)));
+        }
 
         allocator.info();
     }
@@ -227,4 +389,4 @@ namespace infini
         return true;
     }
 
-} // namespace infini
+} // namespace infini

src/operators/concat.cc

@@ -10,13 +10,27 @@ ConcatObj::ConcatObj(GraphObj *graph, TensorVec inputs, Tensor output, int _dim)
 }
 
 optional<vector<Shape>> 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<size_t>(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<int>(axis) != dim && inputDims[axis] != dims[axis])
+                return std::nullopt;
+        }
+        dims[dim] += inputDims[dim];
+    }
 
     return {{dims}};
 }

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<vector<Shape>> 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
+} // namespace infini