quack-kernels 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

quack/tile_scheduler.py

@@ -6,7 +6,7 @@ from enum import IntEnum
 
 import cutlass
 import cutlass.cute as cute
-from cutlass import Int32, Boolean, const_expr
+from cutlass import Int32, Float32, Boolean, const_expr
 
 import quack.utils as utils
 from quack.fast_math import FastDivmod
@@ -135,7 +135,7 @@ class TileScheduler:
         ip=None,
     ):
         self._current_work_linear_idx = current_work_linear_idx
-        self._num_tiles_executed = num_tiles_executed
+        self.num_tiles_executed = num_tiles_executed
         self._tile_count = tile_count
         self._scheduler_pipeline = scheduler_pipeline
         self._pipeline_state = pipeline_state
@@ -251,7 +251,7 @@ class TileScheduler:
         )
         tile_coord_mnkl = (pid_m, pid_n, None, batch_idx)
         if const_expr(not params.is_persistent):
-            is_valid = self._num_tiles_executed == 0
+            is_valid = self.num_tiles_executed == 0
         else:
             is_valid = self._current_work_linear_idx < cute.size(params.problem_shape_ncluster_mnl)
         return cutlass.utils.WorkTileInfo(tile_coord_mnkl, is_valid)
@@ -276,38 +276,6 @@ class TileScheduler:
                 current_work_linear_idx = cute.arch.shuffle_sync(current_work_linear_idx, 0)
             self._current_work_linear_idx = current_work_linear_idx
 
-    # We have to split broadcast_next_work and advance_to_next_work into two functions
-    # due to a bug in cute-dsl 4.2: https://github.com/NVIDIA/cutlass/issues/2647
-    @cute.jit
-    def broadcast_next_work(self, is_scheduler_warp: bool | Boolean = False, *, loc=None, ip=None):
-        """is_scheduler_warp should only be true for one warp in the whole cluster"""
-        params = self.params
-        if const_expr(params.is_persistent and params.tile_count_semaphore is not None):
-            current_work_linear_idx = self._current_work_linear_idx
-            if is_scheduler_warp:
-                self._scheduler_pipeline.producer_acquire(self._pipeline_state)
-                lane_idx = cute.arch.lane_idx()
-                if lane_idx < cute.size(params.cluster_shape_mn):
-                    # cute.printf("Producer bidx = {}, tidx = {}, after empty wait, idx = {}", bidx, tidx, current_work_linear_idx)
-                    if const_expr(cute.size(params.cluster_shape_mn) == 1):
-                        self._tile_count[self._pipeline_state.index] = current_work_linear_idx
-                        self._scheduler_pipeline.producer_commit(self._pipeline_state)
-                    else:
-                        peer_cta_rank_in_cluster = lane_idx
-                        mbar_ptr = self._scheduler_pipeline.producer_get_barrier(
-                            self._pipeline_state
-                        )
-                        cute.arch.mbarrier_arrive_and_expect_tx(
-                            mbar_ptr, 4, peer_cta_rank_in_cluster
-                        )
-                        utils.store_shared_remote(
-                            val=current_work_linear_idx,
-                            smem_ptr=self._tile_count.iterator + self._pipeline_state.index,
-                            mbar_ptr=mbar_ptr,
-                            peer_cta_rank_in_cluster=peer_cta_rank_in_cluster,
-                        )
-                    # cute.printf("Producer bidx = {}, tidx = {}, after full arrive", bidx, tidx)
-
     @cute.jit
     def advance_to_next_work(
         self,
@@ -319,6 +287,7 @@ class TileScheduler:
     ):
         tidx = cute.arch.thread_idx()[0]
         bidx = cute.arch.block_idx()[0]
+        bidz = cute.arch.block_idx()[2]
         params = self.params
         if const_expr(params.is_persistent):
             num_persistent_clusters = cute.arch.grid_dim()[2]
@@ -328,34 +297,60 @@ class TileScheduler:
                 if const_expr(advance_count > 1):
                     self._pipeline_state.advance_iters(advance_count - 1)
                 current_work_linear_idx = self._current_work_linear_idx
