Ai compiler assignment submit 2 from Christina

include/core/allocator.h

@@ -27,7 +27,7 @@ namespace infini {
     // TODO：可能需要设计一个数据结构来存储free block，以便于管理和合并
     // HINT: 可以使用一个 map 来存储 free block，key 为 block 的起始/结尾地址，value 为 block 的大小
     // =================================== 作业 ===================================
-
+    std::map<size_t, size_t> freeBlocksMap;
   public:
     Allocator(Runtime runtime);
 

include/core/graph.h

@@ -50,6 +50,7 @@ namespace infini
          * so the topological sorting fails.
          */
         bool topo_sort();
+        void reconstruct(Operator &op1, Operator &op2, Operator &op3);
 
         void optimize();
 

src/core/allocator.cc

@@ -32,8 +32,31 @@ namespace infini
         // =================================== 作业 ===================================
         // TODO: 设计一个算法来分配内存，返回起始地址偏移量
         // =================================== 作业 ===================================
+        used += size;
+        if(used > peak)
+            peak = used;
 
-        return 0;
+        if(freeBlocksMap.empty()){
+            return peak - size;
+        }
+
+        bool freeBlkOk = false;  //是否存在足够大（>= size）的空闲内存块
+        auto it = freeBlocksMap.begin();
+        for(; it != freeBlocksMap.end(); it++){
+            if(it->second >= size){
+                freeBlkOk = true;
+                if(it->second > size)
+                     freeBlocksMap.insert({it->first + size, it->second - size});
+                break;
+            }
+        }
+        if(!freeBlkOk){
+            it--;
+            size_t moreMemory = size - it->second;
+            peak += moreMemory;
+        }
+        freeBlocksMap.erase(it);
+        return it->first;
     }
 
     void Allocator::free(size_t addr, size_t size)
@@ -44,6 +67,30 @@ namespace infini
         // =================================== 作业 ===================================
         // TODO: 设计一个算法来回收内存
         // =================================== 作业 ===================================
+        auto it = freeBlocksMap.begin();
+        int flag = 0;
+        for(; it != freeBlocksMap.end(); it++){   //遍历空闲内存块，检查是否能与待回收内存合并成一个大内存块
+            if(it->first + it->second == addr){	  //可以与待回收内存合并，该空闲块在前，待回收内存在后
+                flag = -1;
+                break;
+            }
+            else if(it->first == addr + size){	 //可以与待回收内存合并，待回收内存在前，该空闲块在后，
+                flag = 1;
+                break;
+            }
+        }
+        if(flag == -1){
+            freeBlocksMap.insert({it->first, it->second + size}); //因为可以合并，插入合并后的大内存块，删去原来内存块
+            freeBlocksMap.erase(it);
+        }
+        else if(flag == 0){					//标记为0，说明以上遍历时没有找到能合并的空闲块
+            freeBlocksMap.insert({addr, size});
+        }
+        else if(flag == 1){
+            freeBlocksMap.insert({addr, size + it->second});
+            freeBlocksMap.erase(it);
+        }
+        used -= size;
     }
 
     void *Allocator::getPtr()

src/core/graph.cc

@@ -2,10 +2,13 @@
 #include <algorithm>
 #include <numeric>
 #include <queue>
