wehub-resource-sync
543 lines
24 KiB
C++
543 lines
24 KiB
C++
#include "VulkanConv1x1General.hpp"
|
|
#include "VulkanBackend.hpp"
|
|
#include "core/Macro.h"
|
|
#include <algorithm>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <vector>
|
|
|
|
namespace MNN {
|
|
|
|
namespace {
|
|
|
|
struct QuantWeightPrepareParams {
|
|
uint32_t ci;
|
|
uint32_t co;
|
|
uint32_t padN;
|
|
uint32_t weightStride;
|
|
uint32_t srcBytes;
|
|
};
|
|
|
|
struct QuantMetaPrepareParams {
|
|
uint32_t co;
|
|
uint32_t padN;
|
|
uint32_t blockCount;
|
|
uint32_t blockStride;
|
|
uint32_t soSize;
|
|
uint32_t alphaSize;
|
|
};
|
|
|
|
static size_t _alignUp4(size_t size) {
|
|
return (size + 3u) & ~size_t(3u);
|
|
}
|
|
|
|
static bool _prepareQuantBuffersGPU(VulkanBackend* vkBn, const ConvolutionCommon::Int8Common* quantCommon,
|
|
bool useFP16, int ci, int co, uint32_t padN, uint32_t blockStride,
|
|
uint32_t decodeWeightStrideWords, int quantBits,
|
|
std::shared_ptr<VulkanBuffer>& quantWeightBuffer,
|
|
std::shared_ptr<VulkanBuffer>& quantMetaBuffer) {
|
|
if (nullptr == vkBn || nullptr == quantCommon || nullptr == quantCommon->weight.get()) {
|
|
return false;
|
|
}
|
|
|
|
const int soSize = quantCommon->asymmetric ? 2 : 1;
|
|
const int alphaSize = quantCommon->alpha.size();
|
|
const int alphaDenominator = std::max(1, co * soSize);
|
|
const int blockCount = std::max(1, alphaSize / alphaDenominator);
|
|
const int8_t* qWeight = quantCommon->weight.get();
|
|
const size_t rawWeightBytes = static_cast<size_t>(quantCommon->weight.size());
|
|
const size_t alignedWeightBytes = std::max<size_t>(4u, _alignUp4(rawWeightBytes));
|
|
// int3 stores 2 uints per group of 16 weights; others store 1 uint per group.
|
|
const uint32_t wordsPerGroup = (quantBits == 3) ? 2u : 1u;
|
|
const size_t decodeWeightBytes =
|
|
static_cast<size_t>(padN) * static_cast<size_t>(decodeWeightStrideWords) *
|
|
static_cast<size_t>(wordsPerGroup) * sizeof(uint32_t);
|
|
const size_t metaElem = static_cast<size_t>(padN) * static_cast<size_t>(blockStride) * 2u;
|
|
const size_t metaBytes = metaElem * (useFP16 ? sizeof(int16_t) : sizeof(float));
|
|
|
|
const void* rawWeightSrc = qWeight;
|
|
std::vector<uint8_t> weightAlignedHost;
|
|
if (alignedWeightBytes != rawWeightBytes) {
|
|
weightAlignedHost.resize(alignedWeightBytes, 0);
|
|
if (rawWeightBytes > 0u) {
|
|
::memcpy(weightAlignedHost.data(), qWeight, rawWeightBytes);
|
|
}
|
|
rawWeightSrc = weightAlignedHost.data();
|
|
}
|
|
|
|
std::shared_ptr<VulkanBuffer> rawWeightBuffer(new VulkanBuffer(
|
|
vkBn->getMemoryPool(), false, alignedWeightBytes, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_SHARING_MODE_EXCLUSIVE, 0));
|
|
vkBn->copyToGPUBuffer(rawWeightSrc, rawWeightBuffer->buffer(), alignedWeightBytes, 0);
|
|
|
|
quantWeightBuffer.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, decodeWeightBytes, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_SHARING_MODE_EXCLUSIVE, 0));
|
|
|
|
const float* alphaPtr = quantCommon->alpha.get();
|
|
const size_t rawAlphaBytes = static_cast<size_t>(std::max(alphaSize, 1)) * sizeof(float);
|
|
std::shared_ptr<VulkanBuffer> rawAlphaBuffer(new VulkanBuffer(
|
|
vkBn->getMemoryPool(), false, rawAlphaBytes, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_SHARING_MODE_EXCLUSIVE, 0));
|
|
if (alphaSize > 0 && nullptr != alphaPtr) {
|
|
vkBn->copyToGPUBuffer(alphaPtr, rawAlphaBuffer->buffer(), static_cast<size_t>(alphaSize) * sizeof(float), 0);
|
|
} else {
|
|
const float zero = 0.0f;
|
|
vkBn->copyToGPUBuffer(&zero, rawAlphaBuffer->buffer(), sizeof(float), 0);
|
|
}
|
|
|
|
quantMetaBuffer.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, metaBytes, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_SHARING_MODE_EXCLUSIVE, 0));
|
|
|
|
const char* weightShader = nullptr;
|
|
switch (quantBits) {
|
|
case 2: weightShader = "glsl_conv1x1_int2_weight_prepare_comp"; break;
|
|
case 3: weightShader = "glsl_conv1x1_int3_weight_prepare_comp"; break;
|
|
case 4: weightShader = "glsl_conv1x1_int4_weight_prepare_comp"; break;
|
|
default: weightShader = "glsl_conv1x1_int8_weight_prepare_comp"; break;
|
|
}
|
|
const char* metaShader = useFP16 ? "glsl_conv1x1_quant_meta_prepare_FP16_comp"
|
|
: "glsl_conv1x1_quant_meta_prepare_comp";
|
|
|
|
std::vector<VkDescriptorType> prepareTypes = {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
};
|
|
auto weightPipeline = vkBn->getPipeline(weightShader, prepareTypes);
|
|
auto metaPipeline = vkBn->getPipeline(metaShader, prepareTypes);
|
|
if (nullptr == weightPipeline || nullptr == metaPipeline) {
|
|
return false;
|
|
}
|
|
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> weightSet(weightPipeline->createSet());
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> metaSet(metaPipeline->createSet());
|
|
if (nullptr == weightSet.get() || nullptr == metaSet.get()) {
|
|
return false;
|
|
}
|
|
|
|
std::shared_ptr<VulkanCommandPool::Buffer> prepareCmd(vkBn->getPool().allocBuffer());
|
|
prepareCmd->begin(0);
|
|
|
|
{
|
|
QuantWeightPrepareParams pc;
|
|
pc.ci = static_cast<uint32_t>(ci);
|
|
pc.co = static_cast<uint32_t>(co);
|
|
pc.padN = padN;
|
|
pc.weightStride = decodeWeightStrideWords;
|
|
pc.srcBytes = static_cast<uint32_t>(rawWeightBytes);
|
|
|
|
weightSet->writeBuffer(rawWeightBuffer->buffer(), 0, rawWeightBuffer->size());
|
|
weightSet->writeBuffer(quantWeightBuffer->buffer(), 1, quantWeightBuffer->size());
|
|
weightPipeline->bind(prepareCmd->get(), weightSet->get());
|
|
vkCmdPushConstants(prepareCmd->get(), weightPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(pc), &pc);
|
|
vkCmdDispatch(prepareCmd->get(), UP_DIV(decodeWeightStrideWords, 16u), UP_DIV(padN, 16u), 1);
|
|
prepareCmd->barrierSource(quantWeightBuffer->buffer(), 0, quantWeightBuffer->size());
|
|
}
|
|
|
|
{
|
|
QuantMetaPrepareParams pc;
|
|
pc.co = static_cast<uint32_t>(co);
|
|
pc.padN = padN;
|
|
pc.blockCount = static_cast<uint32_t>(blockCount);
|
|
pc.blockStride = blockStride;
|
|
pc.soSize = static_cast<uint32_t>(soSize);
|
|
pc.alphaSize = static_cast<uint32_t>(alphaSize);
|
|
|
|
metaSet->writeBuffer(rawAlphaBuffer->buffer(), 0, rawAlphaBuffer->size());
|
|
metaSet->writeBuffer(quantMetaBuffer->buffer(), 1, quantMetaBuffer->size());
|
|
metaPipeline->bind(prepareCmd->get(), metaSet->get());
|
|
vkCmdPushConstants(prepareCmd->get(), metaPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(pc), &pc);
|
|
vkCmdDispatch(prepareCmd->get(), UP_DIV(blockStride, 16u), UP_DIV(padN, 16u), 1);
|
|
prepareCmd->barrierSource(quantMetaBuffer->buffer(), 0, quantMetaBuffer->size());
|
|
}
|
|
|
|
prepareCmd->end();
|
|
vkBn->getPool().submitAndWait(prepareCmd->get());
|
|
return true;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
VulkanConv1x1General::VulkanConv1x1General(VulkanBackend* backend, const Convolution2DCommon* convOption,
|
|
const float* biasPtr, int ci, int co,
|
|
std::shared_ptr<ConvolutionCommon::Int8Common> quantInfo)
|
|
: VulkanBasicExecution(backend), mCommon(convOption), mCi(ci), mCo(co), mQuantCommon(std::move(quantInfo)) {
|
|
if (!_init(biasPtr, true)) {
|
|
MNN_ERROR("VulkanConv1x1General init failed\n");
|
|
}
|
|
}
|
|
|
|
VulkanConv1x1General::VulkanConv1x1General(VulkanBackend* backend, const Convolution2DCommon* convOption, int ci,
|
|
int co, std::shared_ptr<ConvolutionCommon::Int8Common> quantInfo,
|
|
bool initStaticResource)
|
|
: VulkanBasicExecution(backend), mCommon(convOption), mCi(ci), mCo(co), mQuantCommon(std::move(quantInfo)) {
|
|
if (!_init(nullptr, initStaticResource)) {
|
|
MNN_ERROR("VulkanConv1x1General clone init failed\n");
|
|
}
|
|
}
|
|
|
|
VulkanConv1x1General::~VulkanConv1x1General() {
|
|
}
|
|
|
|
bool VulkanConv1x1General::_init(const float* biasPtr, bool initStaticResource) {
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
if (nullptr == vkBn || nullptr == mQuantCommon.get() || nullptr == mQuantCommon->weight.get()) {
|
|
return false;
|
|
}
|
|
|
|
const bool useFP16 = vkBn->useFP16();
|
|
// Precedence: int2 > int3 > int4 > int8 (default).
|
|
if (mQuantCommon->canUseInt2) {
|
|
mQuantBits = 2;
|
|
} else if (mQuantCommon->canUseInt3) {
|
|
mQuantBits = 3;
|
|
} else if (mQuantCommon->canUseInt4) {
|
|
mQuantBits = 4;
|
|
} else {
|
|
mQuantBits = 8;
|
|
}
|
|
mPadK = ROUND_UP(static_cast<uint32_t>(mCi), 4u);
|
|
mPadN = ROUND_UP(static_cast<uint32_t>(mCo), 32u);
|
|
|
|
if (mPadK == 0u || mPadN == 0u) {
|
|
MNN_ERROR("VulkanConv1x1General invalid shape, ci=%d, co=%d\n", mCi, mCo);
|
|
return false;
|
|
}
|
|
|
|
const int soSize = mQuantCommon->asymmetric ? 2 : 1;
|
|
const int alphaSize = mQuantCommon->alpha.size();
|
|
const int alphaDenominator = std::max(1, mCo * soSize);
|
|
const int blockCount = std::max(1, alphaSize / alphaDenominator);
|
|
mBlockSize = std::max<uint32_t>(1u, static_cast<uint32_t>(UP_DIV(mCi, blockCount)));
|
|
|
|
if ((mBlockSize & 3u) != 0u) {
|
|
MNN_ERROR("VulkanConv1x1General requires blockSize %% 4 == 0, blockSize=%u\n", mBlockSize);
|
|
return false;
|
|
}
|
|
if (mCi % static_cast<int>(mBlockSize) != 0) {
|
|
MNN_ERROR("VulkanConv1x1General requires K %% blockSize == 0, K=%d, blockSize=%u\n", mCi, mBlockSize);
|
|
return false;
|
|
}
|
|
if ((mPadK % mBlockSize) != 0u) {
|
|
MNN_ERROR("VulkanConv1x1General requires padK %% blockSize == 0, padK=%u, blockSize=%u\n", mPadK, mBlockSize);
|
|
return false;
|
|
}
|
|
mBlockStride = mPadK / mBlockSize;
|
|
// For int2/int3, weightStride counts groups-of-16; int3 uses 2 uints per group, int2 uses 1.
|
|
mDecodeWeightStrideWords = (mQuantBits == 2 || mQuantBits == 3) ? UP_DIV(mPadK, 16u)
|
|
: (mQuantBits == 4) ? UP_DIV(mPadK, 8u)
|
|
: (mPadK / 4u);
|
|
|
|
if (initStaticResource) {
|
|
if (!_prepareQuantBuffersGPU(vkBn, mQuantCommon.get(), useFP16, mCi, mCo, mPadN, mBlockStride,
|
|
mDecodeWeightStrideWords, mQuantBits,
|
|
mQuantWeightBuffer, mQuantMetaBuffer)) {
|
|
return false;
|
|
}
|
|
|
|
std::vector<float> biasHost(mPadN, 0.0f);
|
|
if (nullptr != biasPtr) {
|
|
::memcpy(biasHost.data(), biasPtr, static_cast<size_t>(mCo) * sizeof(float));
|
|
}
|
|
if (useFP16) {
|
|
std::vector<int16_t> biasHalf(mPadN);
|
|
FLOAT_TO_HALF(biasHost.data(), biasHalf.data(), static_cast<int>(mPadN));
|
|
mBiasBuffer.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, mPadN * sizeof(int16_t), nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
|
|
VK_BUFFER_USAGE_TRANSFER_DST_BIT,
|
|
VK_SHARING_MODE_EXCLUSIVE, 0));
|
|
vkBn->copyToGPUBuffer(biasHalf.data(), mBiasBuffer->buffer(), mPadN * sizeof(int16_t), 0);
|
|
} else {
|
|
mBiasBuffer.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, mPadN * sizeof(float), nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
|
|
VK_BUFFER_USAGE_TRANSFER_DST_BIT,
|
|
VK_SHARING_MODE_EXCLUSIVE, 0));
|
|
vkBn->copyToGPUBuffer(biasHost.data(), mBiasBuffer->buffer(), mPadN * sizeof(float), 0);
|
|
}
|
|
}
|
|
|
|
int activation = 0;
|
|
if (mCommon->relu()) {
|
|
activation = 1;
|
|
}
|
|
if (mCommon->relu6()) {
|
|
activation = 2;
|
|
}
|
|
|
|
{
|
|
const auto& subgroup = vkBn->getDevice().getSubgroupInfo();
|
|
const VkSubgroupFeatureFlags requiredOps = VK_SUBGROUP_FEATURE_BASIC_BIT | VK_SUBGROUP_FEATURE_ARITHMETIC_BIT;
|
|
mUseSubgroup = subgroup.size > 0 &&
|
|
(subgroup.stages & VK_SHADER_STAGE_COMPUTE_BIT) &&
|
|
((subgroup.ops & requiredOps) == requiredOps);
|
|
}
|
|
|
|
mDecodeSubgroupSize = vkBn->getDevice().getSubgroupSize();
|
|
if (mDecodeSubgroupSize == 0u) {
|
|
mDecodeSubgroupSize = 64u;
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> types = {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
};
|
|
const char* shader = nullptr;
|
|
if (mUseSubgroup) {
|
|
std::vector<uint32_t> spec = {static_cast<uint32_t>(activation)};
|
|
switch (mQuantBits) {
|
|
case 2: shader = useFP16 ? "glsl_gemv_dequant_int2_FP16_comp" : "glsl_gemv_dequant_int2_comp"; break;
|
|
case 3: shader = useFP16 ? "glsl_gemv_dequant_int3_FP16_comp" : "glsl_gemv_dequant_int3_comp"; break;
|
|
case 4: shader = useFP16 ? "glsl_gemv_dequant_int4_FP16_comp" : "glsl_gemv_dequant_int4_comp"; break;
|
|
default: shader = useFP16 ? "glsl_gemv_dequant_int8_FP16_comp" : "glsl_gemv_dequant_int8_comp"; break;
|
|
}
|
|
mDecodePipeline = vkBn->getPipeline(shader, types, {mDecodeSubgroupSize, 1, 1}, spec);
|
|
} else {
|
|
uint32_t localSize = 64u;
|
|
std::vector<uint32_t> spec = {static_cast<uint32_t>(activation), localSize};
|
|
switch (mQuantBits) {
|
|
case 2: shader = useFP16 ? "glsl_gemv_dequant_int2_nosubgroup_FP16_comp" : "glsl_gemv_dequant_int2_nosubgroup_comp"; break;
|
|
case 3: shader = useFP16 ? "glsl_gemv_dequant_int3_nosubgroup_FP16_comp" : "glsl_gemv_dequant_int3_nosubgroup_comp"; break;
|
|
case 4: shader = useFP16 ? "glsl_gemv_dequant_int4_nosubgroup_FP16_comp" : "glsl_gemv_dequant_int4_nosubgroup_comp"; break;
|
|
default: shader = useFP16 ? "glsl_gemv_dequant_int8_nosubgroup_FP16_comp" : "glsl_gemv_dequant_int8_nosubgroup_comp"; break;
|
|
}
|
|
mDecodePipeline = vkBn->getPipeline(shader, types, {localSize, 1, 1}, spec);
|
|
}
|
|
if (nullptr == mDecodePipeline) {
|
|
return false;
|
|
}
|
|
mDecodeSet.reset(mDecodePipeline->createSet());
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> types = {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
};
|
|
const char* shader = useFP16 ? "glsl_pack_a_k4m4_to_m64k4_FP16_comp" : "glsl_pack_a_k4m4_to_m64k4_comp";
|
|
mPackAPipeline = vkBn->getPipeline(shader, types);
|
|
if (nullptr == mPackAPipeline) {
|
|
return false;
|
|
}
|
|
mPackASet.reset(mPackAPipeline->createSet());
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> types = {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
};
|
|
const char* shader = nullptr;
|
|
switch (mQuantBits) {
|
|
case 2: shader = useFP16 ? "glsl_int2_weight_to_pack_FP16_comp" : "glsl_int2_weight_to_pack_comp"; break;
|
|
case 3: shader = useFP16 ? "glsl_int3_weight_to_pack_FP16_comp" : "glsl_int3_weight_to_pack_comp"; break;
|
|
case 4: shader = useFP16 ? "glsl_int4_weight_to_pack_FP16_comp" : "glsl_int4_weight_to_pack_comp"; break;
|
|
default: shader = useFP16 ? "glsl_int8_weight_to_pack_FP16_comp" : "glsl_int8_weight_to_pack_comp"; break;
|
|
}
|
|
mWeightToPackPipeline = vkBn->getPipeline(shader, types);
|
|
if (nullptr == mWeightToPackPipeline) {
|
|
return false;
|
|
}
|
|
mWeightToPackSet.reset(mWeightToPackPipeline->createSet());
|
|
}
|
|
|
|
{
|
|
std::vector<VkDescriptorType> types = {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
};
|
|
std::vector<uint32_t> spec = {static_cast<uint32_t>(activation)};
|
|
const char* shader = useFP16 ? "glsl_gemm_m8n4_FP16_comp" : "glsl_gemm_m8n4_comp";
|
|
mGemmPipeline = vkBn->getPipeline(shader, types, {}, spec);
|
|
if (nullptr == mGemmPipeline) {
|
|
return false;
|
|
}
|
|
mGemmSet.reset(mGemmPipeline->createSet());
|
|
}
|
|
|
|
return mDecodeSet != nullptr && mPackASet != nullptr && mWeightToPackSet != nullptr && mGemmSet != nullptr;
|
|
}
|
|
|
|
bool VulkanConv1x1General::onClone(Backend* bn, const Op* op, VulkanBasicExecution** dst) {
|
|
if (nullptr == dst) {
|
|
return true;
|
|
}
|
|
auto vkBn = static_cast<VulkanBackend*>(bn);
|
|
auto conv2D = op->main_as_Convolution2D();
|
|
if (nullptr == conv2D || nullptr == conv2D->common()) {
|
|
return false;
|
|
}
|
|
auto res = new VulkanConv1x1General(vkBn, conv2D->common(), mCi, mCo, mQuantCommon, false);
|
|
res->mQuantBits = mQuantBits;
|
|
res->mPadK = mPadK;
|
|
res->mPadN = mPadN;
|
|
res->mBlockSize = mBlockSize;
|
|
res->mBlockStride = mBlockStride;
|
|
res->mDecodeWeightStrideWords = mDecodeWeightStrideWords;
|
|
res->mDecodeSubgroupSize = mDecodeSubgroupSize;
|
|
res->mUseSubgroup = mUseSubgroup;
|
|
res->mQuantWeightBuffer = mQuantWeightBuffer;
|
|
res->mQuantMetaBuffer = mQuantMetaBuffer;
|
|
res->mBiasBuffer = mBiasBuffer;
|
|
*dst = res;
|
|
return true;
|
|
}
|
|
|
|
ErrorCode VulkanConv1x1General::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
|
const VulkanCommandPool::Buffer* cmdBuffer) {
|
|
if (inputs.empty() || outputs.empty() || nullptr == mQuantWeightBuffer.get() || nullptr == mQuantMetaBuffer.get() ||
|
|
nullptr == mBiasBuffer.get()) {
|
|
return INVALID_VALUE;
|
|
}
|
|
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
const int M = output->batch() * output->height() * output->width();
|
|
if (M <= 0 || mCi <= 0 || mCo <= 0) {
|
|
return NO_ERROR;
|
|
}
|
|
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
auto srcBuffer = vkBn->getTensorBuffer(input);
|
|
auto dstBuffer = vkBn->getTensorBuffer(output);
|
|
|
|
if (M == 1) {
|
|
struct DecodeParams {
|
|
uint32_t K;
|
|
uint32_t N;
|
|
uint32_t blockSize;
|
|
uint32_t blockStride;
|
|
uint32_t weightStride;
|
|
} pc;
|
|
pc.K = static_cast<uint32_t>(mCi);
|
|
pc.N = static_cast<uint32_t>(mCo);
|
|
pc.blockSize = mBlockSize;
|
|
pc.blockStride = mBlockStride;
|
|
pc.weightStride = mDecodeWeightStrideWords;
|
|
|
|
mDecodeSet->writeBuffer(srcBuffer.first->buffer(), 0, vkBn->getTensorSize(input), srcBuffer.second);
|
|
mDecodeSet->writeBuffer(mQuantWeightBuffer->buffer(), 1, mQuantWeightBuffer->size());
|
|
mDecodeSet->writeBuffer(mQuantMetaBuffer->buffer(), 2, mQuantMetaBuffer->size());
|
|
mDecodeSet->writeBuffer(mBiasBuffer->buffer(), 3, mBiasBuffer->size());
|
|
mDecodeSet->writeBuffer(dstBuffer.first->buffer(), 4, vkBn->getTensorSize(output), dstBuffer.second);
|
|
mDecodePipeline->bind(cmdBuffer->get(), mDecodeSet->get());
|
|
vkCmdPushConstants(cmdBuffer->get(), mDecodePipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0,
|
|
sizeof(DecodeParams), &pc);
|
|
vkCmdDispatch(cmdBuffer->get(), static_cast<uint32_t>(mCo), 1, 1);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
const uint32_t padM = ROUND_UP(static_cast<uint32_t>(M), 64u);
|
|
if ((mPadK & 3u) != 0u || (mPadN & 31u) != 0u || padM == 0u) {
|
|
return INVALID_VALUE;
|
|
}
|
|
|
|
if (vkBn->useFP16()) {
|
|
mTempInputPacked.reset(Tensor::createDevice<int16_t>({static_cast<int>(padM), static_cast<int>(mPadK)}));
|
|
mTempWeightPacked.reset(Tensor::createDevice<int16_t>({static_cast<int>(mPadK), static_cast<int>(mPadN)}));
|
|
} else {
|
|
mTempInputPacked.reset(Tensor::createDevice<float>({static_cast<int>(padM), static_cast<int>(mPadK)}));
|
|
mTempWeightPacked.reset(Tensor::createDevice<float>({static_cast<int>(mPadK), static_cast<int>(mPadN)}));
|
|
}
|
|
|
|
bool acquiredTempA = false;
|
|
bool acquiredTempB = false;
|
|
auto releaseTemp = [&]() {
|
|
if (acquiredTempA && nullptr != mTempInputPacked.get()) {
|
|
vkBn->onReleaseBuffer(mTempInputPacked.get(), Backend::DYNAMIC);
|
|
acquiredTempA = false;
|
|
}
|
|
if (acquiredTempB && nullptr != mTempWeightPacked.get()) {
|
|
vkBn->onReleaseBuffer(mTempWeightPacked.get(), Backend::DYNAMIC);
|
|
acquiredTempB = false;
|
|
}
|
|
};
|
|
|
|
if (!vkBn->onAcquireBuffer(mTempInputPacked.get(), Backend::DYNAMIC)) {
|
|
return OUT_OF_MEMORY;
|
|
}
|
|
acquiredTempA = true;
|
|
|
|
if (!vkBn->onAcquireBuffer(mTempWeightPacked.get(), Backend::DYNAMIC)) {
|
|
releaseTemp();
|
|
return OUT_OF_MEMORY;
|
|
}
|
|
acquiredTempB = true;
|
|
|
|
auto packedABuffer = vkBn->getTensorBuffer(mTempInputPacked.get());
|
|
auto packedBBuffer = vkBn->getTensorBuffer(mTempWeightPacked.get());
|
|
const size_t packedASize = vkBn->getTensorSize(mTempInputPacked.get());
|
|
const size_t packedBSize = vkBn->getTensorSize(mTempWeightPacked.get());
|
|
|
|
{
|
|
struct PackAParams {
|
|
uint32_t M;
|
|
uint32_t K;
|
|
} pc;
|
|
pc.M = static_cast<uint32_t>(M);
|
|
pc.K = mPadK;
|
|
|
|
mPackASet->writeBuffer(srcBuffer.first->buffer(), 0, vkBn->getTensorSize(input), srcBuffer.second);
|
|
mPackASet->writeBuffer(packedABuffer.first->buffer(), 1, packedASize, packedABuffer.second);
|
|
mPackAPipeline->bind(cmdBuffer->get(), mPackASet->get());
|
|
vkCmdPushConstants(cmdBuffer->get(), mPackAPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0,
|
|
sizeof(PackAParams), &pc);
|
|
vkCmdDispatch(cmdBuffer->get(), mPadK / 4u, padM / 64u, 1);
|
|
cmdBuffer->barrierSource(packedABuffer.first->buffer(), packedABuffer.second, packedASize);
|
|
}
|
|
|
|
{
|
|
struct WeightToPackParams {
|
|
uint32_t N;
|
|
uint32_t K;
|
|
uint32_t blockSize;
|
|
uint32_t KBlocks;
|
|
} pc;
|
|
pc.N = mPadN;
|
|
pc.K = mPadK;
|
|
pc.blockSize = mBlockSize;
|
|
pc.KBlocks = mBlockStride;
|
|
|
|
mWeightToPackSet->writeBuffer(mQuantWeightBuffer->buffer(), 0, mQuantWeightBuffer->size());
|
|
mWeightToPackSet->writeBuffer(mQuantMetaBuffer->buffer(), 1, mQuantMetaBuffer->size());
|
|
mWeightToPackSet->writeBuffer(packedBBuffer.first->buffer(), 2, packedBSize, packedBBuffer.second);
|
|
mWeightToPackPipeline->bind(cmdBuffer->get(), mWeightToPackSet->get());
|
|
vkCmdPushConstants(cmdBuffer->get(), mWeightToPackPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0,
|
|
sizeof(WeightToPackParams), &pc);
|
|
vkCmdDispatch(cmdBuffer->get(), UP_DIV(mPadN / 4u, 16u), UP_DIV(mPadK / 4u, 8u), 1);
|
|
cmdBuffer->barrierSource(packedBBuffer.first->buffer(), packedBBuffer.second, packedBSize);
|
|
}
|
|
|
|
{
|
|
struct GemmParams {
|
|
uint32_t M;
|
|
uint32_t N;
|
|
uint32_t K;
|
|
uint32_t padN;
|
|
} pc;
|
|
pc.M = static_cast<uint32_t>(M);
|
|
pc.N = static_cast<uint32_t>(mCo);
|
|
pc.K = mPadK;
|
|
pc.padN = mPadN;
|
|
|
|
mGemmSet->writeBuffer(packedABuffer.first->buffer(), 0, packedASize, packedABuffer.second);
|
|
mGemmSet->writeBuffer(packedBBuffer.first->buffer(), 1, packedBSize, packedBBuffer.second);
|
|
mGemmSet->writeBuffer(mBiasBuffer->buffer(), 2, mBiasBuffer->size());
|
|
mGemmSet->writeBuffer(dstBuffer.first->buffer(), 3, vkBn->getTensorSize(output), dstBuffer.second);
|
|
mGemmPipeline->bind(cmdBuffer->get(), mGemmSet->get());
|
|
vkCmdPushConstants(cmdBuffer->get(), mGemmPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0,
|
|
sizeof(GemmParams), &pc);
|
|
vkCmdDispatch(cmdBuffer->get(), mPadN / 32u, padM / 64u, 1);
|
|
}
|
|
|
|
releaseTemp();
|
|
return NO_ERROR;
|
|
}
|
|
|
|
} // namespace MNN
|