diff --git a/tomobar/cuda_kernels/primal_dual_for_total_variation.cu b/tomobar/cuda_kernels/primal_dual_for_total_variation.cu index 2e286ab63..a1a4d5ec9 100644 --- a/tomobar/cuda_kernels/primal_dual_for_total_variation.cu +++ b/tomobar/cuda_kernels/primal_dual_for_total_variation.cu @@ -122,8 +122,8 @@ __device__ float DivProj3D(float Input, float U_in, float P1, float P2, float P3 return (U_in - tau * div_var + lt * Input) / (1.0f + lt); } -template -__device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *Input, float *U_in, float *U_out, T *P1_in, T *P2_in, T *P3_in, T *P1_out, T *P2_out, T *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +template +__device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *Input, InT *U_in, OutT *U_out, T *P1_in, T *P2_in, T *P3_in, T *P1_out, T *P2_out, T *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) { // calculate each thread global index const long xIndex = blockIdx.x * blockDim.x + threadIdx.x; @@ -166,7 +166,7 @@ __device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *I float P1 = read_as_float(P1_in, index); float P2 = read_as_float(P2_in, index); float P3 = read_as_float(P3_in, index); - float U = U_in[index]; + float U = read_as_float(U_in, index); float Input_value = Input[index]; if (xIndex > 0) @@ -174,7 +174,7 @@ __device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *I P1_prev_x = read_as_float(P1_in, index_prev_x); P2_prev_x = read_as_float(P2_in, index_prev_x); P3_prev_x = read_as_float(P3_in, index_prev_x); - U_prev_x = U_in[index_prev_x]; + U_prev_x = read_as_float(U_in, index_prev_x); } if (yIndex > 0) @@ -182,7 +182,7 @@ __device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *I P1_prev_y = read_as_float(P1_in, index_prev_y); P2_prev_y = read_as_float(P2_in, index_prev_y); P3_prev_y = read_as_float(P3_in, index_prev_y); - U_prev_y = U_in[index_prev_y]; + U_prev_y = read_as_float(U_in, index_prev_y); } if (zIndex > 0) @@ -190,7 +190,7 @@ __device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *I P1_prev_z = read_as_float(P1_in, index_prev_z); P2_prev_z = read_as_float(P2_in, index_prev_z); P3_prev_z = read_as_float(P3_in, index_prev_z); - U_prev_z = U_in[index_prev_z]; + U_prev_z = read_as_float(U_in, index_prev_z); } bool last_x = xIndex == dimX - 1; @@ -199,25 +199,25 @@ __device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *I if (((xIndex > 0) && last_y) || ((yIndex > 0) && last_x)) { - U_prev_x_prev_y = U_in[index - xStride - yStride]; + U_prev_x_prev_y = read_as_float(U_in, index - xStride - yStride); } if (((xIndex > 0) && last_z) || ((zIndex > 0) && last_x)) { - U_prev_x_prev_z = U_in[index - xStride - zStride]; + U_prev_x_prev_z = read_as_float(U_in, index - xStride - zStride); } if (((yIndex > 0) && last_z) || ((zIndex > 0) && last_y)) { - U_prev_y_prev_z = U_in[index - yStride - zStride]; + U_prev_y_prev_z = read_as_float(U_in, index - yStride - zStride); } { float U_values[4] = { U, - last_x ? U_prev_x : U_in[index + xStride], - last_y ? U_prev_y : U_in[index + yStride], - last_z ? U_prev_z : U_in[index + zStride]}; + last_x ? U_prev_x : read_as_float(U_in, index + xStride), + last_y ? U_prev_y : read_as_float(U_in, index + yStride), + last_z ? U_prev_z : read_as_float(U_in, index + zStride)}; dualPD3D(U_values, &P1, &P2, &P3, sigma); } @@ -226,8 +226,8 @@ __device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *I float U_values[4] = { U_prev_x, U, - last_y ? U_prev_x_prev_y : U_in[index - xStride + yStride], - last_z ? U_prev_x_prev_z : U_in[index - xStride + zStride]}; + last_y ? U_prev_x_prev_y : read_as_float(U_in, index - xStride + yStride), + last_z ? U_prev_x_prev_z : read_as_float(U_in, index - xStride + zStride)}; dualPD3D(U_values, &P1_prev_x, &P2_prev_x, &P3_prev_x, sigma); } @@ -235,9 +235,9 @@ __device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *I { float U_values[4] = { U_prev_y, - last_x ? U_prev_x_prev_y : U_in[index + xStride - yStride], + last_x ? U_prev_x_prev_y : read_as_float(U_in, index + xStride - yStride), U, - last_z ? U_prev_y_prev_z : U_in[index - yStride + zStride]}; + last_z ? U_prev_y_prev_z : read_as_float(U_in, index - yStride + zStride)}; dualPD3D(U_values, &P1_prev_y, &P2_prev_y, &P3_prev_y, sigma); } @@ -245,59 +245,103 @@ __device__ __forceinline__ void primal_dual_for_total_variation_3D_impl(float *I { float U_values[4] = { U_prev_z, - last_x ? U_prev_x_prev_z : U_in[index + xStride - zStride], - last_y ? U_prev_y_prev_z : U_in[index + yStride - zStride], + last_x ? U_prev_x_prev_z : read_as_float(U_in, index + xStride - zStride), + last_y ? U_prev_y_prev_z : read_as_float(U_in, index + yStride - zStride), U}; dualPD3D(U_values, &P1_prev_z, &P2_prev_z, &P3_prev_z, sigma); } U = clamp_to_zero(U); float new_U = DivProj3D(Input_value, U, P1, P2, P3, P1_prev_x, P2_prev_y, P3_prev_z, tau, lt); - U_out[index] = new_U + theta * (new_U - U); + write_float(U_out, index, new_U + theta * (new_U - U)); write_float(P1_out, index, P1); write_float(P2_out, index, P2); write_float(P3_out, index, P3); } -extern "C" __global__ void primal_dual_for_total_variation_3D_half(float *Input, float *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +/*** half, half, half combinations ***/ +extern "C" __global__ void primal_dual_for_total_variation_3D_half_half_half(float *Input, __half *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) { - primal_dual_for_total_variation_3D_impl<__half, false, false>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); + primal_dual_for_total_variation_3D_impl<__half, __half, __half, false, false>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); } -extern "C" __global__ void primal_dual_for_total_variation_3D_half_nonneg(float *Input, float *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +extern "C" __global__ void primal_dual_for_total_variation_3D_half_half_half_nonneg(float *Input, __half *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) { - primal_dual_for_total_variation_3D_impl<__half, true, false>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); + primal_dual_for_total_variation_3D_impl<__half, __half, __half, true, false>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); } -extern "C" __global__ void primal_dual_for_total_variation_3D_half_methodTV(float *Input, float *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +extern "C" __global__ void primal_dual_for_total_variation_3D_half_half_half_methodTV(float *Input, __half *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) { - primal_dual_for_total_variation_3D_impl<__half, false, true>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); + primal_dual_for_total_variation_3D_impl<__half, __half, __half, false, true>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); } -extern "C" __global__ void primal_dual_for_total_variation_3D_half_nonneg_methodTV(float *Input, float *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +extern "C" __global__ void primal_dual_for_total_variation_3D_half_half_half_nonneg_methodTV(float *Input, __half *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) { - primal_dual_for_total_variation_3D_impl<__half, true, true>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); + primal_dual_for_total_variation_3D_impl<__half, __half, __half, true, true>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); } -extern "C" __global__ void primal_dual_for_total_variation_3D_float(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P3_in, float *P1_out, float *P2_out, float *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +/*** half, float, half combinations ***/ +extern "C" __global__ void primal_dual_for_total_variation_3D_half_float_half(float *Input, float *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) { - primal_dual_for_total_variation_3D_impl(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); + primal_dual_for_total_variation_3D_impl<__half, float, __half, false, false>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); } -extern "C" __global__ void primal_dual_for_total_variation_3D_float_nonneg(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P3_in, float *P1_out, float *P2_out, float *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +extern "C" __global__ void primal_dual_for_total_variation_3D_half_float_half_nonneg(float *Input, float *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) { - primal_dual_for_total_variation_3D_impl(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); + primal_dual_for_total_variation_3D_impl<__half, float, __half, true, false>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); } -extern "C" __global__ void primal_dual_for_total_variation_3D_float_methodTV(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P3_in, float *P1_out, float *P2_out, float *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +extern "C" __global__ void primal_dual_for_total_variation_3D_half_float_half_methodTV(float *Input, float *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) { - primal_dual_for_total_variation_3D_impl(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); + primal_dual_for_total_variation_3D_impl<__half, float, __half, false, true>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); } -extern "C" __global__ void primal_dual_for_total_variation_3D_float_nonneg_methodTV(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P3_in, float *P1_out, float *P2_out, float *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +extern "C" __global__ void primal_dual_for_total_variation_3D_half_float_half_nonneg_methodTV(float *Input, float *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) { - primal_dual_for_total_variation_3D_impl(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); + primal_dual_for_total_variation_3D_impl<__half, float, __half, true, true>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); +} + +/*** half, half, float combinations ***/ +extern "C" __global__ void primal_dual_for_total_variation_3D_half_half_float(float *Input, __half *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +{ + primal_dual_for_total_variation_3D_impl<__half, __half, float, false, false>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); +} + +extern "C" __global__ void primal_dual_for_total_variation_3D_half_half_float_nonneg(float *Input, __half *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +{ + primal_dual_for_total_variation_3D_impl<__half, __half, float, true, false>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); +} + +extern "C" __global__ void primal_dual_for_total_variation_3D_half_half_float_methodTV(float *Input, __half *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +{ + primal_dual_for_total_variation_3D_impl<__half, __half, float, false, true>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); +} + +extern "C" __global__ void primal_dual_for_total_variation_3D_half_half_float_nonneg_methodTV(float *Input, __half *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P3_in, __half *P1_out, __half *P2_out, __half *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +{ + primal_dual_for_total_variation_3D_impl<__half, __half, float, true, true>(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); +} + +/*** float, float, float combinations ***/ +extern "C" __global__ void primal_dual_for_total_variation_3D_float_float_float(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P3_in, float *P1_out, float *P2_out, float *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +{ + primal_dual_for_total_variation_3D_impl(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); +} + +extern "C" __global__ void primal_dual_for_total_variation_3D_float_float_float_nonneg(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P3_in, float *P1_out, float *P2_out, float *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +{ + primal_dual_for_total_variation_3D_impl(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); +} + +extern "C" __global__ void primal_dual_for_total_variation_3D_float_float_float_methodTV(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P3_in, float *P1_out, float *P2_out, float *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +{ + primal_dual_for_total_variation_3D_impl(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); +} + +extern "C" __global__ void primal_dual_for_total_variation_3D_float_float_float_nonneg_methodTV(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P3_in, float *P1_out, float *P2_out, float *P3_out, float sigma, float tau, float lt, float theta, int dimX, int dimY, int dimZ) +{ + primal_dual_for_total_variation_3D_impl(Input, U_in, U_out, P1_in, P2_in, P3_in, P1_out, P2_out, P3_out, sigma, tau, lt, theta, dimX, dimY, dimZ); } /************************************************/ @@ -357,8 +401,8 @@ __device__ float DivProj2D(float Input, float U_in, float P1, float P2, float P1 return (U_in - tau * div_var + lt * Input) / (1.0f + lt); } -template -__device__ __forceinline__ void primal_dual_for_total_variation_2D_impl(float *Input, float *U_in, float *U_out, T *P1_in, T *P2_in, T *P1_out, T *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +template +__device__ __forceinline__ void primal_dual_for_total_variation_2D_impl(float *Input, InT *U_in, OutT *U_out, T *P1_in, T *P2_in, T *P1_out, T *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) { // calculate each thread global index const long xIndex = blockIdx.x * blockDim.x + threadIdx.x; @@ -388,21 +432,21 @@ __device__ __forceinline__ void primal_dual_for_total_variation_2D_impl(float *I float P1 = read_as_float(P1_in, index); float P2 = read_as_float(P2_in, index); - float U = U_in[index]; + float U = read_as_float(U_in, index); float Input_value = Input[index]; if (xIndex > 0) { P1_prev_x = read_as_float(P1_in, index_prev_x); P2_prev_x = read_as_float(P2_in, index_prev_x); - U_prev_x = U_in[index_prev_x]; + U_prev_x = read_as_float(U_in, index_prev_x); } if (yIndex > 0) { P1_prev_y = read_as_float(P1_in, index_prev_y); P2_prev_y = read_as_float(P2_in, index_prev_y); - U_prev_y = U_in[index_prev_y]; + U_prev_y = read_as_float(U_in, index_prev_y); } bool last_x = xIndex == dimX - 1; @@ -410,14 +454,14 @@ __device__ __forceinline__ void primal_dual_for_total_variation_2D_impl(float *I if (((xIndex > 0) && last_y) || ((yIndex > 0) && last_x)) { - U_prev_x_prev_y = U_in[index - xStride - yStride]; + U_prev_x_prev_y = read_as_float(U_in, index - xStride - yStride); } { float U_values[3] = { U, - last_x ? U_prev_x : U_in[index + xStride], - last_y ? U_prev_y : U_in[index + yStride], + last_x ? U_prev_x : read_as_float(U_in, index + xStride), + last_y ? U_prev_y : read_as_float(U_in, index + yStride), }; dualPD2D(U_values, &P1, &P2, sigma); @@ -428,7 +472,7 @@ __device__ __forceinline__ void primal_dual_for_total_variation_2D_impl(float *I float U_values[3] = { U_prev_x, U, - last_y ? U_prev_x_prev_y : U_in[index - xStride + yStride], + last_y ? U_prev_x_prev_y : read_as_float(U_in, index - xStride + yStride), }; dualPD2D(U_values, &P1_prev_x, &P2_prev_x, sigma); } @@ -437,7 +481,7 @@ __device__ __forceinline__ void primal_dual_for_total_variation_2D_impl(float *I { float U_values[3] = { U_prev_y, - last_x ? U_prev_x_prev_y : U_in[index + xStride - yStride], + last_x ? U_prev_x_prev_y : read_as_float(U_in, index + xStride - yStride), U, }; dualPD2D(U_values, &P1_prev_y, &P2_prev_y, sigma); @@ -445,48 +489,92 @@ __device__ __forceinline__ void primal_dual_for_total_variation_2D_impl(float *I U = clamp_to_zero(U); float new_U = DivProj2D(Input_value, U, P1, P2, P1_prev_x, P2_prev_y, tau, lt); - U_out[index] = new_U + theta * (new_U - U); + write_float(U_out, index, new_U + theta * (new_U - U)); write_float(P1_out, index, P1); write_float(P2_out, index, P2); } -extern "C" __global__ void primal_dual_for_total_variation_2D_half(float *Input, float *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +/*** half, half, half combinations ***/ +extern "C" __global__ void primal_dual_for_total_variation_2D_half_half_half(float *Input, __half *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +{ + primal_dual_for_total_variation_2D_impl<__half, __half, __half, false, false>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); +} + +extern "C" __global__ void primal_dual_for_total_variation_2D_half_half_half_nonneg(float *Input, __half *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +{ + primal_dual_for_total_variation_2D_impl<__half, __half, __half, true, false>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); +} + +extern "C" __global__ void primal_dual_for_total_variation_2D_half_half_half_methodTV(float *Input, __half *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +{ + primal_dual_for_total_variation_2D_impl<__half, __half, __half, false, true>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); +} + +extern "C" __global__ void primal_dual_for_total_variation_2D_half_half_half_nonneg_methodTV(float *Input, __half *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +{ + primal_dual_for_total_variation_2D_impl<__half, __half, __half, true, true>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); +} + +/*** half, float, half combinations ***/ +extern "C" __global__ void primal_dual_for_total_variation_2D_half_float_half(float *Input, float *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +{ + primal_dual_for_total_variation_2D_impl<__half, float, __half, false, false>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); +} + +extern "C" __global__ void primal_dual_for_total_variation_2D_half_float_half_nonneg(float *Input, float *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +{ + primal_dual_for_total_variation_2D_impl<__half, float, __half, true, false>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); +} + +extern "C" __global__ void primal_dual_for_total_variation_2D_half_float_half_methodTV(float *Input, float *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +{ + primal_dual_for_total_variation_2D_impl<__half, float, __half, false, true>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); +} + +extern "C" __global__ void primal_dual_for_total_variation_2D_half_float_half_nonneg_methodTV(float *Input, float *U_in, __half *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +{ + primal_dual_for_total_variation_2D_impl<__half, float, __half, true, true>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); +} + +/*** half, half, float combinations ***/ +extern "C" __global__ void primal_dual_for_total_variation_2D_half_half_float(float *Input, __half *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) { - primal_dual_for_total_variation_2D_impl<__half, false, false>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); + primal_dual_for_total_variation_2D_impl<__half, __half, float, false, false>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); } -extern "C" __global__ void primal_dual_for_total_variation_2D_half_nonneg(float *Input, float *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +extern "C" __global__ void primal_dual_for_total_variation_2D_half_half_float_nonneg(float *Input, __half *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) { - primal_dual_for_total_variation_2D_impl<__half, true, false>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); + primal_dual_for_total_variation_2D_impl<__half, __half, float, true, false>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); } -extern "C" __global__ void primal_dual_for_total_variation_2D_half_methodTV(float *Input, float *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +extern "C" __global__ void primal_dual_for_total_variation_2D_half_half_float_methodTV(float *Input, __half *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) { - primal_dual_for_total_variation_2D_impl<__half, false, true>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); + primal_dual_for_total_variation_2D_impl<__half, __half, float, false, true>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); } -extern "C" __global__ void primal_dual_for_total_variation_2D_half_nonneg_methodTV(float *Input, float *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +extern "C" __global__ void primal_dual_for_total_variation_2D_half_half_float_nonneg_methodTV(float *Input, __half *U_in, float *U_out, __half *P1_in, __half *P2_in, __half *P1_out, __half *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) { - primal_dual_for_total_variation_2D_impl<__half, true, true>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); + primal_dual_for_total_variation_2D_impl<__half, __half, float, true, true>(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); } -extern "C" __global__ void primal_dual_for_total_variation_2D_float(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P1_out, float *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +/*** float, float, float combinations ***/ +extern "C" __global__ void primal_dual_for_total_variation_2D_float_float_float(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P1_out, float *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) { - primal_dual_for_total_variation_2D_impl(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); + primal_dual_for_total_variation_2D_impl(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); } -extern "C" __global__ void primal_dual_for_total_variation_2D_float_nonneg(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P1_out, float *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +extern "C" __global__ void primal_dual_for_total_variation_2D_float_float_float_nonneg(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P1_out, float *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) { - primal_dual_for_total_variation_2D_impl(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); + primal_dual_for_total_variation_2D_impl(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); } -extern "C" __global__ void primal_dual_for_total_variation_2D_float_methodTV(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P1_out, float *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +extern "C" __global__ void primal_dual_for_total_variation_2D_float_float_float_methodTV(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P1_out, float *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) { - primal_dual_for_total_variation_2D_impl(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); + primal_dual_for_total_variation_2D_impl(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); } -extern "C" __global__ void primal_dual_for_total_variation_2D_float_nonneg_methodTV(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P1_out, float *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) +extern "C" __global__ void primal_dual_for_total_variation_2D_float_float_float_nonneg_methodTV(float *Input, float *U_in, float *U_out, float *P1_in, float *P2_in, float *P1_out, float *P2_out, float sigma, float tau, float lt, float theta, int dimX, int dimY) { - primal_dual_for_total_variation_2D_impl(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); + primal_dual_for_total_variation_2D_impl(Input, U_in, U_out, P1_in, P2_in, P1_out, P2_out, sigma, tau, lt, theta, dimX, dimY); } diff --git a/tomobar/regularisersCuPy.py b/tomobar/regularisersCuPy.py index 55a0d9d86..ca59bcf52 100644 --- a/tomobar/regularisersCuPy.py +++ b/tomobar/regularisersCuPy.py @@ -218,19 +218,34 @@ def PD_TV_cupy( lt = cp.float32(tau / regularisation_parameter) # initialise CuPy arrays here: - U_arrays = [data.copy(), cp.zeros(data.shape, dtype=cp.float32, order="C")] + U_arrays = [ + cp.copy(data), + cp.empty_like(data), + ] P1_arrays = [cp.zeros(data.shape, dtype=dtype_of_P, order="C") for _ in range(2)] P2_arrays = [cp.zeros(data.shape, dtype=dtype_of_P, order="C") for _ in range(2)] # loading and compiling CUDA kernels: - kernel_name = f"primal_dual_for_total_variation_{'3D' if data.ndim == 3 else '2D'}_{'half' if half_precision else 'float'}" - if nonneg: - kernel_name += "_nonneg" - if methodTV: - kernel_name += "_methodTV" + intermediate_typename = "half" if half_precision else "float" + kernel_name_base = ( + f"primal_dual_for_total_variation_{'3D' if data.ndim == 3 else '2D'}_" + ) + kernel_names = [ + kernel_name_base + f"{intermediate_typename}_float_{intermediate_typename}", + kernel_name_base + + f"{intermediate_typename}_{intermediate_typename}_{intermediate_typename}", + kernel_name_base + f"{intermediate_typename}_{intermediate_typename}_float", + ] + for kernel_name in kernel_names: + if nonneg: + kernel_name += "_nonneg" + if methodTV: + kernel_name += "_methodTV" module = load_cuda_module("primal_dual_for_total_variation") - primal_dual_for_total_variation = module.get_function(kernel_name) + first_kernel, middle_kernel, last_kernel = [ + module.get_function(kernel_name) for kernel_name in kernel_names + ] dz, dy, dx = (0,) * (3 - data.ndim) + data.shape block_x = 128 @@ -252,7 +267,7 @@ def PD_TV_cupy( input_index = 0 output_index = 1 - for _ in range(iterations): + for i in range(iterations): if input_is_2d: params = ( data, @@ -285,15 +300,23 @@ def PD_TV_cupy( theta, *data_dims, ) + if i == 0: + current_kernel = first_kernel + elif i == iterations - 1: + current_kernel = last_kernel + else: + current_kernel = middle_kernel - primal_dual_for_total_variation(grid_dims, block_dims, params) + current_kernel(grid_dims, block_dims, params) input_index = 1 - input_index output_index = 1 - output_index if input_is_2d: - return cp.expand_dims(U_arrays[input_index], axis=ind_axis) + return cp.expand_dims( + cp.asarray(U_arrays[input_index], dtype=cp.float32), axis=ind_axis + ) else: - return U_arrays[input_index] + return cp.asarray(U_arrays[input_index], dtype=cp.float32) def __check_if_input_2d_or_3d(data: cp.ndarray) -> Tuple[cp.ndarray, bool, int]: