wehub-resource-sync
48 lines
1.4 KiB
C++
48 lines
1.4 KiB
C++
//
|
|
// RopeBufExecution.hpp
|
|
// MNN
|
|
//
|
|
// OpenCL buffer-path implementation of RoPE (Rotary Positional Embedding).
|
|
//
|
|
|
|
#ifdef MNN_SUPPORT_TRANSFORMER_FUSE
|
|
|
|
#ifndef RopeBufExecution_hpp
|
|
#define RopeBufExecution_hpp
|
|
|
|
#include "backend/opencl/execution/image/CommonExecution.hpp"
|
|
|
|
namespace MNN {
|
|
namespace OpenCL {
|
|
|
|
class RopeBufExecution : public CommonExecution {
|
|
public:
|
|
RopeBufExecution(const MNN::Op* op, Backend* backend);
|
|
RopeBufExecution(const MNN::Op* op, Backend* backend, int ropeCutHeadDim, int numHead, int kvNumHead, int headDim,
|
|
std::shared_ptr<cl::Buffer> qGamma, float qEps, std::shared_ptr<cl::Buffer> kGamma, float kEps);
|
|
virtual ~RopeBufExecution() = default;
|
|
|
|
virtual ErrorCode onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) override;
|
|
virtual bool onClone(Backend* bn, const Op* op, Execution** dst) override;
|
|
|
|
private:
|
|
OpenCLBackend* mOpenCLBackend = nullptr;
|
|
uint32_t mMaxWorkGroupSize = 0;
|
|
std::vector<uint32_t> mGlobalWorkSize = {1, 1, 1};
|
|
std::vector<uint32_t> mLocalWorkSize = {1, 1, 1};
|
|
int mRopeCutHeadDim = 0;
|
|
int mNumHead = 0;
|
|
int mKvNumHead = 0;
|
|
int mHeadDim = 0;
|
|
std::shared_ptr<cl::Buffer> mQGamma;
|
|
std::shared_ptr<cl::Buffer> mKGamma;
|
|
float mQEps = 0.0f;
|
|
float mKEps = 0.0f;
|
|
};
|
|
|
|
} // namespace OpenCL
|
|
} // namespace MNN
|
|
|
|
#endif /* RopeBufExecution_hpp */
|
|
#endif /* MNN_SUPPORT_TRANSFORMER_FUSE */
|