Complete tinyinfinitensor assignments

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

src/core/allocator.cc

@@ -29,21 +29,73 @@ 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)
     {
         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()

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,96 @@ namespace infini
 
     void GraphObj::optimize()
     {
-        // =================================== 作业 ===================================
-        // TODO: 设计一个算法来实现指定的图优化规则
-        // 图优化规则如下：
-        // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除）
-        // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去）
-        // =================================== 作业 ===================================
+        auto inversePermutations = [](const vector<int> &a,
+                                      const vector<int> &b) {
+            if (a.size() != b.size())
+                return false;
+            for (size_t i = 0; i < a.size(); ++i)
+                if (a[b[i]] != static_cast<int>(i))
+                    return false;
+            return true;
+        };
+        auto swapsLastTwo = [](const vector<int> &perm) {
+            if (perm.size() < 2)
+                return false;
+            for (size_t i = 0; i + 2 < perm.size(); ++i)
+                if (perm[i] != static_cast<int>(i))
+                    return false;
+            return perm[perm.size() - 2] == static_cast<int>(perm.size() - 1) &&
+                   perm[perm.size() - 1] == static_cast<int>(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<TransposeObj>(second);
+                if (!t2)
+                    continue;
+                auto middle = t2->getInputs(0);
+                auto first = as<TransposeObj>(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<MatmulObj>(op);
+            if (!matmul)
+                continue;
+            for (size_t inputIndex = 0; inputIndex < 2; ++inputIndex) {
+                auto transposed = matmul->getInputs(inputIndex);
+                auto transpose = as<TransposeObj>(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<TensorObj *, size_t> offsets;
+        std::unordered_map<TensorObj *, size_t> 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<uint8_t *>(allocator.getPtr());
+        for (auto &tensor : tensors)
+            tensor->setDataBlob(
+                make_ref<BlobObj>(runtime, base + offsets.at(tensor.get())));
 
         allocator.info();
     }
@@ -227,4 +334,4 @@ namespace infini
         return true;
     }
 
-} // namespace infini
+} // namespace infini

src/operators/concat.cc

@@ -16,8 +16,15 @@ optional<vector<Shape>> 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<int>(axis) != dim)
+                IT_ASSERT(current[axis] == dims[axis]);
+        dims[dim] += current[dim];
+    }
     return {{dims}};
 }
 

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

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<int>(transposePermute.size()) == rank);
+        vector<bool> 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

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<vector<Shape>> 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

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