nabu 2023.2.1__py3-none-any.whl → 2024.1.0rc3__py3-none-any.whl

nabu/cuda/processing.py

@@ -1,16 +1,24 @@
-from ..misc.processing_base import ProcessingBase
+from ..utils import MissingComponentError
+from ..processing.processing_base import ProcessingBase
 from .utils import get_cuda_context, __has_pycuda__
 
 if __has_pycuda__:
     import pycuda.driver as cuda
     import pycuda.gpuarray as garray
+    from ..cuda.kernel import CudaKernel
 
     dev_attrs = cuda.device_attribute
+    GPUArray = garray.GPUArray
+    from pycuda.tools import dtype_to_ctype
+else:
+    GPUArray = MissingComponentError("pycuda")
+    dtype_to_ctype = MissingComponentError("pycuda")
 
 
 # NB: we must detach from a context before creating another context
 class CudaProcessing(ProcessingBase):
-    _array_class = garray.GPUArray if __has_pycuda__ else None
+    array_class = GPUArray if __has_pycuda__ else None
+    dtype_to_ctype = dtype_to_ctype
 
     def __init__(self, device_id=None, ctx=None, stream=None, cleanup_at_exit=True):
         """
@@ -52,3 +60,16 @@ class CudaProcessing(ProcessingBase):
 
     def _allocate_array_mem(self, shape, dtype):
         return garray.zeros(shape, dtype)
+
+    def kernel(
+        self, kernel_name, filename=None, src=None, signature=None, texrefs=None, automation_params=None, **build_kwargs
+    ):
+        return CudaKernel(
+            kernel_name,
+            filename=filename,
+            src=src,
+            signature=signature,
+            texrefs=texrefs,
+            automation_params=automation_params,
+            **build_kwargs,
+        )

nabu/cuda/src/ElementOp.cu

@@ -47,6 +47,33 @@ __global__ void inplace_generic_op_2Dby2D(float* arr2D, float* arr2D_other, int
 }
 
 
+// launched with (Nx, Ny, Nz) threads
+// does array3D[x, y, z] = op(array3D[x, y, z], array1D[x]) (in the "numpy broadcasting" sense)
+__global__ void inplace_generic_op_3Dby1D(
+    float * array3D,
+    float* array1D,
+    int Nx, // input/output number of columns
+    int Ny, // input/output number of rows
+    int Nz  // input/output depth
+) {
+    uint x = blockDim.x * blockIdx.x + threadIdx.x;
+    uint y = blockDim.y * blockIdx.y + threadIdx.y;
+    uint z = blockDim.z * blockIdx.z + threadIdx.z;
+    if ((x >= Nx) || (y >= Ny) || (z >= Nz)) return;
+    size_t idx = ((z * Ny) + y)*Nx + x;
+
+    #if GENERIC_OP == OP_ADD
+    array3D[idx] += array1D[x];
+    #elif GENERIC_OP == OP_SUB
+    array3D[idx] -= array1D[x];
+    #elif GENERIC_OP == OP_MUL
+    array3D[idx] *= array1D[x];
+    #elif GENERIC_OP == OP_DIV
+    array3D[idx] /= array1D[x];
+    #endif
+}
+
+
 // arr3D *= arr1D (along fast dim)
 __global__ void inplace_complex_mul_3Dby1D(complex* arr3D, complex* arr1D, int width, int height, int depth) {
     int x = blockDim.x * blockIdx.x + threadIdx.x;
@@ -161,3 +188,54 @@ __global__ void reverse2D_x(float* array, int Nx, int Ny) {
     array[pos2] = tmp;
 }
 
+
+/**
+
+  Generic mul-add kernel with possibly-complicated indexing.
+
+  dst[DST_IDX] = fac_dst*dst[DST_IDX] + fac_other*other[OTHER_IDX]
+  where
+    DST_IDX = dst_start_row:dst_end_row, dst_start_col:dst_end_col
+    OTHER_IDX = other_start_row:other_end_row, other_start_col:other_end_col
+
+  Usage:
+    mul_add(dst, other, dst_nx, other_nx, a, b, (x1, x2), (y1, y2), (x3, x4), (y3, y4))
+*/
+
+__global__ void  mul_add(
+    float* dst,
+    float* other,
+    int dst_width,
+    int other_width,
+    float fac_dst,
+    float fac_other,
+    int2 dst_x_range,
+    int2 dst_y_range,
+    int2 other_x_range,
+    int2 other_y_range
+    )
+{
+    size_t x = blockDim.x * blockIdx.x + threadIdx.x;
+    size_t y = blockDim.y * blockIdx.y + threadIdx.y;
+
+    int x_start_dst = dst_x_range.x;
+    int x_stop_dst = dst_x_range.y;
+    int y_start_dst = dst_y_range.x;
+    int y_stop_dst = dst_y_range.y;
+
+    int x_start_other = other_x_range.x;
+    int x_stop_other = other_x_range.y;
+    int y_start_other = other_y_range.x;
+    int y_stop_other = other_y_range.y;
+
+    int operation_width = x_stop_dst - x_start_dst; // assumed == x_stop_other - x_start_other
+    int operation_height = y_stop_dst - y_start_dst; // assumed == y_stop_other - y_start_other
+
+    if ((x >= operation_width) || (y >= operation_height)) return;
+
+    size_t idx_in_dst = (y + y_start_dst)*dst_width + (x + x_start_dst);
+    size_t idx_in_other = (y + y_start_other)*other_width + (x + x_start_other);
+
+    dst[idx_in_dst] = fac_dst * dst[idx_in_dst] + fac_other * other[idx_in_other];
+}
+

nabu/cuda/src/backproj.cu

@@ -2,7 +2,9 @@
     #define SHARED_SIZE 256
 #endif
 
+#ifdef USE_TEXTURES
 texture<float, 2, cudaReadModeElementType> tex_projections;
+#endif
 
 #ifdef CLIP_OUTER_CIRCLE
 inline __device__ int is_in_circle(int x, int y, float center_x, float center_y, int radius2) {
@@ -11,6 +13,20 @@ inline __device__ int is_in_circle(int x, int y, float center_x, float center_y,
 #endif
 
 
+/*
+  Linear interpolation on a 2D array, horizontally.
+  This will return arr[y][x] where y is an int (exact access) and x is a float (linear interp horizontally)
+*/
+static inline __device__ float linear_interpolation(float* arr, int Nx, float x, int y) {
+    // check commented to gain a bit of speed - the check was done before function call
+    // if (x < 0 || x >= Nx) return 0.0f; // texture address mode CLAMP_TO_EDGE
+    int xm = (int) floorf(x);
+    int xp = (int) ceilf(x);
+    if ((xm == xp) || (xp >= Nx)) return arr[y*Nx+xm];
+    else return (arr[y*Nx+xm] * (xp - x)) + (arr[y*Nx+xp] * (x - xm));
+}
+
+
 /**
 
 Implementation details
@@ -32,13 +48,16 @@ blocks.
 // One thread handles up to 4 pixels in the output slice
 __global__ void backproj(
     float* d_slice,
+    #ifndef USE_TEXTURES
+    float* d_sino,
+    #endif
     int num_projs,
     int num_bins,
     float axis_position,
     int n_x,
     int n_y,
-    int offset_x,
-    int offset_y,
+    float offset_x,
+    float offset_y,
     float* d_cos,
     float* d_msin,
     #ifdef DO_AXIS_CORRECTION
@@ -107,10 +126,17 @@ __global__ void backproj(
         h4 += axcorr;
         #endif
 
+        #ifdef USE_TEXTURES
         if (h1 >= 0 && h1 < num_bins) sum1 += tex2D(tex_projections, h1 + 0.5f, proj + 0.5f);
         if (h2 >= 0 && h2 < num_bins) sum2 += tex2D(tex_projections, h2 + 0.5f, proj + 0.5f);
         if (h3 >= 0 && h3 < num_bins) sum3 += tex2D(tex_projections, h3 + 0.5f, proj + 0.5f);
         if (h4 >= 0 && h4 < num_bins) sum4 += tex2D(tex_projections, h4 + 0.5f, proj + 0.5f);
+        #else
+        if (h1 >= 0 && h1 < num_bins) sum1 += linear_interpolation(d_sino, num_bins, h1, proj);
+        if (h2 >= 0 && h2 < num_bins) sum2 += linear_interpolation(d_sino, num_bins, h2, proj);
+        if (h3 >= 0 && h3 < num_bins) sum3 += linear_interpolation(d_sino, num_bins, h3, proj);
+        if (h4 >= 0 && h4 < num_bins) sum4 += linear_interpolation(d_sino, num_bins, h4, proj);
+        #endif
     }
 
     int write_topleft = 1, write_topright = 1, write_botleft = 1, write_botright = 1;

nabu/cuda/src/fourier_wavelets.cu

@@ -12,6 +12,6 @@ __global__ void kern_fourierwavelets(float2* sinoF, int Nx, int Ny, float wsigma
 
     int tid = gidy*Nx + gidx;
     // do not forget the scale factor (here Ny)
-    sinoF[tid].x *= factor/Ny;
-    sinoF[tid].y *= factor/Ny;
+    sinoF[tid].x *= factor;
+    sinoF[tid].y *= factor;
 }

nabu/cuda/src/normalization.cu

@@ -45,3 +45,26 @@ __global__ void normalize_chebyshev(
         array[pos] -= ff0*f0/sum0 + ff1*f1/sum1 + ff2*f2/sum2;
     }
 }
+
+
+// launched with (Nx, 1, Nz) threads
+__global__ void vertical_mean(
+    float * array,
+    float* output,
+    int Nx, // input/output number of columns
+    int Ny, // input/output number of rows
+    int Nz  // input/output depth
+) {
+    uint x = blockDim.x * blockIdx.x + threadIdx.x;
+    uint y = blockDim.y * blockIdx.y + threadIdx.y;
+    uint z = blockDim.z * blockIdx.z + threadIdx.z;
+    if ((x >= Nx) || (y >= 1) || (z >= Nz)) return;
+
+    float m = 0.0f;
+    for (uint i = 0; i < Ny; i++) {
+        float s = array[(z * Ny + i) * Nx + x];
+        m += (s - m)/(i+1);
+    }
+    output[z * Nx + x] = (float) m;
+}
+

nabu/cuda/src/padding.cu

@@ -129,8 +129,8 @@ __global__ void coordinate_transform(
     uint y = blockDim.y * blockIdx.y + threadIdx.y;
     if ((x >= Nx_padded) || (y >= Ny_padded)) return;
     uint idx = y*Nx_padded  +  x;
-    int x2 = cols_inds[idx];
-    int y2 = rows_inds[idx];
+    int x2 = cols_inds[x];
+    int y2 = rows_inds[y];
     array_out[idx] = array_in[y2*Nx + x2];
 }
 

nabu/cuda/src/transpose.cu

@@ -0,0 +1,16 @@
+#ifndef SRC_DTYPE
+  #define SRC_DTYPE float
+#endif
+#ifndef DST_DTYPE
+  #define DST_DTYPE float
+#endif
+
+#include <pycuda-complex.hpp>
+
+__global__ void transpose(SRC_DTYPE* src, DST_DTYPE* dst, int src_width, int src_height) {
+    // coordinates for "dst"
+    uint x = blockDim.x * blockIdx.x + threadIdx.x;
+    uint y = blockDim.y * blockIdx.y + threadIdx.y;
+    if ((x >= src_height) || (y >= src_width)) return;
+    dst[y*src_height + x] = (DST_DTYPE) src[x*src_width + y];
+}

nabu/cuda/utils.py

@@ -8,6 +8,7 @@ try:
     import pycuda.driver as cuda
     from pycuda import gpuarray as garray
     from pycuda.tools import clear_context_caches
+    from pycuda.compiler import get_nvcc_version as pycuda_get_nvcc_version
 
     __has_pycuda__ = True
     __pycuda_error_msg__ = None
@@ -24,6 +25,13 @@ try:
 except ImportError:
     __has_cufft__ = False
 
+try:
+    import cupy
+
+    __has_cupy__ = True
+except ImportError:
+    __has_cupy__ = False
+
 
 def get_cuda_context(device_id=None, cleanup_at_exit=True):
     """
