yum-mirror/VkFFT

git clone https://git.yummers.dev/yum-mirror/VkFFT

Tolmachev DmitriiVkFFT v1.3.3 released6f7ded

develop
branchesdevelopmasterall refs →
1.7 KiB65 linesraw
1//general parts
2#include <stdio.h>
3#include <vector>
4#include <memory>
5#include <string.h>
6#include <chrono>
7#include <thread>
8#include <iostream>
9#ifndef __STDC_FORMAT_MACROS
10#define __STDC_FORMAT_MACROS
11#endif
12#include <inttypes.h>
13
14//CUDA parts
15#include "cuda_runtime.h"
16#include "device_launch_parameters.h"
17#include <cufft.h>
18#include <cuda_fp16.h>
19#include <cufftXt.h>
20
21void launch_precision_cuFFT_half(void* inputC, void* output_cuFFT, int device_id, uint64_t* dims)
22{
23	cudaSetDevice(device_id);
24	cufftHandle planHalf;
25	half2* dataC;
26	cudaMalloc((void**)&dataC, sizeof(half2) * dims[0] * dims[1] * dims[2]);
27	cudaMemcpy(dataC, inputC, sizeof(half2) * dims[0] * dims[1] * dims[2], cudaMemcpyHostToDevice);
28	if (cudaGetLastError() != cudaSuccess) {
29		fprintf(stderr, "Cuda error: Failed to allocate\n");
30		return;
31	}
32	uint64_t sizeCUDA;
33	cufftResult res = cufftCreate(&planHalf);
34	size_t ws = 0;
35	long long local_dims[3];
36	switch (dims[4]) {
37	case 1:
38		local_dims[0] = (long long)dims[0];
39		local_dims[1] = (long long)dims[1];
40		local_dims[2] = (long long)dims[2];
41		break;
42	case 2:
43		local_dims[0] = (long long)dims[1];
44		local_dims[1] = (long long)dims[0];
45		local_dims[2] = (long long)dims[2];
46		break;
47	case 3:
48		local_dims[0] = (long long)dims[2];
49		local_dims[1] = (long long)dims[1];
50		local_dims[2] = (long long)dims[0];
51		break;
52	}
53	res = cufftXtMakePlanMany(
54		planHalf, dims[4], local_dims, NULL, 1, 1, CUDA_C_16F,
55		NULL, 1, 1, CUDA_C_16F, 1, &ws, CUDA_C_16F);
56
57	for (int i = 0; i < 1; i++) {
58		res = cufftXtExec(planHalf, dataC, dataC, -1);
59	}
60	cudaDeviceSynchronize();
61	cudaMemcpy(output_cuFFT, dataC, sizeof(half2) * dims[0] * dims[1] * dims[2], cudaMemcpyDeviceToHost);
62	cudaDeviceSynchronize();
63	cufftDestroy(planHalf);
64	cudaFree(dataC);
65}