Repository: NVIDIA/Fuser
Pull request #4215: warp specializied tma persistent kernel, step-1b, add scheduleInnerOuterWarpSpecializedTmaKernel 
Author: login: liqiangxl
Created: 2025-04-09T15:58:06Z | Merged: 2025-04-14T20:35:16Z

This PR follows following https://github.com/NVIDIA/Fuser/pull/4201. It is step-1b of implementing warp specializied tma persistent kernel described in the [design doc](https://docs.google.com/document/d/1gBlTQZTLyWsCR4IzHDBAs8D8uOgppDSSYl9BzvE9mD4/edit?tab=t.0#heading=h.j8ae9cdb7l0y).
**Changes:** minor changes to use a new scheduling function which is a copy of the multi-wave version.I duplicated it here, then the changes in following PRs will be clearly marked.

- Add variable `tma_warp_specialized` to scheduler heuristics, set to `true` when `WarpSpecializedNormalization` is enabled.
- Duplicate `scheduleInnerOuterPersistentKernel()` as `scheduleInnerOuterWarpSpecializedTmaKernel()`, no actual logic change in the code except for clean up.



Diff:
diff --git a/csrc/scheduler/normalization_inner_outer.cpp b/csrc/scheduler/normalization_inner_outer.cpp
index ad0497230e5..fcb7260501d 100644
--- a/csrc/scheduler/normalization_inner_outer.cpp
+++ b/csrc/scheduler/normalization_inner_outer.cpp
@@ -1171,6 +1171,7 @@ std::unique_ptr<ReductionParams> getInnerOuterPersistentHeuristics(
         hp.threads_per_block_max,
         buffer_params.project_to_input,
         runtime_info.getIndexType());
