Support PDL in sm90_gemm_array_tma_warpspecialized_cooperative

Author: HydraQYH

include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp

@@ -610,6 +610,9 @@ class GemmUniversal<
       if (producer_warp_role == ProducerWarpRole::Scheduler) {
         // GroupScheduler requires a producer warp to iterate over the group infos and push
         // the work tile infos to the downstream pipelines.
+        #ifdef CUTLASS_ENABLE_GDC_FOR_SM90
+        cutlass::arch::wait_on_dependent_grids();
+        #endif
         if constexpr (cute::is_same_v<SchedulerTag, GroupScheduler>) {
           do {
             auto [next_work_tile_info, increment_pipe] = scheduler.advance_to_next_work(tile_scheduler_pipeline, tile_scheduler_pipe_producer_state);
@@ -623,6 +626,9 @@ class GemmUniversal<
       }
       // Mainloop Producer Warp
       else if (producer_warp_role == ProducerWarpRole::Mainloop) {
+        #ifdef CUTLASS_ENABLE_GDC_FOR_SM90
+        cutlass::arch::wait_on_dependent_grids();
+        #endif
         int32_t curr_batch = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx;
         int32_t const mock_l_coord = 0;
         int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x);
@@ -787,6 +793,9 @@ class GemmUniversal<
       } // Mainloop Auxiliary Load Producer Warp End
       // Epilogue Producer Warp
       else if (producer_warp_role == ProducerWarpRole::Epilogue && collective_epilogue.is_producer_load_needed()) {
+        #ifdef CUTLASS_ENABLE_GDC_FOR_SM90
+        cutlass::arch::wait_on_dependent_grids();
+        #endif
         int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x);
         int32_t const sm_count = params.hw_info.sm_count;
 
@@ -990,6 +999,11 @@ class GemmUniversal<
 
         // Get next work tile
         auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info, tile_scheduler_pipeline, tile_scheduler_pipe_consumer_state);
+        #ifdef CUTLASS_ENABLE_GDC_FOR_SM90
+        if (!next_work_tile_info.is_valid()) {
+          cutlass::arch::launch_dependent_grids();
+        }
+        #endif
         work_tile_info = next_work_tile_info;
         if (increment_pipe) {
           ++tile_scheduler_pipe_consumer_state;