@@ -133,6 +141,27 @@ def collect_cuda_gpus():
     return cuda_gpus
 
 
+def get_nvcc_version(nvcc_cmd="nvcc"):
+    try:
+        ver = "".join(pycuda_get_nvcc_version(nvcc_cmd)).split("release")[1].strip().split(" ")[0].strip(",")
+    except:
+        ver = None
+    return ver
+
+
+def check_textures_availability():
+    """
+    Check whether Cuda textures can be used.
+    The only limitation is pycuda which does not support texture objects.
+    Textures references were deprecated, and removed from Cuda 12.
+    """
+    nvcc_ver = get_nvcc_version()
+    if nvcc_ver is None:
+        return False  # unknown - can't parse NVCC version for some reason
+    nvcc_major = int(nvcc_ver.split(".")[0])
+    return nvcc_major < 12
+
+
 """
 pycuda/driver.py
     np.complex64: SIGNED_INT32, num_channels = 2
@@ -290,3 +319,13 @@ def replace_array_memory(arr, new_shape):
     arr.shape = new_shape
     # TODO re-compute strides
     return arr
+
+
+def pycuda_to_cupy(arr_pycuda):
+    arr_cupy_mem = cupy.cuda.UnownedMemory(arr_pycuda.ptr, arr_pycuda.size, arr_pycuda)
+    arr_cupy_memptr = cupy.cuda.MemoryPointer(arr_cupy_mem, offset=0)
+    return cupy.ndarray(arr_pycuda.shape, dtype=arr_pycuda.dtype, memptr=arr_cupy_memptr)  # pylint: disable=E1123
+
+
+def cupy_to_pycuda(arr_cupy):
+    return garray.empty(arr_cupy.shape, arr_cupy.dtype, gpudata=arr_cupy.data.ptr)

nabu/estimation/alignment.py

@@ -38,6 +38,7 @@ local_ifftn = scipy.fft.irfftn
 
 class AlignmentBase:
     default_extra_options = {"blocking_plots": False}
+    _default_cor_options = {}
 
     def __init__(
         self,
@@ -47,6 +48,7 @@ class AlignmentBase:
         logger=None,
         data_type=np.float32,
         extra_options=None,
+        cor_options=None,
     ):
         """
         Alignment basic functions.
@@ -71,6 +73,13 @@ class AlignmentBase:
 
         self._init_parameters(vert_fft_width, horz_fft_width, verbose, logger, data_type, extra_options=extra_options)
         self._plot_windows = {}
+        self._set_cor_options(cor_options)
+
+    def _set_cor_options(self, cor_options):
+        default_dict = self._default_cor_options.copy()
+        if cor_options is not None:
+            default_dict.update(cor_options)
+        self.cor_options = default_dict
 
     def _init_parameters(self, vert_fft_width, horz_fft_width, verbose, logger, data_type, extra_options=None):
         self.logger = LoggerOrPrint(logger)
@@ -273,7 +282,7 @@ class AlignmentBase:
                     vertex_x, vertex_min_x, vertex_max_x
                 )
             else:
-                message = "Fitted positions outide the input margins [{}, {}]: {} below and {} above".format(
+                message = "Fitted positions outside the input margins [{}, {}]: {} below and {} above".format(
                     vertex_min_x,
                     vertex_max_x,
                     np.sum(1 - lower_bound_ok),