MNN/source/backend/opengl/execution/GLPermute.cpp at master · DannyP0/MNN · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
//
//  GLPermute.cpp
//  MNN
//
//  Created by MNN on 2019/01/31.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#include "GLPermute.hpp"
#include <sstream>
#include "AllShader.hpp"
#include "GLBackend.hpp"
#include "Macro.h"
#include "TensorUtils.hpp"

namespace MNN {
namespace OpenGL {
GLPermute::GLPermute(const std::vector<Tensor *> &inputs, const Op *op, Backend *bn) : Execution(bn) {

    auto newDim = op->main_as_Permute()->dims();
    for (int i = 0; i < newDim->size(); ++i) {
        mDims.push_back(newDim->data()[i]);
    }
}

GLPermute::~GLPermute() {

}

ErrorCode GLPermute::onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
    auto input = inputs[0];
    auto output = outputs[0];
    std::vector<std::string> prefix;
    setLocalSize(prefix, mLocalSize, 8, 8, 1);
    mSrcBuffer.reset(new GLSSBOBuffer(input->size()));
    mDstBuffer.reset(new GLSSBOBuffer(output->size()));

    mPermuteProgram = ((GLBackend *)backend())->getProgram("permute", glsl_permute_glsl, prefix);
    mSrcProgram = ((GLBackend *)backend())->getProgram("src", glsl_image_to_nchw_buffer_glsl, prefix);
    mDstProgram = ((GLBackend *)backend())->getProgram("dst", glsl_nchw_buffer_to_image_glsl, prefix);

    return NO_ERROR;
}

ErrorCode GLPermute::onExecute(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
    auto input = inputs[0];
    auto output = outputs[0];
    int ih = input->height();
    int iw = input->width();
    int ic = input->channel();
    int ib = input->batch();
    int ic_4 = UP_DIV(ic, 4);

    int oh = output->height();
    int ow = output->width();
    int oc = output->channel();
    int ob = output->batch();
    int oc_4 = UP_DIV(oc, 4);

    //image -> buffer(nchw)
    {
        mSrcProgram->useProgram();
        glBindImageTexture(0, input->deviceId(), 0, GL_TRUE, 0, GL_READ_ONLY, TEXTURE_FORMAT);
        glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, mSrcBuffer->getId());
        glUniform1i(2, iw);
        glUniform1i(3, ih);
        OPENGL_CHECK_ERROR;
        ((GLBackend *)backend())->compute(UP_DIV(iw, mLocalSize[0]), UP_DIV(ih, mLocalSize[1]), UP_DIV(ic_4, mLocalSize[2]));
        OPENGL_CHECK_ERROR;
    }

    //do permute
    {
        mPermuteProgram->useProgram();
        glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, mSrcBuffer->getId());
        glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, mDstBuffer->getId());
        glUniform4i(2, mDims[0], mDims[1], mDims[2], mDims[3]);
        glUniform4i(3, iw, ih, ic, ib);
        glUniform4i(4, ow, oh, oc, ob);
        OPENGL_CHECK_ERROR;
        ((GLBackend *)backend())->compute(UP_DIV(ow, mLocalSize[0]), UP_DIV(oh, mLocalSize[1]), UP_DIV(oc, mLocalSize[2]));
        OPENGL_CHECK_ERROR;
    }

    //buffer(nchw) -> image
    {
        mDstProgram->useProgram();
        glBindImageTexture(0, output->deviceId(), 0, GL_TRUE, 0, GL_WRITE_ONLY, TEXTURE_FORMAT);
        glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, mDstBuffer->getId());
        glUniform1i(2, ow);
        glUniform1i(3, oh);
        OPENGL_CHECK_ERROR;
        ((GLBackend *)backend())->compute(UP_DIV(ow, mLocalSize[0]), UP_DIV(oh, mLocalSize[1]), UP_DIV(oc_4, mLocalSize[2]));
        OPENGL_CHECK_ERROR;
    }

    return NO_ERROR;
}
GLCreatorRegister<TypedCreator<GLPermute>> __permute_op(OpType_Permute);
} // namespace OpenGL
} // namespace MNN