Move tile rotation to top of IPU Jacobi loop body. (#52)

Allows to optimize out one on-tile-copy, saving an additional 10% of cycles.
graphcore-research · Oct 19, 2023 · acabc7d · acabc7d
1 parent b891bd9
commit acabc7d
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/tessellate_ipu/linalg/tile_linalg_jacobi.py b/tessellate_ipu/linalg/tile_linalg_jacobi.py
@@ -221,11 +221,12 @@ def ipu_jacobi_eigh_body(idx: Array, inputs: Tuple[TileShardedArray, ...]) -> Tu
     halfN = Apcols.shape[0]
 
     with jax.named_scope("jacobi_eigh"):
-        # with jax.named_scope("Apqcols_rotation"):
-        #     Apcols, Aqcols = tile_rotate_columns(Apcols, Aqcols)
-        # with jax.named_scope("Vpqcols_rotation"):
-        #     Vpcols, Vqcols = tile_rotate_columns(Vpcols, Vqcols)
-        # Apcols, Aqcols, Vpcols, Vqcols = tile_data_barrier(Apcols, Aqcols, Vpcols, Vqcols)
+        with jax.named_scope("Apqcols_rotation"):
+            Apcols, Aqcols = tile_rotate_columns(Apcols, Aqcols)
+        with jax.named_scope("Vpqcols_rotation"):
+            Vpcols, Vqcols = tile_rotate_columns(Vpcols, Vqcols)
+        # Barrier, to make we sync. both set of tiles A and V and force fused comms.
+        Apcols, Aqcols, Vpcols, Vqcols = tile_data_barrier(Apcols, Aqcols, Vpcols, Vqcols)
 
         # Sharded constant with p/q indices to ignore in second update stage.
         with jax.named_scope("rotset_index_ignored"):
@@ -274,13 +275,12 @@ def ipu_jacobi_eigh_body(idx: Array, inputs: Tuple[TileShardedArray, ...]) -> Tu
             Vqcols,
         )
 
-        # Barrier, to make we sync. both set of tiles A and V
-        Apcols, Aqcols, Vpcols, Vqcols = tile_data_barrier(Apcols, Aqcols, Vpcols, Vqcols)
-        # Move columns between tiles following Jacobi rotation pattern. 2*N commns per tile.
-        with jax.named_scope("Apqcols_rotation"):
-            Apcols, Aqcols = tile_rotate_columns(Apcols, Aqcols)
-        with jax.named_scope("Vpqcols_rotation"):
-            Vpcols, Vqcols = tile_rotate_columns(Vpcols, Vqcols)
+        # Apcols, Aqcols, Vpcols, Vqcols = tile_data_barrier(Apcols, Aqcols, Vpcols, Vqcols)
+        # # Move columns between tiles following Jacobi rotation pattern. 2*N commns per tile.
+        # with jax.named_scope("Apqcols_rotation"):
+        #     Apcols, Aqcols = tile_rotate_columns(Apcols, Aqcols)
+        # with jax.named_scope("Vpqcols_rotation"):
+        #     Vpcols, Vqcols = tile_rotate_columns(Vpcols, Vqcols)
         return Apcols, Aqcols, Vpcols, Vqcols