TinyInfiniTensor/src/core/operator.cc at 7501db8e2fdf5ac027b02d62c1e6dfdfd81b5bdb · InfiniTensor/TinyInfiniTensor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
#include "core/operator.h"
#include "core/graph.h"

namespace infini
{

    OperatorObj::OperatorObj(OpType opType, TensorVec inputs, TensorVec outputs)
        : type(opType), inputs(inputs), outputs(outputs) {}

    void OperatorObj::removePredecessors(const Operator &op)
    {
        for (auto it = predecessors.begin(); it != predecessors.end();)
        {
            if (it->lock() == op)
                it = predecessors.erase(it);
            else
                ++it;
        }
    }

    void OperatorObj::removeSuccessors(const Operator &op)
    {
        for (auto it = successors.begin(); it != successors.end();)
        {
            if (it->lock() == op)
                it = successors.erase(it);
            else
                ++it;
        }
    }

    void OperatorObj::replaceInput(Tensor t1, Tensor t2)
    {
        for (auto itr = inputs.begin(); itr != inputs.end(); ++itr)
        {
            if (*itr == t1)
            {
                *itr = t2;
            }
        }
    }

    bool OperatorObj::checkValid(GraphObj *graph)
    {
        auto optShapes = inferShape();
        if (!optShapes) // shape inference failed
            return false;

        const vector<Shape> &shapes = *optShapes;
        if (shapes.size() != outputs.size())
            return false;
        if (graph)
        { // if graph != nullptr, outputs should be created
            auto dataTypes = inferDataType();
            for (size_t i = 0; i < outputs.size(); i++)
            {
                if (outputs[i] == nullptr) {
                    outputs[i] = graph->addTensor(shapes[i], dataTypes[i]);
                }
            }
        }
        else
        { // if outputs have been created, check their shapes
            for (size_t i = 0; i < shapes.size(); ++i)
            {
                if (!outputs[i]) return false;
                if (shapes[i] != outputs[i]->getDims())
                    return false;
            }
        }
        return true;
    }

    optional<vector<Shape>> OperatorObj::inferShape() { return inferShape(inputs); }

    vector<DataType> OperatorObj::inferDataType(const TensorVec &inputs) const
    {
        auto dataType = inputs[0]->getDType();
        return vector(numOutputs(), dataType);
    }

    vector<DataType> OperatorObj::inferDataType() const
    {
        return inferDataType(inputs);
    }

} // namespace infini