Python: filter FDK/FBP projections in cache-sized blocks, on several threads

Python/tests/test_filtering_workers.py

@@ -0,0 +1,86 @@
+"""Tests for the threaded FDK/FBP filtering step (CPU only, no GPU call)."""
+import numpy as np
+import pytest
+
+import tigre
+import tigre.utilities.filtering as filtering_module
+from tigre.utilities.filtering import default_workers, filter as make_filter, filtering, nextpow2
+from tigre.utilities.filtering import ramp_flat
+
+
+def _geo(nangles, nv, nu, filt=None):
+    geo = tigre.geometry(mode="cone", nVoxel=np.array([8, 8, 8]), default=True)
+    geo.nDetector = np.array([nv, nu])
+    geo.dDetector = np.array([0.8, 0.8])
+    geo.sDetector = geo.nDetector * geo.dDetector
+    angles = np.linspace(0, 2 * np.pi, nangles, endpoint=False, dtype=np.float32)
+    geo.check_geo(angles)
+    geo.filter = filt
+    return geo, angles
+
+
+def _proj(nangles, nv, nu):
+    return np.random.default_rng(nangles + nv + nu).random((nangles, nv, nu), dtype=np.float32)
+
+
+def _plain_filter(proj, geo, angles):
+    """One projection at a time, float64, no packing, no blocks, no threads."""
+    nu = int(geo.nDetector[1])
+    filt_len = max(64, 2 ** nextpow2(2 * nu))
+    filt = make_filter(geo.filter, ramp_flat(filt_len)[0], filt_len, 1).astype(np.float64)
+    padding = (filt_len - nu) // 2
+    scale = (geo.DSD[0] / geo.DSO[0]) * (2 * np.pi / len(angles)) / (4 * geo.dDetector[1])
+    out = np.empty(proj.shape, dtype=np.float64)
+    for i in range(proj.shape[0]):
+        padded = np.zeros((proj.shape[1], filt_len))
+        padded[:, padding:padding + nu] = proj[i]
+        filtered = np.fft.ifft(np.fft.fft(padded, axis=1) * filt, axis=1).real
+        out[i] = filtered[:, padding:padding + nu] * scale
+    return out
+
+
+@pytest.mark.parametrize("shape", [(6, 40, 64), (7, 40, 64), (1, 3, 50), (9, 1, 200), (4, 300, 256)])
+@pytest.mark.parametrize("filt", [None, "ram_lak", "shepp_logan", "cosine", "hamming", "hann"])
+def test_matches_a_plain_per_projection_filter(shape, filt):
+    geo, angles = _geo(*shape, filt=filt)
+    proj = _proj(*shape)
+    expected = _plain_filter(proj, geo, angles)
+    out = filtering(proj.copy(), geo, angles, parker=False)
+    assert out.dtype == np.float32
+    np.testing.assert_allclose(out, expected, rtol=0, atol=2e-6 * np.abs(expected).max())
+
+
+@pytest.mark.parametrize("shape", [(6, 40, 64), (7, 13, 33), (9, 1, 200), (5, 300, 256)])
+def test_same_bits_whatever_the_number_of_workers(shape, monkeypatch):
+    geo, angles = _geo(*shape)
+    proj = _proj(*shape)
+    serial = filtering(proj.copy(), geo, angles, parker=False, workers=1)
+    for workers in (None, 2, 3, 8):
+        assert np.array_equal(serial, filtering(proj.copy(), geo, angles, parker=False,
+                                                workers=workers))
+    # the smallest blocks the function will use: many blocks per projection
+    monkeypatch.setattr(filtering_module, "BLOCK_BYTES", 1)
+    for workers in (1, 4):
+        assert np.array_equal(serial, filtering(proj.copy(), geo, angles, parker=False,
+                                                workers=workers))
+
+
+def test_filters_in_place_and_returns_its_argument():
+    geo, angles = _geo(4, 20, 32)
+    proj = _proj(4, 20, 32)
+    before = proj.copy()
+    out = filtering(proj, geo, angles, parker=False)
+    assert out is proj
+    assert not np.array_equal(before, proj)
+
+
+def test_error_in_a_worker_thread_reaches_the_caller():
+    geo, angles = _geo(8, 600, 512)
+    with pytest.raises(ValueError):
+        filtering(_proj(8, 600, 500), geo, angles, parker=False, workers=4)
+
+
+def test_default_workers(monkeypatch):
+    assert 1 <= default_workers() <= filtering_module.MAX_WORKERS
+    monkeypatch.setattr(filtering_module, "MAX_WORKERS", 1)
+    assert default_workers() == 1

Python/tigre/utilities/filtering.py

@@ -5,14 +5,22 @@
 import numpy as np
 from scipy.fft  import fft, ifft
 
+import os
+import threading
 import warnings
+from concurrent.futures import ThreadPoolExecutor
 
 import numpy as np
 from tigre.utilities.parkerweight import parkerweight
 
