wehub-resource-sync
865 lines
32 KiB
Plaintext
865 lines
32 KiB
Plaintext
//
|
|
// MetalLoop.mm
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2023/12/28.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#import "core/Macro.h"
|
|
#import "MetalCast.hpp"
|
|
#import "MetalBinary.hpp"
|
|
#import "MetalBackend.hpp"
|
|
#import "MNNMetalContext.h"
|
|
#include "MNN_generated.h"
|
|
|
|
#if MNN_METAL_ENABLED
|
|
namespace MNN {
|
|
static const char* gMatMulUnitTemplate = R"metal(
|
|
#include <metal_stdlib>
|
|
#include <simd/simd.h>
|
|
|
|
using namespace metal;
|
|
|
|
struct constBuffer
|
|
{
|
|
int4 size;
|
|
int4 stride_o;
|
|
int4 stride_a;
|
|
int4 stride_b;
|
|
int4 stride_c;
|
|
int4 _step;
|
|
int4 iter;
|
|
};
|
|
|
|
kernel void loop_matmul(device T* uOutput [[buffer(0)]], const device T* uInputA [[buffer(1)]], const device T* uInputB [[buffer(2)]],
|
|
#ifdef HAS_BIAS
|
|
const device T* uInputC [[buffer(3)]],
|
|
const device int* uOOffset [[buffer(4)]],
|
|
const device int* uAOffset [[buffer(5)]],
|
|
const device int* uBOffset [[buffer(6)]],
|
|
const device int* uCOffset [[buffer(7)]],
|
|
constant constBuffer& uConstant [[buffer(8)]],
|
|
#else
|
|
const device int* uOOffset [[buffer(3)]],
|
|
const device int* uAOffset [[buffer(4)]],
|
|
const device int* uBOffset [[buffer(5)]],
|
|
constant constBuffer& uConstant [[buffer(6)]],
|
|
#endif
|
|
uint3 gl_GlobalInvocationID [[thread_position_in_grid]])
|
|
{
|
|
int e = uConstant.size.x;
|
|
int l = uConstant.size.y;
|
|
int h = uConstant.size.z;
|
|
int n = uConstant.size.w;
|
|
|
|
int X0 = gl_GlobalInvocationID.x * 4;
|
|
int Y0 = gl_GlobalInvocationID.y * 4;
|
|
int regionOutsideIndex = gl_GlobalInvocationID.z;
|
|
|
|
if (X0 >= e || Y0 >= h || regionOutsideIndex >= n) {
|
|
return;
|
|
}
|
|
|
|
int4 index = int4(regionOutsideIndex, regionOutsideIndex, regionOutsideIndex, regionOutsideIndex);
|
|
if (uConstant.iter.x >= 0) {
|
|
index.x = uOOffset[regionOutsideIndex];
|
|
}
|
|
if (uConstant.iter.y >= 0) {
|
|
index.y = uAOffset[regionOutsideIndex];
|
|
}
|
|
if (uConstant.iter.z >= 0) {
|
|
index.z = uBOffset[regionOutsideIndex];
|
|
}
|
|
#ifdef HAS_BIAS
|
|
if (uConstant.iter.w >= 0) {
|
|
index.w = uCOffset[regionOutsideIndex];
|
|
}
|
|
#endif
|
|
|
|
int4 offset = index * uConstant._step;
|
|
|
|
T value[4][4];
|
|
for (int y = 0; y < 4; ++y) {
|
|
for (int x = 0; x < 4; ++x) {
|
|
value[x][y] = T(0.0);
|
|
}
|
|
}
|
|
|
|
int aOffset0 = offset.y + uConstant.stride_a.w;
|
|
int bOffset0 = offset.z + uConstant.stride_b.w;
|
|
|
|
int a_idx[4];
|
|
int b_idx[4];
|
|
for (int x = 0; x < 4; ++x) {
|
|
a_idx[x] = min(X0 + x, e - 1) * uConstant.stride_a.x;
|
|
}
|
|
for (int y = 0; y < 4; ++y) {
|
|
b_idx[y] = min(Y0 + y, h - 1) * uConstant.stride_b.z;
|
|
}
|
|
|
|
bool safe = (X0 + 3 < e) && (Y0 + 3 < h);
|
|
|
|
if (safe) {
|
|
for (int i = 0; i < l; i++) {
|
|
T a[4];
|
|
T b[4];
|
|
int a_base = aOffset0 + i * uConstant.stride_a.y;
|
|
int b_base = bOffset0 + i * uConstant.stride_b.y;
|
|
|
|
for(int x = 0; x < 4; ++x) {
|
|
a[x] = uInputA[a_base + a_idx[x]];
|
|
}
|
|
|
|
for(int y = 0; y < 4; ++y) {
|
|
b[y] = uInputB[b_base + b_idx[y]];
|
|
}
|
|
|
|
for(int y = 0; y < 4; ++y) {
|
|
for(int x = 0; x < 4; ++x) {
|
|
value[x][y] += a[x] * b[y];
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
for (int i = 0; i < l; i++) {
|
|
T a[4];
|
|
T b[4];
|
|
int a_base = aOffset0 + i * uConstant.stride_a.y;
|
|
int b_base = bOffset0 + i * uConstant.stride_b.y;
|
|
|
|
// Load A with boundary check
|
|
for(int x = 0; x < 4; ++x) {
|
|
if (X0 + x < e) {
|
|
a[x] = uInputA[a_base + a_idx[x]];
|
|
} else {
|
|
a[x] = T(0.0);
|
|
}
|
|
}
|
|
|
|
// Load B with boundary check
|
|
for(int y = 0; y < 4; ++y) {
|
|
if (Y0 + y < h) {
|
|
b[y] = uInputB[b_base + b_idx[y]];
|
|
} else {
|
|
b[y] = T(0.0);
|
|
}
|
|
}
|
|
|
|
for(int y = 0; y < 4; ++y) {
|
|
for(int x = 0; x < 4; ++x) {
|
|
value[x][y] += a[x] * b[y];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
for (int y = 0; y < 4; ++y) {
|
|
if (Y0 + y >= h) continue;
|
|
for (int x = 0; x < 4; ++x) {
|
|
if (X0 + x >= e) continue;
|
|
T outVal = value[x][y];
|
|
#ifdef HAS_BIAS
|
|
outVal += uInputC[offset.w + (Y0 + y) * uConstant.stride_c.z + (X0 + x) * uConstant.stride_c.x];
|
|
#endif
|
|
uOutput[offset.x + uConstant.stride_o.w + (X0 + x) * uConstant.stride_o.x + (Y0 + y) * uConstant.stride_o.z] = outVal;
|
|
}
|
|
}
|
|
}
|
|
)metal";
|
|
|
|
struct VulkanBatchMatMulInfo {
|
|
int size[4];
|
|
int stride_o[4];
|
|
int stride_a[4];
|
|
int stride_b[4];
|
|
int stride_c[4];
|
|
int step[4];
|
|
int iter[4];
|
|
};
|
|
static void _setTensorStack(std::vector<Tensor*>& result, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const LoopParam* loop) {
|
|
if (loop->inputIndexes() != nullptr) {
|
|
for (int i=0; i<loop->inputIndexes()->size(); ++i) {
|
|
result[loop->inputIndexes()->data()[i]] = inputs[i];
|
|
}
|
|
}
|
|
for (int i=0; i<loop->outputIndexes()->size(); ++i) {
|
|
result[loop->outputIndexes()->data()[i]] = outputs[i];
|
|
}
|
|
}
|
|
|
|
class MetalBatchMatMul : public MetalExecution {
|
|
private:
|
|
const LoopParam* mLoop;
|
|
id<MTLBuffer> mParam;
|
|
id<MTLComputePipelineState> mPipeline;
|
|
std::vector<Tensor*> mTensors;
|
|
bool mHasBias = false;
|
|
int mCmdIndex = 0;
|
|
|
|
public:
|
|
MetalBatchMatMul(const LoopParam* loop, Backend *bn, int index = 0) : MetalExecution(bn) {
|
|
mLoop = loop;
|
|
auto mtbn = static_cast<MetalBackend *>(bn);
|
|
auto context = (__bridge MNNMetalContext *)mtbn->context();
|
|
mParam = [context newDeviceBuffer:sizeof(VulkanBatchMatMulInfo) access:CPUWriteOnly];
|
|
bool useFp16 = mtbn->useFp16InsteadFp32();
|
|
NSString* T = nil;
|
|
if (useFp16) {
|
|
T = @"half";
|
|
} else {
|
|
T = @"float";
|
|
}
|
|
std::vector<std::string> keys = {
|
|
std::string([T UTF8String]),
|
|
"matmulunit"
|
|
};
|
|
auto cmd = loop->commands()->GetAs<RegionCommand>(index);
|
|
mHasBias = cmd->indexes()->size() > 3;
|
|
if (mHasBias) {
|
|
keys.emplace_back("BIAS");
|
|
}
|
|
auto pipeline = mtbn->runtime()->findPipeline(keys);
|
|
if (nil == pipeline) {
|
|
MTLCompileOptions *compileOptions = [[MTLCompileOptions alloc] init];
|
|
if (!mHasBias) {
|
|
compileOptions.preprocessorMacros = @{
|
|
@"T" : T,
|
|
};
|
|
} else {
|
|
compileOptions.preprocessorMacros = @{
|
|
@"T" : T,
|
|
@"HAS_BIAS":@"1",
|
|
};
|
|
}
|
|
pipeline = mtbn->makeComputePipelineWithSourceOption(gMatMulUnitTemplate, "loop_matmul", compileOptions);
|
|
mtbn->runtime()->insertPipeline(keys, pipeline);
|
|
|
|
}
|
|
if (nil == pipeline) {
|
|
MNN_ERROR("Create batch matmul pipeline error\n");
|
|
}
|
|
mPipeline = pipeline;
|
|
mTensors.resize(mLoop->tensorNumber());
|
|
mCmdIndex = index;
|
|
}
|
|
virtual ~MetalBatchMatMul() = default;
|
|
virtual ErrorCode onResize(const std::vector<Tensor *>& inputs, const std::vector<Tensor *>& outputs) override {
|
|
_setTensorStack(mTensors, inputs, outputs, mLoop);
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(mCmdIndex);
|
|
auto size = cmd->size()->data();
|
|
auto AStride = cmd->view()->GetAs<View>(1)->stride()->data();
|
|
auto BStride = cmd->view()->GetAs<View>(2)->stride()->data();
|
|
auto OStride = cmd->view()->GetAs<View>(0)->stride()->data();
|
|
auto param = reinterpret_cast<VulkanBatchMatMulInfo*>([mParam contents]);
|
|
param->size[3] = mLoop->loopNumber();
|
|
for (int i=0; i<3; ++i) {
|
|
param->size[i] = size[i];
|
|
param->stride_o[i] = OStride[i];
|
|
param->stride_a[i] = AStride[i];
|
|
param->stride_b[i] = BStride[i];
|
|
}
|
|
param->stride_o[3] = cmd->view()->GetAs<View>(0)->offset();
|
|
param->stride_a[3] = cmd->view()->GetAs<View>(1)->offset();
|
|
param->stride_b[3] = cmd->view()->GetAs<View>(2)->offset();
|
|
if (mHasBias) {
|
|
param->stride_c[3] = cmd->view()->GetAs<View>(3)->offset();
|
|
}
|
|
::memcpy(param->step, cmd->steps()->data(), cmd->steps()->size() * sizeof(int));
|
|
::memcpy(param->iter, cmd->iterIndexes()->data(), cmd->iterIndexes()->size() * sizeof(int));
|
|
return NO_ERROR;
|
|
}
|
|
virtual void onEncode(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs, id<MTLComputeCommandEncoder> encoder) override {
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(mCmdIndex);
|
|
auto size = cmd->size()->data();
|
|
auto AStride = cmd->view()->GetAs<View>(1)->stride()->data();
|
|
auto BStride = cmd->view()->GetAs<View>(2)->stride()->data();
|
|
auto OStride = cmd->view()->GetAs<View>(0)->stride()->data();
|
|
[encoder setComputePipelineState:mPipeline];
|
|
for (int i=0; i<cmd->indexes()->size(); ++i) {
|
|
MetalBackend::setTensor(mTensors[cmd->indexes()->data()[i]], encoder, i);
|
|
}
|
|
auto iter = cmd->iterIndexes()->data();
|
|
for (int i=0; i<cmd->indexes()->size(); ++i) {
|
|
if (iter[i] >= 0) {
|
|
MetalBackend::setTensor(mTensors[iter[i]], encoder, cmd->indexes()->size() + i);
|
|
} else {
|
|
MetalBackend::setTensor(inputs[0], encoder, cmd->indexes()->size() + i);
|
|
}
|
|
}
|
|
// printf("loop_matmul out dequant BMNK: %d %d %d %d\n", mLoop->loopNumber(), size[0], size[2], size[1]);
|
|
|
|
[encoder setBuffer:mParam offset:0 atIndex:cmd->indexes()->size() * 2];
|
|
|
|
int e = size[0];
|
|
int h = size[2];
|
|
int n = mLoop->loopNumber();
|
|
int threadsX = UP_DIV(e, 4);
|
|
int threadsY = UP_DIV(h, 4);
|
|
[encoder dispatchThreadgroups:MTLSizeMake(UP_DIV(threadsX, 8), UP_DIV(threadsY, 8), n) threadsPerThreadgroup:MTLSizeMake(8, 8, 1)];
|
|
}
|
|
};
|
|
|
|
static const char* gBlitRegion = R"metal(
|
|
#include <metal_stdlib>
|
|
#include <simd/simd.h>
|
|
using namespace metal;
|
|
struct constBuffer
|
|
{
|
|
int4 stride;
|
|
int4 size;
|
|
int4 extent;
|
|
int4 _step;
|
|
int4 iter;
|
|
int4 totalSize;
|
|
};
|
|
|
|
struct s1
|
|
{
|
|
int data[1];
|
|
};
|
|
|
|
struct s2
|
|
{
|
|
int data[1];
|
|
};
|
|
|
|
struct sourceBuffer
|
|
{
|
|
T data[1];
|
|
};
|
|
|
|
struct s0
|
|
{
|
|
T data[1];
|
|
};
|
|
|
|
kernel void gather_blit(device sourceBuffer& uOutput [[buffer(0)]], const device s0& uInput [[buffer(1)]], const device s1& uSrcOffset [[buffer(2)]], const device s2& uDstOffset [[buffer(3)]], constant constBuffer& uConstant [[buffer(4)]], uint3 gl_GlobalInvocationID [[thread_position_in_grid]])
|
|
{
|
|
int3 posTmp = int3(gl_GlobalInvocationID);
|
|
if (posTmp.x < uConstant._step.w)
|
|
{
|
|
int regionInsideIndex = posTmp.x % uConstant.size.w;
|
|
int regionOutsideIndex = posTmp.x / uConstant.size.w;
|
|
int3 pos;
|
|
pos.x = regionInsideIndex / (uConstant.size.y * uConstant.size.z);
|
|
int subIndex = regionInsideIndex % (uConstant.size.y * uConstant.size.z);
|
|
pos.z = subIndex % uConstant.size.z;
|
|
pos.y = subIndex / uConstant.size.z;
|
|
int srcBasicOffset;
|
|
if (uConstant.iter.y > 0)
|
|
{
|
|
srcBasicOffset = uConstant._step.y * int(uSrcOffset.data[regionOutsideIndex]);
|
|
}
|
|
else
|
|
{
|
|
srcBasicOffset = uConstant._step.y * regionOutsideIndex;
|
|
}
|
|
int dstBasicOffset;
|
|
if (uConstant.iter.x > 0)
|
|
{
|
|
dstBasicOffset = uConstant._step.x * int(uDstOffset.data[regionOutsideIndex]);
|
|
}
|
|
else
|
|
{
|
|
dstBasicOffset = uConstant._step.x * regionOutsideIndex;
|
|
}
|
|
int srcOffset = (((srcBasicOffset + uConstant.stride.w) + (uConstant.stride.z * pos.z)) + (uConstant.stride.y * pos.y)) + (uConstant.stride.x * pos.x);
|
|
int dstOffset = (((dstBasicOffset + uConstant.extent.w) + (pos.x * uConstant.extent.x)) + (pos.y * uConstant.extent.y)) + (pos.z * uConstant.extent.z);
|
|
if(srcOffset >= 0 && srcOffset < uConstant.totalSize.x) {
|
|
if(dstOffset >= 0 && dstOffset < uConstant.totalSize.y) {
|
|
uOutput.data[dstOffset] = uInput.data[srcOffset];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
)metal";
|
|
|
|
struct GatherInfo {
|
|
int stride[4];
|
|
int size[4];
|
|
int extent[4];
|
|
int step[4];
|
|
int iter[4];
|
|
int totalSize[4];
|
|
};
|
|
struct InitInfo {
|
|
int srcStride[4];
|
|
int dstStride[4];
|
|
int size[4];
|
|
int totalSize[4];
|
|
};
|
|
static const char* gInitRegion = R"metal(
|
|
#include <metal_stdlib>
|
|
#include <simd/simd.h>
|
|
using namespace metal;
|
|
|
|
struct constBuffer
|
|
{
|
|
int4 srcStride;
|
|
int4 dstStride;
|
|
int4 size;
|
|
int4 totalSize;
|
|
};
|
|
|
|
kernel void set_zero(device T *out [[buffer(0)]],
|
|
const device T *in [[buffer(1)]],
|
|
constant constBuffer &info [[buffer(2)]],
|
|
uint3 gl_GlobalInvocationID [[thread_position_in_grid]]) {
|
|
int3 gid = int3(gl_GlobalInvocationID);
|
|
if (gid.x >= info.size.x || gid.y >= info.size.y || gid.z >= info.size.z) {
|
|
return;
|
|
}
|
|
int dst_offset = (gid.z * info.size.y + gid.y) * info.size.x + gid.x;
|
|
if(dst_offset >= 0 && dst_offset < info.totalSize.y) {
|
|
out[dst_offset] = (T)0;
|
|
}
|
|
}
|
|
|
|
kernel void set_copy(device T *out [[buffer(0)]],
|
|
const device T *in [[buffer(1)]],
|
|
constant constBuffer &info [[buffer(2)]],
|
|
uint3 gl_GlobalInvocationID [[thread_position_in_grid]]) {
|
|
int3 gid = int3(gl_GlobalInvocationID);
|
|
if (gid.x >= info.size.x || gid.y >= info.size.y || gid.z >= info.size.z) {
|
|
return;
|
|
}
|
|
int src_offset = gid.x * info.srcStride.x + gid.y * info.srcStride.y + gid.z * info.srcStride.z;
|
|
int dst_offset = gid.x * info.dstStride.x + gid.y * info.dstStride.y + gid.z * info.dstStride.z;
|
|
if(src_offset >= 0 && src_offset < info.totalSize.x) {
|
|
if(dst_offset >= 0 && dst_offset < info.totalSize.y) {
|
|
out[dst_offset] = in[src_offset];
|
|
}
|
|
}
|
|
}
|
|
)metal";
|
|
|
|
class MetalGather : public MetalExecution {
|
|
private:
|
|
const LoopParam* mLoop;
|
|
int mCmdIndex = 0;
|
|
id<MTLBuffer> mParam;
|
|
id<MTLComputePipelineState> mPipeline;
|
|
std::vector<Tensor*> mTensors;
|
|
public:
|
|
MetalGather(const LoopParam* loop, Backend *bn, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, int index = 0) : MetalExecution(bn) {
|
|
mLoop = loop;
|
|
auto mtbn = static_cast<MetalBackend *>(bn);
|
|
auto context = (__bridge MNNMetalContext *)mtbn->context();
|
|
|
|
mParam = [context newDeviceBuffer:sizeof(GatherInfo) access:CPUWriteOnly];
|
|
bool useFp16 = mtbn->useFp16InsteadFp32();
|
|
mTensors.resize(mLoop->tensorNumber());
|
|
mCmdIndex = index;
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(mCmdIndex);
|
|
_setTensorStack(mTensors, inputs, outputs, mLoop);
|
|
auto dstTensor = mTensors[cmd->indexes()->data()[0]];
|
|
|
|
NSString* T = MetalCast::getScalarType(dstTensor->getType(), useFp16);
|
|
// gather blit command pipeline
|
|
{
|
|
std::vector<std::string> keys = {
|
|
std::string([T UTF8String]),
|
|
"blitregion"
|
|
};
|
|
auto pipeline = mtbn->runtime()->findPipeline(keys);
|
|
if (nil == pipeline) {
|
|
MTLCompileOptions *compileOptions = [[MTLCompileOptions alloc] init];
|
|
compileOptions.preprocessorMacros = @{
|
|
@"T" : T,
|
|
};
|
|
pipeline = mtbn->makeComputePipelineWithSourceOption(gBlitRegion, "gather_blit", compileOptions);
|
|
mtbn->runtime()->insertPipeline(keys, pipeline);
|
|
}
|
|
if (nil == pipeline) {
|
|
MNN_ERROR("Create gather pipeline error\n");
|
|
}
|
|
mPipeline = pipeline;
|
|
}
|
|
}
|
|
virtual ~MetalGather() = default;
|
|
virtual ErrorCode onResize(const std::vector<Tensor *>& inputs, const std::vector<Tensor *>& outputs) override {
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(mCmdIndex);
|
|
_setTensorStack(mTensors, inputs, outputs, mLoop);
|
|
|
|
auto srcStride = cmd->view()->GetAs<View>(1)->stride()->data();
|
|
auto dstStride = cmd->view()->GetAs<View>(0)->stride()->data();
|
|
auto size = cmd->size()->data();
|
|
int totalSize = mLoop->loopNumber() * size[0] * size[1] * size[2];
|
|
auto param = reinterpret_cast<GatherInfo*>([mParam contents]);
|
|
for (int i=0; i<3; ++i) {
|
|
param->size[i] = size[i];
|
|
param->stride[i] = srcStride[i];
|
|
param->extent[i] = dstStride[i];
|
|
}
|
|
param->stride[3] = cmd->view()->GetAs<View>(1)->offset();
|
|
param->extent[3] = cmd->view()->GetAs<View>(0)->offset();
|
|
param->size[3] = size[0] * size[1] * size[2];
|
|
param->step[3] = totalSize;
|
|
param->step[0] = cmd->steps()->data()[0];
|
|
param->step[1] = cmd->steps()->data()[1];
|
|
param->iter[0] = 0;
|
|
param->iter[1] = 0;
|
|
auto iterIndex = cmd->iterIndexes()->data();
|
|
if (iterIndex[0] >= 0) {
|
|
param->iter[0] = 1;
|
|
}
|
|
if (iterIndex[1] >= 0) {
|
|
param->iter[1] = 1;
|
|
}
|
|
|
|
auto dstTensor = mTensors[cmd->indexes()->data()[0]];
|
|
auto srcTensor = mTensors[cmd->indexes()->data()[1]];
|
|
auto inputSize = srcTensor->usize() / srcTensor->buffer().type.bytes();
|
|
auto outputSize = dstTensor->usize() / dstTensor->buffer().type.bytes();
|
|
param->totalSize[0] = inputSize;
|
|
param->totalSize[1] = outputSize;
|
|
|
|
return NO_ERROR;
|
|
}
|
|
virtual void onEncode(const std::vector<Tensor *>& inputs, const std::vector<Tensor *>& outputs,
|
|
id<MTLComputeCommandEncoder> encoder) override {
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(mCmdIndex);
|
|
auto size = cmd->size()->data();
|
|
auto srcStride = cmd->view()->GetAs<View>(1)->stride()->data();
|
|
auto dstStride = cmd->view()->GetAs<View>(0)->stride()->data();
|
|
int totalSize = mLoop->loopNumber() * size[0] * size[1] * size[2];
|
|
|
|
[encoder setComputePipelineState:mPipeline];
|
|
auto dstTensor = mTensors[cmd->indexes()->data()[0]];
|
|
auto srcTensor = mTensors[cmd->indexes()->data()[1]];
|
|
MetalBackend::setTensor(dstTensor, encoder, 0);
|
|
MetalBackend::setTensor(srcTensor, encoder, 1);
|
|
|
|
auto iterIndex = cmd->iterIndexes()->data();
|
|
if (iterIndex[0] >= 0) {
|
|
MetalBackend::setTensor(mTensors[iterIndex[0]], encoder, 3);
|
|
} else {
|
|
MetalBackend::setTensor(dstTensor, encoder, 3);
|
|
}
|
|
if (iterIndex[1] >= 0) {
|
|
MetalBackend::setTensor(mTensors[iterIndex[1]], encoder, 2);
|
|
} else {
|
|
MetalBackend::setTensor(srcTensor, encoder, 2);
|
|
}
|
|
[encoder setBuffer:mParam offset:0 atIndex:4];
|
|
[encoder dispatchThreadgroups:MTLSizeMake(UP_DIV(totalSize, 256), 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
|
|
}
|
|
};
|
|
|
|
|
|
static const char* gBinaryBroadcast = R"metal(
|
|
#include <metal_stdlib>
|
|
#include <simd/simd.h>
|
|
using namespace metal;
|
|
struct constBuffer
|
|
{
|
|
int4 srcview0;
|
|
int4 srcview1;
|
|
int4 dstview;
|
|
int4 size;
|
|
};
|
|
|
|
static inline __attribute__((always_inline))
|
|
int computeVec4dot(thread const int4& a, thread const int4& b)
|
|
{
|
|
return (((a.x * b.x) + (a.y * b.y)) + (a.z * b.z)) + (a.w * b.w);
|
|
}
|
|
|
|
kernel void loop_binary(device T1* uOutput [[buffer(0)]], const device T0* uInput0 [[buffer(1)]], const device T0* uInput1 [[buffer(2)]], constant constBuffer& uConstant [[buffer(3)]], uint3 gl_GlobalInvocationID [[thread_position_in_grid]])
|
|
{
|
|
int3 posTmp = int3(gl_GlobalInvocationID);
|
|
if (posTmp.x < uConstant.size.w)
|
|
{
|
|
int4 pos;
|
|
pos.x = posTmp.x / (uConstant.size.y * uConstant.size.z);
|
|
int subIndex = posTmp.x % (uConstant.size.y * uConstant.size.z);
|
|
pos.z = subIndex % uConstant.size.z;
|
|
pos.y = subIndex / uConstant.size.z;
|
|
pos.w = 1;
|
|
int4 param = uConstant.srcview0;
|
|
int4 param_1 = pos;
|
|
int s0 = computeVec4dot(param, param_1);
|
|
int4 param_2 = uConstant.srcview1;
|
|
int4 param_3 = pos;
|
|
int s1 = computeVec4dot(param_2, param_3);
|
|
int4 param_4 = uConstant.dstview;
|
|
int4 param_5 = pos;
|
|
int d = computeVec4dot(param_4, param_5);
|
|
T0 V0 = uInput0[s0];
|
|
T0 V1 = uInput1[s1];
|
|
uOutput[d] = CUSTOM;
|
|
}
|
|
}
|
|
)metal";
|
|
|
|
struct BinaryBroadCastInfo {
|
|
int srcview0[4];
|
|
int srcview1[4];
|
|
int dstview[4];
|
|
int size[4];
|
|
};
|
|
|
|
class MetalBinaryBroadCast : public MetalExecution {
|
|
public:
|
|
MetalBinaryBroadCast(const LoopParam* loop, Backend *bn, const std::vector<Tensor*>& tensors, NSString* CUSTOM, int index = 0) : MetalExecution(bn) {
|
|
mLoop = loop;
|
|
auto mtbn = static_cast<MetalBackend *>(bn);
|
|
auto context = (__bridge MNNMetalContext *)mtbn->context();
|
|
mParam = mtbn->getConstBuffer(sizeof(BinaryBroadCastInfo));
|
|
mTensors = tensors;
|
|
mCmdIndex = index;
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(mCmdIndex);
|
|
auto dstTensor = mTensors[cmd->indexes()->data()[0]];
|
|
auto srcTensor = mTensors[cmd->indexes()->data()[1]];
|
|
auto srcTensor1 = mTensors[cmd->indexes()->data()[2]];
|
|
|
|
NSString* T1 = MetalCast::getScalarType(dstTensor->getType(), mtbn->useFp16InsteadFp32());
|
|
NSString* T0 = MetalCast::getScalarType(srcTensor->getType(), mtbn->useFp16InsteadFp32());
|
|
std::vector<std::string> keys = {
|
|
std::string([T0 UTF8String]),
|
|
std::string([T1 UTF8String]),
|
|
std::string([CUSTOM UTF8String]),
|
|
"binary_broadcast"
|
|
};
|
|
auto pipeline = mtbn->runtime()->findPipeline(keys);
|
|
if (nil == pipeline) {
|
|
MTLCompileOptions *compileOptions = [[MTLCompileOptions alloc] init];
|
|
compileOptions.preprocessorMacros = @{
|
|
@"T0" : T0,
|
|
@"T1" : T1,
|
|
@"CUSTOM" : CUSTOM,
|
|
};
|
|
pipeline = mtbn->makeComputePipelineWithSourceOption(gBinaryBroadcast, "loop_binary", compileOptions);
|
|
mtbn->runtime()->insertPipeline(keys, pipeline);
|
|
}
|
|
if (nil == pipeline) {
|
|
MNN_ERROR("Create Binary Broadcast pipeline error\n");
|
|
}
|
|
mPipeline = pipeline;
|
|
}
|
|
virtual ~MetalBinaryBroadCast() {
|
|
auto mtbn = static_cast<MetalBackend*>(backend());
|
|
mtbn->returnConstBuffer(mParam);
|
|
}
|
|
virtual ErrorCode onResize(const std::vector<Tensor *>& inputs, const std::vector<Tensor *>& outputs) override {
|
|
_setTensorStack(mTensors, inputs, outputs, mLoop);
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(mCmdIndex);
|
|
auto size = cmd->size()->data();
|
|
auto srcStride0 = cmd->view()->GetAs<View>(1)->stride()->data();
|
|
auto srcStride1 = cmd->view()->GetAs<View>(2)->stride()->data();
|
|
auto dstStride = cmd->view()->GetAs<View>(0)->stride()->data();
|
|
mTotalSize = size[0] * size[1] * size[2];
|
|
auto param = reinterpret_cast<BinaryBroadCastInfo*>([mParam contents]);
|
|
for (int i=0; i<3; ++i) {
|
|
param->size[i] = size[i];
|
|
param->srcview0[i] = srcStride0[i];
|
|
param->srcview1[i] = srcStride1[i];
|
|
param->dstview[i] = dstStride[i];
|
|
}
|
|
param->srcview0[3] = cmd->view()->GetAs<View>(1)->offset();
|
|
param->srcview1[3] = cmd->view()->GetAs<View>(2)->offset();
|
|
param->dstview[3] = cmd->view()->GetAs<View>(0)->offset();
|
|
param->size[3] = size[0] * size[1] * size[2];
|
|
return NO_ERROR;
|
|
}
|
|
virtual void onEncode(const std::vector<Tensor *>& inputs, const std::vector<Tensor *>& outputs,
|
|
id<MTLComputeCommandEncoder> encoder) override {
|
|
|
|
auto cmd = mLoop->commands()->GetAs<RegionCommand>(mCmdIndex);
|
|
auto dstTensor = mTensors[cmd->indexes()->data()[0]];
|
|
auto srcTensor = mTensors[cmd->indexes()->data()[1]];
|
|
auto srcTensor1 = mTensors[cmd->indexes()->data()[2]];
|
|
[encoder setComputePipelineState:mPipeline];
|
|
MetalBackend::setTensor(dstTensor, encoder, 0);
|
|
MetalBackend::setTensor(srcTensor, encoder, 1);
|
|
MetalBackend::setTensor(srcTensor1, encoder, 2);
|
|
[encoder setBuffer:mParam offset:0 atIndex:3];
|
|
[encoder dispatchThreadgroups:MTLSizeMake(UP_DIV(mTotalSize, 256), 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
|
|
}
|
|
private:
|
|
const LoopParam* mLoop;
|
|
id<MTLComputePipelineState> mPipeline;
|
|
id<MTLBuffer> mParam;
|
|
std::vector<Tensor*> mTensors;
|
|
int mTotalSize;
|
|
int mCmdIndex = 0;
|
|
};
|
|
|
|
class MetalLoop : public MetalExecution {
|
|
public:
|
|
MetalLoop(const LoopParam* loop, Backend *bn, const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) : MetalExecution(bn) {
|
|
mLoop = loop;
|
|
auto mtbn = static_cast<MetalBackend *>(bn);
|
|
auto context = (__bridge MNNMetalContext *)mtbn->context();
|
|
mTensors.resize(mLoop->tensorNumber());
|
|
_setTensorStack(mTensors, inputs, outputs, mLoop);
|
|
|
|
// Init
|
|
if(mLoop->initCommand() != nullptr) {
|
|
mNeedInit = true;
|
|
std::string shader = "set_copy";
|
|
auto dstTensor = mTensors[mLoop->initCommand()->GetAs<RegionCommand>(0)->indexes()->data()[0]];
|
|
NSString* T = MetalCast::getScalarType(dstTensor->getType(), mtbn->useFp16InsteadFp32());
|
|
|
|
auto cmd = mLoop->initCommand()->GetAs<RegionCommand>(0);
|
|
if (cmd->op() == nullptr){
|
|
shader = "set_zero";
|
|
} else {
|
|
mInitParam = [context newDeviceBuffer:sizeof(InitInfo) access:CPUWriteOnly];
|
|
}
|
|
std::vector<std::string> keys = {
|
|
std::string([T UTF8String]),
|
|
"init_region",
|
|
shader
|
|
};
|
|
auto pipeline = mtbn->runtime()->findPipeline(keys);
|
|
if (nil == pipeline) {
|
|
MTLCompileOptions *compileOptions = [[MTLCompileOptions alloc] init];
|
|
compileOptions.preprocessorMacros = @{
|
|
@"T" : T,
|
|
};
|
|
pipeline = mtbn->makeComputePipelineWithSourceOption(gInitRegion, shader.c_str(), compileOptions);
|
|
mtbn->runtime()->insertPipeline(keys, pipeline);
|
|
}
|
|
if (nil == pipeline) {
|
|
MNN_ERROR("Create gather init pipeline error\n");
|
|
}
|
|
mInitPipeline = pipeline;
|
|
}
|
|
|
|
bool valid = true;
|
|
for (int i=0; i<loop->commands()->size(); ++i) {
|
|
auto cmd = loop->commands()->GetAs<RegionCommand>(i);
|
|
auto subop = cmd->op();
|
|
if (OpType_UnaryOp == subop->type() && nullptr == subop->main() && cmd->fuse() < 0) {
|
|
mExecutions.emplace_back(std::make_shared<MetalGather>(loop, bn, inputs, outputs, i));
|
|
} else if (OpType_MatMul == subop->type() && loop->parallel()) {
|
|
mExecutions.emplace_back(std::make_shared<MetalBatchMatMul>(loop, bn, i));
|
|
} else if (OpType_BinaryOp == subop->type() && cmd->fuse() < 0 && 1 == loop->loopNumber()) {
|
|
mExecutions.emplace_back(std::make_shared<MetalBinaryBroadCast>(loop, bn, mTensors, MetalBinary::convert(cmd->op()->main_as_BinaryOp()->opType(), mTensors[cmd->indexes()->data()[1]]->getType().code == halide_type_float), i));
|
|
} else {
|
|
valid = false;
|
|
break;
|
|
}
|
|
}
|
|
if (!valid) {
|
|
mExecutions.clear();
|
|
}
|
|
}
|
|
virtual ~MetalLoop() = default;
|
|
|
|
virtual ErrorCode onResize(const std::vector<Tensor *>& inputs, const std::vector<Tensor *>& outputs) override {
|
|
// Init
|
|
if(mNeedInit) {
|
|
_setTensorStack(mTensors, inputs, outputs, mLoop);
|
|
auto initCmd = mLoop->initCommand()->GetAs<RegionCommand>(0);
|
|
auto data = reinterpret_cast<InitInfo*>([mInitParam contents]);
|
|
|
|
auto srcStride = initCmd->view()->GetAs<View>(1)->stride()->data();
|
|
auto dstStride = initCmd->view()->GetAs<View>(0)->stride()->data();
|
|
auto dataSize = initCmd->size()->data();
|
|
for (int i = 0; i < 3; ++i) {
|
|
data->srcStride[i] = srcStride[i];
|
|
data->dstStride[i] = dstStride[i];
|
|
data->size[i] = dataSize[i];
|
|
}
|
|
|
|
auto initDstTensor = mTensors[initCmd->indexes()->data()[0]];
|
|
auto initSrcTensor = mTensors[initCmd->indexes()->data()[1]];
|
|
auto initInputSize = initSrcTensor->usize() / initSrcTensor->buffer().type.bytes();
|
|
auto initOutputSize = initDstTensor->usize() / initDstTensor->buffer().type.bytes();
|
|
data->totalSize[0] = initInputSize;
|
|
data->totalSize[1] = initOutputSize;
|
|
|
|
auto backend = static_cast<MetalBackend *>(this->backend());
|
|
auto context = (__bridge MNNMetalContext *)backend->context();
|
|
mInitThreads = [context computeBestGroupAndLocal:mInitPipeline threads:MTLSizeMake(data->size[0], data->size[1], data->size[2])];
|
|
}
|
|
|
|
// Loop commands
|
|
for (auto& exe : mExecutions) {
|
|
auto code = exe->onResize(inputs, outputs);
|
|
if (NO_ERROR != code) {
|
|
return code;
|
|
}
|
|
}
|
|
return NO_ERROR;
|
|
}
|
|
virtual void onEncode(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs, id<MTLComputeCommandEncoder> encoder) override {
|
|
// Init
|
|
if(mNeedInit) {
|
|
auto initCmd = mLoop->initCommand()->GetAs<RegionCommand>(0);
|
|
[encoder setComputePipelineState:mInitPipeline];
|
|
auto dstTensor = mTensors[initCmd->indexes()->data()[0]];
|
|
auto srcTensor = mTensors[initCmd->indexes()->data()[1]];
|
|
MetalBackend::setTensor(dstTensor, encoder, 0);
|
|
MetalBackend::setTensor(srcTensor, encoder, 1);
|
|
[encoder setBuffer:mInitParam offset:0 atIndex:2];
|
|
[encoder dispatchThreadgroups:mInitThreads.first threadsPerThreadgroup:mInitThreads.second];
|
|
}
|
|
// Loop commands
|
|
for (auto& exe : mExecutions) {
|
|
exe->onEncode(inputs, outputs, encoder);
|
|
}
|
|
}
|
|
bool isValid() {
|
|
return !mExecutions.empty();
|
|
}
|
|
private:
|
|
const LoopParam* mLoop;
|
|
std::vector<std::shared_ptr<MetalExecution>> mExecutions;
|
|
std::vector<Tensor*> mTensors;
|
|
// For Init
|
|
bool mNeedInit = false;
|
|
std::pair<MTLSize, MTLSize> mInitThreads;
|
|
id<MTLComputePipelineState> mInitPipeline;
|
|
id<MTLBuffer> mInitParam;
|
|
|
|
};
|
|
|
|
class MetalLoopCreator : public MetalBackend::Creator {
|
|
public:
|
|
virtual Execution *onCreate(const std::vector<Tensor *> &inputs, const MNN::Op *op, Backend *bn, const std::vector<Tensor *> &outputs) const {
|
|
auto loop = op->main_as_LoopParam();
|
|
if (nullptr == loop || loop->commands() == nullptr) {
|
|
return nullptr;
|
|
}
|
|
|
|
// Make Tensor Stack
|
|
if (1 == loop->commands()->size()) {
|
|
auto cmd = loop->commands()->GetAs<RegionCommand>(0);
|
|
auto subop = cmd->op();
|
|
if (OpType_UnaryOp == subop->type() && nullptr == subop->main() && cmd->fuse() < 0 && nullptr == loop->initCommand()) {
|
|
return new MetalGather(loop, bn, inputs, outputs);
|
|
}
|
|
if (OpType_MatMul == subop->type() && loop->parallel() && nullptr == loop->initCommand()) {
|
|
return new MetalBatchMatMul(loop, bn);
|
|
}
|
|
if (OpType_BinaryOp == subop->type() && cmd->fuse() < 0 && 1 == loop->loopNumber() && nullptr == loop->initCommand()) {
|
|
std::vector<MNN::Tensor*> tensors(loop->tensorNumber());
|
|
_setTensorStack(tensors, inputs, outputs, loop);
|
|
auto srcTensor = tensors[cmd->indexes()->data()[1]];
|
|
|
|
NSString* CUSTOM = MetalBinary::convert(cmd->op()->main_as_BinaryOp()->opType(), srcTensor->getType().code == halide_type_float);
|
|
if (nil == CUSTOM) {
|
|
MNN_ERROR("Metal Don't support binary - %d \n", cmd->op()->main_as_BinaryOp()->opType());
|
|
return nullptr;
|
|
}
|
|
return new MetalBinaryBroadCast(loop, bn, tensors, CUSTOM);
|
|
}
|
|
}
|
|
// General Case
|
|
auto exe = new MetalLoop(loop, bn, inputs, outputs);
|
|
if (exe->isValid()) {
|
|
return exe;
|
|
}
|
|
delete exe;
|
|
return nullptr;
|
|
|
|
}
|
|
};
|
|
REGISTER_METAL_OP_CREATOR(MetalLoopCreator, OpType_While);
|
|
} // namespace MNN
|
|
#endif /* MNN_METAL_ENABLED */
|