-                if not is_scheduler_warp:
-                    # if tidx % 64 == 0: cute.printf("bidx = {},tidx = {}, before full wait, idx = {}", bidx, tidx, current_work_linear_idx)
+                if is_scheduler_warp:
+                    self._scheduler_pipeline.producer_acquire(self._pipeline_state)
+                    lane_idx = cute.arch.lane_idx()
+                    if lane_idx < cute.size(params.cluster_shape_mn):
+                        # cute.printf("Producer bidx = {}, bidz = {}, tidx = {}, after empty wait, idx = {}", bidx, bidz, tidx, current_work_linear_idx)
+                        if const_expr(cute.size(params.cluster_shape_mn) == 1):
+                            self._tile_count[self._pipeline_state.index] = current_work_linear_idx
+                            self._scheduler_pipeline.producer_commit(self._pipeline_state)
+                        else:
+                            peer_cta_rank_in_cluster = lane_idx
+                            mbar_ptr = self._scheduler_pipeline.producer_get_barrier(
+                                self._pipeline_state
+                            )
+                            cute.arch.mbarrier_arrive_and_expect_tx(
+                                mbar_ptr, 4, peer_cta_rank_in_cluster
+                            )
+                            utils.store_shared_remote(
+                                val=current_work_linear_idx,
+                                smem_ptr=self._tile_count.iterator + self._pipeline_state.index,
+                                mbar_ptr=mbar_ptr,
+                                peer_cta_rank_in_cluster=peer_cta_rank_in_cluster,
+                            )
+                        # cute.printf("Producer bidx = {}, bidz = {}, tidx = {}, after full arrive", bidx, bidz, tidx)
+                else:
+                    # if tidx % 32 == 0: cute.printf("bidx = {}, bidz = {}, tidx = {}, before full wait, idx = {}", bidx, bidz, tidx, current_work_linear_idx)
                     self._scheduler_pipeline.consumer_wait(self._pipeline_state)
-                    # if tidx % 64 == 0: cute.printf("bidx = {}, tidx = {}, after full wait, idx = {}", bidx, tidx, current_work_linear_idx)
+                    # if tidx % 32 == 0: cute.printf("bidx = {}, bidz = {}, tidx = {}, after full wait, idx = {}", bidx, bidz, tidx, current_work_linear_idx)
                     current_work_linear_idx = self._tile_count[self._pipeline_state.index]
-                    # if tidx % 64 == 0: cute.printf("bidx = {}, tidx = {}, after smem read, idx = {}", bidx, tidx, current_work_linear_idx)
+                    # if tidx % 32 == 0: cute.printf("bidx = {}, bidz = {}, tidx = {}, after smem read, idx = {}", bidx, bidz, tidx, current_work_linear_idx)
+                    # Need this fence since the STAS from the producer is using the async proxy.
+                    # Without this, we get race condition / deadlock.
+                    if const_expr(cute.size(params.cluster_shape_mn) > 1):
+                        cute.arch.fence_proxy(
+                            cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta
+                        )
                     cute.arch.sync_warp()
                     with cute.arch.elect_one():
-                        # if tidx % 64 == 0: cute.printf("bidx = {}, tidx = {}, before empty arrive", bidx, tidx)
+                        # if tidx % 32 == 0: cute.printf("bidx = {}, bidz = {}, tidx = {}, before empty arrive", bidx, bidz, tidx)
                         self._scheduler_pipeline.consumer_release(self._pipeline_state)
-                        # if tidx % 64 == 0: cute.printf("bidx = {}, tidx = {}, after empty arrive", bidx, tidx)
+                        # if tidx == 320: cute.printf("bidx = {}, bidz = {}, tidx = {}, idx = {}, after empty arrive", bidx, bidz, tidx, current_work_linear_idx)
+                    # if tidx == 320: cute.printf("bidx = {}, bidz = {}, tidx = {}, idx = {}, after empty arrive", bidx, bidz, tidx, current_work_linear_idx)
                 self._current_work_linear_idx = current_work_linear_idx
                 self._pipeline_state.advance()
-        self._num_tiles_executed += Int32(advance_count)
+        self.num_tiles_executed += Int32(advance_count)
 
     def producer_tail(self):
         if const_expr(self.params.is_persistent and self.params.tile_count_semaphore is not None):
             self._scheduler_pipeline.producer_tail(self._pipeline_state)
 
