def buildCache(self, input, type): k = input.ndimension() - 3 n = np.asarray([input.size(k), input.size(k+1)], np.int32) batch = input.nelement() // (2*input.size(k) * input.size(k + 1)) idist = input.size(k) * input.size(k + 1) istride = 1 ostride = istride odist = idist rank = 2 plan = cufft.cufftPlanMany(rank, n.ctypes.data, n.ctypes.data, istride, idist, n.ctypes.data, ostride, odist, type, batch) self.fft_cache[(input.size(), type, input.get_device())] = plan
def buildCufftCache(self, input, type): batch_size, M, N, O, _ = input.size() signal_dims = np.asarray([M, N, O], np.int32) batch = batch_size idist = M * N * O istride = 1 ostride = istride odist = idist rank = 3 print(rank, signal_dims.ctypes.data, signal_dims.ctypes.data, istride, idist, signal_dims.ctypes.data, ostride, odist, type, batch) plan = cufft.cufftPlanMany(rank, signal_dims.ctypes.data, signal_dims.ctypes.data, istride, idist, signal_dims.ctypes.data, ostride, odist, type, batch) self.cufft_cache[(input.size(), type, input.get_device())] = plan