wehub-resource-sync
642 lines
30 KiB
C++
642 lines
30 KiB
C++
#include "VulkanConv1x1CoopA8.hpp"
|
|
#include "core/TensorUtils.hpp"
|
|
#include "core/Macro.h"
|
|
#include "VulkanBackend.hpp"
|
|
#include "backend/vulkan/vulkan/vulkan_wrapper.h"
|
|
#include <algorithm>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
|
|
#ifdef ENABLE_VULKAN_TIME_PROFILE
|
|
#include "backend/vulkan/component/VulkanTimeProfiler.hpp"
|
|
#endif
|
|
|
|
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;
|
|
};
|
|
|
|
struct SumKPrepareParams {
|
|
uint32_t padN;
|
|
uint32_t weightStride;
|
|
};
|
|
|
|
// Builds the static GPU buffers shared by decode + prefill paths:
|
|
// - quantWeightBuffer: INT8: [padN, padK/4] uint32 (4x int8 packed).
|
|
// INT4: [padN, padK/8] uint32 (8x int4 packed,
|
|
// unsigned 0..15 with host +8 offset; padding
|
|
// nibbles are 0x8 -> decoded 0).
|
|
// - quantMetaBuffer: [padN, 2] FP interleaved (scale, offset).
|
|
// - sumWqBuffer: [padN] int32 = sum_k Wq[n, k], padding rows zero.
|
|
// Per-channel asym only — caller must guard.
|
|
static bool _prepareStaticBuffersGPU(VulkanBackend* vkBn, const ConvolutionCommon::Int8Common* quantCommon,
|
|
bool useFP16, int ci, int co, uint32_t padN, uint32_t padK,
|
|
std::shared_ptr<VulkanBuffer>& quantWeightBuffer,
|
|
std::shared_ptr<VulkanBuffer>& quantMetaBuffer,
|
|
std::shared_ptr<VulkanBuffer>& sumWqBuffer) {
|
|
if (nullptr == vkBn || nullptr == quantCommon || nullptr == quantCommon->weight.get()) {
|
|
return false;
|
|
}
|
|
MNN_ASSERT(quantCommon->asymmetric);
|
|
const bool isInt4 = quantCommon->canUseInt4;
|
|
if (isInt4) {
|
|
MNN_ASSERT(padK % 8u == 0u);
|
|
}
|
|
|
|
const int soSize = 2;
|
|
const int alphaSize = quantCommon->alpha.size();
|
|
const int blockCount = std::max(1, alphaSize / std::max(1, co * soSize));
|
|
if (blockCount != 1) {
|
|
return false;
|
|
}
|
|
const uint32_t blockStride = 1u;
|
|
const uint32_t decodeWeightStrideWords = isInt4 ? (padK / 8u) : (padK / 4u);
|
|
|
|
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, ALIGN_UP4(rawWeightBytes));
|
|
const size_t decodeWeightBytes =
|
|
static_cast<size_t>(padN) * static_cast<size_t>(decodeWeightStrideWords) * sizeof(uint32_t);
|
|
const size_t metaBytes = static_cast<size_t>(padN) * static_cast<size_t>(blockStride) * 2u
|
|
* (useFP16 ? sizeof(int16_t) : sizeof(float));
|
|
const size_t sumWqBytes = static_cast<size_t>(padN) * sizeof(int32_t);
|
|
|
|
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));
|
|
sumWqBuffer.reset(new VulkanBuffer(vkBn->getMemoryPool(), false, sumWqBytes, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_SHARING_MODE_EXCLUSIVE, 0));
|
|
|
|
const char* weightShader = isInt4 ? "glsl_conv1x1_int4_weight_prepare_comp"
|
|
: "glsl_conv1x1_int8_weight_prepare_comp";
|
|
const char* metaShader = useFP16 ? "glsl_conv1x1_quant_meta_prepare_FP16_comp"
|
|
: "glsl_conv1x1_quant_meta_prepare_comp";
|
|
const char* sumKShader = isInt4 ? "glsl_conv1x1_int4_weight_sumK_comp"
|
|
: "glsl_conv1x1_int8_weight_sumK_comp";
|
|
|
|
std::vector<VkDescriptorType> twoBufTypes = {
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
|
|
};
|
|
auto weightPipeline = vkBn->getPipeline(weightShader, twoBufTypes);
|
|
auto metaPipeline = vkBn->getPipeline(metaShader, twoBufTypes);
|
|
auto sumKPipeline = vkBn->getPipeline(sumKShader, twoBufTypes);
|
|
if (nullptr == weightPipeline || nullptr == metaPipeline || nullptr == sumKPipeline) {
|
|
return false;
|
|
}
|
|
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> weightSet(weightPipeline->createSet());
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> metaSet(metaPipeline->createSet());
|
|
std::shared_ptr<VulkanLayout::DescriptorSet> sumKSet(sumKPipeline->createSet());
|
|
if (nullptr == weightSet.get() || nullptr == metaSet.get() || nullptr == sumKSet.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());
|
|
}
|
|
|
|
{
|
|
SumKPrepareParams pc;
|
|
pc.padN = padN;
|
|
pc.weightStride = decodeWeightStrideWords;
|
|
|
|
sumKSet->writeBuffer(quantWeightBuffer->buffer(), 0, quantWeightBuffer->size());
|
|
sumKSet->writeBuffer(sumWqBuffer->buffer(), 1, sumWqBuffer->size());
|
|
sumKPipeline->bind(prepareCmd->get(), sumKSet->get());
|
|
vkCmdPushConstants(prepareCmd->get(), sumKPipeline->layout(), VK_SHADER_STAGE_COMPUTE_BIT,
|
|
0, sizeof(pc), &pc);
|
|
vkCmdDispatch(prepareCmd->get(), UP_DIV(padN, 64u), 1, 1);
|
|
}
|
|
|
|
prepareCmd->end();
|
|
vkBn->getPool().submitAndWait(prepareCmd->get());
|
|
return true;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
VulkanConv1x1CoopA8::VulkanConv1x1CoopA8(VulkanBackend* backend, const Convolution2DCommon* convOption,
|
|
const float* biasPtr, int ci, int co,
|
|
VulkanDevice::CoopMatInfo coopMatInfo,
|
|
std::shared_ptr<ConvolutionCommon::Int8Common> quantInfo)
|
|
: VulkanBasicExecution(backend),
|
|
mCommon(convOption),
|
|
mCi(ci),
|
|
mCo(co),
|
|
mQuantCommon(std::move(quantInfo)) {
|
|
MNN_ASSERT(coopMatInfo.supportS8S8S32);
|
|
MNN_ASSERT(coopMatInfo.selectedS8CoopMatShape.size() == 3);
|
|
mCoopM = coopMatInfo.selectedS8CoopMatShape[0];
|
|
mCoopN = coopMatInfo.selectedS8CoopMatShape[1];
|
|
mCoopK = coopMatInfo.selectedS8CoopMatShape[2];
|
|
uint32_t subgroupSize = backend->getDevice().getSubgroupSize();
|
|
MNN_ASSERT(subgroupSize > 0);
|
|
mSubgroupSize = subgroupSize;
|
|
_init(biasPtr, true);
|
|
}
|
|
|
|
VulkanConv1x1CoopA8::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)
|
|
: VulkanBasicExecution(backend),
|
|
mCommon(convOption),
|
|
mCi(ci),
|
|
mCo(co),
|
|
mQuantCommon(std::move(quantInfo)) {
|
|
mCoopM = coopM;
|
|
mCoopN = coopN;
|
|
mCoopK = coopK;
|
|
MNN_ASSERT(subgroupSize > 0);
|
|
mSubgroupSize = subgroupSize;
|
|
_init(nullptr, initStaticResource);
|
|
}
|
|
|
|
VulkanConv1x1CoopA8::~VulkanConv1x1CoopA8() {
|
|
}
|
|
|
|
bool VulkanConv1x1CoopA8::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 VulkanConv1x1CoopA8(vkBn, conv2D->common(), mCi, mCo,
|
|
mCoopM, mCoopN, mCoopK, mSubgroupSize,
|
|
mQuantCommon, false);
|
|
res->mPadK = mPadK;
|
|
res->mPadN = mPadN;
|
|
res->mBiasBuffer = mBiasBuffer;
|
|
res->mQuantWeightBuffer = mQuantWeightBuffer;
|
|
res->mQuantMetaBuffer = mQuantMetaBuffer;
|
|
res->mSumWqBuffer = mSumWqBuffer;
|
|
*dst = res;
|
|
return true;
|
|
}
|
|
|
|
bool VulkanConv1x1CoopA8::_init(const float* biasPtr, bool initStaticResource) {
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
const bool useFP16 = vkBn->useFP16();
|
|
const size_t fpElem = useFP16 ? sizeof(int16_t) : sizeof(float);
|
|
mIsInt4 = mQuantCommon != nullptr && mQuantCommon->canUseInt4;
|
|
|
|
const uint32_t K = mCi;
|
|
const uint32_t N = mCo;
|
|
mPadK = ROUND_UP(K, mCoopK);
|
|
mPadN = ROUND_UP(N, mCoopN);
|
|
|
|
if (initStaticResource) {
|
|
mBiasBuffer = std::make_shared<VulkanBuffer>(vkBn->getMemoryPool(), false, fpElem * mPadN, nullptr,
|
|
VK_BUFFER_USAGE_STORAGE_BUFFER_BIT);
|
|
auto biasMap = mBiasBuffer->map();
|
|
::memset(biasMap, 0, mPadN * fpElem);
|
|
if (biasPtr) {
|
|
if (useFP16) {
|
|
std::vector<int16_t> biasFP16(N);
|
|
FLOAT_TO_HALF(biasPtr, biasFP16.data(), N);
|
|
::memcpy(biasMap, biasFP16.data(), N * sizeof(int16_t));
|
|
} else {
|
|
::memcpy(biasMap, biasPtr, N * sizeof(float));
|
|
}
|
|
}
|
|
mBiasBuffer->unmap();
|
|
|
|
if (!_prepareStaticBuffersGPU(vkBn, mQuantCommon.get(), useFP16, mCi, mCo, mPadN, mPadK,
|
|
mQuantWeightBuffer, mQuantMetaBuffer, mSumWqBuffer)) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
int activation = 0;
|
|
if (mCommon->relu()) {
|
|
activation = 1;
|
|
} else if (mCommon->relu6()) {
|
|
activation = 2;
|
|
}
|
|
|
|
// Decode (M == 1): reuse Coop-A16's gemv_dequant_int{4,8} — per-channel is
|
|
// the single-block subset of block-quant. INT4 / INT8 share the binding
|
|
// layout and push-constant struct; only weightStride and shader name vary
|
|
// (handled in onEncode).
|
|
{
|
|
std::vector<VkDescriptorType> types(5, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
|
|
std::vector<uint32_t> localSize = {mSubgroupSize, 1, 1};
|
|
std::vector<uint32_t> spec = {(uint32_t)activation};
|
|
const char* shader;
|
|
if (mIsInt4) {
|
|
shader = useFP16 ? "glsl_gemv_dequant_int4_FP16_comp" : "glsl_gemv_dequant_int4_comp";
|
|
} else {
|
|
shader = useFP16 ? "glsl_gemv_dequant_int8_FP16_comp" : "glsl_gemv_dequant_int8_comp";
|
|
}
|
|
mDecodePipeline = vkBn->getPipeline(shader, types, localSize, spec);
|
|
mDecodeSet.reset(mDecodePipeline->createSet());
|
|
}
|
|
|
|
// INT4-only: runtime nibble unpack pipeline. Idempotent across INT8 mode
|
|
// (skipped on dispatch side); registered here only when needed.
|
|
if (mIsInt4) {
|
|
std::vector<VkDescriptorType> types(2, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
|
|
std::vector<uint32_t> localSize = {16u, 16u, 1u};
|
|
mInt4UnpackPipeline = vkBn->getPipeline("glsl_dynamic_int4_to_int8_unpack_comp", types, localSize, {});
|
|
mInt4UnpackSet.reset(mInt4UnpackPipeline->createSet());
|
|
}
|
|
|
|
// Prefill stages — spec constant_id starts at 3 (after local_size_x/y/z).
|
|
{
|
|
std::vector<VkDescriptorType> types(3, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
|
|
std::vector<uint32_t> localSize = {mCoopM, 1, 1};
|
|
std::vector<uint32_t> spec = {mCoopM};
|
|
const char* shader = useFP16 ? "glsl_dynamic_quant_minmax_FP16_comp"
|
|
: "glsl_dynamic_quant_minmax_comp";
|
|
mQuantMinMaxPipeline = vkBn->getPipeline(shader, types, localSize, spec);
|
|
mQuantMinMaxSet.reset(mQuantMinMaxPipeline->createSet());
|
|
}
|
|
{
|
|
std::vector<VkDescriptorType> types(4, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
|
|
std::vector<uint32_t> localSize = {64u, 1u, 1u};
|
|
const char* shader = useFP16 ? "glsl_dynamic_quant_reduce_minmax_FP16_comp"
|
|
: "glsl_dynamic_quant_reduce_minmax_comp";
|
|
mQuantReduceMinMaxPipeline = vkBn->getPipeline(shader, types, localSize, {});
|
|
mQuantReduceMinMaxSet.reset(mQuantReduceMinMaxPipeline->createSet());
|
|
}
|
|
{
|
|
std::vector<VkDescriptorType> types(5, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
|
|
const uint32_t packLocalX = std::max(mCoopM, mCoopK);
|
|
std::vector<uint32_t> localSize = {packLocalX, 1u, 1u};
|
|
std::vector<uint32_t> spec = {mCoopM, mCoopK};
|
|
const char* shader = useFP16 ? "glsl_dynamic_quant_pack_FP16_comp"
|
|
: "glsl_dynamic_quant_pack_comp";
|
|
mQuantPackPipeline = vkBn->getPipeline(shader, types, localSize, spec);
|
|
mQuantPackSet.reset(mQuantPackPipeline->createSet());
|
|
}
|
|
{
|
|
std::vector<VkDescriptorType> types(2, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
|
|
std::vector<uint32_t> localSize = {64u, 1u, 1u};
|
|
mQuantReduceSumPipeline = vkBn->getPipeline("glsl_dynamic_quant_reduce_sum_comp", types, localSize, {});
|
|
mQuantReduceSumSet.reset(mQuantReduceSumPipeline->createSet());
|
|
}
|
|
{
|
|
// GEMM spec constants 3..9: COOP_M, COOP_N, COOP_K, A_COL_MAJOR=0,
|
|
// B_COL_MAJOR=1, A_BLOCK_LINEAR=1, B_BLOCK_LINEAR=0.
|
|
std::vector<VkDescriptorType> types(3, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
|
|
std::vector<uint32_t> localSize = {mSubgroupSize, 1u, 1u};
|
|
std::vector<uint32_t> spec = {mCoopM, mCoopN, mCoopK, 0u, 1u, 1u, 0u};
|
|
mGemmS8Pipeline = vkBn->getPipeline("glsl_dynamic_w8a8_coop_gemm_comp", types, localSize, spec);
|
|
mGemmSet.reset(mGemmS8Pipeline->createSet());
|
|
}
|
|
{
|
|
std::vector<VkDescriptorType> types(8, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
|
|
std::vector<uint32_t> localSize = {16u, 16u, 1u};
|
|
std::vector<uint32_t> spec = {(uint32_t)activation};
|
|
const char* shader = useFP16 ? "glsl_dynamic_w8a8_dequant_correction_FP16_comp"
|
|
: "glsl_dynamic_w8a8_dequant_correction_comp";
|
|
mDequantPipeline = vkBn->getPipeline(shader, types, localSize, spec);
|
|
mDequantSet.reset(mDequantPipeline->createSet());
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
ErrorCode VulkanConv1x1CoopA8::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
|
const VulkanCommandPool::Buffer* cmdBuffer) {
|
|
auto input = inputs[0];
|
|
auto output = outputs[0];
|
|
auto vkBn = static_cast<VulkanBackend*>(backend());
|
|
|
|
const int M = input->batch() * input->width() * input->height();
|
|
const int K = mCi;
|
|
const int N = mCo;
|
|
|
|
const uint32_t padK = mPadK;
|
|
const uint32_t padN = mPadN;
|
|
|
|
auto srcBuffer = vkBn->getTensorBuffer(input);
|
|
auto dstBuffer = vkBn->getTensorBuffer(output);
|
|
const bool useFP16 = vkBn->useFP16();
|
|
|
|
// Per-shader timing scope (no-op when ENABLE_VULKAN_TIME_PROFILE is off).
|
|
// The label MUST match the shader registry key passed to getPipeline so
|
|
// the profile output and the shader registry are 1:1 traceable.
|
|
auto dispatchWithProfile = [&](const char* name, const VulkanPipeline* pipeline,
|
|
const std::shared_ptr<VulkanLayout::DescriptorSet>& set,
|
|
uint32_t gx, uint32_t gy, uint32_t gz,
|
|
const void* pc, uint32_t pcSize) {
|
|
#ifdef ENABLE_VULKAN_TIME_PROFILE
|
|
auto* profiler = vkBn->timeProfiler();
|
|
if (nullptr != profiler) {
|
|
VulkanTimeProfileScope scope(profiler, cmdBuffer->get(), name,
|
|
VulkanTimeProfiler::Kind::Shader);
|
|
pipeline->bind(cmdBuffer->get(), set->get());
|
|
if (pc != nullptr) {
|
|
vkCmdPushConstants(cmdBuffer->get(), pipeline->layout(),
|
|
VK_SHADER_STAGE_COMPUTE_BIT, 0, pcSize, pc);
|
|
}
|
|
vkCmdDispatch(cmdBuffer->get(), gx, gy, gz);
|
|
return;
|
|
}
|
|
#else
|
|
(void)name;
|
|
#endif
|
|
pipeline->bind(cmdBuffer->get(), set->get());
|
|
if (pc != nullptr) {
|
|
vkCmdPushConstants(cmdBuffer->get(), pipeline->layout(),
|
|
VK_SHADER_STAGE_COMPUTE_BIT, 0, pcSize, pc);
|
|
}
|
|
vkCmdDispatch(cmdBuffer->get(), gx, gy, gz);
|
|
};
|
|
|
|
if (M == 1) {
|
|
struct DecodeParams {
|
|
uint32_t K;
|
|
uint32_t N;
|
|
uint32_t blockSize;
|
|
uint32_t blockStride;
|
|
uint32_t weightStride;
|
|
} pc;
|
|
pc.K = (uint32_t)K;
|
|
pc.N = (uint32_t)N;
|
|
pc.blockSize = (uint32_t)K;
|
|
pc.blockStride = 1u;
|
|
pc.weightStride = mIsInt4 ? (padK / 8u) : (padK / 4u);
|
|
|
|
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);
|
|
const char* decodeName;
|
|
if (mIsInt4) {
|
|
decodeName = useFP16 ? "glsl_gemv_dequant_int4_FP16_comp" : "glsl_gemv_dequant_int4_comp";
|
|
} else {
|
|
decodeName = useFP16 ? "glsl_gemv_dequant_int8_FP16_comp" : "glsl_gemv_dequant_int8_comp";
|
|
}
|
|
dispatchWithProfile(decodeName, mDecodePipeline, mDecodeSet,
|
|
(uint32_t)N, 1, 1, &pc, sizeof(pc));
|
|
return NO_ERROR;
|
|
}
|
|
|
|
const uint32_t padM = ROUND_UP(M, mCoopM);
|
|
const uint32_t k4Count = UP_DIV((uint32_t)K, 4u);
|
|
// 64 below mirrors K4_BLOCK in dynamic_quant_minmax.comp — keep in sync.
|
|
const uint32_t partialBlocks = std::max(1u, UP_DIV(k4Count, 64u));
|
|
const uint32_t tilesK = padK / mCoopK;
|
|
const halide_type_t fpType = useFP16 ? halide_type_of<int16_t>() : halide_type_of<float>();
|
|
|
|
std::shared_ptr<Tensor> tPartialMin(Tensor::createDevice({(int)partialBlocks, (int)padM}, fpType));
|
|
std::shared_ptr<Tensor> tPartialMax(Tensor::createDevice({(int)partialBlocks, (int)padM}, fpType));
|
|
std::shared_ptr<Tensor> tScaleA(Tensor::createDevice({(int)padM}, fpType));
|
|
std::shared_ptr<Tensor> tOffsetA(Tensor::createDevice({(int)padM}, fpType));
|
|
std::shared_ptr<Tensor> tAq(Tensor::createDevice<int8_t>({(int)padM, (int)padK}));
|
|
std::shared_ptr<Tensor> tPartialSumAq(Tensor::createDevice<int32_t>({(int)tilesK, (int)padM}));
|
|
std::shared_ptr<Tensor> tSumAq(Tensor::createDevice<int32_t>({(int)padM}));
|
|
std::shared_ptr<Tensor> tAcc(Tensor::createDevice<int32_t>({(int)padM, (int)padN}));
|
|
// INT4 only: unpacked int8 weight buffer that the GEMM reads in place of
|
|
// mQuantWeightBuffer. Same lifetime as the other DYNAMIC tensors.
|
|
std::shared_ptr<Tensor> tWqInt8;
|
|
if (mIsInt4) {
|
|
tWqInt8.reset(Tensor::createDevice<int8_t>({(int)padN, (int)padK}));
|
|
}
|
|
|
|
std::vector<Tensor*> dyns = {tPartialMin.get(), tPartialMax.get(), tScaleA.get(), tOffsetA.get(),
|
|
tAq.get(), tPartialSumAq.get(), tSumAq.get(), tAcc.get()};
|
|
if (mIsInt4) {
|
|
dyns.push_back(tWqInt8.get());
|
|
}
|
|
size_t acquired = 0;
|
|
for (Tensor* t : dyns) {
|
|
if (!vkBn->onAcquireBuffer(t, Backend::DYNAMIC)) {
|
|
for (size_t j = 0; j < acquired; ++j) {
|
|
vkBn->onReleaseBuffer(dyns[j], Backend::DYNAMIC);
|
|
}
|
|
return OUT_OF_MEMORY;
|
|
}
|
|
++acquired;
|
|
}
|
|
|
|
auto bPartialMin = vkBn->getTensorBuffer(tPartialMin.get());
|
|
auto bPartialMax = vkBn->getTensorBuffer(tPartialMax.get());
|
|
auto bScaleA = vkBn->getTensorBuffer(tScaleA.get());
|
|
auto bOffsetA = vkBn->getTensorBuffer(tOffsetA.get());
|
|
auto bAq = vkBn->getTensorBuffer(tAq.get());
|
|
auto bPartialSum = vkBn->getTensorBuffer(tPartialSumAq.get());
|
|
auto bSumAq = vkBn->getTensorBuffer(tSumAq.get());
|
|
auto bAcc = vkBn->getTensorBuffer(tAcc.get());
|
|
const size_t szPartialMin = vkBn->getTensorSize(tPartialMin.get());
|
|
const size_t szPartialMax = vkBn->getTensorSize(tPartialMax.get());
|
|
const size_t szScaleA = vkBn->getTensorSize(tScaleA.get());
|
|
const size_t szOffsetA = vkBn->getTensorSize(tOffsetA.get());
|
|
const size_t szAq = vkBn->getTensorSize(tAq.get());
|
|
const size_t szPartialSum = vkBn->getTensorSize(tPartialSumAq.get());
|
|
const size_t szSumAq = vkBn->getTensorSize(tSumAq.get());
|
|
const size_t szAcc = vkBn->getTensorSize(tAcc.get());
|
|
|
|
{
|
|
struct PC { uint32_t M, K, padM, partialBlocks; } pc;
|
|
pc.M = (uint32_t)M; pc.K = (uint32_t)K; pc.padM = padM; pc.partialBlocks = partialBlocks;
|
|
|
|
mQuantMinMaxSet->writeBuffer(srcBuffer.first->buffer(), 0, vkBn->getTensorSize(input), srcBuffer.second);
|
|
mQuantMinMaxSet->writeBuffer(bPartialMin.first->buffer(), 1, szPartialMin, bPartialMin.second);
|
|
mQuantMinMaxSet->writeBuffer(bPartialMax.first->buffer(), 2, szPartialMax, bPartialMax.second);
|
|
dispatchWithProfile(useFP16 ? "glsl_dynamic_quant_minmax_FP16_comp"
|
|
: "glsl_dynamic_quant_minmax_comp",
|
|
mQuantMinMaxPipeline, mQuantMinMaxSet,
|
|
padM / mCoopM, partialBlocks, 1, &pc, sizeof(pc));
|
|
cmdBuffer->barrierSource(bPartialMin.first->buffer(), bPartialMin.second, szPartialMin);
|
|
cmdBuffer->barrierSource(bPartialMax.first->buffer(), bPartialMax.second, szPartialMax);
|
|
}
|
|
|
|
{
|
|
struct PC { uint32_t M, padM, partialBlocks, reserved; } pc;
|
|
pc.M = (uint32_t)M; pc.padM = padM; pc.partialBlocks = partialBlocks; pc.reserved = 0;
|
|
|
|
mQuantReduceMinMaxSet->writeBuffer(bPartialMin.first->buffer(), 0, szPartialMin, bPartialMin.second);
|
|
mQuantReduceMinMaxSet->writeBuffer(bPartialMax.first->buffer(), 1, szPartialMax, bPartialMax.second);
|
|
mQuantReduceMinMaxSet->writeBuffer(bScaleA.first->buffer(), 2, szScaleA, bScaleA.second);
|
|
mQuantReduceMinMaxSet->writeBuffer(bOffsetA.first->buffer(), 3, szOffsetA, bOffsetA.second);
|
|
dispatchWithProfile(useFP16 ? "glsl_dynamic_quant_reduce_minmax_FP16_comp"
|
|
: "glsl_dynamic_quant_reduce_minmax_comp",
|
|
mQuantReduceMinMaxPipeline, mQuantReduceMinMaxSet,
|
|
UP_DIV(padM, 64u), 1, 1, &pc, sizeof(pc));
|
|
cmdBuffer->barrierSource(bScaleA.first->buffer(), bScaleA.second, szScaleA);
|
|
cmdBuffer->barrierSource(bOffsetA.first->buffer(), bOffsetA.second, szOffsetA);
|
|
}
|
|
|
|
{
|
|
struct PC { uint32_t M, K, padM, padK; } pc;
|
|
pc.M = (uint32_t)M; pc.K = (uint32_t)K; pc.padM = padM; pc.padK = padK;
|
|
|
|
mQuantPackSet->writeBuffer(srcBuffer.first->buffer(), 0, vkBn->getTensorSize(input), srcBuffer.second);
|
|
mQuantPackSet->writeBuffer(bScaleA.first->buffer(), 1, szScaleA, bScaleA.second);
|
|
mQuantPackSet->writeBuffer(bOffsetA.first->buffer(), 2, szOffsetA, bOffsetA.second);
|
|
mQuantPackSet->writeBuffer(bAq.first->buffer(), 3, szAq, bAq.second);
|
|
mQuantPackSet->writeBuffer(bPartialSum.first->buffer(), 4, szPartialSum, bPartialSum.second);
|
|
dispatchWithProfile(useFP16 ? "glsl_dynamic_quant_pack_FP16_comp"
|
|
: "glsl_dynamic_quant_pack_comp",
|
|
mQuantPackPipeline, mQuantPackSet,
|
|
padM / mCoopM, tilesK, 1, &pc, sizeof(pc));
|
|
cmdBuffer->barrierSource(bAq.first->buffer(), bAq.second, szAq);
|
|
cmdBuffer->barrierSource(bPartialSum.first->buffer(), bPartialSum.second, szPartialSum);
|
|
}
|
|
|
|
{
|
|
struct PC { uint32_t M, padM, tilesK, reserved; } pc;
|
|
pc.M = (uint32_t)M; pc.padM = padM; pc.tilesK = tilesK; pc.reserved = 0;
|
|
|
|
mQuantReduceSumSet->writeBuffer(bPartialSum.first->buffer(), 0, szPartialSum, bPartialSum.second);
|
|
mQuantReduceSumSet->writeBuffer(bSumAq.first->buffer(), 1, szSumAq, bSumAq.second);
|
|
dispatchWithProfile("glsl_dynamic_quant_reduce_sum_comp",
|
|
mQuantReduceSumPipeline, mQuantReduceSumSet,
|
|
UP_DIV(padM, 64u), 1, 1, &pc, sizeof(pc));
|
|
cmdBuffer->barrierSource(bSumAq.first->buffer(), bSumAq.second, szSumAq);
|
|
}
|
|
|
|
// INT4 only: nibble-packed weight -> int8 [padN, padK] before GEMM. The
|
|
// O(padN*padK) write/read here trades INT4 runtime bandwidth back to INT8
|
|
// levels, in exchange for fully reusing the W8A8 GEMM body.
|
|
if (mIsInt4) {
|
|
struct PC { uint32_t padN, padK, halfK; } pc;
|
|
pc.padN = padN;
|
|
pc.padK = padK;
|
|
pc.halfK = padK / 2u;
|
|
|
|
auto bWqInt8 = vkBn->getTensorBuffer(tWqInt8.get());
|
|
const size_t szWqInt8 = vkBn->getTensorSize(tWqInt8.get());
|
|
mInt4UnpackSet->writeBuffer(mQuantWeightBuffer->buffer(), 0, mQuantWeightBuffer->size());
|
|
mInt4UnpackSet->writeBuffer(bWqInt8.first->buffer(), 1, szWqInt8, bWqInt8.second);
|
|
dispatchWithProfile("glsl_dynamic_int4_to_int8_unpack_comp",
|
|
mInt4UnpackPipeline, mInt4UnpackSet,
|
|
UP_DIV(pc.halfK, 16u), UP_DIV(padN, 16u), 1, &pc, sizeof(pc));
|
|
cmdBuffer->barrierSource(bWqInt8.first->buffer(), bWqInt8.second, szWqInt8);
|
|
}
|
|
|
|
{
|
|
struct PC { uint32_t M, N, K; } pc;
|
|
pc.M = padM; pc.N = padN; pc.K = padK;
|
|
|
|
mGemmSet->writeBuffer(bAq.first->buffer(), 0, szAq, bAq.second);
|
|
if (mIsInt4) {
|
|
auto bWqInt8 = vkBn->getTensorBuffer(tWqInt8.get());
|
|
const size_t szWqInt8 = vkBn->getTensorSize(tWqInt8.get());
|
|
mGemmSet->writeBuffer(bWqInt8.first->buffer(), 1, szWqInt8, bWqInt8.second);
|
|
} else {
|
|
mGemmSet->writeBuffer(mQuantWeightBuffer->buffer(), 1, mQuantWeightBuffer->size());
|
|
}
|
|
mGemmSet->writeBuffer(bAcc.first->buffer(), 2, szAcc, bAcc.second);
|
|
dispatchWithProfile("glsl_dynamic_w8a8_coop_gemm_comp",
|
|
mGemmS8Pipeline, mGemmSet,
|
|
padN / mCoopN, padM / mCoopM, 1, &pc, sizeof(pc));
|
|
cmdBuffer->barrierSource(bAcc.first->buffer(), bAcc.second, szAcc);
|
|
}
|
|
|
|
{
|
|
struct PC { uint32_t M, N, K, padM, padN; } pc;
|
|
pc.M = (uint32_t)M; pc.N = (uint32_t)N; pc.K = (uint32_t)K; pc.padM = padM; pc.padN = padN;
|
|
|
|
const uint32_t n4_valid = UP_DIV((uint32_t)N, 4u);
|
|
mDequantSet->writeBuffer(bAcc.first->buffer(), 0, szAcc, bAcc.second);
|
|
mDequantSet->writeBuffer(bScaleA.first->buffer(), 1, szScaleA, bScaleA.second);
|
|
mDequantSet->writeBuffer(bOffsetA.first->buffer(), 2, szOffsetA, bOffsetA.second);
|
|
mDequantSet->writeBuffer(bSumAq.first->buffer(), 3, szSumAq, bSumAq.second);
|
|
mDequantSet->writeBuffer(mQuantMetaBuffer->buffer(), 4, mQuantMetaBuffer->size());
|
|
mDequantSet->writeBuffer(mSumWqBuffer->buffer(), 5, mSumWqBuffer->size());
|
|
mDequantSet->writeBuffer(mBiasBuffer->buffer(), 6, mBiasBuffer->size());
|
|
mDequantSet->writeBuffer(dstBuffer.first->buffer(), 7, vkBn->getTensorSize(output), dstBuffer.second);
|
|
dispatchWithProfile(useFP16 ? "glsl_dynamic_w8a8_dequant_correction_FP16_comp"
|
|
: "glsl_dynamic_w8a8_dequant_correction_comp",
|
|
mDequantPipeline, mDequantSet,
|
|
UP_DIV((uint32_t)M, 16u), UP_DIV(n4_valid, 16u), 1, &pc, sizeof(pc));
|
|
}
|
|
|
|
for (Tensor* t : dyns) {
|
|
vkBn->onReleaseBuffer(t, Backend::DYNAMIC);
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
} // namespace MNN
|