-    @property
-    def num_tiles_executed(self) -> Int32:
-        return self._num_tiles_executed
-
     def __extract_mlir_values__(self):
         values, self._values_pos = [], []
         for obj in [
             self._current_work_linear_idx,
-            self._num_tiles_executed,
+            self.num_tiles_executed,
             self._tile_count,
             self._scheduler_pipeline,
             self._pipeline_state,
@@ -371,7 +366,7 @@ class TileScheduler:
         for obj, n_items in zip(
             [
                 self._current_work_linear_idx,
-                self._num_tiles_executed,
+                self.num_tiles_executed,
                 self._tile_count,
                 self._scheduler_pipeline,
                 self._pipeline_state,
@@ -390,7 +385,7 @@ def triangular_idx_to_coord(idx: Int32) -> Tuple[Int32, Int32]:
     Convert a triangular index to 2D coordinates.
     This is used to convert the linear index to 2D coordinates for triangular matrices.
     """
-    row = utils.ceil((cute.math.sqrt(2 * idx + 2.25, fastmath=True) - 0.5)) - 1
+    row = utils.ceil((utils.sqrt(2 * idx + 2.25) - 0.5)) - 1
     col = idx - (row * (row + 1)) // 2
     return row, col
 
@@ -402,7 +397,7 @@ class TriangularTileScheduler(TileScheduler):
     class Params(ParamsBase):
         problem_shape_ncluster_mnl: cute.Shape
         num_clusters_per_problem_divmod: FastDivmod
-        group_size_inv_f32: cutlass.Float32
+        group_size_inv_f32: Float32
         num_groups_regular: Int32
         group_size_divmod: FastDivmod
         group_size_tail_divmod: FastDivmod
@@ -433,7 +428,7 @@ class TriangularTileScheduler(TileScheduler):
             return TriangularTileScheduler.Params(
                 problem_shape_ncluster_mnl,
                 FastDivmod.create(num_clusters_per_problem),
-                cutlass.Float32(1.0 / group_size),
+                Float32(1.0 / group_size),
                 num_groups_regular,
                 FastDivmod.create(group_size),
                 # Don't divide by 0
@@ -524,8 +519,7 @@ class TriangularTileScheduler(TileScheduler):
         group_size = params.group_size_divmod.divisor
         group_id = (
             utils.ceil(
-                (cute.math.sqrt(2 * cluster_id_in_problem + 2.25, fastmath=True) - 0.5)
-                * params.group_size_inv_f32
+                (utils.sqrt(2 * cluster_id_in_problem + 2.25) - 0.5) * params.group_size_inv_f32
             )
             - 1
         )
@@ -562,7 +556,7 @@ class TriangularTileScheduler(TileScheduler):
         pid_n = cid_n * params.cluster_shape_mn[1] + bidx_in_cluster[1]
         tile_coord_mnkl = (pid_m, pid_n, None, bidz)
         if const_expr(not params.is_persistent):
-            is_valid = self._num_tiles_executed == 0
+            is_valid = self.num_tiles_executed == 0
         else:
             is_valid = (
                 self._current_work_linear_idx
@@ -681,7 +675,7 @@ class VarlenMTileScheduler(TileScheduler):
         ip=None,
     ):
         self._current_work_linear_idx = current_work_linear_idx
-        self._num_tiles_executed = num_tiles_executed
+        self.num_tiles_executed = num_tiles_executed
         self._current_batch_idx = current_batch_idx
         self._num_work_idx_before_cur_batch = num_work_idx_before_cur_batch
         self._tile_count = tile_count
@@ -878,25 +872,25 @@ class VarlenMTileScheduler(TileScheduler):
         pid_n = cid_n * params.cluster_shape_mn[1] + bidx_in_cluster[1]
         tile_coord_mnkl = (pid_m, pid_n, None, batch_idx)
         if const_expr(not params.is_persistent):
-            is_valid = self._num_tiles_executed == 0 and batch_idx < num_batch
+            is_valid = self.num_tiles_executed == 0 and batch_idx < num_batch
         else:
             is_valid = batch_idx < num_batch
         return cutlass.utils.WorkTileInfo(tile_coord_mnkl, is_valid)
 
     @cute.jit
-    def fetch_next_work(self, is_scheduler_warp: bool | Boolean, *, loc=None, ip=None):
+    def fetch_next_work(self, is_scheduler_warp: bool | Boolean = False, *, loc=None, ip=None):
         """is_scheduler_warp should only be true for one warp in the whole cluster"""
         if const_expr(self.params.tile_count_semaphore is not None):
             params = self.params
             current_work_linear_idx = self._current_work_linear_idx
             if is_scheduler_warp:
                 if cute.arch.lane_idx() == 0:
-                    # cute.printf("before atomicadd, tidx = {}, idx = {}", cute.arch.thread_idx()[0], current_work_linear_idx)
+                    # cute.printf("before atomicadd, tidx = {}, bidz = {}, idx = {}", cute.arch.thread_idx()[0], cute.arch.block_idx()[2], current_work_linear_idx)
                     num_persistent_clusters = cute.arch.grid_dim()[2]
                     current_work_linear_idx = num_persistent_clusters + utils.atomic_add_i32(
                         1, params.tile_count_semaphore
                     )
-                    # cute.printf("after atomicadd, tidx = {}, idx = {}", cute.arch.thread_idx()[0], current_work_linear_idx)
+                    # cute.printf("after atomicadd, tidx = {}, bidz = {}, idx = {}", cute.arch.thread_idx()[0], cute.arch.block_idx()[2], current_work_linear_idx)
                 # lane 0 already has the right tile_idx, just need to broadcast
                 current_work_linear_idx = cute.arch.shuffle_sync(current_work_linear_idx, 0)
             self._current_work_linear_idx = current_work_linear_idx
@@ -905,7 +899,7 @@ class VarlenMTileScheduler(TileScheduler):
         values, self._values_pos = [], []
         for obj in [
             self._current_work_linear_idx,
-            self._num_tiles_executed,
+            self.num_tiles_executed,
             self._current_batch_idx,
             self._num_work_idx_before_cur_batch,
             self._tile_count,
@@ -923,7 +917,7 @@ class VarlenMTileScheduler(TileScheduler):
         for obj, n_items in zip(
             [
                 self._current_work_linear_idx,
-                self._num_tiles_executed,
+                self.num_tiles_executed,
                 self._current_batch_idx,
                 self._num_work_idx_before_cur_batch,
                 self._tile_count,