+    rparams->tma_warp_specialized = true;
   } else {
     rparams = innerOuterPersistentHeuristic(
         properties.total_iteration_numel,
@@ -1487,6 +1488,250 @@ void scheduleInnerOuterPersistentKernel(
   inlineMost();
 }
 
+void scheduleTmaWarpSpecializedOuter(
+    Fusion* fusion,
+    const ReductionParams* rparams,
+    const std::vector<TensorView*>& outer_reduction_tvs,
+    std::vector<TensorView*>& cached_gmem,
+    std::vector<TensorView*>& cached_gmem_reload,
+    std::vector<TensorView*>& outer_reference_tvs,
+    std::unordered_set<TensorView*>& boundaryNodesSet) {
+  auto mergeReductionOrIterDomains = [](TensorView* tv, bool mergeReduction) {
+    int prev_i = -1;
+    for (int i = static_cast<int>(tv->nDims()) - 1; i >= 0; i--) {
+      if (mergeReduction == tv->axis(i)->isReduction()) {
+        if (prev_i == -1) {
+          prev_i = i;
+        } else {
+          tv->merge(i, prev_i);
+          prev_i = i;
+        }
+      }
+    }
+  };
+  for (auto& outer_reduction_tv : outer_reduction_tvs) {
+    // Similar to the inner reduction, we need to reorder the outer reduction tv
+    // when there are view operations.
+    if (!ir_utils::getViewOps(fusion).empty()) {
+      // Reorder reference_tv after propagating the view operation. This will
+      // reorder for better merging.
+      outer_reduction_tv->reorder(
+          scheduler_utils::domainReorderAsLogicalMap(outer_reduction_tv));
+    }
+
+    // merge tensorview to [reduction, iteraiton] domains
+    mergeReductionOrIterDomains(outer_reduction_tv, true);
+    mergeReductionOrIterDomains(outer_reduction_tv, false);
+
+    // First-stage of outer reduction
+    outer_reduction_tv->split(0, rparams->lparams.gdimy());
+
+    TensorView* partialResult = outer_reduction_tv->rFactor({0});
+    partialResult->cacheBefore();
+    partialResult->setMemoryType(MemoryType::Global);
+    TensorView* partialResultReload = partialResult->cacheAfter();
+
+    boundaryNodesSet.insert(partialResultReload);
+    cached_gmem.emplace_back(partialResult);
+    cached_gmem_reload.emplace_back(partialResultReload);
+
+    // Second-stage of outer reduction
+    // reduction domain, [I1/TIDy, TIDy]
+    outer_reduction_tv->split(0, rparams->lparams.bdimy());
+    outer_reduction_tv->axis(1)->parallelize(ParallelType::TIDy);
+    // iteration domain, [BIDy, TIDx, Vect]
+    int axisID = -1;
+    if (rparams->vectorization_factor_outer > 1) {
+      outer_reduction_tv->split(axisID, rparams->vectorization_factor_outer);
+      outer_reduction_tv->axis(axisID--)->parallelize(ParallelType::Vectorize);
+    }
+
+    if (rparams->lparams.bdimx() > 1) {
+      outer_reduction_tv->split(axisID, rparams->lparams.bdimx());
+      outer_reduction_tv->axis(axisID--)->parallelize(ParallelType::TIDx);
+    }
+
+    if (rparams->combined_split_grid_inner_dim) {
+      outer_reduction_tv->split(
+          axisID, NamedScalar::getParallelDim(ParallelType::BIDy));
+    }
+
+    outer_reduction_tv->axis(axisID--)->parallelize(ParallelType::BIDy);
+
+    auto outer_reference_tv =
+        reduction_scheduler_utils::sortAndRFactor(outer_reduction_tv);
+    outer_reference_tvs.emplace_back(outer_reference_tv);
+  }
+}
+
+void scheduleTmaWarpSpecializedInnerOuter(
+    Fusion* fusion,
+    const ReductionParams* rparams) {
+  FusionGuard fg(fusion);
+
+  // Grab the reduction, input, and output tensor views. dummy_outputs are
+  // helper tensors for persistent buffer projection.
+  std::vector<TensorView*> dummy_outputs, cached_inputs, reduction_tvs,
+      smem_consumers;
+  std::vector<std::pair<TensorView*, TensorView*>> cached_outputs;
+  normalization_scheduler_utils::beforeSchedule(
+      fusion,
+      rparams,
+      dummy_outputs,
+      cached_inputs,
+      reduction_tvs,
+      smem_consumers,
+      cached_outputs);
+
+  // split reduction_tvs into inner and outer reduction_tvs
+  std::vector<TensorView*> inner_reduction_tvs, outer_reduction_tvs;
+  for (auto tv : reduction_tvs) {
+    if (scheduler_utils::isFastestDimReduction(tv)) {
+      inner_reduction_tvs.emplace_back(tv);
+    } else {
+      outer_reduction_tvs.emplace_back(tv);
+    }
+  }
+  NVF_ERROR(
+      !inner_reduction_tvs.empty(),
+      "schedulePersistentKernelInnerOuter is called but no inner reduction is found.");
+  NVF_ERROR(
+      !outer_reduction_tvs.empty(),
+      "schedulePersistentKernelInnerOuter is called but no outer reduction is found.");
+
+  // schedule inner reduction, only schedule the first inner reduction tv,
+  // then will be propagated to other inner reduction tvs.
+  TensorView* inner_reference_tv =
+      normalization_scheduler_utils::scheduleReductionGeneral(
+          fusion,
+          rparams,
+          inner_reduction_tvs,
+          InnerOuterPersistentKernelScheduler::schedulerType());
+
+  // schedule outer reduction, schedule all the outer reduction tvs since we
+  // need to store the intermediate results.
+  std::vector<TensorView*> cached_gmem;
+  std::vector<TensorView*> cached_gmem_reload;
+  std::vector<TensorView*> outer_reference_tvs;
+  std::unordered_set<TensorView*> boundaryNodesSet;
+  scheduleTmaWarpSpecializedOuter(
+      fusion,
+      rparams,
+      outer_reduction_tvs,
+      cached_gmem,
+      cached_gmem_reload,
+      outer_reference_tvs,
+      boundaryNodesSet);
+
+  // Propagate inner reduction and outer reductions
+  for (auto output : dummy_outputs) {
+    fusion->addOutput(output);
+  }
+
+  const bool is_unroll_or_vectorization = rparams->isUnrolled();
+  const bool is_vectorize =
+      rparams->vectorize_inner_reduction || rparams->vectorize_iter_dom;
+  const bool is_outer_grid_persistence = rparams->persistent_kernel &&
+      rparams->cross_grid_inner_reduction && !rparams->fastest_dim;
+
+  // Propagate inner reduction. There is a cutoff at boundaryNodesSet, so this
+  // propagation will not propagate to the final outer reduction.
+  reduction_scheduler_utils::propagateTransformation(
+      inner_reference_tv, boundaryNodesSet);
+  reduction_scheduler_utils::propagateRFactor(
+      inner_reference_tv, inner_reduction_tvs[0], inner_reduction_tvs);
+
+  // Don't allow parallelization propagation goes through boundaryNodesSet
+  const auto& selected_tvs_inner =
+      scheduler_utils::getAllTvsFrom(inner_reduction_tvs, boundaryNodesSet);
+  const auto& unroll_vectorizable_cached_tvs =
+      reduction_scheduler_utils::getCachedTvsToUnrollOrVectorize(
+          inner_reference_tv, is_vectorize, cached_inputs, cached_outputs);
+  reduction_scheduler_utils::propagateParallelization(
+      inner_reduction_tvs[0],
+      inner_reference_tv,
+      is_unroll_or_vectorization,
+      is_outer_grid_persistence,
+      inner_reduction_tvs,
+      unroll_vectorizable_cached_tvs,
+      {selected_tvs_inner.begin(), selected_tvs_inner.end()});
+
+  // Propagate outer reduction. Each outer reduction is connected with its
+  // cached_gmem and output, since we added all the cached_gmem to the
+  // boundaryNodesSet, the transformation from one outer reduction can't
+  // propagate to other outer reductions due to the cutoff at
+  // boundaryNodesSet. Thus, we need a loop to initiate the propagation from
+  // each outer reduction. Don't allow parallelization propagation goes
+  // through cached_gmem, see issue 246.
+  for (long unsigned int i = 0; i < outer_reference_tvs.size(); i++) {
+    const auto& selected_tvs_outer = scheduler_utils::getAllTvsFrom(
+        {outer_reduction_tvs[i]}, {cached_gmem[i]});
+    reduction_scheduler_utils::propagateTransformation(
+        outer_reference_tvs[i], boundaryNodesSet);
+    const auto& unroll_vectorizable_cached_tvs =
+        reduction_scheduler_utils::getCachedTvsToUnrollOrVectorize(
+            outer_reference_tvs[i],
+            is_vectorize,
+            cached_inputs,
+            cached_outputs);
+    reduction_scheduler_utils::propagateParallelization(
+        outer_reduction_tvs[i],
+        outer_reference_tvs[i],
+        is_unroll_or_vectorization,
+        is_outer_grid_persistence,
+        outer_reduction_tvs,
+        unroll_vectorizable_cached_tvs,
+        {selected_tvs_outer.begin(), selected_tvs_outer.end()});
+  }
+
+  // special vectorization of temp gmem, vectorization_factor_tmp_gmem_write
+  // is guaranteed to be smaller or equal to input vectorization factor.
+  if (rparams->vectorization_factor_tmp_gmem_write > 1) {
+    for (auto tv : cached_gmem) {
+      NVF_ERROR(
+          rparams->vectorization_factor_tmp_gmem_write <=
+              rparams->unroll_factor_inner_reduction,
+          "vectorization factor of temp gmem write should be smaller than that of inner reduction.")
+      if (rparams->vectorization_factor_tmp_gmem_write <
+          rparams->unroll_factor_inner_reduction) {
+        tv->split(-1, rparams->vectorization_factor_tmp_gmem_write);
+      }
+      tv->axis(-1)->parallelize(ParallelType::Vectorize);
+    }
+  }
+  // vectorization propagate through propagateParallelization only works for
+  // input and output tensors. propagate vectorization to cached_gmem_reload
+  // directly from output tv using parallelizeAllLike. must propagate
+  // seperaely for different tvs as outer reductions are transformed
+  // seperately.
+  if (rparams->vectorization_factor_outer > 1) {
+    for (auto tv : cached_gmem_reload) {
+      auto output_tvs = ir_utils::outputTvsOf(tv);
+      NVF_ERROR(
+          !output_tvs.empty(),
+          "cached_gmem_reload should have at least one output tensor.")
+      scheduler_utils::parallelizeAllLike(
+          output_tvs[0],
+          -1,
+          {cached_gmem_reload.begin(), cached_gmem_reload.end()},
+          {ParallelType::Vectorize});
+    }
+  }
+
+  // Needs special handling of vectorized loading from shared memory due to
+  // potential different data types of inputs and shared memory tensor.
+  if (is_vectorize) {
+    reduction_scheduler_utils::sharedMemoryConsumerVectorization(
+        smem_consumers, rparams->unroll_factor_inner_reduction);
+  }
+
+  // Remove dummy outputs as they can inadvertently affect CA positions
+  for (auto output : dummy_outputs) {
+    fusion->removeOutput(output);
+  }
+  inlineMost();
+}
+
 } // namespace
 
 bool InnerOuterPersistentKernelScheduler::canScheduleCompileTime(
@@ -1768,6 +2013,10 @@ void InnerOuterPersistentKernelScheduler::schedule(
       rparams != nullptr && rparams->scheduler_type == schedulerType(),
       "Incorrect parameters sent to InnerOuterPersistentKernelScheduler::schedule",
       params);
-  scheduleInnerOuterPersistentKernel(fusion, rparams);
+  if (rparams->tma_warp_specialized) {
+    scheduleTmaWarpSpecializedInnerOuter(fusion, rparams);
+  } else {
+    scheduleInnerOuterPersistentKernel(fusion, rparams);
+  }
 }
 } // namespace nvfuser
