chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,260 @@
|
||||
//
|
||||
// GeometryBinary.cpp
|
||||
// MNN
|
||||
//
|
||||
// Created by MNN on 2020/05/07.
|
||||
// Copyright © 2018, Alibaba Group Holding Limited
|
||||
//
|
||||
|
||||
#include "ConvertUtils.hpp"
|
||||
#include "geometry/GeometryComputer.hpp"
|
||||
#include "geometry/GeometryComputerUtils.hpp"
|
||||
#include "shape/SizeComputer.hpp"
|
||||
#ifndef MNN_REDUCE_SIZE
|
||||
#define MNN_BINARY_LOOP_OPT
|
||||
#endif
|
||||
namespace MNN {
|
||||
class GeometryBinary : public GeometryComputer {
|
||||
public:
|
||||
virtual bool onRecompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
||||
Context& context, CommandBuffer& res) const override {
|
||||
if (res.command.size() != 1) {
|
||||
return false;
|
||||
}
|
||||
auto input0 = inputs[0];
|
||||
auto input1 = inputs[1];
|
||||
auto output = outputs[0];
|
||||
auto inputL0 = TensorUtils::getRawSize(input0);
|
||||
auto inputL1 = TensorUtils::getRawSize(input1);
|
||||
auto outputSize = TensorUtils::getRawSize(output);
|
||||
auto inp0format = TensorUtils::getDescribe(inputs[0])->dimensionFormat;
|
||||
auto inp1format = TensorUtils::getDescribe(inputs[1])->dimensionFormat;
|
||||
auto outFormat = TensorUtils::getDescribe(output)->dimensionFormat;
|
||||
auto cmdP = res.command[0];
|
||||
if (cmdP->op->type() != OpType_BinaryOp) {
|
||||
return false;
|
||||
}
|
||||
MNN_ASSERT(0 != inputL1 && 0 != inputL0 && 0 != outputSize);
|
||||
//MNN_PRINT("On compute geometry: %d - %d - %d\n", inputL0, inputL1, outputSize);
|
||||
if (1 == inputL0 || 1 == inputL1) {
|
||||
// Can directly compute
|
||||
cmdP->inputs[0] = input0;
|
||||
cmdP->inputs[1] = input1;
|
||||
return true;
|
||||
}
|
||||
// Need Broadcast or same shape
|
||||
bool input0Broadcast = false;
|
||||
bool input1Broadcast = false;
|
||||
// MNN: also force the broadcast on Vulkan. The Vulkan binary kernel
|
||||
// does not handle non-equal-rank inputs (e.g. {4} broadcast onto
|
||||
// {1,1,4}) — without this, AddBroast on Vulkan reads the wrong
|
||||
// image plane and outputs `input0` instead of `input0+input1`.
|
||||
bool gpuNeedsBroadcast = (MNN_DATA_FORMAT_NC4HW4 == outFormat
|
||||
|| context.forwardType() == MNN_FORWARD_OPENCL
|
||||
|| context.forwardType() == MNN_FORWARD_VULKAN);
|
||||
if (outputSize != inputL0 || inp0format != outFormat ||
|
||||
(output->dimensions() != input0->dimensions() && gpuNeedsBroadcast)) {
|
||||
input0Broadcast = true;
|
||||
}
|
||||
if (outputSize != inputL1 || inp1format != outFormat ||
|
||||
(output->dimensions() != input1->dimensions() && gpuNeedsBroadcast)) {
|
||||
input1Broadcast = true;
|
||||
}
|
||||
auto cacheTensor = std::move(res.extras);
|
||||
if (input0Broadcast) {
|
||||
std::shared_ptr<Tensor> newTensor;
|
||||
if (!cacheTensor.empty()) {
|
||||
newTensor = cacheTensor[cacheTensor.size() - 1];
|
||||
cacheTensor.erase(cacheTensor.begin() + cacheTensor.size() - 1);
|
||||
} else {
|
||||
newTensor.reset(new Tensor);
|
||||
}
|
||||
TensorUtils::copyShape(output, newTensor.get(), true);
|
||||
newTensor->buffer().type = output->buffer().type;
|
||||
ConvertUtils::broadcastto(input0, newTensor.get());
|
||||
input0 = newTensor.get();
|
||||
res.extras.emplace_back(newTensor);
|
||||
}
|
||||
if (input1Broadcast) {
|
||||
std::shared_ptr<Tensor> newTensor;
|
||||
if (!cacheTensor.empty()) {
|
||||
newTensor = cacheTensor[cacheTensor.size() - 1];
|
||||
cacheTensor.erase(cacheTensor.begin() + cacheTensor.size() - 1);
|
||||
} else {
|
||||
newTensor.reset(new Tensor);
|
||||
}
|
||||
TensorUtils::copyShape(output, newTensor.get(), true);
|
||||
newTensor->buffer().type = output->buffer().type;
|
||||
ConvertUtils::broadcastto(input1, newTensor.get());
|
||||
input1 = newTensor.get();
|
||||
res.extras.emplace_back(newTensor);
|
||||
}
|
||||
cmdP->inputs[0] = input0;
|
||||
cmdP->inputs[1] = input1;
|
||||
return true;
|
||||
}
|
||||
|
||||
virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
|
||||
Context& context, CommandBuffer& res) const override {
|
||||
auto input0 = inputs[0];
|
||||
auto input1 = inputs[1];
|
||||
auto output = outputs[0];
|
||||
auto inputL0 = TensorUtils::getRawSize(input0);
|
||||
auto inputL1 = TensorUtils::getRawSize(input1);
|
||||
auto outputSize = TensorUtils::getRawSize(output);
|
||||
auto inp0format = TensorUtils::getDescribe(inputs[0])->dimensionFormat;
|
||||
auto inp1format = TensorUtils::getDescribe(inputs[1])->dimensionFormat;
|
||||
auto outFormat = TensorUtils::getDescribe(output)->dimensionFormat;
|
||||
MNN_ASSERT(0 != inputL1 && 0 != inputL0 && 0 != outputSize);
|
||||
//MNN_PRINT("On compute geometry: %d - %d - %d\n", inputL0, inputL1, outputSize);
|
||||
if (1 == inputL0 || 1 == inputL1 || context.forwardType() == MNN_FORWARD_NN) {
|
||||
// Can directly compute
|
||||
std::shared_ptr<Command> cmdP(new Command);
|
||||
auto& cmd = *cmdP;
|
||||
cmd.op = op;
|
||||
cmd.inputs = {input0, input1};
|
||||
cmd.outputs = std::move(outputs);
|
||||
res.command.emplace_back(std::move(cmdP));
|
||||
return true;
|
||||
}
|
||||
// Need Broadcast or same shape
|
||||
bool input0Broadcast = false;
|
||||
bool input1Broadcast = false;
|
||||
// MNN: also force the broadcast on Vulkan. The Vulkan binary kernel
|
||||
// does not handle non-equal-rank inputs (e.g. {4} broadcast onto
|
||||
// {1,1,4}) — without this, AddBroast on Vulkan reads the wrong
|
||||
// image plane and outputs `input0` instead of `input0+input1`.
|
||||
bool gpuNeedsBroadcast = (MNN_DATA_FORMAT_NC4HW4 == outFormat
|
||||
|| context.forwardType() == MNN_FORWARD_OPENCL
|
||||
|| context.forwardType() == MNN_FORWARD_VULKAN);
|
||||
if (outputSize != inputL0 || inp0format != outFormat ||
|
||||
(output->dimensions() != input0->dimensions() && gpuNeedsBroadcast)) {
|
||||
input0Broadcast = true;
|
||||
}
|
||||
if (outputSize != inputL1 || inp1format != outFormat ||
|
||||
(output->dimensions() != input1->dimensions() && gpuNeedsBroadcast)) {
|
||||
input1Broadcast = true;
|
||||
}
|
||||
#ifdef MNN_BINARY_LOOP_OPT
|
||||
// One input need broadcast, the other needn't
|
||||
bool singleBroadCast = (!(input0Broadcast && input1Broadcast)) && (input0Broadcast || input1Broadcast);
|
||||
bool forwardSupportLoop = inp0format == outFormat && inp1format == outFormat && outFormat != MNN_DATA_FORMAT_NC4HW4 && input0->getType().code == halide_type_float && op->main_as_BinaryOp()->activationType() == 0;
|
||||
bool openLoop = context.support(Interpreter::GeometryComputeMask::GEOMETRCOMPUTEMASK_USELOOP);
|
||||
if (singleBroadCast && forwardSupportLoop && openLoop) {
|
||||
// Use Loop instead of broadcast
|
||||
std::shared_ptr<Tensor> newTensor(new Tensor);
|
||||
TensorUtils::copyShape(output, newTensor.get(), true);
|
||||
newTensor->buffer().type = output->buffer().type;
|
||||
int srcIndex = 1;
|
||||
int dstIndex = 2;
|
||||
if (input0Broadcast) {
|
||||
ConvertUtils::broadcastto(input0, newTensor.get());
|
||||
} else {
|
||||
srcIndex = 2;
|
||||
dstIndex = 1;
|
||||
ConvertUtils::broadcastto(input1, newTensor.get());
|
||||
}
|
||||
auto des = TensorUtils::getDescribe(newTensor.get());
|
||||
flatbuffers::FlatBufferBuilder builder;
|
||||
BinaryOpBuilder binaryOpParamBuilder(builder);
|
||||
binaryOpParamBuilder.add_opType(op->main_as_BinaryOp()->opType());
|
||||
auto binaryOpParamOffset = binaryOpParamBuilder.Finish();
|
||||
OpBuilder cmdOpBuilder(builder);
|
||||
cmdOpBuilder.add_type(OpType_BinaryOp);
|
||||
cmdOpBuilder.add_main(binaryOpParamOffset.Union());
|
||||
cmdOpBuilder.add_main_type(OpParameter_BinaryOp);
|
||||
auto cmdOpOffset = cmdOpBuilder.Finish();
|
||||
auto iterIndexesOffset = builder.CreateVector(std::vector<int>{-1, -1, -1});
|
||||
auto stepOffset = builder.CreateVector(std::vector<int>{0, 0, 0});
|
||||
auto indexesOffset = builder.CreateVector(std::vector<int>{2, 0, 1});
|
||||
std::vector<flatbuffers::Offset<RegionCommand>> regionCommands;
|
||||
|
||||
for (int i=0; i<des->regions.size(); ++i) {
|
||||
auto& reg = des->regions[i];
|
||||
auto sizeOffset = builder.CreateVector(reg.size, 3);
|
||||
auto dstStride = builder.CreateVector(reg.dst.stride, 3);
|
||||
auto srcStride = builder.CreateVector(reg.src.stride, 3);
|
||||
std::vector<flatbuffers::Offset<View>> views(3);
|
||||
{
|
||||
ViewBuilder dstBuilder(builder);
|
||||
dstBuilder.add_offset(reg.dst.offset);
|
||||
dstBuilder.add_stride(dstStride);
|
||||
views[0] = dstBuilder.Finish();
|
||||
views[dstIndex] = views[0];
|
||||
ViewBuilder srcBuilder(builder);
|
||||
srcBuilder.add_offset(reg.src.offset);
|
||||
srcBuilder.add_stride(srcStride);
|
||||
views[srcIndex] = srcBuilder.Finish();
|
||||
}
|
||||
auto viewsOffset = builder.CreateVector<flatbuffers::Offset<View>>(views);
|
||||
RegionCommandBuilder cmdBuilder(builder);
|
||||
cmdBuilder.add_op(cmdOpOffset);
|
||||
cmdBuilder.add_view(viewsOffset);
|
||||
cmdBuilder.add_size(sizeOffset);
|
||||
cmdBuilder.add_steps(stepOffset);
|
||||
cmdBuilder.add_iterIndexes(iterIndexesOffset);
|
||||
cmdBuilder.add_indexes(indexesOffset);
|
||||
|
||||
regionCommands.emplace_back(cmdBuilder.Finish());
|
||||
}
|
||||
auto rcmdAllOffset = builder.CreateVector<flatbuffers::Offset<RegionCommand>>(regionCommands);
|
||||
auto inputIndexesOffset = builder.CreateVector(std::vector<int>{0, 1});
|
||||
auto outputIndexesOffset = builder.CreateVector(std::vector<int>{2});
|
||||
LoopParamBuilder loopBuilder(builder);
|
||||
loopBuilder.add_commands(rcmdAllOffset);
|
||||
loopBuilder.add_loopNumber(1);
|
||||
loopBuilder.add_tensorNumber(3);
|
||||
loopBuilder.add_inputIndexes(inputIndexesOffset);
|
||||
loopBuilder.add_outputIndexes(outputIndexesOffset);
|
||||
auto loopOffset = loopBuilder.Finish();
|
||||
flatbuffers::Offset<flatbuffers::String> nameOffset;
|
||||
if (nullptr != op->name()) {
|
||||
nameOffset = builder.CreateString(op->name()->c_str());
|
||||
}
|
||||
OpBuilder finishBuilder(builder);
|
||||
finishBuilder.add_main(loopOffset.Union());
|
||||
finishBuilder.add_main_type(OpParameter_LoopParam);
|
||||
finishBuilder.add_type(OpType_While);
|
||||
if (nullptr != op->name()) {
|
||||
finishBuilder.add_name(nameOffset);
|
||||
}
|
||||
builder.Finish(finishBuilder.Finish());
|
||||
auto cmd = GeometryComputerUtils::makeCommand(builder, {input0, input1}, outputs);
|
||||
res.command.emplace_back(std::move(cmd));
|
||||
return true;
|
||||
}
|
||||
#endif
|
||||
if (input0Broadcast) {
|
||||
std::shared_ptr<Tensor> newTensor(new Tensor);
|
||||
TensorUtils::copyShape(output, newTensor.get(), true);
|
||||
newTensor->buffer().type = output->buffer().type;
|
||||
ConvertUtils::broadcastto(input0, newTensor.get());
|
||||
input0 = newTensor.get();
|
||||
res.extras.emplace_back(newTensor);
|
||||
}
|
||||
if (input1Broadcast) {
|
||||
std::shared_ptr<Tensor> newTensor(new Tensor);
|
||||
TensorUtils::copyShape(output, newTensor.get(), true);
|
||||
newTensor->buffer().type = output->buffer().type;
|
||||
ConvertUtils::broadcastto(input1, newTensor.get());
|
||||
input1 = newTensor.get();
|
||||
res.extras.emplace_back(newTensor);
|
||||
}
|
||||
std::shared_ptr<Command> cmdP(new Command);
|
||||
auto& cmd = *cmdP;
|
||||
cmd.op = op;
|
||||
cmd.inputs = {input0, input1};
|
||||
cmd.outputs = std::move(outputs);
|
||||
res.command.emplace_back(std::move(cmdP));
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
static void _create() {
|
||||
std::shared_ptr<GeometryComputer> comp(new GeometryBinary);
|
||||
GeometryComputer::registerGeometryComputer(comp, {OpType_BinaryOp});
|
||||
}
|
||||
|
||||
REGISTER_GEOMETRY(GeometryBinary, _create);
|
||||
|
||||
} // namespace MNN
|
||||
Reference in New Issue
Block a user