fix mem fence (#2030)

Co-authored-by: yuzhai <[email protected]>
NVIDIA · Jan 8, 2025 · c506e16 · c506e16
1 parent 7494a18
commit c506e16
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/include/cutlass/pipeline/sm90_pipeline.hpp b/include/cutlass/pipeline/sm90_pipeline.hpp
@@ -316,6 +316,7 @@ class PipelineTmaAsync {
       cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
           storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, multicast_consumer_arrival_count);
     }
+    cutlass::arch::fence_barrier_init();
   }
 
   template<class ClusterShape, class InitBarriers, class InitMasks>
@@ -757,6 +758,7 @@ class PipelineTransactionAsync {
       cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(full_barrier_ptr), decltype(empty_barrier_ptr), Stages>(
           full_barrier_ptr, empty_barrier_ptr, params.producer_arv_count, params.consumer_arv_count);
     }
+    cutlass::arch::fence_barrier_init();
   }
 
   // Constructor
@@ -993,6 +995,7 @@ class PipelineAsync {
       cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
           storage.full_barrier_, storage.empty_barrier_, params.producer_arv_count, params.consumer_arv_count);
     }
+    cutlass::arch::fence_barrier_init();
   }
 
   template<class InitBarriers>
@@ -1249,6 +1252,7 @@ class OrderedSequenceBarrier {
         }
       }
     }
+    cutlass::arch::fence_barrier_init();
   }
 
   // Wait on a stage to be unlocked