diff --git a/csrc/scheduler/reduction_heuristic.h b/csrc/scheduler/reduction_heuristic.h
index e4e08e2f28c..833c1efe618 100644
--- a/csrc/scheduler/reduction_heuristic.h
+++ b/csrc/scheduler/reduction_heuristic.h
@@ -136,6 +136,9 @@ class ReductionParams : public HeuristicParams {
   // in outer reduction part of inner-outer persistent scheduler, may further
   // split inner dim by grid
   bool combined_split_grid_inner_dim = false;
+
+  // TMA warp specialized, only used in inner-outer persistent scheduler
+  bool tma_warp_specialized = false;
   // partial result of outer reduction is written to gmem then read back in a
   // different parallel pattern set the vectorization factor of its read and
   // write


Reviews:
user: login: liqiangxl
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T18:51:45Z
user: login: liqiangxl
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T18:54:34Z
user: login: jjsjann123
body: stamp for code prep
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-14T17:20:27Z
user: login: rdspring1
body: Is there a reason to use `ReductionParams` for so many schedulers?
Should `scheduleTmaWarpSpecializedOuter` go in a new cpp file?

Creating a new warp specialized scheduler seems like an opportunity to reduce technical debt.

state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T17:33:26Z

Review comments:
diff_hunk: @@ -1768,6 +2013,10 @@ void InnerOuterPersistentKernelScheduler::schedule(
       rparams != nullptr && rparams->scheduler_type == schedulerType(),
       "Incorrect parameters sent to InnerOuterPersistentKernelScheduler::schedule",
       params);
-  scheduleInnerOuterPersistentKernel(fusion, rparams);
+  if (rparams->tma_warp_specialized) {
+    scheduleInnerOuterWarpSpecializedTmaKernel(fusion, rparams);
path: csrc/scheduler/normalization_inner_outer.cpp
user: login: liqiangxl
body: `scheduleInnerOuterWarpSpecializedTmaKernel` is a duplication of `scheduleInnerOuterPersistentKernel`, will be modified in following PRs.
created_at: 2025-04-09T18:51:45Z
author_association: COLLABORATOR
diff_hunk: @@ -1487,6 +1488,250 @@ void scheduleInnerOuterPersistentKernel(
   inlineMost();
 }
 
+void scheduleTmaWarpSpecializedOuter(
path: csrc/scheduler/normalization_inner_outer.cpp
user: login: liqiangxl
body: `scheduleTmaWarpSpecializedOuter()` is a duplication of `scheduleReductionCombinedOuter()` with some clean ups e.g. removing code handling `rparams->multiple_reds_per_blk` since it will not be used for TMA warp specialized kernel. 
created_at: 2025-04-09T18:54:34Z
author_association: COLLABORATOR
diff_hunk: @@ -1768,6 +2013,10 @@ void InnerOuterPersistentKernelScheduler::schedule(
       rparams != nullptr && rparams->scheduler_type == schedulerType(),
       "Incorrect parameters sent to InnerOuterPersistentKernelScheduler::schedule",
       params);
-  scheduleInnerOuterPersistentKernel(fusion, rparams);
+  if (rparams->tma_warp_specialized) {
+    scheduleInnerOuterWarpSpecializedTmaKernel(fusion, rparams);
path: csrc/scheduler/normalization_inner_outer.cpp
user: login: jjsjann123
body: nitpick: can we keep the name of the function to be more consistent with the variable in params?
something like `scheduleTmaWarpSpecializedInnerOuter`, similar to how the other function is named as `scheduleTmaWarpSpecializedOuter`?
created_at: 2025-04-14T17:19:54Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-09T15:58:58Z
body: <!-- llm action combined persistent comment, dont modify this line -->

*Review updated until commit https://github.com/NVIDIA/Fuser/commit/222bbc6e37f577b40718b5ad75e606e8dedfcc2c*


### **Description**
- Added `tma_warp_specialized` flag to `ReductionParams`.

- Duplicated `scheduleInnerOuterPersistentKernel` as `scheduleTmaWarpSpecializedInnerOuter`.

- Implemented `scheduleTmaWarpSpecializedOuter` for specialized scheduling.

- Conditional scheduling based on `tma_warp_specialized` flag.


___



### **Changes walkthrough** 📝
<table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Enhancement</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>normalization_inner_outer.cpp</strong><dd><code>Add Warp Specialized TMA Persistent Kernel Scheduling</code>&nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/normalization_inner_outer.cpp

<li>Added <code>tma_warp_specialized</code> to <code>getInnerOuterPersistentHeuristics</code>.<br> <li> Duplicated <code>scheduleInnerOuterPersistentKernel</code> as <br><code>scheduleTmaWarpSpecializedInnerOuter</code>.<br> <li> Implemented <code>scheduleTmaWarpSpecializedOuter</code> for specialized <br>scheduling.<br> <li> Conditional scheduling based on <code>tma_warp_specialized</code> flag.


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4215/files#diff-1b9c932dfa154070bef4f9c75d30bb331ff35abd54063e020bf9f35d901e224a">+250/-1</a>&nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>reduction_heuristic.h</strong><dd><code>Add `tma_warp_specialized` Flag</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/reduction_heuristic.h

- Added `tma_warp_specialized` flag to `ReductionParams`.


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4215/files#diff-7166deb38b09ff7b2babe4bd629bc77e1dfe30b0ed29f4db0e73a440847fdf00">+3/-0</a>&nbsp; &nbsp; &nbsp; </td>

</tr>
</table></td></tr></tr></tbody></table>

---

### PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

<table>
<tr><td>🧪&nbsp;<strong>No relevant tests</strong></td></tr>
<tr><td>⚡&nbsp;<strong>Recommended focus areas for review</strong><br><br>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4215/files#diff-1b9c932dfa154070bef4f9c75d30bb331ff35abd54063e020bf9f35d901e224aR1491-R1511'><strong>Possible Issue</strong></a>

The new function `scheduleTmaWarpSpecializedOuter` introduces a lambda function `mergeReductionOrIterDomains` that modifies tensor views. Ensure that this function correctly handles all edge cases, especially when dealing with tensor views that have no reduction or iteration domains.</summary>

```c++
void scheduleTmaWarpSpecializedOuter(
    Fusion* fusion,
    const ReductionParams* rparams,
    const std::vector<TensorView*>& outer_reduction_tvs,
    std::vector<TensorView*>& cached_gmem,
    std::vector<TensorView*>& cached_gmem_reload,
    std::vector<TensorView*>& outer_reference_tvs,
    std::unordered_set<TensorView*>& boundaryNodesSet) {
  auto mergeReductionOrIterDomains = [](TensorView* tv, bool mergeReduction) {
    int prev_i = -1;
    for (int i = static_cast<int>(tv->nDims()) - 1; i >= 0; i--) {
      if (mergeReduction == tv->axis(i)->isReduction()) {
        if (prev_i == -1) {
          prev_i = i;
        } else {
          tv->merge(i, prev_i);
          prev_i = i;
        }
      }
    }
  };
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4215/files#diff-1b9c932dfa154070bef4f9c75d30bb331ff35abd54063e020bf9f35d901e224aR1567-R1733'><strong>Code Duplication</strong></a>

The new function `scheduleTmaWarpSpecializedInnerOuter` is very similar to `scheduleInnerOuterPersistentKernel`. Consider refactoring to avoid code duplication and ensure consistency.</summary>

```c++
void scheduleTmaWarpSpecializedInnerOuter(
    Fusion* fusion,
    const ReductionParams* rparams) {
  FusionGuard fg(fusion);

  // Grab the reduction, input, and output tensor views. dummy_outputs are
  // helper tensors for persistent buffer projection.
  std::vector<TensorView*> dummy_outputs, cached_inputs, reduction_tvs,
      smem_consumers;
  std::vector<std::pair<TensorView*, TensorView*>> cached_outputs;
  normalization_scheduler_utils::beforeSchedule(
      fusion,
      rparams,
      dummy_outputs,
      cached_inputs,
      reduction_tvs,
      smem_consumers,
      cached_outputs);

  // split reduction_tvs into inner and outer reduction_tvs
  std::vector<TensorView*> inner_reduction_tvs, outer_reduction_tvs;
  for (auto tv : reduction_tvs) {
    if (scheduler_utils::isFastestDimReduction(tv)) {
      inner_reduction_tvs.emplace_back(tv);
    } else {
      outer_reduction_tvs.emplace_back(tv);
    }
  }
  NVF_ERROR(
      !inner_reduction_tvs.empty(),
      "schedulePersistentKernelInnerOuter is called but no inner reduction is found.");
  NVF_ERROR(
      !outer_reduction_tvs.empty(),
      "schedulePersistentKernelInnerOuter is called but no outer reduction is found.");

  // schedule inner reduction, only schedule the first inner reduction tv,
  // then will be propagated to other inner reduction tvs.
  TensorView* inner_reference_tv =
      normalization_scheduler_utils::scheduleReductionGeneral(
          fusion,
          rparams,
          inner_reduction_tvs,
          InnerOuterPersistentKernelScheduler::schedulerType());

  // schedule outer reduction, schedule all the outer reduction tvs since we
  // need to store the intermediate results.
  std::vector<TensorView*> cached_gmem;
  std::vector<TensorView*> cached_gmem_reload;
  std::vector<TensorView*> outer_reference_tvs;
  std::unordered_set<TensorView*> boundaryNodesSet;
  scheduleTmaWarpSpecializedOuter(
      fusion,
      rparams,
      outer_reduction_tvs,
      cached_gmem,
      cached_gmem_reload,
      outer_reference_tvs,
      boundaryNodesSet);

  // Propagate inner reduction and outer reductions
  for (auto output : dummy_outputs) {
    fusion->addOutput(output);
  }

  const bool is_unroll_or_vectorization = rparams->isUnrolled();
  const bool is_vectorize =
      rparams->vectorize_inner_reduction || rparams->vectorize_iter_dom;
  const bool is_outer_grid_persistence = rparams->persistent_kernel &&
      rparams->cross_grid_inner_reduction && !rparams->fastest_dim;

  // Propagate inner reduction. There is a cutoff at boundaryNodesSet, so this
  // propagation will not propagate to the final outer reduction.
  reduction_scheduler_utils::propagateTransformation(
      inner_reference_tv, boundaryNodesSet);
  reduction_scheduler_utils::propagateRFactor(
      inner_reference_tv, inner_reduction_tvs[0], inner_reduction_tvs);

  // Don't allow parallelization propagation goes through boundaryNodesSet
  const auto& selected_tvs_inner =
      scheduler_utils::getAllTvsFrom(inner_reduction_tvs, boundaryNodesSet);
  const auto& unroll_vectorizable_cached_tvs =
      reduction_scheduler_utils::getCachedTvsToUnrollOrVectorize(
          inner_reference_tv, is_vectorize, cached_inputs, cached_outputs);
  reduction_scheduler_utils::propagateParallelization(
      inner_reduction_tvs[0],
      inner_reference_tv,
      is_unroll_or_vectorization,
      is_outer_grid_persistence,
      inner_reduction_tvs,
      unroll_vectorizable_cached_tvs,
      {selected_tvs_inner.begin(), selected_tvs_inner.end()});

  // Propagate outer reduction. Each outer reduction is connected with its
  // cached_gmem and output, since we added all the cached_gmem to the
  // boundaryNodesSet, the transformation from one outer reduction can't
  // propagate to other outer reductions due to the cutoff at
  // boundaryNodesSet. Thus, we need a loop to initiate the propagation from
  // each outer reduction. Don't allow parallelization propagation goes
  // through cached_gmem, see issue 246.
  for (long unsigned int i = 0; i < outer_reference_tvs.size(); i++) {
    const auto& selected_tvs_outer = scheduler_utils::getAllTvsFrom(
        {outer_reduction_tvs[i]}, {cached_gmem[i]});
    reduction_scheduler_utils::propagateTransformation(
        outer_reference_tvs[i], boundaryNodesSet);
    const auto& unroll_vectorizable_cached_tvs =
        reduction_scheduler_utils::getCachedTvsToUnrollOrVectorize(
            outer_reference_tvs[i],
            is_vectorize,
            cached_inputs,
            cached_outputs);
    reduction_scheduler_utils::propagateParallelization(
        outer_reduction_tvs[i],
        outer_reference_tvs[i],
        is_unroll_or_vectorization,
        is_outer_grid_persistence,
        outer_reduction_tvs,
        unroll_vectorizable_cached_tvs,
        {selected_tvs_outer.begin(), selected_tvs_outer.end()});
  }

  // special vectorization of temp gmem, vectorization_factor_tmp_gmem_write
  // is guaranteed to be smaller or equal to input vectorization factor.
  if (rparams->vectorization_factor_tmp_gmem_write > 1) {
    for (auto tv : cached_gmem) {
      NVF_ERROR(
          rparams->vectorization_factor_tmp_gmem_write <=
              rparams->unroll_factor_inner_reduction,
          "vectorization factor of temp gmem write should be smaller than that of inner reduction.")
      if (rparams->vectorization_factor_tmp_gmem_write <
          rparams->unroll_factor_inner_reduction) {
        tv->split(-1, rparams->vectorization_factor_tmp_gmem_write);
      }
      tv->axis(-1)->parallelize(ParallelType::Vectorize);
    }
  }
  // vectorization propagate through propagateParallelization only works for
  // input and output tensors. propagate vectorization to cached_gmem_reload
  // directly from output tv using parallelizeAllLike. must propagate
  // seperaely for different tvs as outer reductions are transformed
  // seperately.
  if (rparams->vectorization_factor_outer > 1) {
    for (auto tv : cached_gmem_reload) {
      auto output_tvs = ir_utils::outputTvsOf(tv);
      NVF_ERROR(
          !output_tvs.empty(),
          "cached_gmem_reload should have at least one output tensor.")
      scheduler_utils::parallelizeAllLike(
          output_tvs[0],
          -1,
          {cached_gmem_reload.begin(), cached_gmem_reload.end()},
          {ParallelType::Vectorize});
    }
  }

  // Needs special handling of vectorized loading from shared memory due to
  // potential different data types of inputs and shared memory tensor.
  if (is_vectorize) {
    reduction_scheduler_utils::sharedMemoryConsumerVectorization(
        smem_consumers, rparams->unroll_factor_inner_reduction);
  }

  // Remove dummy outputs as they can inadvertently affect CA positions
  for (auto output : dummy_outputs) {
    fusion->removeOutput(output);
  }
  inlineMost();
}
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4215/files#diff-7166deb38b09ff7b2babe4bd629bc77e1dfe30b0ed29f4db0e73a440847fdf00R140-R141'><strong>Documentation</strong></a>

The new member variable `tma_warp_specialized` in `ReductionParams` should be documented with a clear explanation of its purpose and usage.</summary>

```c
// TMA warp specialized, only used in inner-outer persistent scheduler
bool tma_warp_specialized = false;
```

</details>

</td></tr>
</table>

author_association: NONE
user: login: liqiangxl
created_at: 2025-04-09T18:59:11Z
body: !test
author_association: COLLABORATOR
user: login: liqiangxl
created_at: 2025-04-14T19:57:03Z
body: > Is there a reason to use `ReductionParams` for so many schedulers?

No solid reason, we need to split it into different schedulers. It is just not prioritized in our work plan.

> Should `scheduleTmaWarpSpecializedOuter` go in a new cpp file?
> Creating a new warp specialized scheduler seems like an opportunity to reduce technical debt.

Since it shares some functions with the muti-wave version, we can keep them in the same file and re-consider after the whole dev is done. 


author_association: COLLABORATOR
user: login: liqiangxl
created_at: 2025-04-14T19:58:04Z
body: !build
author_association: COLLABORATOR