Updates for 3.1 (#932)

2023-04-29 09:34:27 -04:00
parent 6f8596ce3f
commit 7c04f95415
51 changed files with 1796 additions and 328 deletions
@@ -450,7 +450,7 @@ private :
  CUTLASS_DEVICE
  void consumer_wait(uint32_t stage, uint32_t phase, ConsumerToken barrier_token) {
    if (barrier_token == BarrierStatus::WaitAgain) {
-      consumer_wait(stage, phase);
+      full_barrier_ptr_[stage].wait(phase);
    }
  }

@@ -654,7 +654,7 @@ public :
    consumer_release(state.index());
  }

-protected:
+private:
  FullBarrier *full_barrier_ptr_ = nullptr;
  EmptyBarrier *empty_barrier_ptr_ = nullptr;
  Params params_;
@@ -976,6 +976,11 @@ public:
    ++stage_;
  }

+  CUTLASS_DEVICE
+  void advance() {
+    ++stage_;
+  }
+
 private:

  CUTLASS_DEVICE