graph.cc

#include "core/graph.h"
#include <algorithm>
#include <numeric>
#include <queue>
#include "operators/transpose.h" 
#include "operators/matmul.h"

namespace infini
{

    void GraphObj::addOperatorAndConnect(const Operator &op)
    {
        sorted = false;
        ops.push_back(op);
        for (auto &input : op->getInputs())
        {
            if (input)
            {
                input->addTarget(op);
                if (auto pred = input->getSource())
                {
                    pred->addSuccessors(op);
                    op->addPredecessors(pred);
                }
            }
        }
        for (auto &output : op->getOutputs())
        {
            if (output)
            {
                output->setSource(op);
                for (auto &succ : output->getTargets())
                {
                    succ->addPredecessors(op);
                    op->addSuccessors(succ);
                }
            }
        }
    }

    string GraphObj::toString() const
    {
        std::ostringstream oss;
        oss << "Graph Tensors:\n";
        for (const auto &tensor : tensors)
            oss << tensor << "\n";
        printf("完成tensor访问\n");
        oss << "Graph operators:\n";
        for (const auto &op : ops)
        {
            vector<UidBaseType> preds, succs;
            for (auto &o : op->getPredecessors())
                preds.emplace_back(o->getGuid());
            for (auto &o : op->getSuccessors())
                succs.emplace_back(o->getGuid());
            oss << "OP " << op->getGuid();
            oss << ", pred " << vecToString(preds);
            oss << ", succ " << vecToString(succs);
            oss << ", " << op << "\n";
        }
        printf("完成op访问\n");
        
        return oss.str();
    }

    bool GraphObj::topo_sort()
    {
        if (this->sorted)
        {
            return true;
        }
        std::vector<Operator> sorted;
        std::unordered_set<OperatorObj *> flags;
        sorted.reserve(ops.size());
        flags.reserve(ops.size());
        while (sorted.size() < ops.size())
        {
            // Any node is move to sorted in this loop.
            auto modified = false;
            for (auto const &op : ops)
            {
                if (auto const &inputs = op->getInputs();
                    flags.find(op.get()) == flags.end() &&
                    std::all_of(inputs.begin(), inputs.end(),
                                [&flags](auto const &input)
                                {
                                    auto ptr = input->getSource().get();
                                    return !ptr || flags.find(ptr) != flags.end();
                                }))
                {
                    modified = true;
                    sorted.emplace_back(op);
                    flags.insert(op.get());
                }
            }
            if (!modified)
            {
                return false;
            }
        }
        this->ops = std::move(sorted);
        return this->sorted = true;
    }

