MNN:Bugfix: Fix crash bug for ConvolutionPackWinograd when threadNumb…

…er is changed
alibaba · Jul 23, 2024 · 6ba9f9f · 6ba9f9f
1 parent c6f25ca
commit 6ba9f9f
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 13 deletions.
diff --git a/source/backend/cpu/compute/ConvolutionPackWinograd.cpp b/source/backend/cpu/compute/ConvolutionPackWinograd.cpp
@@ -215,6 +215,10 @@ WinogradConfig ConvolutionPackWinograd::bestWinogradUnit(const Convolution2DComm
 
 ErrorCode ConvolutionPackWinograd::onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
     CPUConvolution::onResize(inputs, outputs);
+    int threadNumber = ((CPUBackend*)(backend()))->threadNumber();
+    mTempBuffer->setLength(0, threadNumber);
+    mGemmMidBuffer->setLength(0, threadNumber);
+    mTransformMidBuffer->setLength(0, threadNumber);
     // FUNC_PRINT(mA->length(1));
     bool success = backend()->onAcquireBuffer(mTempBuffer.get(), Backend::DYNAMIC);
     success      = success && backend()->onAcquireBuffer(mGemmMidBuffer.get(), Backend::DYNAMIC);
@@ -256,7 +260,6 @@ ErrorCode ConvolutionPackWinograd::onResize(const std::vector<Tensor *> &inputs,
 
     auto totalCount   = wUnit * hUnit * batch;
     // MNN_PRINT("ow=%d, oh=%d\n", ow, oh);
-    int threadNumber = ((CPUBackend*)(backend()))->threadNumber();
 
     std::vector<int> divides(threadNumber+1);
     static_cast<const CPURuntime*>( static_cast<CPUBackend*>(backend())->getRuntime())->computeDivideSizes(totalCount, divides.data()+1);

diff --git a/source/core/BufferAllocator.cpp b/source/core/BufferAllocator.cpp
@@ -301,6 +301,10 @@ DeferBufferAllocator::DeferBufferAllocator(std::shared_ptr<Allocator> parent, si
 
 //------------------------------- DeferBufferAllocator -----------------------------------//
 MemChunk DeferBufferAllocator::alloc(size_t size, bool separate, size_t align) {
+    if (0 == align) {
+        align = mAlign;
+    }
+    size = UP_DIV(size, align) * align;
     if (mFreeList.empty() || separate) {
         auto newChunk = createMemNode(size);
         insert_after(newChunk);

diff --git a/test/expr/ModuleTest.cpp b/test/expr/ModuleTest.cpp
@@ -482,6 +482,11 @@ class ModuleTestSpeed : public MNNTestCase {
         Module::Config config;
         config.shapeMutable = false;
         config.rearrange = true;
+        auto x = _Input({1, 3, 224, 224}, NC4HW4, halide_type_of<float>());
+        auto xPtr = x->writeMap<float>();
+        ::memset(xPtr, 0, 1*3*224*224*sizeof(float));
+        x->unMap();
+        int runTime = 10;
         std::shared_ptr<Module> interp0;
         {
             MNN::ScheduleConfig sconfig;
@@ -490,6 +495,13 @@ class ModuleTestSpeed : public MNNTestCase {
             std::shared_ptr<Executor::RuntimeManager> rtMgr(Executor::RuntimeManager::createRuntimeManager(sconfigs));
             interp0.reset(Module::load({"Input"}, {"Prob"}, bufferOutput, sizeOutput, rtMgr, &config), Module::destroy);
         }
+        {
+            Timer _l;
+            for (int i=0; i<runTime; ++i) {
+                auto y0 = interp0->onForward({x});
+            }
+            MNN_PRINT("Thread 1 avg cost: %f ms\n", (float)_l.durationInUs() / 1000.0f / runTime);
+        }
         std::shared_ptr<Module> interp1;
         {
             MNN::ScheduleConfig sconfig;
@@ -499,18 +511,6 @@ class ModuleTestSpeed : public MNNTestCase {
             rtMgr->setHint(Interpreter::STRICT_CHECK_MODEL, 0);
             interp1.reset(Module::load({"Input"}, {"Prob"}, bufferOutput, sizeOutput, rtMgr, &config), Module::destroy);
         }
-        auto x = _Input({1, 3, 224, 224}, NC4HW4, halide_type_of<float>());
-        auto xPtr = x->writeMap<float>();
-        ::memset(xPtr, 0, 1*3*224*224*sizeof(float));
-        x->unMap();
-        int runTime = 10;
-        {
-            Timer _l;
-            for (int i=0; i<runTime; ++i) {
-                auto y0 = interp0->onForward({x});
-            }
-            MNN_PRINT("Thread 1 avg cost: %f ms\n", (float)_l.durationInUs() / 1000.0f / runTime);
-        }
         {
             Timer _l;
             for (int i=0; i<runTime; ++i) {