Utilizing cuFFT library on the backend source code


I am trying to editing pytorch backend to test an improved conv algorithm. All my work was on ConvolutionMM2d.cu. cuFFT was used for all FFTs in my case. However, I just find out creating individual plan for each batch extremely hurt the performance. Is there any possibility that I could create a plan for all batches, and only need to transfer the created plan to each batch(or, to the ConvMM2d.cu file)?



Inspired by cufft plan LRU cache in pytorch(source code in CuFFTPlanCache.h, SpectralOps.cpp), I wrote a static harsh map to store the address of plans, and it worked perfectly.