Fix workspace allocation and reuse in matmul_all_reduce

benchmark/ops/bench_matmul_all_reduce.py

@@ -33,18 +33,10 @@ def matmul_all_reduce(state, ctx):
     state.set_flops(2 * M * N * K)
 
     def _run():
-        workspace.prepared = False
-        ctx.ops.matmul_all_reduce(C, A, B, config=config, workspace=workspace)
-
-    def _preamble():
-        C.zero_()
-        if workspace.locks is not None:
-            workspace.locks.zero_()
-        if workspace.aux_buffer is not None:
-            workspace.aux_buffer.zero_()
-        workspace.prepared = True
-
-    state.exec(_run, preamble_fn=_preamble)
+        nonlocal workspace
+        workspace = ctx.ops.matmul_all_reduce(C, A, B, config=config, workspace=workspace)
+
+    state.exec(_run)
 
 
 if __name__ == "__main__":

benchmark/ops/bench_matmul_all_reduce_stages.py

@@ -89,18 +89,10 @@ def fused_gemm_allreduce(state, ctx):
     state.set_flops(2 * M * N * K_local)
 
     def _run():
-        workspace.prepared = False
-        ctx.ops.matmul_all_reduce(C, A, B, config=config, workspace=workspace)
-
-    def _preamble():
-        C.zero_()
-        if workspace.locks is not None:
-            workspace.locks.zero_()
-        if workspace.aux_buffer is not None:
-            workspace.aux_buffer.zero_()
-        workspace.prepared = True
-
-    state.exec(_run, preamble_fn=_preamble)
+        nonlocal workspace
+        workspace = ctx.ops.matmul_all_reduce(C, A, B, config=config, workspace=workspace)
+
+    state.exec(_run)
 
 
 if __name__ == "__main__":

iris/mem/triton/context.py

@@ -731,7 +731,7 @@ def all_reduce_spinlock(self, tile: Tile, dst_view: TensorView, locks):
             self.atomic_xchg(locks + tile_id, 0, to_rank=dest_rank, sem="release", scope="sys")
 
     @triton.jit
-    def all_reduce_one_shot(self, tile: Tile, src_view: TensorView, dst_view: TensorView, locks):
+    def all_reduce_one_shot(self, tile: Tile, src_view: TensorView, dst_view: TensorView, locks, generation=1):
         """
         Tile-level all-reduce using one-shot algorithm.
 
@@ -755,7 +755,10 @@ def all_reduce_one_shot(self, tile: Tile, src_view: TensorView, dst_view: Tensor
         for remote_rank in range(self.world_size):
             if remote_rank != self.rank:
                 lock_ptr = locks + tile_id
-                while self.atomic_add(lock_ptr, 0, to_rank=remote_rank, sem="acquire", scope="sys") != 1:
+                while (
+                    self.atomic_cas(lock_ptr, generation, generation, to_rank=remote_rank, sem="acquire", scope="sys")
+                    != generation
+                ):
                     pass
                 partial = self.load(src_tile_ptr, from_rank=remote_rank, mask=mask)
                 acc += partial.to(acc_dtype)
@@ -799,7 +802,7 @@ def all_reduce_ring(self, tile: Tile, src_view: TensorView, dst_view: TensorView
                 tl.store(dst_tile_ptr, remote_result, mask=mask)
 
     @triton.jit
-    def all_reduce_two_shot(self, tile: Tile, src_view: TensorView, dst_view: TensorView, locks):
+    def all_reduce_two_shot(self, tile: Tile, src_view: TensorView, dst_view: TensorView, locks, generation=1):
         """
         Tile-level all-reduce using two-shot algorithm with work distribution.
 
@@ -826,7 +829,12 @@ def all_reduce_two_shot(self, tile: Tile, src_view: TensorView, dst_view: Tensor
             for remote_rank in range(self.world_size):
                 if remote_rank != self.rank:
                     lock_ptr = locks + tile_id
-                    while self.atomic_add(lock_ptr, 0, to_rank=remote_rank, sem="acquire", scope="sys") != 1:
+                    while (
+                        self.atomic_cas(
+                            lock_ptr, generation, generation, to_rank=remote_rank, sem="acquire", scope="sys"
+                        )
+                        != generation
+                    ):
                         pass
                     partial = self.load(src_tile_ptr, from_rank=remote_rank, mask=mask)
                     acc += partial.to(acc_dtype)