wehub-resource-sync
148 lines
4.8 KiB
C++
148 lines
4.8 KiB
C++
//
|
|
// TopKV2Execution.cpp
|
|
// MNN
|
|
//
|
|
// OpenCL image-path implementation of TopKV2.
|
|
//
|
|
|
|
#include "TopKV2Execution.hpp"
|
|
#include "core/TensorUtils.hpp"
|
|
#include "MNN_generated.h"
|
|
|
|
namespace MNN {
|
|
namespace OpenCL {
|
|
|
|
static const int kTopKThreadNumber = 128;
|
|
static const int kTopKLocalK = 8;
|
|
static const int kTopKCandidateNumber = kTopKThreadNumber * kTopKLocalK;
|
|
|
|
TopKV2Execution::TopKV2Execution(const MNN::Op* op, Backend* backend, int k) : CommonExecution(backend, op) {
|
|
mOpenCLBackend = static_cast<OpenCLBackend*>(backend);
|
|
mK = k;
|
|
mLargest = true;
|
|
auto param = op->main_as_TopKV2();
|
|
if (nullptr != param) {
|
|
mLargest = param->largest();
|
|
}
|
|
}
|
|
|
|
ErrorCode TopKV2Execution::onResize(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
auto input = inputs[0];
|
|
const int rowSize = input->length(input->dimensions() - 1);
|
|
if (rowSize <= 0) {
|
|
mNumRows = 0;
|
|
return NO_ERROR;
|
|
}
|
|
mNumRows = input->elementSize() / rowSize;
|
|
CommonExecution::onResize(inputs, outputs);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode TopKV2Execution::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
if (mNumRows <= 0) {
|
|
return NO_ERROR;
|
|
}
|
|
|
|
MNN_ASSERT(inputs.size() >= 1);
|
|
MNN_ASSERT(outputs.size() == 2);
|
|
|
|
auto input = inputs[0];
|
|
auto outputValue = outputs[0];
|
|
auto outputIndex = outputs[1];
|
|
|
|
const int rowSize = input->length(input->dimensions() - 1);
|
|
const int k = mK;
|
|
|
|
if (k > kTopKCandidateNumber) {
|
|
MNN_ERROR("TopKV2: k is too large, current implementation supports k <= %d\n", kTopKCandidateNumber);
|
|
return NOT_SUPPORT;
|
|
}
|
|
|
|
auto runtime = mOpenCLBackend->getOpenCLRuntime();
|
|
|
|
// Get shape info: tensorShapeFormat returns {N, H, W, C}
|
|
std::vector<int> inputShape = tensorShapeFormat(input);
|
|
const int width = inputShape[2]; // W
|
|
const int channels = inputShape[3]; // C = rowSize (for 1D/2D)
|
|
const int channelBlocks = UP_DIV(channels, 4);
|
|
|
|
// Build kernel with appropriate options
|
|
std::set<std::string> buildOptions;
|
|
if (mLargest) {
|
|
buildOptions.insert("-DSORT_DESC=1");
|
|
}
|
|
auto inputType = input->getType();
|
|
if (inputType.code == halide_type_int && inputType.bits == 32) {
|
|
buildOptions.insert("-DIS_INT=1");
|
|
}
|
|
|
|
mUnits.resize(1);
|
|
Unit& unit = mUnits[0];
|
|
unit.kernel = runtime->buildKernel("topkv2", "topkv2_channel", buildOptions, mOpenCLBackend->getPrecision());
|
|
OPENCL_CHECK_KERNEL(unit.kernel);
|
|
|
|
std::vector<uint32_t> gws = {
|
|
static_cast<uint32_t>(kTopKThreadNumber),
|
|
static_cast<uint32_t>(mNumRows),
|
|
static_cast<uint32_t>(1),
|
|
};
|
|
std::vector<uint32_t> lws = {
|
|
static_cast<uint32_t>(kTopKThreadNumber),
|
|
static_cast<uint32_t>(1),
|
|
static_cast<uint32_t>(1),
|
|
};
|
|
|
|
uint32_t idx = 0;
|
|
cl_int ret = CL_SUCCESS;
|
|
ret |= unit.kernel->get().setArg(idx++, gws[0]);
|
|
ret |= unit.kernel->get().setArg(idx++, gws[1]);
|
|
ret |= unit.kernel->get().setArg(idx++, gws[2]);
|
|
ret |= unit.kernel->get().setArg(idx++, openCLImage(input));
|
|
ret |= unit.kernel->get().setArg(idx++, openCLImage(outputValue));
|
|
ret |= unit.kernel->get().setArg(idx++, openCLImage(outputIndex));
|
|
ret |= unit.kernel->get().setArg(idx++, rowSize);
|
|
ret |= unit.kernel->get().setArg(idx++, k);
|
|
ret |= unit.kernel->get().setArg(idx++, width);
|
|
ret |= unit.kernel->get().setArg(idx++, channelBlocks);
|
|
MNN_CHECK_CL_SUCCESS(ret, "setArg TopKV2 topkv2_channel");
|
|
|
|
mOpenCLBackend->recordKernel3d(unit.kernel, gws, lws);
|
|
|
|
unit.globalWorkSize = {gws[0], gws[1], gws[2]};
|
|
unit.localWorkSize = {lws[0], lws[1], lws[2]};
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class TopKV2Creator : public OpenCLBackend::Creator {
|
|
public:
|
|
virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
|
const MNN::Op* op, Backend* backend) const override {
|
|
if (inputs.size() < 2 || outputs.size() != 2) {
|
|
return nullptr;
|
|
}
|
|
if (TensorUtils::getDescribe(inputs[0])->dimensionFormat == MNN_DATA_FORMAT_NC4HW4) {
|
|
return nullptr;
|
|
}
|
|
// Only support sort along channel (last dim mapped to C) for dims <= 2
|
|
if (inputs[0]->dimensions() > 2) {
|
|
return nullptr;
|
|
}
|
|
|
|
const int k = inputs[1]->host<int32_t>()[0];
|
|
|
|
for (int i = 0; i < inputs.size(); ++i) {
|
|
TensorUtils::setTensorSupportPack(inputs[i], false);
|
|
}
|
|
for (int i = 0; i < outputs.size(); ++i) {
|
|
TensorUtils::setTensorSupportPack(outputs[i], false);
|
|
}
|
|
OPENCL_CREATOR_CHECK(new TopKV2Execution(op, backend, k));
|
|
}
|
|
};
|
|
|
|
REGISTER_OPENCL_OP_CREATOR(TopKV2Creator, OpType_TopKV2, IMAGE);
|
|
|
|
} // namespace OpenCL
|
|
} // namespace MNN
|