    void GraphObj::optimize()
    {
        // =================================== 作业 ===================================
        // TODO: 设计一个算法来实现指定的图优化规则
        // 图优化规则如下：
        // 1. 去除冗余的算子（例如，两个相邻的算子都是 transpose 算子，且做的是相反的操作，可以将其全部删除）
        // 2. 合并算子（例如，矩阵乘算子中含有属性transA、transB，如果其输入存在transpose，且对最后两个维度做交换，就可以将transpose融入到矩阵乘算子的属性中去）
        // =================================== 作业 ===================================
        bool changed = true;
        while(changed){
            changed = false;
            auto ops = this->ops;
            for(auto op:ops){
                //case1::检查如果是两个连续的转置
                
                if(op->getOpType()==OpType::Transpose){
                    auto trans1 = as<TransposeObj>(op);
                    auto tensorA = trans1->getInputs()[0];//输入，不可丢弃
                    auto tensorB = trans1->getOutput();//中间变量，可丢弃
                    auto prevop = tensorA->getSource();//A的前置op
                    //如果只有一个去向
                    if (tensorB->getTargets().size() == 1){
                        auto nextOp = tensorB->getTargets()[0];
                        if (nextOp->getOpType() == OpType::Transpose){//如果这个刚好是transpose
                            auto trans2 = as<TransposeObj>(nextOp);
                            auto tensorC = trans2->getOutput();//最终输出，可丢弃
                            if (trans1->getPermute() == trans2->getPermute()){
                                auto targets = tensorC->getTargets();
                                //把C流向的算子们的张量来源设置为A，算子来源设置为prevop；
                                for (auto target : targets) {
                                    target->removePredecessors(trans2);
                                    if (prevop) {
                                        prevop->addSuccessors(target);
                                        target->addPredecessors(prevop); // 只有当 A 有生产者时，才连接
                                    }
                                    target->replaceInput(tensorC, tensorA);
                                    tensorA->addTarget(target);
                                }
                                if (prevop) {
                                    prevop->removeSuccessors(trans1);
                                }
                                tensorA->removeTarget(trans1);
                                this->removeOperator(trans1);
                                this->removeOperator(trans2);
                                this->removeTensor(tensorB);
                                this->removeTensor(tensorC);
                                //printf("优化1,成功移除\n");
                                changed = true;
                                break;
                            }
                        }
                    }
                }
                if (op->getOpType() == OpType::MatMul){
                    auto matmul = as<MatmulObj>(op);
                    auto inputs = matmul->getInputs();
                    for (size_t i = 0; i < 2; ++i){
                        auto tensorB = inputs[i];
                        auto prevOp = tensorB->getSource();
                        if (prevOp && prevOp->getOpType() == OpType::Transpose){
                            auto trans = as<TransposeObj>(prevOp);
                            auto tensorA = trans->getInputs()[0];
                            auto perm = trans->getPermute();

                            int rank = perm.size();
                            bool isSwapLastTwo = true;
                            if (rank < 2) isSwapLastTwo = false;
                            else {
                                if (perm[rank-1] != rank-2 || perm[rank-2] != rank-1) isSwapLastTwo = false;
                                for (int k = 0; k < rank - 2; ++k) {
                                    if (perm[k] != k) { isSwapLastTwo = false; break; }
                                }
                            }
                            if(isSwapLastTwo){
                                if (i == 0) matmul->setTransA(!matmul->getTransA());
                                else        matmul->setTransB(!matmul->getTransB());

                                matmul->replaceInput(tensorB, tensorA);
                                auto sourceA = tensorA->getSource();
                                if(sourceA){
                                    sourceA->addSuccessors(matmul);
                                    matmul->addPredecessors(sourceA);
                                    
                                }
                                matmul->removePredecessors(trans);
                                trans->removeSuccessors(matmul);
                                tensorA->addTarget(matmul);
                                tensorB->removeTarget(matmul);
                                if(tensorB->getTargets().empty()){
                                    if(sourceA){
                                        trans->removePredecessors(sourceA);
                                        sourceA->removeSuccessors(trans);
                                    }
                                    tensorA->removeTarget(trans);
                                    this->removeOperator(trans);
                                    this->removeTensor(tensorB);
                                }
                                //printf("优化2，成功\n");
                                changed = true;
                            }
                        }
                    }
                    if(changed) break;
                }
            }
        }
        //printf("结束循环\n");
    }

    Tensor GraphObj::getTensor(int fuid) const
    {
        for (auto tensor : tensors)
        {
            if (tensor->getFuid() == fuid)
            {
                return tensor;
            }
        }
        return nullptr;
    }

    void GraphObj::shape_infer()
    {
        for (auto &op : ops)
        {
            auto ans = op->inferShape();
            IT_ASSERT(ans.has_value());
            auto oldOutputs = op->getOutputs();
            IT_ASSERT(ans.value().size() == oldOutputs.size());
            // replace the old outputshape and size with new one
            for (int i = 0; i < (int)ans.value().size(); ++i)
            {
                auto newShape = ans.value()[i];
                auto oldShape = oldOutputs[i]->getDims();
                auto fuid = oldOutputs[i]->getFuid();
                if (newShape != oldShape)
                {
                    auto tensor = this->getTensor(fuid);
                    tensor->setShape(newShape);
                }
            }
        }
    }

