wehub-resource-sync
248 lines
8.8 KiB
C++
248 lines
8.8 KiB
C++
#include <MNN/expr/ExprCreator.hpp>
|
|
#include <MNN/expr/Module.hpp>
|
|
|
|
#include "MNNTestSuite.h"
|
|
#include "TestUtils.h"
|
|
#include "core/TensorUtils.hpp"
|
|
|
|
using namespace MNN;
|
|
using namespace MNN::Express;
|
|
|
|
static VARPS makeComplexGraph(VARP x) {
|
|
// Input: NCHW float, shape {1, 4, 32, 32}
|
|
// Graph intent:
|
|
// - Introduce multiple ops (convert/conv/pool/concat/transpose) to increase
|
|
// intermediate allocations.
|
|
// - Keep an early large tensor as an output (Aux) to make later output tensor
|
|
// more likely to be allocated with non-zero offset on Metal.
|
|
|
|
auto x4 = _Convert(x, NC4HW4);
|
|
|
|
auto c0 = _Conv(0.01f, 0.0f, x4, {4, 8}, {3, 3}, SAME, {1, 1}, {1, 1}, 1);
|
|
c0 = _Relu(c0);
|
|
|
|
auto maxP = _MaxPool(c0, {2, 2}, {2, 2}, VALID);
|
|
auto aveP = _AvePool(c0, {2, 2}, {2, 2}, VALID);
|
|
|
|
auto aux = _Concat({maxP, aveP}, 1);
|
|
aux->setName("Aux");
|
|
|
|
auto c1 = _Conv(0.02f, 0.01f, aux, {16, 4}, {1, 1}, SAME, {1, 1}, {1, 1}, 1);
|
|
c1 = _Relu6(c1);
|
|
|
|
auto y = _Convert(c1, NCHW);
|
|
y = _Transpose(y, {0, 2, 3, 1});
|
|
y = _Transpose(y, {0, 3, 1, 2});
|
|
y = y + _Scalar<float>(1.0f);
|
|
y = _ReduceSum(y, {2}, true);
|
|
y->setName("Output");
|
|
|
|
auto s = _Shape(y);
|
|
s->setName("Shape");
|
|
|
|
return {aux, y, s};
|
|
}
|
|
|
|
static VARP makeInput(float base) {
|
|
auto x = _Input({1, 4, 32, 32}, NCHW, halide_type_of<float>());
|
|
auto ptr = x->writeMap<float>();
|
|
for (int i = 0; i < x->getInfo()->size; ++i) {
|
|
ptr[i] = base + (float)i * 0.001f;
|
|
}
|
|
x->unMap();
|
|
return x;
|
|
}
|
|
|
|
static bool isMetalRuntime() {
|
|
auto rtInfo = Express::ExecutorScope::Current()->getRuntime();
|
|
if (rtInfo.first.empty()) {
|
|
return false;
|
|
}
|
|
return rtInfo.first.begin()->first == MNN_FORWARD_METAL;
|
|
}
|
|
|
|
class StaticModuleOutputReuseTest : public MNNTestCase {
|
|
public:
|
|
bool run(int precision) override {
|
|
auto executor = cloneCurrentExecutor();
|
|
ExecutorScope scope(executor);
|
|
|
|
std::vector<int8_t> buffer;
|
|
{
|
|
auto x = _Input({1, 4, 32, 32}, NCHW, halide_type_of<float>());
|
|
x->setName("Input");
|
|
auto outputs = makeComplexGraph(x);
|
|
buffer = Variable::save(outputs);
|
|
}
|
|
|
|
Module::Config config;
|
|
config.shapeMutable = true;
|
|
|
|
std::shared_ptr<Module> module(
|
|
Module::load({"Input"}, {"Aux", "Output", "Shape"}, (const uint8_t*)buffer.data(), buffer.size(), &config),
|
|
Module::destroy);
|
|
if (nullptr == module) {
|
|
return false;
|
|
}
|
|
{
|
|
auto x = _Input({1, 4, 320, 320}, NCHW, halide_type_of<float>());
|
|
auto ptr = x->writeMap<float>();
|
|
module->onForward({x});
|
|
}
|
|
|
|
auto input0 = makeInput(0.0f);
|
|
auto outputs0 = module->onForward({input0});
|
|
if (outputs0.size() != 3) {
|
|
return false;
|
|
}
|
|
|
|
auto aux0Info = outputs0[0]->getInfo();
|
|
auto output0Info = outputs0[1]->getInfo();
|
|
auto shape0Info = outputs0[2]->getInfo();
|
|
if (nullptr == aux0Info || nullptr == output0Info || nullptr == shape0Info) {
|
|
return false;
|
|
}
|
|
if (aux0Info->dim != std::vector<int>({1, 16, 16, 16})) {
|
|
return false;
|
|
}
|
|
if (output0Info->dim != std::vector<int>({1, 4, 1, 16})) {
|
|
return false;
|
|
}
|
|
if (shape0Info->dim != std::vector<int>({4})) {
|
|
return false;
|
|
}
|
|
|
|
int expectedShape[4] = {1, 4, 1, 16};
|
|
auto shapeResult0 = outputs0[2]->readMap<int>();
|
|
if (!checkVector(shapeResult0, expectedShape, 4, 0)) {
|
|
return false;
|
|
}
|
|
|
|
// Snapshot first inference results. We'll verify they remain unchanged after
|
|
// the next inference (to catch output buffer reuse issues).
|
|
std::vector<float> aux0Snapshot(aux0Info->size);
|
|
std::vector<float> output0Snapshot(output0Info->size);
|
|
std::vector<int> shape0Snapshot(4);
|
|
{
|
|
auto auxPtr = outputs0[0]->readMap<float>();
|
|
auto outPtr = outputs0[1]->readMap<float>();
|
|
for (int i = 0; i < aux0Info->size; ++i) {
|
|
aux0Snapshot[i] = auxPtr[i];
|
|
}
|
|
for (int i = 0; i < output0Info->size; ++i) {
|
|
output0Snapshot[i] = outPtr[i];
|
|
}
|
|
for (int i = 0; i < 4; ++i) {
|
|
shape0Snapshot[i] = shapeResult0[i];
|
|
}
|
|
}
|
|
|
|
auto aux0Tensor = outputs0[0]->getTensor();
|
|
auto output0Tensor = outputs0[1]->getTensor();
|
|
auto shape0Tensor = outputs0[2]->getTensor();
|
|
if (nullptr == aux0Tensor || nullptr == output0Tensor || nullptr == shape0Tensor) {
|
|
return false;
|
|
}
|
|
|
|
auto aux0DescribeOrigin = TensorUtils::getDescribeOrigin(aux0Tensor);
|
|
auto output0DescribeOrigin = TensorUtils::getDescribeOrigin(output0Tensor);
|
|
auto shape0DescribeOrigin = TensorUtils::getDescribeOrigin(shape0Tensor);
|
|
if (nullptr == aux0DescribeOrigin || nullptr == output0DescribeOrigin || nullptr == shape0DescribeOrigin) {
|
|
return false;
|
|
}
|
|
|
|
int aux0Offset = aux0DescribeOrigin->offset;
|
|
int output0Offset = output0DescribeOrigin->offset;
|
|
int shape0Offset = shape0DescribeOrigin->offset;
|
|
|
|
if (isMetalRuntime() && output0Offset <= 0) {
|
|
// Not a hard assert: offset depends on allocator strategy / platform.
|
|
// But the model is designed to make output buffer more likely to be a
|
|
// non-zero slice inside a shared MTLBuffer.
|
|
MNN_PRINT("[StaticModuleOutputReuseTest] Metal output offset=%d (aux=%d, shape=%d)\n", output0Offset, aux0Offset,
|
|
shape0Offset);
|
|
}
|
|
|
|
auto input1 = makeInput(10.0f);
|
|
auto outputs1 = module->onForward({input1});
|
|
for (int i=0; i<10; ++i) {
|
|
outputs1 = module->onForward({input1});
|
|
}
|
|
if (outputs1.size() != 3) {
|
|
return false;
|
|
}
|
|
|
|
auto aux1Info = outputs1[0]->getInfo();
|
|
auto output1Info = outputs1[1]->getInfo();
|
|
auto shape1Info = outputs1[2]->getInfo();
|
|
if (nullptr == aux1Info || nullptr == output1Info || nullptr == shape1Info) {
|
|
return false;
|
|
}
|
|
if (aux1Info->dim != std::vector<int>({1, 16, 16, 16})) {
|
|
return false;
|
|
}
|
|
if (output1Info->dim != std::vector<int>({1, 4, 1, 16})) {
|
|
return false;
|
|
}
|
|
if (shape1Info->dim != std::vector<int>({4})) {
|
|
return false;
|
|
}
|
|
|
|
auto shapeResult1 = outputs1[2]->readMap<int>();
|
|
if (!checkVector(shapeResult1, expectedShape, 4, 0)) {
|
|
return false;
|
|
}
|
|
|
|
|
|
// Ensure the previous forward's outputs are still valid, and their offsets
|
|
// stay unchanged after the next forward.
|
|
auto aux0InfoAfter = outputs0[0]->getInfo();
|
|
auto output0InfoAfter = outputs0[1]->getInfo();
|
|
auto shape0InfoAfter = outputs0[2]->getInfo();
|
|
if (nullptr == aux0InfoAfter || nullptr == output0InfoAfter || nullptr == shape0InfoAfter) {
|
|
return false;
|
|
}
|
|
if (aux0InfoAfter->dim != std::vector<int>({1, 16, 16, 16})) {
|
|
return false;
|
|
}
|
|
if (output0InfoAfter->dim != std::vector<int>({1, 4, 1, 16})) {
|
|
return false;
|
|
}
|
|
if (shape0InfoAfter->dim != std::vector<int>({4})) {
|
|
return false;
|
|
}
|
|
|
|
// Ensure the previous inference results remain unchanged after the next forward.
|
|
if (!checkVector(outputs0[0]->readMap<float>(), aux0Snapshot.data(), (int)aux0Snapshot.size(), 1e-4f)) {
|
|
return false;
|
|
}
|
|
if (!checkVector(outputs0[1]->readMap<float>(), output0Snapshot.data(), (int)output0Snapshot.size(), 1e-4f)) {
|
|
return false;
|
|
}
|
|
if (!checkVector(outputs0[2]->readMap<int>(), shape0Snapshot.data(), (int)shape0Snapshot.size(), 0)) {
|
|
return false;
|
|
}
|
|
|
|
auto aux0DescribeAfterOrigin = TensorUtils::getDescribeOrigin(outputs0[0]->getTensor());
|
|
auto output0DescribeAfterOrigin = TensorUtils::getDescribeOrigin(outputs0[1]->getTensor());
|
|
auto shape0DescribeAfterOrigin = TensorUtils::getDescribeOrigin(outputs0[2]->getTensor());
|
|
if (nullptr == aux0DescribeAfterOrigin || nullptr == output0DescribeAfterOrigin || nullptr == shape0DescribeAfterOrigin) {
|
|
return false;
|
|
}
|
|
|
|
if (aux0Offset != aux0DescribeAfterOrigin->offset) {
|
|
return false;
|
|
}
|
|
if (output0Offset != output0DescribeAfterOrigin->offset) {
|
|
return false;
|
|
}
|
|
if (shape0Offset != shape0DescribeAfterOrigin->offset) {
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
};
|
|
|
|
MNNTestSuiteRegister(StaticModuleOutputReuseTest, "expr/StaticModuleOutputReuseTest");
|