wehub-resource-sync
620 lines
25 KiB
Plaintext
620 lines
25 KiB
Plaintext
//
|
||
// MetalConvolutionCommon.mm
|
||
// MNN
|
||
//
|
||
// Created by MNN on 2019/02/25.
|
||
// Copyright © 2018, Alibaba Group Holding Limited
|
||
//
|
||
|
||
#import "backend/metal/MetalConvolutionCommon.hpp"
|
||
#import "core/Macro.h"
|
||
#import "backend/metal/MetalBackend.hpp"
|
||
#import "backend/metal/MetalConvolution1x1.hpp"
|
||
#import "backend/metal/MetalConvolutionWinograd.hpp"
|
||
#import "core/TensorUtils.hpp"
|
||
#include "core/OpCommonUtils.hpp"
|
||
|
||
#if MNN_METAL_ENABLED
|
||
namespace MNN {
|
||
|
||
static const char* gWeightTrans = R"metal(
|
||
#include <metal_stdlib>
|
||
#include <simd/simd.h>
|
||
using namespace metal;
|
||
struct weight_shape {
|
||
int group;
|
||
int goc;
|
||
int goc_4;
|
||
int gic;
|
||
int gic_4;
|
||
int kh;
|
||
int kw;
|
||
};
|
||
kernel void int4_weight_transform_fast(const device uint8_t* src [[buffer(0)]],
|
||
device uint8_t* dst [[buffer(1)]],
|
||
constant weight_shape &uConstant [[buffer(2)]],
|
||
uint3 gid [[thread_position_in_grid]]
|
||
) {
|
||
if ((int)gid.x < uConstant.goc && (int)gid.y < uConstant.gic) {
|
||
auto zo = gid.x / 4, ro = gid.x % 4;
|
||
auto zi = gid.y / 4, ri = gid.y % 4;
|
||
|
||
dst[((zo * uConstant.gic_4 + zi) * 16 + ro * 4 + ri) / 2] = src[(gid.x * uConstant.gic + gid.y) / 2];
|
||
}
|
||
}
|
||
|
||
kernel void int4_weight_transform_c4_fast(const device uint16_t* src [[buffer(0)]],
|
||
device uint16_t* dst [[buffer(1)]],
|
||
constant weight_shape &uConstant [[buffer(2)]],
|
||
uint3 gid [[thread_position_in_grid]]
|
||
) {
|
||
if ((int)gid.x < uConstant.goc && (int)gid.y < uConstant.gic_4) {
|
||
auto zo = gid.x / 4, ro = gid.x % 4;
|
||
|
||
dst[(zo * uConstant.gic_4 + gid.y) * 4 + ro] = src[gid.x * uConstant.gic_4 + gid.y];
|
||
}
|
||
}
|
||
|
||
kernel void weight_transform_common(const device IType* src [[buffer(0)]],
|
||
device OType* dst [[buffer(1)]],
|
||
constant weight_shape &uConstant [[buffer(2)]],
|
||
uint3 gid [[thread_position_in_grid]]
|
||
) {
|
||
if ((int)gid.x < uConstant.group * uConstant.goc && (int)gid.y < uConstant.gic && (int)gid.z < uConstant.kh * uConstant.kw) {
|
||
|
||
auto g = gid.x / uConstant.goc;
|
||
auto goc = gid.x % uConstant.goc;
|
||
auto zo = goc / 4, ro = goc % 4;
|
||
auto zi = gid.y / 4, ri = gid.y % 4;
|
||
auto h = gid.z / uConstant.kw;
|
||
auto w = gid.z % uConstant.kw;
|
||
|
||
// to [g][o/4][i/4][h][w][16]
|
||
// from [g][o][i][h][w]
|
||
int dx = g * uConstant.goc_4 * uConstant.gic_4 * uConstant.kh * uConstant.kw * 16 + zo * uConstant.gic_4 * uConstant.kh * uConstant.kw * 16 + ro * 4 + zi * uConstant.kh * uConstant.kw * 16 + ri + (h * uConstant.kw + w) * 16;
|
||
int sx = (gid.x * uConstant.gic + gid.y) * uConstant.kh * uConstant.kw + gid.z;
|
||
|
||
dst[dx] = (OType)src[sx];
|
||
}
|
||
}
|
||
)metal";
|
||
|
||
static std::shared_ptr<MNN::Tensor> biasForConv(Backend *bn, const Op* op, const Convolution2D *conv, bool fp16) {
|
||
auto bias = conv->bias();
|
||
auto oc = conv->common()->outputCount();
|
||
int bytes = fp16 ? 2 : 4;
|
||
|
||
auto bias_size_unit = UP_DIV(oc, 16) * 16;
|
||
std::shared_ptr<MNN::Tensor> t(MNN::Tensor::createDevice<float>({bias_size_unit}));
|
||
auto res = bn->onAcquireBuffer(t.get(), Backend::STATIC);
|
||
if (!res) {
|
||
return nullptr;
|
||
}
|
||
|
||
const bool useCachedMmap = (bn->getRuntime() && bn->getRuntime()->hint().useCachedMmap > 1);
|
||
if (useCachedMmap) {
|
||
return t;
|
||
}
|
||
|
||
auto buffer = MetalBackend::getBuffer(t.get());
|
||
auto dstOrigin = (uint8_t*)[buffer.first contents] + buffer.second;
|
||
::memset(dstOrigin, 0, bias_size_unit * bytes);
|
||
|
||
const float* src = nullptr;
|
||
std::unique_ptr<float[]> externalBias;
|
||
if (nullptr != bias && bias->size() >= oc) {
|
||
src = bias->data();
|
||
} else if (nullptr != op && nullptr != op->externalPath() && USE_EXTERNAL_DATA(conv) && nullptr != conv->external() && conv->external()->size() >= 3 && conv->external()->data()[2] > 0) {
|
||
auto externalInfo = conv->external()->data();
|
||
size_t biasBytes = externalInfo[2];
|
||
size_t expectedBytes = (size_t)oc * sizeof(float);
|
||
if (biasBytes < expectedBytes) {
|
||
return t;
|
||
}
|
||
externalBias.reset(new float[oc]);
|
||
std::unique_ptr<FileLoader> external(new FileLoader(op->externalPath()->c_str()));
|
||
external->offset(externalInfo[0] + externalInfo[1]);
|
||
external->read((char*)externalBias.get(), expectedBytes);
|
||
src = externalBias.get();
|
||
}
|
||
|
||
if (nullptr == src) {
|
||
return t;
|
||
}
|
||
|
||
if (fp16) {
|
||
auto dst = (__fp16 *)dstOrigin;
|
||
#pragma clang loop vectorize(enable) unroll(enable)
|
||
for (int i = 0; i < oc; i++) {
|
||
dst[i] = src[i];
|
||
}
|
||
} else {
|
||
::memcpy(dstOrigin, src, oc * sizeof(float));
|
||
}
|
||
|
||
return t;
|
||
}
|
||
|
||
MetalConvolutionCommon::MetalConvolutionCommon(Backend *backend, const MNN::Op *op, std::shared_ptr<MNN::Tensor> bias) : MetalExecution(backend) {
|
||
auto mtbn = static_cast<MetalBackend*>(backend);
|
||
auto conv = op->main_as_Convolution2D();
|
||
auto common = conv->common();
|
||
mOp = op;
|
||
mKernelX = common->kernelX();
|
||
mKernelY = common->kernelY();
|
||
mStrideX = common->strideX();
|
||
mStrideY = common->strideY();
|
||
mDilateX = common->dilateX();
|
||
mDilateY = common->dilateY();
|
||
if (nullptr != bias) {
|
||
mBias = bias;
|
||
} else {
|
||
mBias = biasForConv(backend, op, conv, mtbn->useFp16InsteadFp32());
|
||
}
|
||
mActivationType = common->relu() ? 1 : (common->relu6() ? 2 : 0);
|
||
if (nullptr == mBias) {
|
||
mValid = false;
|
||
}
|
||
}
|
||
|
||
template <typename FType, typename TType>
|
||
void MetalConvolutionCommon::convertWeightFormat(int group, int oc, int ic, int kh, int kw, const FType *src, TType* dstOrigion, Tensor* dstTensor, id<MTLBuffer> srcGpuBuffer) {
|
||
auto goc = oc / group;
|
||
auto gic = ic / group;
|
||
auto goc_4 = UP_DIV(goc, 4);
|
||
auto gic_4 = UP_DIV(gic, 4);
|
||
|
||
auto backend = static_cast<MetalBackend *>(this->backend());
|
||
auto context = (__bridge MNNMetalContext *)static_cast<MetalBackend *>(backend)->context();
|
||
if(srcGpuBuffer == nil) {
|
||
srcGpuBuffer = [context newDeviceBuffer:group * goc * gic * kh * kw * sizeof(FType) access:CPUReadWrite];
|
||
::memcpy((void *)srcGpuBuffer.contents, (void *)src, group * goc * gic * kh * kw * sizeof(FType));
|
||
}
|
||
std::string IType = "float";
|
||
std::string OType = "float";
|
||
if(std::is_same<FType, int8_t>::value) {
|
||
IType = "int8_t";
|
||
}
|
||
if(std::is_same<TType, __fp16>::value) {
|
||
OType = "half";
|
||
} else if(std::is_same<TType, int8_t>::value) {
|
||
OType = "int8_t";
|
||
}
|
||
MTLCompileOptions *option = [[MTLCompileOptions alloc] init];
|
||
auto dic = [NSMutableDictionary dictionaryWithCapacity:0];
|
||
[dic setValue:@(IType.c_str()) forKey:@"IType"];
|
||
[dic setValue:@(OType.c_str()) forKey:@"OType"];
|
||
option.preprocessorMacros = dic;
|
||
// create const buffer
|
||
int constants[] = {group,
|
||
goc,
|
||
goc_4,
|
||
gic,
|
||
gic_4,
|
||
kh,
|
||
kw};
|
||
auto constBuffer = backend->getConstBuffer(sizeof(constants));
|
||
::memcpy(constBuffer.contents, constants, sizeof(constants));
|
||
|
||
auto encoder = [backend->getCommandBufferForBufferCopy() computeCommandEncoder];
|
||
auto pipeline = backend->makeComputePipelineWithSourceOption(gWeightTrans, "weight_transform_common", option);
|
||
[encoder setComputePipelineState:pipeline];
|
||
[encoder setBuffer:srcGpuBuffer offset:0 atIndex:0];
|
||
MetalBackend::setTensor(dstTensor, encoder, 1);
|
||
[encoder setBuffer:constBuffer offset:0 atIndex:2];
|
||
auto gl = [context computeBestGroupAndLocal:pipeline threads:MTLSizeMake((NSInteger)group * goc, (NSInteger)gic, (NSInteger)kh * kw)];
|
||
[encoder dispatchThreadgroups:gl.first threadsPerThreadgroup:gl.second];
|
||
[encoder endEncoding];
|
||
// just commit, don‘t wait for not block
|
||
backend->commit();
|
||
}
|
||
|
||
template<typename DType>
|
||
static std::pair<std::shared_ptr<MNN::Tensor>, float> getDequantScale(const float* scale, int size, MetalBackend *backend, bool asymmetric, int oc) {
|
||
int totalCount = 0;
|
||
if (asymmetric) {
|
||
totalCount = size / 2;
|
||
} else {
|
||
totalCount = size;
|
||
}
|
||
int blockSize = totalCount / oc;
|
||
int alignOutputCount = ALIGN_UP4(oc);
|
||
std::shared_ptr<MNN::Tensor> dequantScale(MNN::Tensor::createDevice<uint8_t>({alignOutputCount * blockSize * (int)(sizeof(DType) * 2) + (int)sizeof(float)}));
|
||
bool res = backend->onAcquireBuffer(dequantScale.get(), Backend::STATIC);
|
||
if (!res) {
|
||
MNN_ERROR("Buffer allocated error!\n");
|
||
return std::make_pair(nullptr, 1.0);
|
||
}
|
||
auto buffer0 = MetalBackend::getBuffer(dequantScale.get());
|
||
DType* dst_scale = (DType*)((uint8_t*)[buffer0.first contents] + buffer0.second);
|
||
auto coefPtr = (float*)((uint8_t*)dst_scale + alignOutputCount * blockSize * (int)(sizeof(DType) * 2));
|
||
if (backend->getRuntime()->hint().useCachedMmap > 1) {
|
||
return std::make_pair(dequantScale, *coefPtr);
|
||
}
|
||
::memset(dst_scale, 0, dequantScale->usize());
|
||
|
||
float coef = 1.0;
|
||
if(std::is_same<DType, __fp16>::value) {
|
||
float max_data = 0.0;
|
||
if(asymmetric) {
|
||
for (int z=0; z<oc; ++z) {
|
||
auto srcZ = scale + z * blockSize * 2;
|
||
for (int bi=0; bi<blockSize; ++bi) {
|
||
float s = fabs(srcZ[2*bi+1]);
|
||
float b = fabs(srcZ[2*bi+0]);
|
||
float temp = ALIMAX(s, b);
|
||
if(temp > max_data) {
|
||
max_data = temp;
|
||
}
|
||
}
|
||
}
|
||
} else {
|
||
for (int z=0; z<oc; ++z) {
|
||
auto srcZ = scale + z * blockSize;
|
||
for (int bi=0; bi<blockSize; ++bi) {
|
||
float s = srcZ[bi];
|
||
if(s > max_data) {
|
||
max_data = s;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// too big scale may cause half precision loss
|
||
coef = 1000.0 / max_data;
|
||
}
|
||
if (asymmetric) {
|
||
for (int z=0; z<oc; ++z) {
|
||
int zo = z / 4;
|
||
int zi = z % 4;
|
||
auto srcZ = scale + z * blockSize * 2;
|
||
auto dstSZ = dst_scale + zo * blockSize * 8 + zi;
|
||
auto dstBZ = dst_scale + zo * blockSize * 8 + zi + 4;
|
||
for (int bi=0; bi<blockSize; ++bi) {
|
||
float s = srcZ[2*bi+1];
|
||
float b = srcZ[2*bi+0];
|
||
dstSZ[bi * 8] = (DType)(s * coef);
|
||
dstBZ[bi * 8] = (DType)(b * coef);
|
||
}
|
||
}
|
||
} else {
|
||
for (int z=0; z<oc; ++z) {
|
||
int zo = z / 4;
|
||
int zi = z % 4;
|
||
auto srcZ = scale + z * blockSize;
|
||
auto dstSZ = dst_scale + zo * blockSize * 8 + zi;
|
||
auto dstBZ = dst_scale + zo * blockSize * 8 + zi + 4;
|
||
for (int bi=0; bi<blockSize; ++bi) {
|
||
float s = srcZ[bi];
|
||
float b = 0.0f;
|
||
dstSZ[bi * 8] = (DType)(s * coef);
|
||
dstBZ[bi * 8] = b;
|
||
}
|
||
}
|
||
}
|
||
*coefPtr = coef;
|
||
return std::make_pair(dequantScale, coef);
|
||
}
|
||
void MetalConvolutionCommon::loadWeight(const MNN::Op *op, bool loadWeightInt8) {
|
||
auto conv = op->main_as_Convolution2D();
|
||
auto common = conv->common();
|
||
auto kw = common->kernelX();
|
||
auto kh = common->kernelY();
|
||
auto group = common->group();
|
||
auto oc = common->outputCount();
|
||
int ic = common->inputCount();
|
||
|
||
void* weightMemPtr = nullptr;
|
||
id<MTLBuffer> srcGpuBuffer = nil;
|
||
auto useOriginMmap = backend()->getRuntime()->hint().useCachedMmap > 1;
|
||
bool preAllocGpuMem = ic != 0 && conv->quanParameter();
|
||
int quantBit;
|
||
// only for weight int4/int8 now.
|
||
if(loadWeightInt8) {
|
||
quantBit = conv->quanParameter()->aMaxOrBits();
|
||
// 3.1.2 and after has aMaxOrBits for quant bits
|
||
if (quantBit == 0) {
|
||
// support old model for external weight file with int4/int8 quant
|
||
quantBit = ConvolutionCommon::getQuantBitFromExternalFile(op);
|
||
}
|
||
if(quantBit != 4 && quantBit != 8) {
|
||
preAllocGpuMem = false;
|
||
}
|
||
}
|
||
if (preAllocGpuMem && (!useOriginMmap)) {
|
||
size_t size = oc * ic * kh * kw / group;
|
||
|
||
if (loadWeightInt8) {
|
||
if(quantBit == 4) {
|
||
size = UP_DIV(size, 2);
|
||
}
|
||
} else {
|
||
size *= sizeof(float);
|
||
}
|
||
|
||
auto backend = static_cast<MetalBackend *>(this->backend());
|
||
auto context = (__bridge MNNMetalContext *)static_cast<MetalBackend *>(backend)->context();
|
||
|
||
srcGpuBuffer = [context newDeviceBuffer:size access:CPUReadWrite];
|
||
}
|
||
|
||
std::shared_ptr<ConvolutionCommon::Int8Common> qnt = NULL;
|
||
if (loadWeightInt8) {
|
||
qnt = ConvolutionCommon::load(op, backend(), false, true, (void *)srcGpuBuffer.contents);
|
||
} else if (conv->quanParameter()) {
|
||
qnt = ConvolutionCommon::load(op, backend(), true, false, (void *)srcGpuBuffer.contents);
|
||
}
|
||
// param
|
||
size_t size = 0;
|
||
if (ic > 0) {
|
||
size = oc * ic * kh * kw / group;
|
||
} else {
|
||
size = qnt ? MAX(qnt->weight.size(), qnt->weightFloat.size()) : conv->weight()->size();
|
||
ic = size / kw / kh / (oc / group);
|
||
}
|
||
// convert
|
||
if (loadWeightInt8) {
|
||
auto backend = static_cast<MetalBackend *>(this->backend());
|
||
bool useInt2 = qnt->canUseInt2;
|
||
bool useInt3 = qnt->canUseInt3;
|
||
bool int4Path = qnt->canUseInt4 && !useInt2 && !useInt3;
|
||
bool int8Path = !int4Path && !useInt2 && !useInt3;
|
||
int subBits = useInt2 ? 2 : (useInt3 ? 3 : 0);
|
||
mWeight = weightTransform(group, oc, ic, kh, kw, (float*)qnt->weight.get(), int8Path, int4Path, srcGpuBuffer, subBits);
|
||
if(backend->useFp16InsteadFp32()) {
|
||
auto dequantParams = getDequantScale<__fp16>(qnt->alpha.get(), qnt->alphaSize, backend, qnt->asymmetric, oc);
|
||
mDequantScaleBias = dequantParams.first;
|
||
mScaleCoef = dequantParams.second;
|
||
} else {
|
||
auto dequantParams = getDequantScale<float>(qnt->alpha.get(), qnt->alphaSize, backend, qnt->asymmetric, oc);
|
||
mDequantScaleBias = dequantParams.first;
|
||
mScaleCoef = dequantParams.second;
|
||
}
|
||
|
||
mDequantBits = useInt2 ? 2 : (useInt3 ? 3 : (int4Path ? 4 : 8));
|
||
} else if (qnt && qnt->weightFloat.get()) {
|
||
mWeight = weightTransform(group, oc, ic, kh, kw, qnt->weightFloat.get(), false, false, srcGpuBuffer);
|
||
} else {
|
||
const float* src = nullptr;
|
||
std::unique_ptr<float[]> externalWeight;
|
||
if (nullptr != conv->weight() && conv->weight()->size() > 0) {
|
||
src = conv->weight()->data();
|
||
} else {
|
||
const bool useCachedMmap = (backend()->getRuntime() && backend()->getRuntime()->hint().useCachedMmap > 1);
|
||
if (!useCachedMmap && nullptr != op->externalPath() && USE_EXTERNAL_DATA(conv) && nullptr != conv->external() && conv->external()->size() >= 2 && conv->external()->data()[1] > 0) {
|
||
auto externalInfo = conv->external()->data();
|
||
size_t weightBytes = externalInfo[1];
|
||
size_t expectedBytes = size * sizeof(float);
|
||
if (weightBytes < expectedBytes) {
|
||
mValid = false;
|
||
return;
|
||
}
|
||
externalWeight.reset(new float[size]);
|
||
std::unique_ptr<FileLoader> external(new FileLoader(op->externalPath()->c_str()));
|
||
external->offset(externalInfo[0]);
|
||
external->read((char*)externalWeight.get(), expectedBytes);
|
||
src = externalWeight.get();
|
||
}
|
||
}
|
||
mWeight = weightTransform(group, oc, ic, kh, kw, src, false, false, srcGpuBuffer);
|
||
}
|
||
}
|
||
|
||
std::shared_ptr<MNN::Tensor> MetalConvolutionCommon::weightTransform(int group, int oc, int ic, int kh, int kw, const float *src, bool int8Weight, bool int4Weight, id<MTLBuffer> srcGpuBuffer, int subBits) {
|
||
if(srcGpuBuffer != nil) {
|
||
MNN_ASSERT((void*)src == (void*)srcGpuBuffer.contents);
|
||
}
|
||
auto backend = static_cast<MetalBackend *>(this->backend());
|
||
auto context = (__bridge MNNMetalContext *)static_cast<MetalBackend *>(backend)->context();
|
||
auto goc = oc / group;
|
||
auto gic = ic / group;
|
||
auto goc_4 = UP_DIV(goc, 4);
|
||
auto gic_4 = UP_DIV(gic, 4);
|
||
auto weight_len = group * ROUND_UP(goc_4, 4) * gic_4 * kw * kh * 16;
|
||
auto ori_len = group * goc * gic * kh * kw;
|
||
bool needMemset = (goc % 4 != 0 || gic % 4 != 0);
|
||
#ifdef MNN_LOW_MEMORY
|
||
if (subBits == 3) {
|
||
// 3-bit packed: 6 bytes / (4 OC, 4 IC) tile.
|
||
size_t weight_bytes = (size_t)group * goc_4 * gic_4 * kh * kw * 6;
|
||
std::shared_ptr<MNN::Tensor> weightLow(MNN::Tensor::createDevice<int8_t>({(int)weight_bytes}));
|
||
if (!backend->onAcquireBuffer(weightLow.get(), Backend::STATIC)) {
|
||
MNN_ERROR("Memory alloc error!\n");
|
||
return nullptr;
|
||
}
|
||
if (nil == src) {
|
||
return weightLow;
|
||
}
|
||
auto buf = MetalBackend::getBuffer(weightLow.get());
|
||
auto dstPtr = (uint8_t*)[buf.first contents] + buf.second;
|
||
::memset(dstPtr, 0, weight_bytes);
|
||
auto srcPtr = (const int8_t*)src;
|
||
for (int g = 0; g < group; g++) {
|
||
for (int o = 0; o < goc; o++) {
|
||
int zo = o / 4, ro = o % 4;
|
||
for (int i = 0; i < gic; i++) {
|
||
int zi = i / 4, ri = i % 4;
|
||
for (int h = 0; h < kh; h++) {
|
||
for (int w = 0; w < kw; w++) {
|
||
int srcIdx = ((g * goc + o) * gic + i) * kh * kw + h * kw + w;
|
||
int sv = (int)srcPtr[srcIdx] + 4;
|
||
int tileBase = (((g * goc_4 + zo) * gic_4 + zi) * kh + h) * kw * 6 + w * 6;
|
||
dstPtr[tileBase + ro] |= (uint8_t)((sv & 3) << (6 - ri * 2));
|
||
int hiByte = tileBase + 4 + (ro / 2);
|
||
int hiShift = (ro % 2 == 0 ? 4 : 0) + (3 - ri);
|
||
dstPtr[hiByte] |= (uint8_t)(((sv >> 2) & 1) << hiShift);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return weightLow;
|
||
}
|
||
if (subBits == 2) {
|
||
// 2-bit packed: 4 bytes / (4 OC, 4 IC) tile.
|
||
size_t weight_bytes = (size_t)group * goc_4 * gic_4 * kh * kw * 4;
|
||
std::shared_ptr<MNN::Tensor> weightLow(MNN::Tensor::createDevice<int8_t>({(int)weight_bytes}));
|
||
if (!backend->onAcquireBuffer(weightLow.get(), Backend::STATIC)) {
|
||
MNN_ERROR("Memory alloc error!\n");
|
||
return nullptr;
|
||
}
|
||
if (nil == src) {
|
||
return weightLow;
|
||
}
|
||
auto buf = MetalBackend::getBuffer(weightLow.get());
|
||
auto dstPtr = (uint8_t*)[buf.first contents] + buf.second;
|
||
::memset(dstPtr, 0, weight_bytes);
|
||
auto srcPtr = (const int8_t*)src;
|
||
for (int g = 0; g < group; g++) {
|
||
for (int o = 0; o < goc; o++) {
|
||
int zo = o / 4, ro = o % 4;
|
||
for (int i = 0; i < gic; i++) {
|
||
int zi = i / 4, ri = i % 4;
|
||
for (int h = 0; h < kh; h++) {
|
||
for (int w = 0; w < kw; w++) {
|
||
int srcIdx = ((g * goc + o) * gic + i) * kh * kw + h * kw + w;
|
||
int sv = (int)srcPtr[srcIdx] + 2;
|
||
int tileBase = (((g * goc_4 + zo) * gic_4 + zi) * kh + h) * kw * 4 + w * 4;
|
||
dstPtr[tileBase + ro] |= (uint8_t)((sv & 3) << (6 - ri * 2));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return weightLow;
|
||
}
|
||
if (int4Weight) {
|
||
weight_len = UP_DIV(weight_len, 2);
|
||
std::shared_ptr<MNN::Tensor> weightLow(MNN::Tensor::createDevice<int8_t>({weight_len}));
|
||
auto res = backend->onAcquireBuffer(weightLow.get(), Backend::STATIC);
|
||
if (!res) {
|
||
MNN_ERROR("Memory alloc error!\n");
|
||
return nullptr;
|
||
}
|
||
if (nil == src) {
|
||
// Use mmap weight. No need to compute
|
||
return weightLow;
|
||
}
|
||
|
||
auto buf = MetalBackend::getBuffer(weightLow.get());
|
||
auto dstPtr = (uint8_t*)[buf.first contents] + buf.second;
|
||
if(needMemset) {
|
||
::memset(dstPtr, 0, weight_len);
|
||
}
|
||
bool fastBlit = (group == 1 && kh == 1 && kw == 1 && ic % 2 == 0);
|
||
|
||
auto oc_4 = UP_DIV(oc, 4);
|
||
auto ic_4 = UP_DIV(ic, 4);
|
||
// fast int4 reorder
|
||
if (fastBlit) {
|
||
if(srcGpuBuffer == nil) {
|
||
srcGpuBuffer = [context newDeviceBuffer:UP_DIV(ori_len, 2) access:CPUReadWrite];
|
||
::memcpy((void *)srcGpuBuffer.contents, (void *)src, UP_DIV(ori_len, 2));
|
||
}
|
||
|
||
MTLCompileOptions *option = [[MTLCompileOptions alloc] init];
|
||
auto dic = [NSMutableDictionary dictionaryWithCapacity:0];
|
||
[dic setValue:@"uint8_t" forKey:@"IType"];
|
||
[dic setValue:@"uint8_t" forKey:@"OType"];
|
||
option.preprocessorMacros = dic;
|
||
// create const buffer
|
||
int constants[] = {group,
|
||
goc,
|
||
goc_4,
|
||
gic,
|
||
gic_4,
|
||
kh,
|
||
kw};
|
||
auto constBuffer = backend->getConstBuffer(sizeof(constants));
|
||
::memcpy(constBuffer.contents, constants, sizeof(constants));
|
||
|
||
auto encoder = [backend->getCommandBufferForBufferCopy() computeCommandEncoder];
|
||
id<MTLComputePipelineState> pipeline;
|
||
bool c4_fast = ic % 4 == 0;
|
||
if(c4_fast) {
|
||
pipeline = backend->makeComputePipelineWithSourceOption(gWeightTrans, "int4_weight_transform_c4_fast", option);
|
||
} else {
|
||
pipeline = backend->makeComputePipelineWithSourceOption(gWeightTrans, "int4_weight_transform_fast", option);
|
||
}
|
||
[encoder setComputePipelineState:pipeline];
|
||
[encoder setBuffer:srcGpuBuffer offset:0 atIndex:0];
|
||
MetalBackend::setTensor(weightLow.get(), encoder, 1);
|
||
[encoder setBuffer:constBuffer offset:0 atIndex:2];
|
||
MTLSize totalThread;
|
||
if(c4_fast) {
|
||
totalThread = MTLSizeMake((NSInteger)goc, (NSInteger)gic_4, (NSInteger)1);
|
||
} else {
|
||
totalThread = MTLSizeMake((NSInteger)goc, (NSInteger)gic, (NSInteger)1);
|
||
}
|
||
auto gl = [context computeBestGroupAndLocal:pipeline threads:totalThread];
|
||
[encoder dispatchThreadgroups:gl.first threadsPerThreadgroup:gl.second];
|
||
[encoder endEncoding];
|
||
// just commit, don‘t wait for not block
|
||
backend->commit();
|
||
} else {
|
||
auto srcPtr = (int8_t*)src;
|
||
// slow int4 reorder
|
||
int sx = 0;
|
||
auto goc_4 = UP_DIV(goc, 4);
|
||
auto gic_4 = UP_DIV(gic, 4);
|
||
for (int g = 0; g < group; g++) {
|
||
for (int o = 0; o < goc; o++) {
|
||
auto zo = o / 4, ro = o % 4;
|
||
for (int i = 0; i < gic; i++) {
|
||
auto zi = i / 4, ri = i % 4;
|
||
for (int h = 0; h < kh; h++) {
|
||
for (int w = 0; w < kw; w++) {
|
||
// to [g][o/4][i/4][h][w][16]
|
||
// from [g][o][i][h][w]
|
||
int dx = g * goc_4 * gic_4 * kh * kw * 16 + zo * gic_4 * kh * kw * 16 + ro * 4 + zi * kh * kw * 16 + ri + (h * kw + w) * 16;
|
||
uint8_t s = srcPtr[sx/2];
|
||
s = (sx % 2) ? (s & 0xf) : (s >> 4);
|
||
s = (dx % 2) ? s : (s << 4);
|
||
dstPtr[dx/2] |= s;
|
||
sx++;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return weightLow;
|
||
}
|
||
#endif
|
||
std::shared_ptr<MNN::Tensor> t(MNN::Tensor::createDevice<float>({weight_len}));
|
||
if (int8Weight || int4Weight) {
|
||
t.reset(MNN::Tensor::createDevice<int8_t>({weight_len}));
|
||
}
|
||
bool res = backend->onAcquireBuffer(t.get(), Backend::STATIC);
|
||
if (!res) {
|
||
return nullptr;
|
||
}
|
||
if (nullptr == src) {
|
||
// No need to compute
|
||
return t;
|
||
}
|
||
auto buffer = MetalBackend::getBuffer(t.get());
|
||
auto dst = (uint8_t*)[buffer.first contents] + buffer.second;
|
||
if (int8Weight) {
|
||
if(needMemset) {
|
||
::memset(dst, 0, weight_len);
|
||
}
|
||
convertWeightFormat<int8_t, int8_t>(group, oc, ic, kh, kw, (const int8_t*)src, (int8_t *)dst, t.get(), srcGpuBuffer);
|
||
} else if (backend->useFp16InsteadFp32()) {
|
||
if(needMemset) {
|
||
::memset(dst, 0, weight_len * sizeof(__fp16));
|
||
}
|
||
convertWeightFormat<float, __fp16>(group, oc, ic, kh, kw, (const float*)src, (__fp16 *)dst, t.get(), srcGpuBuffer);
|
||
} else {
|
||
if(needMemset) {
|
||
::memset(dst, 0, weight_len * sizeof(float));
|
||
}
|
||
convertWeightFormat<float, float>(group, oc, ic, kh, kw, (const float*)src, (float *)dst, t.get(), srcGpuBuffer);
|
||
}
|
||
|
||
return t;
|
||
}
|
||
|
||
} // namespace MNN
|
||
|
||
#endif
|