    void GraphObj::dataMalloc()
    {
        // topological sorting first
        IT_ASSERT(topo_sort() == true);
        // =================================== 作业 ===================================
        // TODO：利用 allocator 给计算图分配内存
        // HINT: 获取分配好的内存指针后，可以调用 tensor 的 setDataBlob 函数给 tensor 绑定内存
        // =================================== 作业 ===================================
        std::unordered_map<int, size_t> ref_counts;//记录引用数量，如果引用归0了就直接free
        for(auto &tensor: this->tensors){
            ref_counts[tensor->getFuid()] = tensor->getTargets().size();
        }
        //对需要最终输出的output进行人为+1让它可以被永久保存下来
        for (auto &tensor : this->getOutputs()) {
            ref_counts[tensor->getFuid()]++; // 人为加1，保证不被回收
       }
        std::unordered_map<int, size_t> offsets;
        for (auto &tensor : this->tensors) {
            if (!tensor->getSource()) { // 如果没有来源算子，说明它是图的输入
                size_t size = tensor->getBytes();
                size_t offset = allocator.alloc(size);
                offsets[tensor->getFuid()] = offset;
            }
        }

        //现在开始遍历Ops模拟内存分配情况：
        for(auto &op:this->ops){
            //这里我们检查输出，因为要为每个输出分配空间
            for(auto& tensor:op->getOutputs()){
                size_t size = tensor->getBytes();//拿到大小
                size_t offset = allocator.alloc(size);//申请内存
                offsets[tensor->getFuid()] = offset;//把分配的空间偏移保存下来
            }
            //检查输入看看能不能把输入释放了
            for (auto &tensor : op->getInputs()){
                int fuid = tensor->getFuid();
                ref_counts[fuid]--;
                if(ref_counts[fuid]==0){
                    //首先我们要理解的是，我们这里的分配只是为中间结果分配，
                    // 外部输入的向量是不包含在这里的，所以我们要检查这个是不是外部的输入。
                    if (offsets.find(fuid) != offsets.end()) {
                        //free掉
                        allocator.free(offsets[fuid], tensor->getBytes());
                    }
                }
            }
        }
        //现在我们为每个fuid对应的tensor找到了合适的offset。
        void *basePtr = allocator.getPtr();
        //开始分配Blob
        for(auto&tensor:this->tensors){
            int fuid = tensor->getFuid();
            //首先我们要检查这个是不是外部的输入，或者说万一是计算图多余的
            if(offsets.find(fuid)!=offsets.end()){
                size_t offset = offsets[fuid];
                //void*不能进行算数运算，所以先转成char*
                void *ptr = static_cast<char *>(basePtr) + offset;
                //首先我们要创建一个blob：BlobObj(Runtime runtime, void *ptr)
                auto blob = make_ref<BlobObj>(this->runtime, ptr);
                tensor->setDataBlob(blob);
            }
        }
        allocator.info();
    }

    Tensor GraphObj::addTensor(Shape dim, DataType dtype)
    {
        return tensors.emplace_back(make_ref<TensorObj>(dim, dtype, runtime));
    }

    Tensor GraphObj::addTensor(const Tensor &tensor)
    {
        IT_ASSERT(tensor->getRuntime() == runtime,
                  std::string("Tensor runtime mismatch: cannot add a tenosr in ") +
                      tensor->getRuntime()->toString() + " to " +
                      runtime->toString());
        tensors.emplace_back(tensor);
        return tensor;
    }

    TensorVec GraphObj::addTensor(const TensorVec &tensors)
    {
        for (auto &t : tensors)
            addTensor(t);
        return tensors;
    }

    // tensor's "source" and "target" must be in "ops".
    // tensor has no "source" and no "target" must not exist.
    // "inputs" or "outputs" of operators must be in "tensors"
    // "predecessors" and "successors" of an operator of "ops" must be in "ops".
    bool GraphObj::checkValid() const
    {
        for (auto tensor : tensors)
        {
            IT_ASSERT(!(tensor->getTargets().size() == 0 &&
                        nullptr == tensor->getSource()));
            for (auto op : tensor->getTargets())
            {
                IT_ASSERT(std::find(ops.begin(), ops.end(), op) != ops.end());
            }
            auto op = tensor->getSource();
            IT_ASSERT(!(op && std::find(ops.begin(), ops.end(), op) == ops.end()));
        }
        for (auto op : ops)
        {
            for (auto tensor : op->getInputs())
            {
                IT_ASSERT(std::find(tensors.begin(), tensors.end(), tensor) !=
                          tensors.end());
            }
            for (auto tensor : op->getOutputs())
            {
                IT_ASSERT(std::find(tensors.begin(), tensors.end(), tensor) !=
                          tensors.end());
            }
            for (auto pre : op->getPredecessors())
            {
                IT_ASSERT(std::find(ops.begin(), ops.end(), pre) != ops.end());
            }
            for (auto suc : op->getSuccessors())
            {
                IT_ASSERT(std::find(ops.begin(), ops.end(), suc) != ops.end());
            }
        }
        std::set<UidBaseType> s;
        // check whether two tensors with the same FUID exist
        for (auto tensor : tensors)
        {
            int cnt = s.count(tensor->getFuid());
            IT_ASSERT(cnt == 0, std::to_string(tensor->getFuid()));
            s.insert(tensor->getFuid());
        }
        return true;
    }

} // namespace infini