Repository: NVIDIA/Fuser
Pull request #4269: Forward full op
Author: login: naoyam
Labels: name: rope
Created: 2025-04-17T19:58:33Z | Merged: 2025-04-24T00:00:03Z

Fusion segmenter sets aside a certain sequence of unary ops starting with fusion inputs, which we call forwarding. It effectively works as an optimization by recomputing (cheap) unary ops instead of passing tensors from one segment to another.

This PR extends the forwarding optimization to those starting with factory methods. Here's a motivating example (Litgpt Llama 3 RoPE backward):

![llama_bwd](https://github.com/user-attachments/assets/84f83b2e-d7c6-4fad-9dee-6cc17578285d)

The `T81` tensor is the output a full op. The tensor is used inside both yellow and gray segments. The op itself is in the yellow segment, so it's created inside the yellow segment, and that is passed, through gmem, to the gray segment. Obviously, cheap ops like this should be just replicated in the gray segment instead of passing a full tensor. Here's another way to see it:

```
g{(resize)
group id: 4
inputs:
  T1_g___bfloat[bS3{1}, iS4{32}, iS5{8192}, iS6{128}] __bfloat
  T3_g___bfloat[bS11{1}, iS12{8192}, iS13{128}] __bfloat
  T9_g___bfloat[bS38{1}, bS39{1 ex 32}, iS40{8192}, iS41{128}] __bfloat
outputs:
  T25_g___bfloat[bS107{1}, bS108{1 ex 32}, iS109{8192}, iS110{128}] __bfloat
  T54_g___bfloat[bS233{1}, iS238{8}rf, iS239{4}rf, iS235{8192}, iS236{128}] __bfloat
  T81_g___bfloat[bS366{1}, iS367{32}, iS368{8192}, iS369{128}] __bfloat


T81_g___bfloat[bS366{1}, iS367{32}, iS368{8192}, iS369{128}]
   = full({1, 32, 8192, 128}, __bfloat(0));
(121)
...
```

And `T81` is used in the next segment of:

```
g{(resize)
group id: 3
inputs:
  T18_g___bfloat[bS79{1}, iS80{32}, iS81{8192}, iS82{128}] __bfloat
  T25_g___bfloat[bS107{1}, bS108{1 ex 32}, iS109{8192}, iS110{128}] __bfloat
  T34_g___bfloat[bS144{1}, iS145{32}, iS146{8192}, iS147{128}] __bfloat
  T81_g___bfloat[bS366{1}, iS367{32}, iS368{8192}, iS369{128}] __bfloat
outputs:
  T75_g___bfloat[bS328{1}, iS329{8}, iS331{6}rf, iS332{8192}, iS333{128}] __bfloat


T50_l___bfloat[bS212{1}, iS213{32}, iS214{8192}, iS216{64}rf]
   = slice( T34_g___bfloat[bS144{1}, iS145{32}, iS146{8192}, iS147{128}], { {0, 1, 1} {0, 32, 1} {0, 8192, 1} {64, 128, 1} } )
(52)
T55_g___bfloat[bS240{1}, iS241{32}, iS242{8192}, iS244{128}rf]
   = pad( T50_l___bfloat[bS212{1}, iS213{32}, iS214{8192}, iS216{64}rf], {0, 0, 0, 0, 0, 0, 0, 64} )
(61)
T39_g___bfloat[bS166{1}, iS167{32}, iS168{8192}, iS170{64}rf]
   = slice( T34_g___bfloat[bS144{1}, iS145{32}, iS146{8192}, iS147{128}], { {0, 1, 1} {0, 32, 1} {0, 8192, 1} {0, 64, 1} } )
(39)
T43_l_float[bS184{1}, iS185{32}, iS186{8192}, iS187{64}]
   = __bfloat2float(T39_g___bfloat[bS166{1}, iS167{32}, iS168{8192}, iS170{64}rf]);
(44)
T46_l_float[bS196{1}, iS197{32}, iS198{8192}, iS199{64}]
   = -T43_l_float[bS184{1}, iS185{32}, iS186{8192}, iS187{64}];
(47)
T48_g___bfloat[bS204{1}, iS205{32}, iS206{8192}, iS207{64}]
   = __float2bfloat(T46_l_float[bS196{1}, iS197{32}, iS198{8192}, iS199{64}]);
(49)
T51_g___bfloat[bS217{1}, iS218{32}, iS219{8192}, iS221{128}rf]
   = pad( T48_g___bfloat[bS204{1}, iS205{32}, iS206{8192}, iS207{64}], {0, 0, 0, 0, 0, 0, 64, 0} )
(54)
T38_l_float[bS162{1}, iS163{32}, iS164{8192}, iS165{128}]
   = __bfloat2float(T81_g___bfloat[bS366{1}, iS367{32}, iS368{8192}, iS369{128}]);
(101)
...
```

There are multiple ways to achieve that. What seems to most make sense to me is to extend the existing forwarding method to handle cases like this. The existing method only considers ops starting with fusion inputs, which do not include factory-created tensors. 

This PR applies a small change to the forwarding logic to include factory ops as well. The end result of this change with the above example case is that the full result is no longer passed around. Here's the first segment:

```
g{(resize)
group id: 3
inputs:
  T0_g___bfloat[bS0{1}, iS1{8192}, iS2{128}] __bfloat
  T1_g___bfloat[bS3{1}, iS4{32}, iS5{8192}, iS6{128}] __bfloat
  T3_g___bfloat[bS11{1}, iS12{8192}, iS13{128}] __bfloat
outputs:
  T49_g___bfloat[bS208{1}, iS209{32}, iS210{8192}, iS211{128}] __bfloat


T20_l___bfloat[bS87{1}, bS88{1}, iS89{8192}, iS90{128}]
   = broadcast( T3_g___bfloat[bS11{1}, iS12{8192}, iS13{128}], flags = {false, true, false, false} )
(16)
T25_g___bfloat[bS107{1}, bS108{1 ex 32}, iS109{8192}, iS110{128}] = expand( T20_l___bfloat[bS87{1}, bS88{1}, iS89{8192}, iS90{128}], {1, 32, 8192, 128} )
(129)
T5_l___bfloat[bS18{1}, bS19{1}, iS20{8192}, iS21{128}]
   = broadcast( T0_g___bfloat[bS0{1}, iS1{8192}, iS2{128}], flags = {false, true, false, false} )
(0)
T9_g___bfloat[bS38{1}, bS39{1 ex 32}, iS40{8192}, iS41{128}] = expand( T5_l___bfloat[bS18{1}, bS19{1}, iS20{8192}, iS21{128}], {1, 32, 8192, 128} )
(128)
T81_g___bfloat[bS366{1}, iS367{32}, iS368{8192}, iS369{128}]
   = full({1, 32, 8192, 128}, __bfloat(0));
...
```

Notice that `T81` is no longer a segment output. And the second segment is:

```
g{(resize)
group id: 4
inputs:
  T0_g___bfloat[bS0{1}, iS1{8192}, iS2{128}] __bfloat
  T2_g___bfloat[bS7{1}, iS8{32}, iS9{8192}, iS10{128}] __bfloat
  T3_g___bfloat[bS11{1}, iS12{8192}, iS13{128}] __bfloat
outputs:
  T74_g___bfloat[bS321{1}, iS326{8}rf, iS327{4}rf, iS323{8192}, iS324{128}] __bfloat


T20_l___bfloat[bS87{1}, bS88{1}, iS89{8192}, iS90{128}]
   = broadcast( T3_g___bfloat[bS11{1}, iS12{8192}, iS13{128}], flags = {false, true, false, false} )
(16)
T25_g___bfloat[bS107{1}, bS108{1 ex 32}, iS109{8192}, iS110{128}] = expand( T20_l___bfloat[bS87{1}, bS88{1}, iS89{8192}, iS90{128}], {1, 32, 8192, 128} )
(129)
T5_l___bfloat[bS18{1}, bS19{1}, iS20{8192}, iS21{128}]
   = broadcast( T0_g___bfloat[bS0{1}, iS1{8192}, iS2{128}], flags = {false, true, false, false} )
(0)
T9_g___bfloat[bS38{1}, bS39{1 ex 32}, iS40{8192}, iS41{128}] = expand( T5_l___bfloat[bS18{1}, bS19{1}, iS20{8192}, iS21{128}], {1, 32, 8192, 128} )
(128)
T81_g___bfloat[bS366{1}, iS367{32}, iS368{8192}, iS369{128}]
   = full({1, 32, 8192, 128}, __bfloat(0));
(121)
...
```




Diff:
diff --git a/csrc/fusion_segmenter.cpp b/csrc/fusion_segmenter.cpp
index bf7693c6028..7242054864f 100644
--- a/csrc/fusion_segmenter.cpp
+++ b/csrc/fusion_segmenter.cpp
@@ -2156,6 +2156,21 @@ SegmentedGroup* SegmentCandidateFinder::mergeAllGivenGroups(
       !groups_to_merge.empty(),
       "fusion segment :(mergeAllGivenGroups) tried to merge no groups");
 
+  // The fusion input auxiliary groups should never be merged.
+  const auto& aux_input_groups = getAuxiliaryInputGroups();
+  std::vector<SegmentedGroup*> aux_groups_to_merge;
+  std::ranges::copy_if(
+      groups_to_merge,
+      std::back_inserter(aux_groups_to_merge),
+      [&](SegmentedGroup* group) {
+        return std::ranges::find(aux_input_groups, group) !=
+            aux_input_groups.end();
+      });
+  NVF_ERROR(
+      aux_groups_to_merge.empty(),
+      "Trying to merge auxiliary input groups: ",
+      toDelimitedString(aux_groups_to_merge));
+
   // Make a set to detect internal edges
   std::unordered_set<SegmentedGroup*> group_set(
       groups_to_merge.begin(), groups_to_merge.end());
@@ -3298,7 +3313,6 @@ class CombineReductions {
                       return groups_to_merge_set.has(group);
                     }),
                 groups_with_reductions_.end());