+# bytes of padded complex data a thread filters at a time, small enough to stay in cache
+BLOCK_BYTES = 2**21
+# threads used when the caller does not choose; more than this was not faster where measured
+MAX_WORKERS = 16
+
 
 # TODO: Fix parker
-def filtering(proj, geo, angles, parker, verbose=False):
+def filtering(proj, geo, angles, parker, verbose=False, workers=None):
     if parker:
         proj=parkerweight(proj.transpose(0,2,1),geo,angles,parker).transpose(0,2,1)
 
@@ -21,36 +29,65 @@ def filtering(proj, geo, angles, parker, verbose=False):
 
     d=1
     filt=filter(geo.filter,ramp_kernel[0],filt_len,d,verbose=verbose)
-    filt=np.kron(np.ones((np.int64(geo.nDetector[0]),1),dtype=np.float32),filt)
 
     padding = int((filt_len-geo.nDetector[1])//2 )
     scale_factor = (geo.DSD[0]/geo.DSO[0]) * (2 * np.pi/ len(angles)) / ( 4 * geo.dDetector[1] ) 
 
-    #filter 2 projection at a time packing in to complex container
-    fproj=np.empty((geo.nDetector[0],filt_len),dtype=np.complex64)
-    for i in range(0,angles.shape[0]-1,2):
-        fproj.fill(0)
-        fproj.real[:,padding:padding+geo.nDetector[1]]=proj[i]
-        fproj.imag[:,padding:padding+geo.nDetector[1]]=proj[i+1]
-
-        fproj=fft(fproj,axis=1)
-        fproj=fproj*filt
-        fproj=ifft(fproj,axis=1)
-
-        proj[i]=fproj.real[:,padding:padding+geo.nDetector[1]] * scale_factor
-        proj[i+1]=fproj.imag[:,padding:padding+geo.nDetector[1]] * scale_factor
+    nangles=angles.shape[0]
+    nrows=int(geo.nDetector[0])
+    cols=slice(padding,padding+geo.nDetector[1])
+    if workers is None:
+        workers=default_workers()
+
+    #detector rows are filtered independently of each other: cut them in blocks that fit
+    #in cache and share the blocks among threads (numpy and scipy.fft release the GIL).
+    #A multiple of 16 rows, so that the FFT packs the same rows in a SIMD vector as it
+    #does on a whole projection, and rounds the same way
+    block=max(16,BLOCK_BYTES//(8*filt_len)//16*16)
+    local=threading.local()
+
+    def filter_blocks(tasks):
+        if not hasattr(local,"fproj"):
+            local.fproj=np.empty((block,filt_len),dtype=np.complex64)
+
+        for i,r0,r1 in tasks:
+            #filter 2 projection at a time packing in to complex container
+            fproj=local.fproj[:r1-r0]
+            fproj.fill(0)
+            fproj.real[:,cols]=proj[i,r0:r1]
+            #if odd number of projections filter last solo
+            if i+1<nangles:
+                fproj.imag[:,cols]=proj[i+1,r0:r1]
+
+            fproj=fft(fproj,axis=1,overwrite_x=True)
+            fproj*=filt
+            fproj=ifft(fproj,axis=1,overwrite_x=True)
+
+            proj[i,r0:r1]=fproj.real[:,cols] * scale_factor
+            if i+1<nangles:
+                proj[i+1,r0:r1]=fproj.imag[:,cols] * scale_factor
+
+    tasks=[(i,r0,min(r0+block,nrows)) for i in range(0,nangles,2) for r0 in range(0,nrows,block)]
+    #a detector with few rows gives small tasks: hand them to the threads several at a time
+    batch=max(1,block//nrows)
+    batches=[tasks[k:k+batch] for k in range(0,len(tasks),batch)]
+    workers=min(workers,len(batches))
+    if workers>1:
+        with ThreadPoolExecutor(workers) as pool:
+            #list() so that an exception raised in a thread is raised here
+            list(pool.map(filter_blocks,batches))
+    else:
+        filter_blocks(tasks)
 
-    #if odd number of projections filter last solo
-    if angles.shape[0] % 2:
-        fproj.fill(0)
-        fproj.real[:,padding:padding+geo.nDetector[1]]=proj[angles.shape[0]-1]
+    return proj
 
-        fproj=fft(fproj,axis=1)
-        fproj=fproj*filt
-        fproj=np.real(ifft(fproj,axis=1))     
-        proj[angles.shape[0]-1]=fproj[:,padding:padding+geo.nDetector[1]] * scale_factor
 
-    return proj
+def default_workers():
+    try:
+        ncpu = len(os.sched_getaffinity(0))  # the cores this process is allowed to use
+    except AttributeError:  # not available on Windows and macOS
+        ncpu = os.cpu_count() or 1
+    return min(ncpu, MAX_WORKERS)
 
 
 def ramp_flat(n, verbose=False):