Bound mapped memory during sampling and block comparison

pilot/analyze.py

@@ -58,7 +58,10 @@ def sample_object(obj, tensors, sample_size=2048):
 
 def exact_pair(a, b, tensors, deadline):
     if a['kind'] != 'tensor':
-        return metrics(read_object(a, tensors), read_object(b, tensors))
+        result = metrics(read_object(a, tensors), read_object(b, tensors))
+        tensors[a['tensor']].evict()
+        tensors[b['tensor']].evict()
+        return result
     ta, tb = tensors[a['tensor']], tensors[b['tensor']]
     xx = yy = xy = delta = sx = sy = 0.
     for start in range(0, ta.numel, 262144):
@@ -227,7 +230,10 @@ def structural_probe(a, b, tensors):
 
 def analyze_mode(mode, objects, tensors, root, deadline, cap, seed, baseline_cap=100):
     print(f'{mode}: sampling {len(objects)} objects', flush=True)
-    features = [sample_object(o, tensors) for o in objects]
+    features = []
+    for obj in objects:
+        features.append(sample_object(obj, tensors))
+        tensors[obj['tensor']].evict()
     candidates, coverage = nearest_candidates(objects, features, cap)
     write_json(root / f'{mode}_objects.json', objects)
     pairs, nulls, structures = [], [], []
@@ -273,6 +279,8 @@ def analyze_mode(mode, objects, tensors, root, deadline, cap, seed, baseline_cap
         ndim = len(objects[a]['shape'])
         if ndim >= 2 and structural_counts[ndim] < 20:
             structures.append({'a': a, 'b': b, 'modes': structural_probe(objects[a], objects[b], tensors)})
+            tensors[objects[a]['tensor']].evict()
+            tensors[objects[b]['tensor']].evict()
             structural_counts[ndim] += 1
         if (i + 1) % 100 == 0:
             write_csv(root / f'{mode}_pairs.csv', pairs)