-
             return joined_group;
           }
         }
@@ -3603,6 +3617,10 @@ class MergeUpAndDownCast {
       SegmentedGroup* group = to_visit.front();
       to_visit.pop_front();
 
+      if (group->exprs().empty()) {
+        continue;
+      }
+
       if (groups_to_merge_set.count(group)) {
         continue;
       }
@@ -4356,12 +4374,28 @@ void SegmentCandidateFinder::forwardInputs() {
   excluded_inp_unary_exprs_ = {};
   input2group_.clear();
 
+  std::vector<Val*> extended_fusion_inputs = completeFusion()->inputs();
+
+  // Grab factory ops that should be forwarded. Add created tensors to
+  // the fusion input list to make them handled like fusion inputs
+  // TODO: Handle more factory methods such as IotaOp, EyeOp,
+  // TensorConstruct. Probably should not include relatively expensive
+  // ops like RNGOp.
+  for (auto expr : completeFusion()->exprs()) {
+    if (expr->isA<FullOp>() &&
+        // Don't bother if it's a fusion output
+        !expr->output(0)->isFusionOutput()) {
+      extended_fusion_inputs.push_back(expr->output(0));
+      excluded_inp_unary_exprs_.pushBack(expr);
+    }
+  }
+
   // "Terminating" outputs from the excluded input unary exprs, these will be
   // treated as complete fusion inputs.
   VectorOfUniqueEntries<Val*> forwarded_inputs;
   {
     std::deque<UnaryOp*> to_visit;
-    for (Val* inp : completeFusion()->inputs()) {
+    for (Val* inp : extended_fusion_inputs) {
       if (UnaryOp* unary_use = shouldForward(inp)) {
         to_visit.push_back(unary_use);
       }
@@ -4384,11 +4418,13 @@ void SegmentCandidateFinder::forwardInputs() {
     }
   }
 
-  auto excluded_fusion_inputs = IterVisitor::getInputsTo(
-      {forwarded_inputs.begin(), forwarded_inputs.end()});
+  // Stop traversing back at factory vals (and fusion inputs)
+  auto excluded_fusion_inputs = InputsOf::getInputsTo(
+      {forwarded_inputs.begin(), forwarded_inputs.end()},
+      extended_fusion_inputs);
 
   // List of vals to treat as complete fusion inputs for segmentation
-  forwarded_fusion_inputs_ = completeFusion()->inputs();
+  forwarded_fusion_inputs_ = extended_fusion_inputs;
 
   forwarded_fusion_inputs_.erase(
       std::remove_if(
@@ -4427,6 +4463,16 @@ void SegmentCandidateFinder::cleanupForwardedInputs() {
   input2group_.clear();
 }
 
+std::vector<SegmentedGroup*> SegmentCandidateFinder::getAuxiliaryInputGroups()
+    const {
+  std::vector<SegmentedGroup*> aux_groups;
+  aux_groups.reserve(input2group_.size());
+  std::ranges::transform(input2group_, aux_groups.begin(), [](const auto& kv) {
+    return kv.second;
+  });
+  return aux_groups;
+}
+
 void SegmentCandidateFinder::finalMerge() {
   FUSER_PERF_SCOPE("SegmentCandidateFinder::finalMerge");
   auto producer_check = getGroupDependency();
@@ -4639,7 +4685,14 @@ void SegmentCandidateFinder::resolveScalarsInGroup(SegmentedGroup* group) {
 
 SegmentedGroup* SegmentCandidateFinder::createInputGroup(Val* forwarded_input) {
   SegmentedGroup* group = segmented_fusion_->newGroup();
-  group->input_vals_ = IterVisitor::getInputsTo({forwarded_input});
+  for (auto inp : IterVisitor::getInputsTo({forwarded_input})) {
+    // inp may be a factory-created tensor, which is not an input to
+    // the group.
+    if (std::ranges::find(completeFusion()->inputs(), inp) !=
+        completeFusion()->inputs().end()) {
+      group->input_vals_.pushBack(inp);
+    }
+  }
   group->exprs_ = StmtSort::getExprsTo({forwarded_input});
   return group;
 }
diff --git a/csrc/fusion_segmenter.h b/csrc/fusion_segmenter.h
index ace8a0bfb02..cfcbf75fd5d 100644
--- a/csrc/fusion_segmenter.h
+++ b/csrc/fusion_segmenter.h
@@ -693,6 +693,9 @@ class SegmentCandidateFinder {
                val) != forwarded_fusion_inputs_.end();
   };
 
+  // Get all auxiliary groups created for fusion inputs
+  std::vector<SegmentedGroup*> getAuxiliaryInputGroups() const;
+
  protected:
   //! These are the merge node heuristic passes, should
   //!  eventually should have a dedicated interface
diff --git a/tests/cpp/test_segmentation.cpp b/tests/cpp/test_segmentation.cpp
index 24db8d7793a..acb65bff7f6 100644
--- a/tests/cpp/test_segmentation.cpp
+++ b/tests/cpp/test_segmentation.cpp
@@ -695,7 +695,7 @@ TEST_F(SegmentationTest, ForwardInputsToSegmenterSetIssue2658) {
 }
 
 // Test to verify an upcast is replicated between different segments
-TEST_F(NVFuserTest, PrivatizeUpcast) {
+TEST_F(SegmentationTest, PrivatizeUpcast) {
   auto fusion_ptr = std::make_unique<Fusion>();
   auto& fusion = *fusion_ptr;
   FusionGuard fg(fusion_ptr.get());
@@ -741,7 +741,7 @@ TEST_F(NVFuserTest, PrivatizeUpcast) {
 
 // Unlike PrivatizeUpcast, verify replicated upcast ops are
 // consolidated back as they are grouped into the same segment
-TEST_F(NVFuserTest, RevertPrivatizedUpcast) {
+TEST_F(SegmentationTest, RevertPrivatizedUpcast) {
   auto fusion_ptr = std::make_unique<Fusion>();
   auto& fusion = *fusion_ptr;
   FusionGuard fg(fusion_ptr.get());
@@ -807,4 +807,59 @@ TEST_F(NVFuserTest, RevertPrivatizedUpcast) {
   }
 }
 
+TEST_F(SegmentationTest, ForwardFull) {
+  auto fusion_ptr = std::make_unique<Fusion>();
+  auto& fusion = *fusion_ptr;
+  FusionGuard fg(fusion_ptr.get());
+
+  auto tv0 = makeSymbolicTensor(1);
+  fusion.addInput(tv0);
+
+  // FullOp that is used in two segments
+  auto tv1 = full({tv0->axis(0)->extent()}, fusion.oneVal(), DataType::Float);
+
+  auto tv2 = add(tv0, tv1);
+  auto tv3 = segment_set(tv2);
+
+  auto tv4 = add(tv3, tv1);
+  fusion.addOutput(tv4);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  auto t0 = at::randn({1024}, options);
+
+  FusionExecutorCache executor_cache(std::move(fusion_ptr));
+  auto outputs = executor_cache.runFusionWithInputs({t0});
+  testValidate(&fusion, outputs, {t0}, __LINE__, __FILE__);
+
+  FusionKernelRuntime* runtime = executor_cache.getMostRecentKernelRuntime();
+  EXPECT_THAT(runtime->fusionSegments()->groups(), SizeIs(2));
+
+  // Make sure the full output should not be a segment input
+  for (const auto& executor : runtime->executors()) {
+    auto ke = dynamic_cast<KernelExecutor*>(executor.get());
+    ASSERT_NE(ke, nullptr);
+    kir::Kernel* kernel = ke->compiledKernel()->kernel();
+    bool full_op_found = false;
+    for (auto expr : KernelExprVisitor::getAllExprs(kernel)) {
+      auto out_tv = ir_utils::getTvOutput(expr);
+      if (out_tv == nullptr) {
+        continue;
+      }
+      auto full_op = dynamic_cast<FullOp*>(out_tv->definition());
+      if (full_op == nullptr) {
+        continue;
+      }
+      full_op_found = true;
+      auto output_it =
+          std::ranges::find_if(kernel->outputs(), [&](Val* output) {
+            return output->isA<TensorView>() &&
+                output->name() == out_tv->name();
+          });
+      EXPECT_EQ(output_it, kernel->outputs().end())
+          << "FullOp ouput should not be a segment output";
+    }
+    EXPECT_TRUE(full_op_found) << "Each segment has its own FullOp";
+  }
+}
+
 } // namespace nvfuser


Reviews:
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-23T02:29:11Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-23T02:30:03Z
user: login: jjsjann123
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-23T22:24:50Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-23T23:01:31Z
user: login: jjsjann123
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-23T23:40:43Z

Review comments:
diff_hunk: @@ -3603,6 +3617,10 @@ class MergeUpAndDownCast {
       SegmentedGroup* group = to_visit.front();
       to_visit.pop_front();
 
+      if (group->exprs().empty()) {
path: csrc/fusion_segmenter.cpp
user: login: naoyam
body: Not strictly related, but while testing this PR, I was hitting a bug that should been easily spotted if this was in place.
created_at: 2025-04-23T02:29:11Z
author_association: COLLABORATOR
diff_hunk: @@ -4356,12 +4374,28 @@ void SegmentCandidateFinder::forwardInputs() {
   excluded_inp_unary_exprs_ = {};
   input2group_.clear();
 
+  std::vector<Val*> extended_fusion_inputs = completeFusion()->inputs();
+
+  // Grab factory ops that should be forwarded. Add created tensors to
+  // the fusion input list to make them handled like fusion inputs
+  // TODO: Handle more factory methods such as IotaOp, EyeOp,
path: csrc/fusion_segmenter.cpp
user: login: naoyam
body: I was originally planning to add support for other factory ops but decided to leave them as follow-up extensions.
created_at: 2025-04-23T02:30:03Z
author_association: COLLABORATOR
diff_hunk: @@ -4356,12 +4374,28 @@ void SegmentCandidateFinder::forwardInputs() {
   excluded_inp_unary_exprs_ = {};
   input2group_.clear();
 
+  std::vector<Val*> extended_fusion_inputs = completeFusion()->inputs();
+
+  // Grab factory ops that should be forwarded. Add created tensors to
+  // the fusion input list to make them handled like fusion inputs
+  // TODO: Handle more factory methods such as IotaOp, EyeOp,
+  // TensorConstruct. Probably should not include relatively expensive
+  // ops like RNGOp.
+  for (auto expr : completeFusion()->exprs()) {
+    if (expr->isA<FullOp>() &&
+        // Don't bother if it's a fusion output
+        !expr->output(0)->isFusionOutput()) {
path: csrc/fusion_segmenter.cpp
user: login: jjsjann123
body: should we not forward even though they are fusion outputs, that would still save bandwidth on reading gmem in other consumers right?!
created_at: 2025-04-23T22:17:49Z
author_association: COLLABORATOR
diff_hunk: @@ -4356,12 +4374,28 @@ void SegmentCandidateFinder::forwardInputs() {
   excluded_inp_unary_exprs_ = {};
   input2group_.clear();
 
+  std::vector<Val*> extended_fusion_inputs = completeFusion()->inputs();
+
+  // Grab factory ops that should be forwarded. Add created tensors to
+  // the fusion input list to make them handled like fusion inputs
+  // TODO: Handle more factory methods such as IotaOp, EyeOp,
+  // TensorConstruct. Probably should not include relatively expensive
+  // ops like RNGOp.
+  for (auto expr : completeFusion()->exprs()) {
+    if (expr->isA<FullOp>() &&
+        // Don't bother if it's a fusion output
+        !expr->output(0)->isFusionOutput()) {
path: csrc/fusion_segmenter.cpp
user: login: naoyam
body: That's true in theory, but in practice it'll be likely cached anyway, so the impact would be minimal. 

Implementation wise, we would need to further extend the forwarding logic since it's likely to result in causing something unexpected if multiple segments had the same op producing the same fusion output.

I just don't think it's worthwhile supporting such a case.
created_at: 2025-04-23T23:01:31Z
author_association: COLLABORATOR

PR comments:
user: login: naoyam
created_at: 2025-04-17T19:59:13Z
body: !test --diff
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-04-17T19:59:24Z
body: <!-- llm action combined persistent comment, dont modify this line -->

*Review updated until commit https://github.com/NVIDIA/Fuser/commit/329f19a580eef17841226dbf20b151ba4bae6538*


### **Description**
- Extend forwarding optimization to factory methods like FullOp

- Prevent merging of auxiliary input groups

- Add test for forwarding FullOp across segments


___



### **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>fusion_segmenter.cpp</strong><dd><code>Extend forwarding to factory methods</code>&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/fusion_segmenter.cpp

<li>Prevent merging of auxiliary input groups<br> <li> Extend forwarding optimization to include FullOp<br> <li> Update logic to handle factory-created tensors


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4269/files#diff-e2f2ad44a6dc03e4ad8e5f0f047be25eb1c142add431d48c1e046c968a577f3b">+59/-6</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>fusion_segmenter.h</strong><dd><code>Add method for auxiliary groups</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/fusion_segmenter.h

- Add method to get auxiliary input groups


</details>


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

</tr>
</table></td></tr><tr><td><strong>Tests</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>test_segmentation.cpp</strong><dd><code>Add test for FullOp forwarding</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; &nbsp; </dd></summary>
<hr>

tests/cpp/test_segmentation.cpp

<li>Rename test cases to use SegmentationTest<br> <li> Add test for forwarding FullOp across segments


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4269/files#diff-256c87b67ef1f18e8078fe44d35392afc11ab21056f855d206917d0411a26840">+57/-2</a>&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>PR contains 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/4269/files#diff-e2f2ad44a6dc03e4ad8e5f0f047be25eb1c142add431d48c1e046c968a577f3bR4373-R4415'><strong>Performance Evaluation</strong></a>

Ensure thorough performance evaluation is provided to demonstrate the benefits of extending forwarding optimization to factory methods.</summary>

```c++
void SegmentCandidateFinder::forwardInputs() {
  excluded_inp_unary_exprs_ = {};
  input2group_.clear();

  std::vector<Val*> extended_fusion_inputs = completeFusion()->inputs();

  // Grab factory ops that should be forwarded. Add created tensors to
  // the fusion input list to make them handled like fusion inputs
  // TODO: Handle more factory methods such as IotaOp, EyeOp,
  // TensorConstruct. Probably should not include relatively expensive
  // ops like RNGOp.
  for (auto expr : completeFusion()->exprs()) {
    if (expr->isA<FullOp>() &&
        // Don't bother if it's a fusion output
        !expr->output(0)->isFusionOutput()) {
      extended_fusion_inputs.push_back(expr->output(0));
      excluded_inp_unary_exprs_.pushBack(expr);
    }
  }

  // "Terminating" outputs from the excluded input unary exprs, these will be
  // treated as complete fusion inputs.
  VectorOfUniqueEntries<Val*> forwarded_inputs;
  {
    std::deque<UnaryOp*> to_visit;
    for (Val* inp : extended_fusion_inputs) {
      if (UnaryOp* unary_use = shouldForward(inp)) {
        to_visit.push_back(unary_use);
      }
    }

    while (!to_visit.empty()) {
      UnaryOp* uop = to_visit.front();
      to_visit.pop_front();

      if (UnaryOp* unary_use = shouldForward(uop->out())) {
        to_visit.push_back(unary_use);
      } else {
        // We cannot extend the chain of unary ops, so we finalize this chain by
        // saving its output as a forwarded input.
        forwarded_inputs.pushBack(uop->out());
      }
      // Either way, `uop` is excluded from merging until
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4269/files#diff-e2f2ad44a6dc03e4ad8e5f0f047be25eb1c142add431d48c1e046c968a577f3bR4385-R4391'><strong>Factory Method Handling</strong></a>

Consider handling more factory methods such as IotaOp, EyeOp, TensorConstruct, and ensure that relatively expensive ops like RNGOp are not included.</summary>

```c++
  if (expr->isA<FullOp>() &&
      // Don't bother if it's a fusion output
      !expr->output(0)->isFusionOutput()) {
    extended_fusion_inputs.push_back(expr->output(0));
    excluded_inp_unary_exprs_.pushBack(expr);
  }
}
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4269/files#diff-256c87b67ef1f18e8078fe44d35392afc11ab21056f855d206917d0411a26840R810-R863'><strong>Test Coverage</strong></a>

Ensure that the new test cases cover a variety of scenarios and edge cases to validate the correctness of the forwarding optimization for factory methods.</summary>

```c++
TEST_F(SegmentationTest, ForwardFull) {
  auto fusion_ptr = std::make_unique<Fusion>();
  auto& fusion = *fusion_ptr;
  FusionGuard fg(fusion_ptr.get());

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  // FullOp that is used in two segments
  auto tv1 = full({tv0->axis(0)->extent()}, fusion.oneVal(), DataType::Float);

  auto tv2 = add(tv0, tv1);
  auto tv3 = segment_set(tv2);

  auto tv4 = add(tv3, tv1);
  fusion.addOutput(tv4);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn({1024}, options);

  FusionExecutorCache executor_cache(std::move(fusion_ptr));
  auto outputs = executor_cache.runFusionWithInputs({t0});
  testValidate(&fusion, outputs, {t0}, __LINE__, __FILE__);

  FusionKernelRuntime* runtime = executor_cache.getMostRecentKernelRuntime();
  EXPECT_THAT(runtime->fusionSegments()->groups(), SizeIs(2));

  // Make sure the full output should not be a segment input
  for (const auto& executor : runtime->executors()) {
    auto ke = dynamic_cast<KernelExecutor*>(executor.get());
    ASSERT_NE(ke, nullptr);
    kir::Kernel* kernel = ke->compiledKernel()->kernel();
    bool full_op_found = false;
    for (auto expr : KernelExprVisitor::getAllExprs(kernel)) {
      auto out_tv = ir_utils::getTvOutput(expr);
      if (out_tv == nullptr) {
        continue;
      }
      auto full_op = dynamic_cast<FullOp*>(out_tv->definition());
      if (full_op == nullptr) {
        continue;
      }
      full_op_found = true;
      auto output_it =
          std::ranges::find_if(kernel->outputs(), [&](Val* output) {
            return output->isA<TensorView>() &&
                output->name() == out_tv->name();
          });
      EXPECT_EQ(output_it, kernel->outputs().end())
          << "FullOp ouput should not be a segment output";
    }
    EXPECT_TRUE(full_op_found) << "Each segment has its own FullOp";
  }
}
```

</details>

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

author_association: NONE
user: login: naoyam
created_at: 2025-04-17T20:52:00Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-18T03:44:03Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-23T02:37:44Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-04-23T03:31:36Z
body: ## Codediff results

### TestNvFuserFrontend::test_issue1872

This seems expected because the full-op result is used in two slices.

### jit_codegen_diff_20_6/7, jit_codegen_diff_20_7/7

Something seems off in these two cases. They are showing lots of new and removed tests. Maybe the script somehow fails to match two results? CC: @jacobhinkle 


author_association: COLLABORATOR