Merge branch 'topic/default/memory-leak-pythran' into 'branch/default'

Fix a memory leak by Pythran

See merge request fluiddyn/fluidimage!97

.gitattributes

@@ -7,4 +7,5 @@ pixi.lock linguist-language=YAML
 old export-ignore
 try export-ignore
 bench export-ignore
+dev export-ignore
 image_samples export-ignore

dev/memory-leak/Makefile

@@ -0,0 +1,6 @@
+
+all:
+	pythran mod_pythran.py
+
+clean:
+	rm -f *.so

dev/memory-leak/README.md

@@ -0,0 +1,21 @@
+# Investigating a leak in fluidimage
+
+There was a leak in {func}`fluidimage.calcul.subpix.compute_subpix_2d_gaussian2`,
+which is a pythranized function.
+
+A simple reproducer is
+
+```python
+
+# pythran export simpler_leak(float32[:, :], int, int)
+# pythran export simpler_no_leak(float32[:, :], int, int)
+
+def simpler_leak(correl, ix, iy):
+    # returning this view leads to a leak!
+    correl_crop = correl[iy - 1 : iy + 2, ix - 1 : ix + 2]
+    return correl_crop
+
+def simpler_no_leak(correl, ix, iy):
+    correl_crop = np.ascontiguousarray(correl[iy - 1 : iy + 2, ix - 1 : ix + 2])
+    return correl_crop
+```

dev/memory-leak/mod_pythran.py

@@ -0,0 +1,95 @@
+from math import log
+
+import numpy as np
+
+# pythran export compute_subpix_2d_gaussian2(float32[:, :], int, int)
+# pythran export compute_subpix_2d_gaussian3(float32[:, :], int, int)
+
+# pythran export simpler_leak(float32[:, :], int, int)
+# pythran export simpler_no_leak(float32[:, :], int, int)
+
+
+def compute_subpix_2d_gaussian2(correl, ix, iy):
+
+    # returning this view leads to a leak!
+    correl_crop = correl[iy - 1 : iy + 2, ix - 1 : ix + 2]
+
+    tmp = np.where(correl_crop < 0)
+    for i0, i1 in zip(tmp[0], tmp[1]):
+        correl_crop[i0, i1] = 1e-6
+
+    c10 = 0
+    c01 = 0
+    c11 = 0
+    c20 = 0
+    c02 = 0
+    for i in range(3):
+        for j in range(3):
+            c10 += (i - 1) * np.log(correl_crop[j, i])
+            c01 += (j - 1) * np.log(correl_crop[j, i])
+            c11 += (i - 1) * (j - 1) * np.log(correl_crop[j, i])
+            c20 += (3 * (i - 1) ** 2 - 2) * np.log(correl_crop[j, i])
+            c02 += (3 * (j - 1) ** 2 - 2) * np.log(correl_crop[j, i])
+            c00 = (5 - 3 * (i - 1) ** 2 - 3 * (j - 1) ** 2) * np.log(
+                correl_crop[j, i]
+            )
+
+    c00, c10, c01, c11, c20, c02 = (
+        c00 / 9,
+        c10 / 6,
+        c01 / 6,
+        c11 / 4,
+        c20 / 6,
+        c02 / 6,
+    )
+    deplx = (c11 * c01 - 2 * c10 * c02) / (4 * c20 * c02 - c11**2)
+    deply = (c11 * c10 - 2 * c01 * c20) / (4 * c20 * c02 - c11**2)
+    return deplx, deply, correl_crop
+
+
+def compute_subpix_2d_gaussian3(correl, ix, iy):
+    correl_crop = np.ascontiguousarray(correl[iy - 1 : iy + 2, ix - 1 : ix + 2])
+
+    for i0 in range(-1, 2):
+        for i1 in range(-1, 2):
+            if correl_crop[i0, i1] < 0:
+                correl_crop[i0, i1] = 1e-6
+
+    c10 = 0
+    c01 = 0
+    c11 = 0
+    c20 = 0
+    c02 = 0
+    for i0 in range(3):
+        for i1 in range(3):
+            c10 += (i1 - 1) * log(correl_crop[i0, i1])
+            c01 += (i0 - 1) * log(correl_crop[i0, i1])
+            c11 += (i1 - 1) * (i0 - 1) * log(correl_crop[i0, i1])
+            c20 += (3 * (i1 - 1) ** 2 - 2) * log(correl_crop[i0, i1])
+            c02 += (3 * (i0 - 1) ** 2 - 2) * log(correl_crop[i0, i1])
+            c00 = (5 - 3 * (i1 - 1) ** 2 - 3 * (i0 - 1) ** 2) * log(
+                correl_crop[i0, i1]
+            )
+
+    c00, c10, c01, c11, c20, c02 = (
+        c00 / 9,
+        c10 / 6,
+        c01 / 6,
+        c11 / 4,
+        c20 / 6,
+        c02 / 6,
+    )
+    deplx = (c11 * c01 - 2 * c10 * c02) / (4 * c20 * c02 - c11**2)
+    deply = (c11 * c10 - 2 * c01 * c20) / (4 * c20 * c02 - c11**2)
+    return deplx, deply, correl_crop
+
+
+def simpler_leak(correl, ix, iy):
+    # returning this view leads to a leak!
+    correl_crop = correl[iy - 1 : iy + 2, ix - 1 : ix + 2]
+    return correl_crop
+
+
+def simpler_no_leak(correl, ix, iy):
+    correl_crop = np.ascontiguousarray(correl[iy - 1 : iy + 2, ix - 1 : ix + 2])
+    return correl_crop

