CpuKernels.h

浏览该文件的文档.

 
#include "common/Export.h"
#ifndef EVE_TENSOR_CPUKERNELS_H
#define EVE_TENSOR_CPUKERNELS_H
 
#include "tensor/Graph.h"
 
#include <cstddef>
 
namespace eve::tensor {
 
namespace kernels {
 
float applyUnary(OpType type, float x, float s0, float s1);
 
bool isElementwiseOp(OpType type);
 
bool broadcastShape(const int *aDims, int aRank, const int *bDims, int bRank, int *outDims,
                    int &outRank);
 
void binaryOp(OpType type, const float *a, const int *aDims, int aRank, const float *b,
              const int *bDims, int bRank, float *out, const int *outDims, int outRank);
 
void unaryOp(OpType type, const float *in, int count, float *out, float s0, float s1);
 
EVENGINE_API_DOMAINS void softmax(const float *in, const int *dims, int rank, int axis, bool logMode, float *out);
void layernorm(const float *in, int rows, int cols, const float *scale, const float *bias,
               float eps, float *out);
void rmsnorm(const float *in, int rows, int cols, const float *scale, float eps, float *out);
 
void conv1d(const float *x, const int *xDims, const float *w, const int *wDims,
            const float *bias, int stride, int pad, float *out);
void conv2d(const float *x, const int *xDims, const float *w, const int *wDims,
            const float *bias, int stride, int pad, float *out);
void maxpool2d(const float *in, const int *dims, int ksize, int stride, int pad, float *out);
void avgpool2d(const float *in, const int *dims, int ksize, int stride, int pad, float *out);
 
void embedding(const float *table, int vocab, int dim, const float *indices, int count,
               float *out);
 
void concat(const float *const *ins, const int *const *inDims, const int *inRanks, int n,
            int axis, float *out, const int *outDims, int outRank);
 
void sliceOp(const float *in, const int *inDims, int inRank, int axis, int begin, int end,
             float *out, const int *outDims, int outRank);
void permute(const float *in, const int *inDims, int rank, const int *order, float *out,
             const int *outDims);
 
void reduceAxis(OpType type, const float *in, const int *dims, int rank, int axis, float *out,
                const int *outDims, int outRank);
void argmax(const float *in, const int *dims, int rank, int axis, float *out,
            const int *outDims, int outRank);
 
void sdpa(const float *q, const float *k, const float *v, const float *mask, int B, int H,
          int T, int S, int D, float scale, float *out);
 
void resize2d(const float *in, const int *inDims, int outW, int outH, int mode, float *out);
 
}  // namespace kernels
}  // namespace eve::tensor
 
#endif  // EVE_TENSOR_CPUKERNELS_H