/* -------------------------------------------------------------------------- * * OpenMM * * -------------------------------------------------------------------------- * * This is part of the OpenMM molecular simulation toolkit originating from * * Simbios, the NIH National Center for Physics-Based Simulation of * * Biological Structures at Stanford, funded under the NIH Roadmap for * * Medical Research, grant U54 GM072970. See https://simtk.org. * * * * Portions copyright (c) 2009-2015 Stanford University and the Authors. * * Authors: Peter Eastman * * Contributors: * * * * This program is free software: you can redistribute it and/or modify * * it under the terms of the GNU Lesser General Public License as published * * by the Free Software Foundation, either version 3 of the License, or * * (at your option) any later version. * * * * This program is distributed in the hope that it will be useful, * * but WITHOUT ANY WARRANTY; without even the implied warranty of * * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the * * GNU Lesser General Public License for more details. * * * * You should have received a copy of the GNU Lesser General Public License * * along with this program. If not, see . * * -------------------------------------------------------------------------- */ #include "CudaFFT3D.h" #include "CudaContext.h" #include "CudaKernelSources.h" #include "SimTKOpenMMRealType.h" #include #include #include using namespace OpenMM; using namespace std; CudaFFT3D::CudaFFT3D(CudaContext& context, int xsize, int ysize, int zsize, bool realToComplex) : context(context), xsize(xsize), ysize(ysize), zsize(zsize) { packRealAsComplex = false; int packedXSize = xsize; int packedYSize = ysize; int packedZSize = zsize; if (realToComplex) { // If any axis size is even, we can pack the real values into a complex grid that is only half as large. // Look for an appropriate axis. packRealAsComplex = true; int packedAxis, bufferSize; if (xsize%2 == 0) { packedAxis = 0; packedXSize /= 2; bufferSize = packedXSize; } else if (ysize%2 == 0) { packedAxis = 1; packedYSize /= 2; bufferSize = packedYSize; } else if (zsize%2 == 0) { packedAxis = 2; packedZSize /= 2; bufferSize = packedZSize; } else packRealAsComplex = false; if (packRealAsComplex) { // Build the kernels for packing and unpacking the data. map defines; defines["XSIZE"] = context.intToString(xsize); defines["YSIZE"] = context.intToString(ysize); defines["ZSIZE"] = context.intToString(zsize); defines["PACKED_AXIS"] = context.intToString(packedAxis); defines["PACKED_XSIZE"] = context.intToString(packedXSize); defines["PACKED_YSIZE"] = context.intToString(packedYSize); defines["PACKED_ZSIZE"] = context.intToString(packedZSize); CUmodule module = context.createModule(CudaKernelSources::vectorOps+CudaKernelSources::fftR2C, defines); packForwardKernel = context.getKernel(module, "packForwardData"); unpackForwardKernel = context.getKernel(module, "unpackForwardData"); packBackwardKernel = context.getKernel(module, "packBackwardData"); unpackBackwardKernel = context.getKernel(module, "unpackBackwardData"); } } bool inputIsReal = (realToComplex && !packRealAsComplex); zkernel = createKernel(packedXSize, packedYSize, packedZSize, zthreads, 0, true, inputIsReal); xkernel = createKernel(packedYSize, packedZSize, packedXSize, xthreads, 1, true, inputIsReal); ykernel = createKernel(packedZSize, packedXSize, packedYSize, ythreads, 2, true, inputIsReal); invzkernel = createKernel(packedXSize, packedYSize, packedZSize, zthreads, 0, false, inputIsReal); invxkernel = createKernel(packedYSize, packedZSize, packedXSize, xthreads, 1, false, inputIsReal); invykernel = createKernel(packedZSize, packedXSize, packedYSize, ythreads, 2, false, inputIsReal); } void CudaFFT3D::execFFT(CudaArray& in, CudaArray& out, bool forward) { CUfunction kernel1 = (forward ? zkernel : invzkernel); CUfunction kernel2 = (forward ? xkernel : invxkernel); CUfunction kernel3 = (forward ? ykernel : invykernel); void* args1[] = {&in.getDevicePointer(), &out.getDevicePointer()}; void* args2[] = {&out.getDevicePointer(), &in.getDevicePointer()}; if (packRealAsComplex) { CUfunction packKernel = (forward ? packForwardKernel : packBackwardKernel); CUfunction unpackKernel = (forward ? unpackForwardKernel : unpackBackwardKernel); int gridSize = xsize*ysize*zsize/2; // Pack the data into a half sized grid. context.executeKernel(packKernel, args1, gridSize, 128); // Perform the FFT. context.executeKernel(kernel1, args2, gridSize, zthreads); context.executeKernel(kernel2, args1, gridSize, xthreads); context.executeKernel(kernel3, args2, gridSize, ythreads); // Unpack the data. context.executeKernel(unpackKernel, args1, gridSize, 128); } else { context.executeKernel(kernel1, args1, xsize*ysize*zsize, zthreads); context.executeKernel(kernel2, args2, xsize*ysize*zsize, xthreads); context.executeKernel(kernel3, args1, xsize*ysize*zsize, ythreads); } } int CudaFFT3D::findLegalDimension(int minimum) { if (minimum < 1) return 1; while (true) { // Attempt to factor the current value. int unfactored = minimum; for (int factor = 2; factor < 8; factor++) { while (unfactored > 1 && unfactored%factor == 0) unfactored /= factor; } if (unfactored == 1) return minimum; minimum++; } } static int getSmallestRadix(int size) { int minRadix = 1; int unfactored = size; while (unfactored%7 == 0) { minRadix = 7; unfactored /= 7; } while (unfactored%5 == 0) { minRadix = 5; unfactored /= 5; } while (unfactored%4 == 0) { minRadix = 4; unfactored /= 4; } while (unfactored%3 == 0) { minRadix = 3; unfactored /= 3; } while (unfactored%2 == 0) { minRadix = 2; unfactored /= 2; } return minRadix; } CUfunction CudaFFT3D::createKernel(int xsize, int ysize, int zsize, int& threads, int axis, bool forward, bool inputIsReal) { int maxThreads = 256; // while (maxThreads > 128 && maxThreads-64 >= zsize) // maxThreads -= 64; int threadsPerBlock = zsize/getSmallestRadix(zsize); stringstream source; int blocksPerGroup = max(1, maxThreads/threadsPerBlock); int stage = 0; int L = zsize; int m = 1; // Factor zsize, generating an appropriate block of code for each factor. while (L > 1) { int input = stage%2; int output = 1-input; int radix; if (L%7 == 0) radix = 7; else if (L%5 == 0) radix = 5; else if (L%4 == 0) radix = 4; else if (L%3 == 0) radix = 3; else if (L%2 == 0) radix = 2; else throw OpenMMException("Illegal size for FFT: "+context.intToString(zsize)); source<<"{\n"; L = L/radix; source<<"// Pass "<<(stage+1)<<" (radix "< replacements; replacements["XSIZE"] = context.intToString(xsize); replacements["YSIZE"] = context.intToString(ysize); replacements["ZSIZE"] = context.intToString(zsize); replacements["BLOCKS_PER_GROUP"] = context.intToString(blocksPerGroup); replacements["THREADS_PER_BLOCK"] = context.intToString(threadsPerBlock); replacements["M_PI"] = context.doubleToString(M_PI); replacements["COMPUTE_FFT"] = source.str(); replacements["SIGN"] = (forward ? "1" : "-1"); replacements["INPUT_TYPE"] = (inputIsReal && axis == 0 && forward ? "real" : "real2"); replacements["OUTPUT_TYPE"] = (outputIsReal ? "real" : "real2"); replacements["INPUT_IS_REAL"] = (inputIsReal && axis == 0 && forward ? "1" : "0"); replacements["INPUT_IS_PACKED"] = (inputIsReal && axis == 0 && !forward ? "1" : "0"); replacements["OUTPUT_IS_PACKED"] = (outputIsPacked ? "1" : "0"); CUmodule module = context.createModule(CudaKernelSources::vectorOps+context.replaceStrings(CudaKernelSources::fft, replacements)); CUfunction kernel = context.getKernel(module, "execFFT"); threads = blocksPerGroup*threadsPerBlock; return kernel; }