+#include "operators/transpose.h"
+#include "operators/matmul.h"
+#include <unordered_set>
+#include <unordered_map>
 
 namespace infini
 {
-
     void GraphObj::addOperatorAndConnect(const Operator &op)
     {
         sorted = false;
@@ -98,15 +101,128 @@ namespace infini
         return this->sorted = true;
     }
 
-    void GraphObj::optimize()
-    {
+    bool transposeOpsCancel(vector<int> a, vector<int> b){
+    	if(a.size() != b.size())
+    		return false;
+    	for(int i = 0; i < (int)a.size(); i++){
+    		if(b[a[i]] != i)
+    			return false;
+    	}
+    	return true;
+    }
+
+    bool transOpCanIntegrateToMatmul(vector<int> perm){
+        int size = perm.size();
+        for(int i = 0; i < size - 2; i++){
+            if(perm[i] != i)
+                return false;
+        }
+        if(perm[size - 2] != size - 1 || perm[size - 1] != size - 2)
+            return false;
+        else
+            return true;
+    }
+
+    void GraphObj::reconstruct(Operator &op1, Operator &op2, Operator &op3){
+        Tensor input = op1->getInputs(0);
+        if(input) {input->addTarget(op3);}
+        if(op2 == nullptr){
+            op3->replaceInput(op1->getOutput(), input);
+            op1->removeSuccessors(op3);
+            op3->removePredecessors(op1);
+            op1->getOutput()->removeTarget(op3);
+        }
+        else{
+            op3->replaceInput(op2->getOutput(), input);
+            op3->removePredecessors(op2);
+            op1->removeSuccessors(op2);
+            op1->getOutput()->removeTarget(op2);
+        }
+        for(auto &pred: op1->getPredecessors()){
+            op3->addPredecessors(pred);
+            pred->addSuccessors(op3);
+        }
+    }
+
+void GraphObj::optimize()
+{
         // =================================== 作业 ===================================
         // TODO: 设计一个算法来实现指定的图优化规则
         // 图优化规则如下：
         // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除）
         // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去）
         // =================================== 作业 ===================================
+    std::unordered_set<OperatorObj *> toDelete;    //待删除的算子
+    std::shared_ptr<OperatorObj> emptyPtr;         //有的情况下重构计算图时需要
+    bool modified = true;
+    while(modified){
+        modified = false;    
+        for(auto &op: ops){
+            if(toDelete.find(op.get()) != toDelete.end()){  //如果遍历到之前发现的待删除算子，略过
+                continue;
+            }
+            OpType opType = op->getOpType();
+            if(opType == OpType::Transpose){           
+                TransposeObj* transOp = dynamic_cast<TransposeObj*>(op.get());               
+                Tensor input = op->getInputs(0);
+                for(auto &succ: op->getSuccessors()){
+                    if(succ->getOpType() == OpType::Transpose){  //优化情况1 去除冗余transpose算子:当两个相邻算子都是transpose且相反操作
+                        TransposeObj* transSucc = dynamic_cast<TransposeObj*>(succ.get());
+                        bool cancelOut = false;                 //判断两个相邻transpose算子是否相反的操作，在transposeOpsCancel函数中进行
+                        if(transOp && transSucc){
+                            cancelOut = transposeOpsCancel(transOp->getPermute(), transSucc->getPermute());
+                        }                       
+                        if(cancelOut){
+                            toDelete.insert(succ.get());
+                            removeTensor(succ->getOutput());
+                            for(auto &succ_succ: succ->getSuccessors()){
+                                reconstruct(op, succ, succ_succ);
+                            }
+                            if(op->getSuccessors().size() == 0){    //如果第1个transpose只有一个后继算子(即第2个transpose),那么可以删除第1个transpose
+                                toDelete.insert(op.get());
+                                //删除第1个transpose时需要进行以下重构
+                                removeTensor(op->getOutput());
+                                if(input) {input->removeTarget(op);}
+                                for(auto &pred: op->getPredecessors())
+                                    pred->removeSuccessors(op);
+                            }
+                            modified = true;
+                        }
+                    }
+                    else if(succ->getOpType() == OpType::MatMul){    //优化情况2 合并transpose与MatMul:当MatMul算子中含有属性transA、transB，且transpose对最后两个维度做交换
+                        bool ok = transOpCanIntegrateToMatmul(transOp->getPermute());  //判断能否合并
+                        if(ok){                                      //即使能够合并，也不一定就能删除transpose，除非transpose只有matMul唯一一个后继结点
+                            MatmulObj* matmulSucc = dynamic_cast<MatmulObj*>(succ.get());
+                            if (matmulSucc){
+                                if(succ->getInputs(0)->getGuid() == op->getOutput()->getGuid())
+                                    matmulSucc->setTransA(!(matmulSucc->getTransA()));
+                                else
+                                    matmulSucc->setTransB(!(matmulSucc->getTransB()));
+                            }
+                            reconstruct(op, emptyPtr, succ);        //合并transpose与matMul时需要进行重构：将transpose的输入作为matMul的输入    
+                            if(op->getSuccessors().size() == 0){    //如果transpose只有一个后继算子(即matMul)那么可以删除transpose。 size()==0是因为前面一步重构时已删去了两算子间的边
+                                toDelete.insert(op.get());          
+                                removeTensor(op->getOutput());
+                                if(input) {input->removeTarget(op);}
+                                for(auto &pred: op->getPredecessors())
+                                    pred->removeSuccessors(op);
+                            }
+                            modified = true;
+                        }
+                    }
+                }
+            }
+        }
+    }
+    //删除之前遍历计算图时发现的冗余算子
+    for(int i = 0; i < (int)ops.size();){
+        if(toDelete.find(ops[i].get()) != toDelete.end())
+            ops.erase(ops.begin() + i);
+        else
+            i++;
     }
+    this->sorted = false;
+}
 
     Tensor GraphObj::getTensor(int fuid) const
     {
@@ -152,10 +268,50 @@ namespace infini
         // TODO：利用 allocator 给计算图分配内存
         // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存
         // =================================== 作业 ===================================
+        std::unordered_map<int, size_t> tensorOffset;  //tensor在内存中的偏移量。 tensor id作为key,偏移量作为value
+        std::unordered_map<int, int> tensorRefNum;    //tensor被引用次数，即被多少个算子使用。tensor id作为key,被引用次数作为value
+        //记录每个张量的被引用次数
+        for(auto &tensor: tensors){
+            tensorRefNum[tensor->getFuid()] = static_cast<int>(tensor->getTargets().size());
+        }
+        //为input分配内存
+        for(auto &g_input : getInputs()){
+            tensorOffset[g_input->getFuid()] = allocator.alloc(g_input->getBytes());  //if(!tensor->getSource())
+        }
+        //遍历算子
+        for(auto &op: ops){
+            for(auto &output: op->getOutputs()){
+                int id = output->getFuid();
+                if(tensorOffset.find(id) == tensorOffset.end())   //确保不会重复分配
+                    tensorOffset[id] = allocator.alloc(output->getBytes());
+            }
+            for(auto &input: op->getInputs()){
+                int id = input->getFuid();
+                if(tensorRefNum.find(id) != tensorRefNum.end()){
+                    tensorRefNum[id]--;
+                    if(tensorRefNum[id] == 0){              //如果一个张量不再被使用((被引用次数是0)，可以释放其内存
+                        auto it = tensorOffset.find(id);    //释放前找到它在内存中的地址（偏移量）
+                        if(it != tensorOffset.end()){
+                            allocator.free(tensorOffset[id], input->getBytes());
+                            //tensorOffset.erase(it);
+                        }
+                        tensorRefNum.erase(id);
+                    }
+                }
+            }
+        }
+        void *memPtr = allocator.getPtr();
+        for(auto &tensor: tensors){
+            auto it = tensorOffset.find(tensor->getFuid());
+            IT_ASSERT(it != tensorOffset.end());
+            void *tensorPtr = static_cast<void *>(static_cast<char *>(memPtr) + it->second);
+            tensor->setDataBlob(make_ref<BlobObj>(runtime, tensorPtr));
+        }
 
         allocator.info();
     }
 
+
     Tensor GraphObj::addTensor(Shape dim, DataType dtype)
     {
         return tensors.emplace_back(make_ref<TensorObj>(dim, dtype, runtime));

src/operators/concat.cc

@@ -17,7 +17,10 @@ optional<vector<Shape>> ConcatObj::inferShape(const TensorVec &inputs) {
     // TODO：修改 dims，返回正确的 concat 后的 shape
     // REF: https://onnx.ai/onnx/operators/onnx__Concat.html#concat-13
     // =================================== 作业 ===================================
-
+    int n = inputs.size();
+    for(int i = 1; i < n; i++){
+        dims[dim] += inputs[i]->getDims()[dim];
+    }
     return {{dims}};
 }
 

src/operators/matmul.cc

@@ -27,7 +27,27 @@ namespace infini
         // TODO：返回经过 matmul 操作后的 shape
         // REF: https://github.com/onnx/onnx/blob/main/docs/Operators.md#gemm
         // =================================== 作业 ===================================
-        return std::nullopt;
+        const auto A = inputs[0];
+        auto a_dim = A->getDims();
+        const auto B = inputs[1];
+        auto b_dim = B->getDims();
+        int size = a_dim.size();
+        Shape res = a_dim;
+        //multi-dimensional broadcasting except for the last two dimensions
+        for(int i = 0; i < size - 2; i++){
+            if(a_dim[i] < b_dim[i])
+                res[i] = b_dim[i];
+        }
+
+        if(transA)
+            res[size - 2] = a_dim[size - 1];
+        if(transB)
+            res[size - 1] = b_dim[size - 2];
+        else
+            res[size - 1] = b_dim[size - 1];
+
+        std::vector<Shape> vec {res};
+        return vec;
     }
 
 } // namespace infini

src/operators/transpose.cc

@@ -33,8 +33,11 @@ namespace infini
         // TODO：修改 output_dim，返回正确的 transpose 后的 shape
         // REF: https://onnx.ai/onnx/operators/onnx__Transpose.html#transpose-21
         // =================================== 作业 ===================================
-
-        return std::nullopt;
+        for (int i = 0; i < rank; i++){
+           output_dim[i] = input_dim[transposePermute[i]];
+        }
+        std::vector<Shape> res {output_dim};
+        return res;
     }
 
     std::string TransposeObj::toString() const