diff --git a/include/BM3D.h b/include/BM3D.h index 0c729a9..9bdec6b 100644 --- a/include/BM3D.h +++ b/include/BM3D.h @@ -94,6 +94,8 @@ struct BM3D_FilterData return *this; } + + static std::mutex s_fftw_planner_mutex; }; diff --git a/include/BM3D_Base.h b/include/BM3D_Base.h index 219d1c2..b8b741c 100644 --- a/include/BM3D_Base.h +++ b/include/BM3D_Base.h @@ -26,6 +26,7 @@ #define BM3D_BASE_H_ +#include #include #include #include "BM3D.h" @@ -54,6 +55,7 @@ class BM3D_Data_Base _Mypara para; std::vector f; + std::mutex mutex0, mutex1, mutex2; std::unordered_map buffer0, buffer1, buffer2; public: @@ -72,17 +74,26 @@ class BM3D_Data_Base { if (rdef && rnode) vsapi->freeNode(rnode); - for (auto &e : buffer0) { - AlignedFree(e.second); + std::lock_guard guard0(mutex0); + for (auto &e : buffer0) + { + AlignedFree(e.second); + } } - for (auto &e : buffer1) { - AlignedFree(e.second); + std::lock_guard guard1(mutex1); + for (auto &e : buffer1) + { + AlignedFree(e.second); + } } - for (auto &e : buffer2) { - AlignedFree(e.second); + std::lock_guard guard2(mutex2); + for (auto &e : buffer2) + { + AlignedFree(e.second); + } } } diff --git a/source/BM3D.cpp b/source/BM3D.cpp index 676e816..f7e757e 100644 --- a/source/BM3D.cpp +++ b/source/BM3D.cpp @@ -132,6 +132,7 @@ void BM3D_Para::thMSE_Default() //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // Functions of struct BM3D_FilterData +std::mutex BM3D_FilterData::s_fftw_planner_mutex; BM3D_FilterData::BM3D_FilterData(bool wiener, double sigma, PCType GroupSize, PCType BlockSize, double lambda) : fp(GroupSize), bp(GroupSize), finalAMP(GroupSize), thrTable(wiener ? 0 : GroupSize), @@ -143,6 +144,9 @@ BM3D_FilterData::BM3D_FilterData(bool wiener, double sigma, PCType GroupSize, PC FLType *temp = nullptr; + // Executing a plan is thread-safe, but creating a plan is not because it mutates the planner's "wisdom". + // Consider that any number of instances of the (V)BM3D filter may exist in a single VapourSynth pipeline. + std::lock_guard guard(s_fftw_planner_mutex); for (PCType i = 1; i <= GroupSize; ++i) { AlignedMalloc(temp, i * BlockSize * BlockSize); diff --git a/source/BM3D_Base.cpp b/source/BM3D_Base.cpp index 1513bdd..cb3be0e 100644 --- a/source/BM3D_Base.cpp +++ b/source/BM3D_Base.cpp @@ -296,14 +296,17 @@ void BM3D_Process_Base::Kernel(FLType *dst, const FLType *src, const FLType *ref std::thread::id threadId = std::this_thread::get_id(); FLType *ResNum = dst, *ResDen = nullptr; - if (!d.buffer0.count(threadId)) { - AlignedMalloc(ResDen, dst_pcount[0]); - d.buffer0.emplace(threadId, ResDen); - } - else - { - ResDen = d.buffer0.at(threadId); + std::lock_guard guard(d.mutex0); + if (!d.buffer0.count(threadId)) + { + AlignedMalloc(ResDen, dst_pcount[0]); + d.buffer0.emplace(threadId, ResDen); + } + else + { + ResDen = d.buffer0.at(threadId); + } } memset(ResNum, 0, sizeof(FLType) * dst_pcount[0]); @@ -363,14 +366,17 @@ void BM3D_Process_Base::Kernel(FLType *dstY, FLType *dstU, FLType *dstV, if (d.process[0]) { - if (!d.buffer0.count(threadId)) - { - AlignedMalloc(ResDenY, dst_pcount[0]); - d.buffer0.emplace(threadId, ResDenY); - } - else { - ResDenY = d.buffer0.at(threadId); + std::lock_guard guard(d.mutex0); + if (!d.buffer0.count(threadId)) + { + AlignedMalloc(ResDenY, dst_pcount[0]); + d.buffer0.emplace(threadId, ResDenY); + } + else + { + ResDenY = d.buffer0.at(threadId); + } } memset(ResNumY, 0, sizeof(FLType) * dst_pcount[0]); @@ -379,30 +385,36 @@ void BM3D_Process_Base::Kernel(FLType *dstY, FLType *dstU, FLType *dstV, if (d.process[1]) { - if (!d.buffer1.count(threadId)) - { - AlignedMalloc(ResDenU, dst_pcount[1]); - d.buffer1.emplace(threadId, ResDenU); - } - else { - ResDenU = d.buffer1.at(threadId); + std::lock_guard guard(d.mutex1); + if (!d.buffer1.count(threadId)) + { + AlignedMalloc(ResDenU, dst_pcount[1]); + d.buffer1.emplace(threadId, ResDenU); + } + else + { + ResDenU = d.buffer1.at(threadId); + } } - + memset(ResNumU, 0, sizeof(FLType) * dst_pcount[1]); memset(ResDenU, 0, sizeof(FLType) * dst_pcount[1]); } if (d.process[2]) { - if (!d.buffer2.count(threadId)) - { - AlignedMalloc(ResDenV, dst_pcount[2]); - d.buffer2.emplace(threadId, ResDenV); - } - else { - ResDenV = d.buffer2.at(threadId); + std::lock_guard guard(d.mutex2); + if (!d.buffer2.count(threadId)) + { + AlignedMalloc(ResDenV, dst_pcount[2]); + d.buffer2.emplace(threadId, ResDenV); + } + else + { + ResDenV = d.buffer2.at(threadId); + } } memset(ResNumV, 0, sizeof(FLType) * dst_pcount[2]);