cherry-pick: fix stmatrix indexing

Author: rdspring1

csrc/device_lower/pass/index.cpp

@@ -2115,13 +2115,13 @@ void IndexLowering::handle(const LoadStoreOp* ldst) {
       switch (swizzle) {
         case MmaInputSmemSwizzle::None:
           out = hardCodedIndexGenerationForStMatrix(
-              ldst, for_loops_[0], m_tile, n_tile, m, n);
+              ldst, for_loops_[for_loops_.size() - 3], m_tile, n_tile, m, n);
           break;
         case MmaInputSmemSwizzle::B128:
         case MmaInputSmemSwizzle::B64:
         case MmaInputSmemSwizzle::B32:
           out = hardCodedIndexGenerationForStMatrixSwizzle(
-              ldst, for_loops_[0], m_tile, n_tile, m, n);
+              ldst, for_loops_[for_loops_.size() - 3], m_tile, n_tile, m, n);
           break;
         default:
           NVF_ERROR("Unsupported Swizzle Type for StMatrix");

csrc/scheduler/mma_utils.cpp

@@ -1315,21 +1315,23 @@ void scheduleStMatrixForMmaOutput(
       dataTypeSize(tv->dtype()) == 2,
       "we only support 16-bit types in stmatrix");
 
+  // NOTE: There can be iterDomains to left of the mma output if there is cta
+  // or warp tiling.
   if (tile_m == 16 && tile_n == 16) {
     // Let [M, N] be [64, 32]
     // After scheduleMmaOutputAllocation: [128(TIDx), 4, 2, 2]
     // [128(TIDx), 4(n), 2, 2] ->  [128(TIDx), 2(no), 2(ni), 2, 2]
     tv->split(-3, 2);
     // [128(TIDx), 2(no), 2(ni), 2, 2] -> [2(no), 128(TIDx), 2(ni), 2, 2]
-    tv->reorder({{-4, 0}});
+    tv->reorder({{-4, -5}});
     // [128(TIDx), 2(no), 2(ni), 2, 2] -> [2(no), 128(TIDx), 8 (vectorize)]
     tv->merge(-3);
     tv->merge(-2);
   } else if (tile_m == 16 && tile_n == 8) {
     // Let [M, N] be [64, 16]
     // After scheduleMmaOutputAllocation: [128(TIDx), 2, 2, 2]
     // [128(TIDx), 2, 2, 2] -> [2, 128(TIDx), 2, 2]
-    tv->reorder({{-3, 0}});
+    tv->reorder({{-3, -4}});
     // [2, 128(TIDx), 2, 2] -> [2, 128(TIDx), 4(vectorize)]
     tv->merge(-2);
   }