first complete assignments

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<size_t, size_t> freeBlocks;
+    // The current end of the simulated heap. `peak` is its high-water mark.
+    size_t cursor;
 
   public:
     Allocator(Runtime runtime);

src/core/allocator.cc

@@ -1,4 +1,5 @@
 #include "core/allocator.h"
+#include <algorithm>
 #include <utility>
 
 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,21 +31,72 @@ 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)
     {
         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;
     }
 

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,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<TransposeObj>(candidate);
+                if (!second)
+                    continue;
+                const Tensor intermediate = second->getInputs(0);
+                auto first = as<TransposeObj>(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<int>(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<MatmulObj>(candidate);
+                if (!matmul)
+                    continue;
+                for (size_t inputId = 0; inputId < 2; ++inputId) {
+                    const Tensor transposed = matmul->getInputs(inputId);
+                    auto transpose = as<TransposeObj>(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<int>(i);
+                    swapsLastTwo = swapsLastTwo &&
+                                       perm[rank - 2] == static_cast<int>(rank - 1) &&
+                                       perm[rank - 1] == static_cast<int>(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<TensorObj *, size_t> offsets;
+        std::unordered_map<TensorObj *, size_t> 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<unsigned char *>(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<BlobObj>(runtime, base + it->second));
+        }
 
         allocator.info();
     }
@@ -227,4 +382,4 @@ namespace infini
         return true;
     }
 
-} // namespace infini
+} // namespace infini

src/operators/concat.cc

@@ -10,13 +10,22 @@ ConcatObj::ConcatObj(GraphObj *graph, TensorVec inputs, Tensor output, int _dim)
 }
 
 optional<vector<Shape>> 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<size_t>(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<size_t>(dim))
+                dims[axis] += current[axis];
+            else
+                IT_ASSERT(dims[axis] == current[axis],
+                          "Concat dimensions do not match");
+        }
+    }
 
     return {{dims}};
 }