Files
2026-07-13 13:33:03 +08:00

79 lines
3.0 KiB
C++

#ifndef VulkanConv1x1Coop_hpp
#define VulkanConv1x1Coop_hpp
#include "VulkanBasicExecution.hpp"
#include "VulkanConvolution.hpp"
#include "core/ConvolutionCommon.hpp"
namespace MNN {
class VulkanConv1x1Coop : public VulkanBasicExecution {
public:
VulkanConv1x1Coop(VulkanBackend* backend, const Convolution2DCommon* convOption, const float* weightPtr, const float* biasPtr, int ci, int co, VulkanDevice::CoopMatInfo coopMatInfo, std::shared_ptr<ConvolutionCommon::Int8Common> quantInfo = nullptr);
virtual ~VulkanConv1x1Coop();
virtual ErrorCode onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
const VulkanCommandPool::Buffer* cmdBuffer) override;
virtual bool onClone(Backend* bn, const Op* op, VulkanBasicExecution** dst) override;
private:
VulkanConv1x1Coop(VulkanBackend* backend, const Convolution2DCommon* convOption, int ci, int co,
uint32_t coopM, uint32_t coopN, uint32_t coopK, uint32_t subgroupSize,
std::shared_ptr<ConvolutionCommon::Int8Common> quantInfo,
bool initStaticResource);
bool _init(const float* weightPtr, const float* biasPtr, bool initStaticResource);
const Convolution2DCommon* mCommon;
int mCi;
int mCo;
bool mIsQuant = false;
std::shared_ptr<ConvolutionCommon::Int8Common> mQuantCommon;
uint32_t mBlockSize = 0; // quant block size along K
uint32_t mPadK = 0;
uint32_t mPadN = 0;
bool mQuantConverted = false;
// Weight and Bias Buffer
std::shared_ptr<VulkanBuffer> mWeightBuffer; // float/half weight (float path) or dequantized coop weight (quant path)
std::shared_ptr<VulkanBuffer> mBiasBuffer;
// Quant resources (static)
std::shared_ptr<VulkanBuffer> mQuantWeightBuffer; // int8/int4 packed in [padN, padK]
std::shared_ptr<VulkanBuffer> mQuantMetaBuffer; // [scale, offset] per (N, K-block)
// Pipelines
const VulkanPipeline* mPackPipeline;
const VulkanPipeline* mMatMulPipeline;
const VulkanPipeline* mUnpackPipeline;
const VulkanPipeline* mPrefillDequantPipeline = nullptr;
const VulkanPipeline* mDecodePipeline = nullptr;
// Descriptor Sets
std::shared_ptr<VulkanLayout::DescriptorSet> mPackSet;
std::shared_ptr<VulkanLayout::DescriptorSet> mMatMulSet;
std::shared_ptr<VulkanLayout::DescriptorSet> mUnpackSet;
std::shared_ptr<VulkanLayout::DescriptorSet> mPrefillDequantSet;
std::shared_ptr<VulkanLayout::DescriptorSet> mDecodeSet;
// Constant Buffers
std::shared_ptr<VulkanBuffer> mPackConst;
std::shared_ptr<VulkanBuffer> mMatMulConst;
std::shared_ptr<VulkanBuffer> mUnpackConst;
// Temp Tensors
std::shared_ptr<Tensor> mTempInput;
std::shared_ptr<Tensor> mTempOutput;
std::shared_ptr<Tensor> mTempWeight;
// COOP Shape
uint32_t COOP_M = 64;
uint32_t COOP_N = 64;
uint32_t COOP_K = 16;
// subgroupSize
uint32_t mSubgroupSize = 64;
};
} // namespace MNN
#endif