Optimizer.h

浏览该文件的文档.

 
#include "common/Export.h"
#ifndef EVE_TENSOR_OPTIMIZER_H
#define EVE_TENSOR_OPTIMIZER_H
 
#include "tensor/Graph.h"
 
#include <vector>
 
namespace eve::tensor {
 
enum class GroupKind : uint8_t {
    Elementwise,
    MatMul,
    Conv1d,
    Conv2d,
    MaxPool2d,
    AvgPool2d,
    Softmax,
    LayerNorm,
    RMSNorm,
    Reduce,
    ArgMax,
    Embedding,
    Concat,
    Slice,
    Permute,
    Sdpa,
    Resize2d,
    Alias,
};
 
struct FusedGroup {
    GroupKind kind = GroupKind::Elementwise;
    OpType op = OpType::Const;
    std::vector<int> nodes;
    std::vector<int> inputs;
    int  outputNode = -1;
 
    // attributes copied from the root node for convenience
    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    int   i0 = 0, i1 = 0, i2 = 0, i3 = 0;
    int   perm[Tensor::kMaxRank] = {0, 1, 2, 3, 4, 5};
    int   permRank = 0;
    int   dtype = static_cast<int>(DType::Float32);
 
    // MatMul / Conv epilogue: bias node id (or -1) + elementwise chain node ids
    int              biasNode = -1;
    std::vector<int> epilogue;
    bool             hasScale = false;
    bool             hasBias  = false;
    bool             masked   = false;
    bool             logMode  = false;
};
 
struct OptimizedGraph {
    std::vector<int> order;
    std::vector<FusedGroup> groups;
    std::vector<int> groupOrder;
    std::vector<int> nodeSlot;
    std::vector<int> slotSize;
    std::vector<int> persistentSlots;
    int outputNode = -1;
};
 
EVENGINE_API_DOMAINS OptimizedGraph optimizeGraph(const Graph &graph, int outputNode);
 
int groupKernelCount(const OptimizedGraph &opt);
 
}  // namespace eve::tensor
 
#endif  // EVE_TENSOR_OPTIMIZER_H