complete the proj

include/core/allocator.h

@@ -27,7 +27,10 @@ namespace infini {
     // TODO：可能需要设计一个数据结构来存储free block，以便于管理和合并
     // HINT: 可以使用一个 map 来存储 free block，key 为 block 的起始/结尾地址，value 为 block 的大小
     // =================================== 作业 ===================================
-
+    //<address, blocksize>
+    std::map<size_t, size_t> free_blocks;
+
+
   public:
     Allocator(Runtime runtime);
 
@@ -55,5 +58,11 @@ namespace infini {
     // function: memory alignment, rouned up
     // return: size of the aligned memory block
     size_t getAlignedSize(size_t size);
+
+    // function: merge adjacent free blocks
+    // arguments:
+    //     addr: address of the newly freed block
+    //     size: size of the newly freed block
+    void mergeAdjacentBlocks(size_t addr, size_t size);
   };
 }

include/core/graph.h

@@ -2,11 +2,14 @@
 #include "core/allocator.h"
 #include "core/operator.h"
 #include "core/tensor.h"
+#include "operators/transpose.h"
 #include <algorithm>
 #include <cstdint>
 
 namespace infini
 {
+    // 前向声明
+    class MatmulObj;
 
     class GraphObj : public Object
     {
@@ -27,6 +30,36 @@ namespace infini
         TensorVec addTensor(const TensorVec &tensors);
         void removeOperator(Operator op)
         {
+            // 清理操作符与张量的连接关系
+            for (auto& input : op->getInputs()) {
+                if (input) {
+                    input->removeTarget(op);
+                }
+            }
+            for (auto& output : op->getOutputs()) {
+                if (output) {
+                    output->setSource(nullptr);
+                }
+            }
+
+            // 清理操作符之间的连接关系
+            // 从所有前驱操作符中移除对当前操作符的引用
+            auto predecessors = op->getPredecessors();
+            for (const auto& pred : predecessors) {
+                if (pred) {
+                    pred->removeSuccessors(op);
+                }
+            }
+
+            // 从所有后继操作符中移除对当前操作符的引用
+            auto successors = op->getSuccessors();
+            for (const auto& succ : successors) {
+                if (succ) {
+                    succ->removePredecessors(op);
+                }
+            }
+
+            // 从操作符列表中删除
             auto it = std::find(ops.begin(), ops.end(), op);
             if (it != ops.end())
                 ops.erase(it);
@@ -116,6 +149,36 @@ namespace infini
          * @brief If the nodes is sorted in topological order.
          */
         bool sorted;
+
+        /**
+         * @brief Add check function for inverse transpose
+         */
+        bool areInverseTransposes(const TransposeObj *transpose1, const TransposeObj *transpose2);
+
+        /**
+         * @brief Add check function for same transpose
+         */
+        bool areSameTransposes(const TransposeObj *transpose1, const TransposeObj *transpose2);
+
+        /**
+         * @brief Check if transpose swaps last two dimensions
+         */
+        bool isLastTwoDimsSwap(const TransposeObj *transpose);
+
+        /**
+         * @brief Merge transpose into matmul operator
+         */
+        void mergeTransposeToMatmul(const Operator& transpose, const Operator& matmul);
+
+        /**
+         * @brief Reconnect graph after removing operators
+         */
+        void reconnectGraph(const Operator& op1, const Operator& op2);
+
+        /**
+         * @brief Clean up unused tensors
+         */
+        void cleanupUnusedTensors();
     };
 
 } // namespace infini

src/core/allocator.cc

@@ -28,12 +28,48 @@ namespace infini
         IT_ASSERT(this->ptr == nullptr);
         // pad the size to the multiple of alignment
         size = this->getAlignedSize(size);
-
+
+
         // =================================== 作业 ===================================
         // TODO: 设计一个算法来分配内存，返回起始地址偏移量
         // =================================== 作业 ===================================
 
-        return 0;
+        // 使用First Fit算法查找合适的空闲块
+        for (auto it = free_blocks.begin(); it != free_blocks.end(); ++it) {
+            size_t block_addr = it->first;
+            size_t block_size = it->second;
+
+            if (block_size >= size) {
+                // 找到合适的块，进行分配
+                size_t remaining_size = block_size - size;
+
+                if (remaining_size > 0) {
+                    // 如果剩余空间足够大，创建新的空闲块
+                    free_blocks[block_addr + size] = remaining_size;
+                }
+
+                // 移除或更新当前块
+                free_blocks.erase(it);
+
+                // 更新使用统计
+                used += size;
+                if (used > peak) {
+                    peak = used;
+                }
+
+                return block_addr;
+            }
+        }
+
+        // 如果没有找到合适的空闲块，需要扩展内存
+        // 这里简化处理，直接返回当前已使用的大小作为新地址
+        size_t new_addr = used;
+        used += size;
+        if (used > peak) {
+            peak = used;
+        }
+
+        return new_addr;
     }
 
     void Allocator::free(size_t addr, size_t size)
@@ -44,6 +80,15 @@ namespace infini
         // =================================== 作业 ===================================
         // TODO: 设计一个算法来回收内存
         // =================================== 作业 ===================================
+
+        // 将释放的内存块添加到空闲块映射中
+        free_blocks[addr] = size;
+
+        // 尝试与相邻的空闲块合并
+        mergeAdjacentBlocks(addr, size);
+
+        // 更新使用统计
+        used -= size;
     }
 
     void *Allocator::getPtr()
@@ -61,6 +106,38 @@ namespace infini
         return ((size - 1) / this->alignment + 1) * this->alignment;
     }
 
+    void Allocator::mergeAdjacentBlocks(size_t addr, size_t size)
+    {
+        // 查找前一个相邻的空闲块
+        auto prev_it = free_blocks.find(addr - 1);
+        if (prev_it != free_blocks.end()) {
+            // 找到前一个块，检查是否真的相邻
+            size_t prev_addr = prev_it->first;
+            size_t prev_size = prev_it->second;
+
+            if (prev_addr + prev_size == addr) {
+                // 可以与前一个块合并
+                size_t new_size = prev_size + size;
+                free_blocks[prev_addr] = new_size;
+                free_blocks.erase(addr);  // 移除当前块
+
+                // 更新参数，继续检查是否可以与后一个块合并
+                addr = prev_addr;
+                size = new_size;
+            }
+        }
+
+        // 查找后一个相邻的空闲块
+        auto next_it = free_blocks.find(addr + size);
+        if (next_it != free_blocks.end()) {
+            // 找到后一个块，可以合并
+            size_t next_size = next_it->second;
+            size_t new_size = size + next_size;
+            free_blocks[addr] = new_size;
+            free_blocks.erase(addr + size);  // 移除后一个块
+        }
+    }
+
     void Allocator::info()
     {
         std::cout << "Used memory: " << this->used