gradient3.cuh 2.97 KB
#ifndef STIM_CUDA_GRADIENT3_H
#define STIM_CUDA_GRADIENT3_H

#include <iostream>
#include <cuda.h>
#include <stim/cuda/cudatools.h>
#include <stim/cuda/cudatools/error.h>

template<typename T>
__global__ void gradient3(T* out, T* in, float anisotropy, int x, int y, int z){

	//calculate x,y,z coordinates for this thread
	int xi = blockIdx.x * blockDim.x + threadIdx.x;
	//find the grid size along y
	int grid_y = y / blockDim.y;
	int blockidx_y = blockIdx.y % grid_y;
	int yi = blockidx_y * blockDim.y + threadIdx.y;
	int zi = blockIdx.y / grid_y;
	int i = zi * x * y + yi * x + xi;
	
	//return if the pixel is outside of the image
	if(xi >= x || yi >= y || zi>=z) return;

	//calculate indices for the forward difference
	int i_xp = zi * x * y + yi * x + (xi + 1);
	int i_yp = zi * x * y + (yi + 1) * x + xi;
	int i_zp = (zi + 1) * x * y + yi * x + xi;

	//calculate indices for the backward difference
	int i_xn = zi * x * y + yi * x + (xi - 1);
	int i_yn = zi * x * y + (yi - 1) * x + xi;
	int i_zn = (zi - 1) * x * y + yi * x + xi;

	//use forward differences if a coordinate is zero
	if(xi == 0)
		out[i * 3 + 0] = in[i_xp] - in[i];
	if(yi == 0)
		out[i * 3 + 1] = in[i_yp] - in[i];
	if (zi==0)
		out[i * 3 + 2] = in[i_zp] - in[i];

	//use backward differences if the coordinate is at the maximum edge
	if(xi == x-1)
		out[i * 3 + 0] = in[i] - in[i_xn];
	if(yi == y-1)
		out[i * 3 + 1] = in[i] - in[i_yn];
	if(zi == z-1)
		out[i * 3 + 2] = in[i] - in[i_zn];

	//otherwise use central differences
	if(xi > 0 && xi < x-1)
		out[i * 3 + 0] = (in[i_xp] - in[i_xn]) / 2;
	if(yi > 0 && yi < y-1)
		out[i * 3 + 1] = (in[i_yp] - in[i_yn]) / 2;
	if(zi > 0 && zi < z-1)
		out[i * 3 + 2] = (in[i_zp] - in[i_zn]) / 2;

	out[i * 3 + 2] *= 1/anisotropy;

}

template<typename T>

void gpu_gradient3(T* gpuGrad, T* gpuI, float anisotropy, unsigned int x, unsigned int y, unsigned int z){

			
	int max_threads = stim::maxThreadsPerBlock();
	dim3 threads(sqrt (max_threads),sqrt (max_threads));
	dim3 blocks(x / threads.x + 1, (y / threads.y + 1) * z);
	
	//call the GPU kernel to determine the gradient
	gradient3<T> <<< blocks, threads >>>(gpuGrad, gpuI, anisotropy, x, y, z);

}

template<typename T>
void cpu_gradient3(T* out, T* in, float anisotropy, unsigned int x, unsigned int y, unsigned int z){

	//get the number of pixels in the image
	unsigned int pixels = x * y * z;
	unsigned int bytes = pixels * sizeof(T);

	//allocate space on the GPU for the input image
	T* gpuIn;
	HANDLE_ERROR(cudaMalloc(&gpuIn, bytes));

	//copy the image data to the GPU
	HANDLE_ERROR(cudaMemcpy(gpuIn, in, bytes, cudaMemcpyHostToDevice));

	//allocate space on the GPU for the output gradient image
	T* gpuOut;
	cudaMalloc(&gpuOut, bytes * 3);		//the output image will have two channels (x, y)

	//call the GPU version of this function
	gpu_gradient3(gpuOut, gpuIn, anisotropy, x, y, z);	

	//copy the results to the CPU
	cudaMemcpy(out, gpuOut, bytes * 3, cudaMemcpyDeviceToHost);

	//free allocated memory
	cudaFree(gpuIn);
	cudaFree(gpuOut);
}



#endif