dev/memory-leak/profiler.py

@@ -0,0 +1,42 @@
+import tracemalloc
+
+# list to store memory snapshots
+snaps = []
+
+
+def snapshot():
+    snaps.append(tracemalloc.take_snapshot())
+
+
+def display_stats():
+    stats = snaps[0].statistics("filename")
+    print("\n*** top 5 stats grouped by filename ***")
+    for s in stats[:5]:
+        print(s)
+
+
+def compare():
+    first = snaps[0]
+    for snapshot in snaps[1:]:
+        stats = snapshot.compare_to(first, "lineno")
+        print("\n*** top 10 stats ***")
+        for s in stats[:10]:
+            print(s)
+
+
+def print_trace():
+    # pick the last saved snapshot, filter noise
+    snapshot = snaps[-1].filter_traces(
+        (
+            tracemalloc.Filter(False, "<frozen importlib._bootstrap>"),
+            tracemalloc.Filter(False, "<frozen importlib._bootstrap_external>"),
+            tracemalloc.Filter(False, "<unknown>"),
+        )
+    )
+    largest = snapshot.statistics("traceback")[0]
+
+    print(
+        f"\n*** Trace for largest memory block - ({largest.count} blocks, {largest.size/1024} Kb) ***"
+    )
+    for line in largest.traceback.format():
+        print(line)

dev/memory-leak/trace_gaussian2.py

@@ -0,0 +1,34 @@
+
+import gc
+import tracemalloc
+
+import numpy as np
+
+from fluidimage.calcul.subpix import compute_subpix_2d_gaussian2
+
+import profiler
+
+
+n0, n1 = 32, 32
+
+correl = np.zeros((n0, n1), dtype=np.float32)
+
+i_max = n0//2
+
+correl[i_max-1:i_max+2, i_max-1:i_max+2] = 0.6
+correl[i_max, i_max] = 1.0
+correl[i_max, i_max-1] = 0.7
+
+print(compute_subpix_2d_gaussian2(correl, i_max, i_max))
+
+tracemalloc.start()
+
+for _ in range(5):
+    for idx in range(1000):
+        compute_subpix_2d_gaussian2(correl, i_max, i_max)
+    gc.collect()
+    profiler.snapshot()
+
+profiler.display_stats()
+profiler.compare()
+profiler.print_trace()

dev/memory-leak/trace_piv_work.py

@@ -0,0 +1,56 @@
+import gc
+import os
+import tracemalloc
+
+from fluidimage.piv import Work
+
+import profiler
+
+os.environ["OMP_NUM_THREADS"] = "1"
+
+params = Work.create_default_params()
+
+params.series.path = "../../image_samples/wake_legi/images/B*.png"
+
+params.piv0.shape_crop_im0 = 40
+params.piv0.displacement_max = 14
+
+params.piv0.nb_peaks_to_search = 1
+params.piv0.particle_radius = 3
+
+params.mask.strcrop = ":, :1500"
+
+params.multipass.number = 2
+
+# params.multipass.use_tps = "last"
+params.multipass.use_tps = False
+params.multipass.subdom_size = 200
+params.multipass.smoothing_coef = 10.0
+params.multipass.threshold_tps = 0.5
+
+params.fix.correl_min = 0.15
+params.fix.threshold_diff_neighbour = 3
+
+work = Work(params=params)
+
+# tracemalloc.start()
+
+# piv = work.process_1_serie()
+
+# snapshot = tracemalloc.take_snapshot()
+# top_stats = snapshot.statistics("lineno")
+
+# print("[ Top 10 ]")
+# for stat in top_stats[:10]:
+#     print(stat)
+
+tracemalloc.start(10)
+
+for _ in range(5):
+    piv = work.process_1_serie()
+    gc.collect()
+    profiler.snapshot()
+
+profiler.display_stats()
+profiler.compare()
+profiler.print_trace()

