Bring back old scratch sorting as default

microsoft · Sep 29, 2021 · 90da877 · 90da877
1 parent dd05b77
commit 90da877
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/sccl/ncclize.py b/sccl/ncclize.py
@@ -336,6 +336,10 @@ def allocate_scratch(gpu, addr):
         _remap_scratch_into_input_output(liveness, gpus, logging)
     elif greedy_scratch_sorting:
         _greedy_scratch_sort(algorithm, gpus)
+    else:
+        # Sort scratch mappings in an attempt to make more of them contiguous (this is of course a heuristic).
+        for gpu in gpus.values():
+            gpu.scratch = { addr: idx for idx, addr in enumerate(sorted(gpu.scratch)) }
 
     # Add any copies from input to output that weren't already added
     for rank, gpu in gpus.items():