Matmul scheduler - apply changes from #2488

- apply improvement in matmul instruction scheduling with loop rotation

Author: drzejan2

third_party/nvfuser/csrc/scheduler/matmul.cpp

@@ -604,12 +604,12 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
   acr->axis(-1)->parallelize(ParallelType::Vectorize);
   bcr->axis(-1)->parallelize(ParallelType::Vectorize);
 
-  //  0   1  2  3    4   5  6  7  8  9  10
-  // [Mo No Ko Mwo  Nwo Kw Mw Nw (Mi Ni Ki)]
+  //  0   1  2  3   4   5   6  7  8  9  10
+  // [Mo No Ko Kwo Mwo Nwo Mw Nw (Mi Ni Ki)]
   cc->axis(0)->parallelize(ParallelType::BIDx);
   cc->axis(1)->parallelize(ParallelType::BIDy);
-  cc->axis(3)->parallelize(ParallelType::TIDz);
-  cc->axis(4)->parallelize(ParallelType::TIDy);
+  cc->axis(4)->parallelize(ParallelType::TIDz);
+  cc->axis(5)->parallelize(ParallelType::TIDy);
 
   // Propagate mma output swizzle and parallelization down the DAG
   if (params.double_buffer_options.double_buffer_smem_write) {
@@ -640,6 +640,11 @@ void scheduleMatmul(Fusion* fusion, const MatmulParams& params) {
       scheduler_utils::BoundedDirectionalTransformPropagator::Options()
           .propagateParallelType()
           .propagateToBoundary());
+
+  if (params.double_buffer_options.double_buffer_smem_read &&
+      params.double_buffer_options.double_buffer_smem_write) {
+    scheduler_utils::rotateLoop(cc, 2, {acr, bcr});
+  }
 }
 
 } // namespace nvfuser

third_party/nvfuser/csrc/scheduler/matmul_heuristic.h

@@ -17,6 +17,9 @@ class MatmulParams : public HeuristicParams {
     int smem_double_buffer_stage = 2;
   };
 
+  //! Whether to rotate the ldmatrix out of the main loop
+  bool rotate_ldmatrix_out_of_main_loop = true;
+
   //! (Ampere+) Use cp.async to load operands.
   bool async_gmem_load_operands = false;