wehub-resource-sync
86 lines
3.7 KiB
C++
86 lines
3.7 KiB
C++
#ifndef VulkanConv1x1CoopA8_hpp
|
|
#define VulkanConv1x1CoopA8_hpp
|
|
|
|
#include "VulkanBasicExecution.hpp"
|
|
#include "VulkanConvolution.hpp"
|
|
#include "core/ConvolutionCommon.hpp"
|
|
|
|
namespace MNN {
|
|
|
|
// W8A8 cooperative-matrix Conv1x1 path (per-channel asym int8 OR int4 +
|
|
// Adreno + S8S8->S32 coop matrix). M==1 reuses Coop-A16's gemv_dequant decode
|
|
// (int4/int8 variant); M>1 runs a 6-stage dynamic-quant -> S8 GEMM ->
|
|
// dequant-correction pipeline. INT4 mode adds a runtime int4->int8 unpack
|
|
// stage so the GEMM body is shared with INT8.
|
|
class VulkanConv1x1CoopA8 : public VulkanBasicExecution {
|
|
public:
|
|
VulkanConv1x1CoopA8(VulkanBackend* backend, const Convolution2DCommon* convOption,
|
|
const float* biasPtr, int ci, int co,
|
|
VulkanDevice::CoopMatInfo coopMatInfo,
|
|
std::shared_ptr<ConvolutionCommon::Int8Common> quantInfo);
|
|
virtual ~VulkanConv1x1CoopA8();
|
|
|
|
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:
|
|
VulkanConv1x1CoopA8(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* biasPtr, bool initStaticResource);
|
|
|
|
const Convolution2DCommon* mCommon;
|
|
int mCi;
|
|
int mCo;
|
|
std::shared_ptr<ConvolutionCommon::Int8Common> mQuantCommon;
|
|
uint32_t mPadK = 0;
|
|
uint32_t mPadN = 0;
|
|
|
|
// Static buffers — decode + prefill share, onClone forwards them.
|
|
// mQuantWeightBuffer: INT8: [padN, padK/4] uint32 (4x int8 packed).
|
|
// INT4: [padN, padK/8] uint32 (8x int4 packed,
|
|
// unsigned 0..15 with host +8 offset).
|
|
// mQuantMetaBuffer: [padN, 2] FP interleaved (scale, offset).
|
|
// mSumWqBuffer: [padN] int32, sum_k Wq[n,k]; padding rows zero.
|
|
std::shared_ptr<VulkanBuffer> mBiasBuffer;
|
|
std::shared_ptr<VulkanBuffer> mQuantWeightBuffer;
|
|
std::shared_ptr<VulkanBuffer> mQuantMetaBuffer;
|
|
std::shared_ptr<VulkanBuffer> mSumWqBuffer;
|
|
|
|
const VulkanPipeline* mDecodePipeline = nullptr;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mDecodeSet;
|
|
|
|
// INT4 prefill only: runtime nibble unpack from mQuantWeightBuffer to
|
|
// a dynamic [padN, padK] int8 tensor (tWqInt8 in onEncode), then the
|
|
// shared dynamic_w8a8_coop_gemm consumes it as if it were INT8.
|
|
const VulkanPipeline* mInt4UnpackPipeline = nullptr;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mInt4UnpackSet;
|
|
|
|
const VulkanPipeline* mQuantMinMaxPipeline = nullptr;
|
|
const VulkanPipeline* mQuantReduceMinMaxPipeline = nullptr;
|
|
const VulkanPipeline* mQuantPackPipeline = nullptr;
|
|
const VulkanPipeline* mQuantReduceSumPipeline = nullptr;
|
|
const VulkanPipeline* mGemmS8Pipeline = nullptr;
|
|
const VulkanPipeline* mDequantPipeline = nullptr;
|
|
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mQuantMinMaxSet;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mQuantReduceMinMaxSet;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mQuantPackSet;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mQuantReduceSumSet;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mGemmSet;
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> mDequantSet;
|
|
|
|
uint32_t mCoopM = 64;
|
|
uint32_t mCoopN = 64;
|
|
uint32_t mCoopK = 32;
|
|
uint32_t mSubgroupSize = 64;
|
|
bool mIsInt4 = false;
|
|
};
|
|
|
|
} // namespace MNN
|
|
|
|
#endif
|