dev/memory-leak/trace_pythran.py

@@ -0,0 +1,35 @@
+import gc
+import tracemalloc
+
+import numpy as np
+
+# from mod_pythran import compute_subpix_2d_gaussian2 as func
+from mod_pythran import compute_subpix_2d_gaussian3 as func
+# from mod_pythran import simpler_leak as func
+# from mod_pythran import simpler_no_leak as func
+
+
+import profiler
+
+
+n0, n1 = 3, 3
+correl = np.zeros((n0, n1), dtype=np.float32)
+
+i_max = n0 // 2
+correl[i_max - 1 : i_max + 2, i_max - 1 : i_max + 2] = 0.6
+correl[i_max, i_max] = 1.0
+correl[i_max, i_max - 1] = 0.7
+
+print(func(correl, i_max, i_max))
+
+tracemalloc.start()
+
+for _ in range(5):
+    for idx in range(1000):
+        func(correl, i_max, i_max)
+    gc.collect()
+    profiler.snapshot()
+
+profiler.display_stats()
+profiler.compare()
+profiler.print_trace()

src/fluidimage/calcul/subpix.py

@@ -15,29 +15,37 @@
 
 @boost
 def compute_subpix_2d_gaussian2(correl: "float32[][]", ix: int, iy: int):
-    correl_crop = correl[iy - 1 : iy + 2, ix - 1 : ix + 2]
-    # hoops, pythran crashes because of this line
-    # correl_crop[correl_crop < 0] = 1e-6
+    # without np.ascontiguousarray => memory leak
+    correl_crop = np.ascontiguousarray(correl[iy - 1 : iy + 2, ix - 1 : ix + 2])
 
-    # we write it like this to please pythran
-    tmp = np.where(correl_crop < 0)
-    for i0, i1 in zip(tmp[0], tmp[1]):
-        correl_crop[i0, i1] = 1e-6
+    correl_crop_ravel = correl_crop.ravel()
+
+    correl_min = correl_crop.min()
+    for idx in range(9):
+        correl_crop_ravel[idx] -= correl_min
+
+    correl_max = correl_crop.max()
+    for idx in range(9):
+        correl_crop_ravel[idx] /= correl_max
+
+    for idx in range(9):
+        if correl_crop_ravel[idx] == 0:
+            correl_crop_ravel[idx] = 1e-8
 
     c10 = 0
     c01 = 0
     c11 = 0
     c20 = 0
     c02 = 0
-    for i in range(3):
-        for j in range(3):
-            c10 += (i - 1) * np.log(correl_crop[j, i])
-            c01 += (j - 1) * np.log(correl_crop[j, i])
-            c11 += (i - 1) * (j - 1) * np.log(correl_crop[j, i])
-            c20 += (3 * (i - 1) ** 2 - 2) * np.log(correl_crop[j, i])
-            c02 += (3 * (j - 1) ** 2 - 2) * np.log(correl_crop[j, i])
-            c00 = (5 - 3 * (i - 1) ** 2 - 3 * (j - 1) ** 2) * np.log(
-                correl_crop[j, i]
+    for i0 in range(3):
+        for i1 in range(3):
+            c10 += (i1 - 1) * np.log(correl_crop[i0, i1])
+            c01 += (i0 - 1) * np.log(correl_crop[i0, i1])
+            c11 += (i1 - 1) * (i0 - 1) * np.log(correl_crop[i0, i1])
+            c20 += (3 * (i1 - 1) ** 2 - 2) * np.log(correl_crop[i0, i1])
+            c02 += (3 * (i0 - 1) ** 2 - 2) * np.log(correl_crop[i0, i1])
+            c00 = (5 - 3 * (i1 - 1) ** 2 - 3 * (i0 - 1) ** 2) * np.log(
+                correl_crop[i0, i1]
             )
 
     c00, c10, c01, c11, c20, c02 = (