Merge remote-tracking branch 'origin/pw_scheduler_reference_find_patc…

…h' into HEAD
NVIDIA · Dec 11, 2024 · 1d021c7 · 1d021c7
2 parents 4ba5baa + b7f628f
commit 1d021c7
Show file tree

Hide file tree

Showing 12 changed files with 636 additions and 656 deletions.
diff --git a/csrc/codegen.cpp b/csrc/codegen.cpp
@@ -273,7 +273,16 @@ class CudaKernelGenerator : private kir::ConstIrVisitor {
 
   // Generates the kernel function declaration
   void genDeclaration(const std::string& kernel_name) {
-    code_ << "__global__ void " << kernel_name << "(";
+    code_ << "__global__ void ";
+    if (kernel_->hasManaged("cluster_dims")) {
+      auto cluster_dims =
+          kernel_->getManaged<std::tuple<int64_t, int64_t, int64_t>>(
+              "cluster_dims");
+      code_ << "__cluster_dims__(" << std::get<0>(cluster_dims) << ", "
+            << std::get<1>(cluster_dims) << ", " << std::get<2>(cluster_dims)
+            << ") ";
+    }
+    code_ << kernel_name << "(";
 
     std::unordered_set<Val*> unique_args;
 

diff --git a/csrc/scheduler/hopper_multi_matmul.cpp b/csrc/scheduler/hopper_multi_matmul.cpp
diff --git a/csrc/scheduler/hopper_multi_matmul.h b/csrc/scheduler/hopper_multi_matmul.h
@@ -178,20 +178,6 @@ class HopperMultipleMatmulScheduler : public MultipleMatmulScheduler {
 
   void setUpCircularBuffering();
 
-  //! Schedules the copy operation of output of a Mma op which resided in the
-  //! registers to shared memory.
-  void scheduleStMatrixForMmaOutput(
-      TensorView* tv,
-      MmaInputSmemSwizzle swizzle,
-      int64_t tile_m,
-      int64_t tile_n);
-
-  //! Schedules the copy operation of output of a Mma op which resided in the
-  //! shared memory to global memory.
-  void scheduleTMAStoreForMmaOutput(
-      TensorView* tv,
-      MmaInputSmemSwizzle swizzle);
-
   // Map TensorView's iterDomain to its ValGroup.
   // Then, find the MatmulDimRole for the ValGroup.
   // Return MatmulDimRole for IterDomain