[Mosaic GPU] Add a synchronization point at the end of the constructor for BarrierArray.

Without such a synchronization point, calls to `mbarrier_init`s may end up
happening after uses of the `mbarrier` being initialized---which is undefined
behaviour and leads to deadlocks.

This allows us to reenable the previously broken test cases.

PiperOrigin-RevId: 638246527

Author: bchetioui

jax/experimental/mosaic/gpu/utils.py

@@ -502,6 +502,7 @@ def __init__(self, num_barriers: int, arrival_count: int = 1):
     with once():
       for i in range(num_barriers):
         nvgpu.mbarrier_init(self.value, c(arrival_count, index), c(i, index))
+    gpu.barrier()
 
   def __iter__(self) -> Iterator["Barrier"]:
     for offset in range(self.num_barriers):

tests/mosaic/matmul_test.py

@@ -60,12 +60,6 @@ def test_matmul(self, m, k, n, stages, tile_m, tile_n, in_dtype):
     if n < tile_n:
       self.skipTest(f"No use in running a test with {n=} < {tile_n=}.")
 
-    # TODO(bchetioui): investigate why this test case fails with error
-    #  Illegal barrier arrive operation
-    # under memcheck.
-    if tile_m == 64 and tile_n == 64 and stages == 2:
-      self.skipTest("Broken test case---skipping.")
-
     try:
       matmul.verify(
           m,
@@ -102,12 +96,6 @@ def test_matmul_f32(self, m, k, n, stages, tile_m, tile_n, high_precision):
     if n < tile_n:
       self.skipTest(f"No use in running a test with {n=} < {tile_n=}.")
 
-    # TODO(bchetioui): investigate why this test case fails with error
-    #  Illegal barrier arrive operation
-    # under memcheck.
-    if tile_m == 64 and tile_n == 64 and stages == 2:
-      self.skipTest("Broken test case---skipping.")
-
     try:
       matmul.verify(
           m,