Repository: NVIDIA/Fuser
Pull request #4051: project persistent buffer to it producer if the buffer is the output of an upcast
Author: login: liqiangxl
Created: 2025-03-09T02:32:18Z | Merged: 2025-03-19T11:39:27Z

**What?**
Enhance persistent buffer projection to allow project to its immediate producer if the buffer is the output of an upcast op.
This is another method to solve segmenter issue discussed at https://github.com/NVIDIA/Fuser/pull/4040

**Why?**

- This saves register usage. It is similar to project to inputs, not abosutely necessary but we always do it to save register usage.
- It is required in cases where the segmenter added privatized upcast (featured added in [PR-3776](https://github.com/NVIDIA/Fuser/pull/3776)) and then reverted at the finalization stage. This may lead to different persistent buffer size and cause the persistent scheudler fails due to not enough registers or shared memory. An actual case is `mpirun -np 1 pytest tests/python/test_multidevice.py --only-mpi -s -k transformer_backward` on RTX-6000.


Diff:
diff --git a/csrc/scheduler/reduction_utils.cpp b/csrc/scheduler/reduction_utils.cpp
index 0cbc76a8b96..b312b73514f 100644
--- a/csrc/scheduler/reduction_utils.cpp
+++ b/csrc/scheduler/reduction_utils.cpp
@@ -798,7 +798,9 @@ class PersistentBufferProjector {
               persistent_buffers[a], persistent_buffers[b]);
         });
 
-    // try to project buffer to its producers
+    // try to project buffer to its producers when
+    // (1) all producers are persistent buffers
+    // (2) or, the buffer is the input to an upcast op
     std::unordered_set<TensorView*> persistent_buffer_set(
         persistent_buffers.begin(), persistent_buffers.end());
     for (auto buffer_i : visiting_order) {
@@ -809,6 +811,23 @@ class PersistentBufferProjector {
         projectToInputOrImmediatePersistentProducer(
             (int)buffer_i,
             std::vector<Val*>(producers.begin(), producers.end()));
+      } else if (
+          auto upcast_input = scheduler_utils::getUpCastInputOf(buffer)) {
+        // Similar to projecting to inputs and persistent producers, this logic
+        // projects the buffer to its producer when the buffer is the output of
+        // an upcast op. This optimization reduces buffer size and can be
+        // extended to project to low-precision intermediate tensors, even if
+        // the recomputation involves non-cast ops. However, this should be
+        // avoided when the recomputation cost outweighs the benefits of reduced
+        // register usage.
+        // TODO: extend to allow non-cast ops in the recomputation.
+        auto consumers = ir_utils::consumerTvsOf(buffer);
+        for (auto i : c10::irange(1, consumers.size())) {
+          ir_utils::replaceValInExprInputs(
+              consumers.at(i)->definition(),
+              buffer,
+              RecomputeTv::recompute(buffer, {upcast_input}));
+        }
       }
     }
   }
diff --git a/csrc/scheduler/utils.cpp b/csrc/scheduler/utils.cpp
index afe039e063f..c6510336989 100644
--- a/csrc/scheduler/utils.cpp
+++ b/csrc/scheduler/utils.cpp
@@ -558,6 +558,24 @@ std::pair<bool, std::vector<TensorView*>> canProjectToInputsWithoutReduction(
   return std::make_pair(true, target_broadcast_tvs);
 }
 
+TensorView* getUpCastInputOf(const TensorView* tv) {
+  // skip if definition is not a unary op
+  if (auto uop = dynamic_cast<UnaryOp*>(tv->definition())) {
+    // skip if the input is a fusion input or the op is not a cast
+    if (uop->input(0)->isFusionInput() ||
+        uop->getUnaryOpType() != UnaryOpType::Cast) {
+      return nullptr;
+    }
+    // skip if the cast is not upcast
+    auto precisions = ir_utils::getPrecisionOfProducerConsumerTensors(uop);
+    if (!precisions.has_value() || precisions->first >= precisions->second) {
+      return nullptr;
+    }
+    return uop->input(0)->as<TensorView>();
+  }
+  return nullptr;
+}
+
 PersistentBufferInfo persistentBuffers(Fusion* fusion) {
   FusionGuard fg(fusion);
   PersistentBufferInfo persistent_buffer_info;
@@ -964,11 +982,21 @@ int64_t getPersistentBufferSizeOfTensor(
       buffer_bytes *= id_size.as<int64_t>();
     }
   }
+  // If the persistent buffer is the output of an upcast op, scheduler will
+  // project it back to the input to save register usage. This is similar to
+  // project to inputs, not abosutely necessary but we always do it to
+  // save register usage. So, need to compute the buffer size using the data
+  // type before upcast.
+  int64_t dtype_size = 1;
+  if (auto upcast_input = getUpCastInputOf(buffer)) {
+    dtype_size = dataTypeSize(
+        upcast_input->getDataType().value(), runtime_info.getIndexType());
+  } else {
+    dtype_size = dataTypeSize(
+        buffer->getDataType().value(), runtime_info.getIndexType());
+  }
 
-  buffer_bytes = buffer_bytes == -1 ? 0
-                                    : buffer_bytes *
-          (int64_t)dataTypeSize(buffer->getDataType().value(),
-                                runtime_info.getIndexType());
+  buffer_bytes = buffer_bytes == -1 ? 0 : buffer_bytes * dtype_size;
   return buffer_bytes;
 }
 
diff --git a/csrc/scheduler/utils.h b/csrc/scheduler/utils.h
index 4e95d854c28..5f5bf45e446 100644
--- a/csrc/scheduler/utils.h
+++ b/csrc/scheduler/utils.h
@@ -803,5 +803,10 @@ inline int64_t nLogicalDims(const TensorView* tv) {
 // reference IDs. Non-matching loop IDs are placed outermost positions.
 void reorderTensorLike(TensorView* tv, const std::vector<IterDomain*>& ref);
 
+// If buffer_tv's definition is an upcast and the input to the cast is not a
+// fusion input, return input to the cast. Otherwise, return nullptr. Used to
+// recompute buffer_tv from its producer to save register/smem usage. Fusion
+// input is skipped as it is handled by project to inputs.
+TensorView* getUpCastInputOf(const TensorView* buffer_tv);
 } // namespace scheduler_utils
 } // namespace nvfuser
diff --git a/tests/cpp/test_persistent_buffer.cpp b/tests/cpp/test_persistent_buffer.cpp
index 3fe22606a88..07ecd913a0d 100644
--- a/tests/cpp/test_persistent_buffer.cpp
+++ b/tests/cpp/test_persistent_buffer.cpp
@@ -1557,6 +1557,63 @@ INSTANTIATE_TEST_SUITE_P(
       return sanitizeTestName(ss.str());
     });
 
+// If the persistent buffer is the output of an upcast op, project
+// it back to the input to save register usage. This is similar to
+// project to inputs, not abosutely necessary but we always do it to
+// save register usage.
+TEST_F(PersistentBufferTest, ProjectToUpcastInput) {
+  std::unique_ptr<Fusion> fusion_ptr = std::make_unique<Fusion>();
+  Fusion& fusion = *fusion_ptr.get();
+  FusionGuard fg(&fusion);
+
+  int64_t dim0 = 128, dim1 = 1024;
+  auto tv0 = makeSymbolicTensor(2);
+  fusion.addInput(tv0);
+  auto tv1 = set(tv0);
+  auto tv2 = gt(tv1, IrBuilder::create<Val>(0.5));
+  auto tv3 = castOp(DataType::Float, tv2);
+  auto tv4 = sum(tv3, {1});
+  auto tv5 = broadcast(tv4, {false, true});
+  auto tv6 = add(tv5, tv3);
+  fusion.addOutput(tv6);
+  auto fusion_copy = fusion;
+
+  // tv3 is the persistent tensor
+  auto persistent_buffer_info = scheduler_utils::persistentBuffers(&fusion);
+  EXPECT_EQ(
+      persistent_buffer_info.persistent_buffers, std::vector<TensorView*>{tv3});
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  at::Tensor aten_input = at::randn({dim0, dim1}, options);
+
+  // The persistent buffer size is dim1 * sizeof(bool) not sizeof(float)
+  // becase tv3 is the output of an upcast op, the scheduler will project it
+  // back to the input which is tv2 and its data type is bool.
+  SchedulerRuntimeInfo runtime_info(&fusion, {aten_input});
+  auto persistent_buffer_size = scheduler_utils::persistentBufferSize(
+      &fusion, runtime_info, persistent_buffer_info);
+  EXPECT_EQ(
+      persistent_buffer_size.persistent_buffer_size,
+      dim1 * dataTypeSize(DataType::Bool));
+
+  // Check the compute position of the bool tensor, tv2, is at the top of the
+  // kernel.
+  FusionExecutorCache executor_cache(std::move(fusion_ptr));
+  auto cg_outputs = executor_cache.runFusionWithInputs({aten_input});
+  auto runtime = executor_cache.getMostRecentKernelRuntime();
+  Fusion* scheduled_fusion = runtime->executors()
+                                 .back()
+                                 ->as<KernelExecutor>()
+                                 ->compiledKernel()
+                                 ->kernel();
+  for (auto tv : scheduled_fusion->allTvs()) {
+    if (tv->getDataType() == DataType::Bool) {
+      EXPECT_EQ(tv->getComputeAtPosition(), 1);
+    }
+  }
+  testValidate(&fusion_copy, cg_outputs, {aten_input}, __LINE__, __FILE__, "");
+}
+
 TEST_F(NVFuserTest, FalsePersistentBuffer) {
   auto fusion_ptr = std::make_unique<Fusion>();
   auto& fusion = *fusion_ptr;
@@ -1587,5 +1644,4 @@ TEST_F(NVFuserTest, FalsePersistentBuffer) {
       scheduler_utils::persistentBuffers(&fusion);
   EXPECT_TRUE(info.persistent_buffers.empty());
 }
-
 } // namespace nvfuser


Reviews:
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T01:13:59Z
user: login: liqiangxl
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T01:25:48Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T01:36:42Z
user: login: liqiangxl
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T01:48:08Z
user: login: naoyam
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-14T04:58:58Z

Review comments:
diff_hunk: @@ -557,6 +557,18 @@ std::pair<bool, std::vector<TensorView*>> canProjectToInputsWithoutReduction(
   return std::make_pair(true, target_broadcast_tvs);
 }
 
+TensorView* getUpCastInputOf(const TensorView* tv) {
+  if (auto uop = dynamic_cast<UnaryOp*>(tv->definition())) {
+    if (uop->getUnaryOpType() == UnaryOpType::Cast &&
+        !uop->input(0)->isFusionInput() &&
+        dataTypeSize(uop->input(0)->getDataType().value()) <
path: csrc/scheduler/utils.cpp
user: login: naoyam
body: Please use `getPrecisionOfProducerConsumerTensors`
created_at: 2025-03-13T01:02:26Z
author_association: COLLABORATOR
diff_hunk: @@ -809,6 +811,15 @@ class PersistentBufferProjector {
         projectToInputOrImmediatePersistentProducer(
             (int)buffer_i,
             std::vector<Val*>(producers.begin(), producers.end()));
+      } else if (
+          auto upcast_input = scheduler_utils::getUpCastInputOf(buffer)) {
+        auto consumers = ir_utils::consumerTvsOf(buffer);
+        for (auto i : c10::irange(1, consumers.size())) {
+          ir_utils::replaceValInExprInputs(
+              consumers.at(i)->definition(),
+              buffer,
+              RecomputeTv::recompute(buffer, {upcast_input}));
+        }
path: csrc/scheduler/reduction_utils.cpp
user: login: naoyam
body: Why is this necessary? Why not just using `upcast_input` as the persistent buffer?
created_at: 2025-03-13T01:13:40Z
author_association: COLLABORATOR
diff_hunk: @@ -809,6 +811,15 @@ class PersistentBufferProjector {
         projectToInputOrImmediatePersistentProducer(
             (int)buffer_i,
             std::vector<Val*>(producers.begin(), producers.end()));
+      } else if (
+          auto upcast_input = scheduler_utils::getUpCastInputOf(buffer)) {
+        auto consumers = ir_utils::consumerTvsOf(buffer);
+        for (auto i : c10::irange(1, consumers.size())) {
+          ir_utils::replaceValInExprInputs(
+              consumers.at(i)->definition(),
+              buffer,
+              RecomputeTv::recompute(buffer, {upcast_input}));
+        }
path: csrc/scheduler/reduction_utils.cpp
user: login: liqiangxl
body: That's what I did at the [first approach (first 3 commits)](https://github.com/NVIDIA/Fuser/pull/4051/files/d4bdfde3afd2574ac7a40ce1efcc66be1a041f7a). With that approach, we need to recompute `upcast output`, similar to cast insertion. Since we already have `project to inputs` and `project to other persistent buffers`, I prefer implementing it as `project to low-precision intermediate tensors` to maintain consistency with our existing structure and suitable for futhure extensions.
created_at: 2025-03-13T01:25:48Z
author_association: COLLABORATOR
diff_hunk: @@ -809,6 +811,15 @@ class PersistentBufferProjector {
         projectToInputOrImmediatePersistentProducer(
             (int)buffer_i,
             std::vector<Val*>(producers.begin(), producers.end()));
+      } else if (
+          auto upcast_input = scheduler_utils::getUpCastInputOf(buffer)) {
+        auto consumers = ir_utils::consumerTvsOf(buffer);
+        for (auto i : c10::irange(1, consumers.size())) {
+          ir_utils::replaceValInExprInputs(
+              consumers.at(i)->definition(),
+              buffer,
+              RecomputeTv::recompute(buffer, {upcast_input}));
+        }
path: csrc/scheduler/reduction_utils.cpp
user: login: naoyam
body: Please make sure to have such a background context as code comments.
created_at: 2025-03-13T01:36:42Z
author_association: COLLABORATOR
diff_hunk: @@ -809,6 +811,15 @@ class PersistentBufferProjector {
         projectToInputOrImmediatePersistentProducer(
             (int)buffer_i,
             std::vector<Val*>(producers.begin(), producers.end()));
+      } else if (
+          auto upcast_input = scheduler_utils::getUpCastInputOf(buffer)) {
+        auto consumers = ir_utils::consumerTvsOf(buffer);
+        for (auto i : c10::irange(1, consumers.size())) {
+          ir_utils::replaceValInExprInputs(
+              consumers.at(i)->definition(),
+              buffer,
+              RecomputeTv::recompute(buffer, {upcast_input}));
+        }
path: csrc/scheduler/reduction_utils.cpp
user: login: liqiangxl
body: Added comments:
```
        // Similar to projecting to inputs and persistent producers, this logic
        // projects the buffer to its producer when the buffer is the output of
        // an upcast op. This optimization reduces buffer size and can be
        // extended to project to low-precision intermediate tensors, even if
        // the recomputation involves non-cast ops. However, this should be
        // avoided when the recomputation cost outweighs the benefits of reduced
        // register usage.
        // TODO: extend to allow non-cast ops in the recomputation.
```
created_at: 2025-03-13T01:48:08Z
author_association: COLLABORATOR

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

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


### **Description**
- Enhance persistent buffer projection for upcast outputs

- Save register usage by projecting buffers to upcast inputs

- Adjust buffer size calculation for upcast outputs

- Add test case for upcast input projection


___



### **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>reduction_utils.cpp</strong><dd><code>Project buffer to upcast input</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>

csrc/scheduler/reduction_utils.cpp

<li>Modify logic to project buffer to producer if buffer is upcast output<br> <li> Add comments explaining the new logic


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4051/files#diff-94d037e9fcb9914df95a45c2b1e5966bb04f10842beebcdb08a17abaa8b62301">+20/-1</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>utils.cpp</strong><dd><code>Add upcast input detection</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; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/scheduler/utils.cpp

<li>Implement <code>getUpCastInputOf</code> function to identify upcast inputs<br> <li> Adjust buffer size calculation for upcast outputs


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4051/files#diff-078826ec098dc190fe9aef215a8130a7f5d268dc36f9537b2bb9e6f255d589d6">+32/-4</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>utils.h</strong><dd><code>Declare upcast input function</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; &nbsp; </dd></summary>
<hr>

csrc/scheduler/utils.h

- Declare `getUpCastInputOf` function in header


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4051/files#diff-ee71c807e8995e215fce4d916b5de8ee271959c631f37a4c368c03636bd1337c">+5/-0</a>&nbsp; &nbsp; &nbsp; </td>

</tr>
</table></td></tr><tr><td><strong>Tests</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>test_persistent_buffer.cpp</strong><dd><code>Add upcast projection test</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; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

tests/cpp/test_persistent_buffer.cpp

- Add test case for projecting buffer to upcast input


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4051/files#diff-5b1a11694088bdbdaf11a6d4b6cc6b3399aef6e73f477de44a06eda733d04bb5">+57/-1</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/4051/files#diff-94d037e9fcb9914df95a45c2b1e5966bb04f10842beebcdb08a17abaa8b62301R814-R831'><strong>Performance Impact</strong></a>

The new logic for projecting buffers to their producers when they are the output of an upcast operation should be evaluated for performance impact. Ensure that the benefits of reduced register usage outweigh the potential costs of recomputation.</summary>

```c++
} else if (
    auto upcast_input = scheduler_utils::getUpCastInputOf(buffer)) {
  // Similar to projecting to inputs and persistent producers, this logic
  // projects the buffer to its producer when the buffer is the output of
  // an upcast op. This optimization reduces buffer size and can be
  // extended to project to low-precision intermediate tensors, even if
  // the recomputation involves non-cast ops. However, this should be
  // avoided when the recomputation cost outweighs the benefits of reduced
  // register usage.
  // TODO: extend to allow non-cast ops in the recomputation.
  auto consumers = ir_utils::consumerTvsOf(buffer);
  for (auto i : c10::irange(1, consumers.size())) {
    ir_utils::replaceValInExprInputs(
        consumers.at(i)->definition(),
        buffer,
        RecomputeTv::recompute(buffer, {upcast_input}));
  }
}
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4051/files#diff-078826ec098dc190fe9aef215a8130a7f5d268dc36f9537b2bb9e6f255d589d6R561-R577'><strong>Code Clarity</strong></a>

The logic in `getUpCastInputOf` could be made clearer by breaking it into smaller functions or adding more comments to explain each step.</summary>

```c++
TensorView* getUpCastInputOf(const TensorView* tv) {
  // skip if definition is not a unary op
  if (auto uop = dynamic_cast<UnaryOp*>(tv->definition())) {
    // skip if the input is a fusion input or the op is not a cast
    if (uop->input(0)->isFusionInput() ||
        uop->getUnaryOpType() != UnaryOpType::Cast) {
      return nullptr;
    }
    // skip if the cast is not upcast
    auto precisions = ir_utils::getPrecisionOfProducerConsumerTensors(uop);
    if (!precisions.has_value() || precisions->first >= precisions->second) {
      return nullptr;
    }
    return uop->input(0)->as<TensorView>();
  }
  return nullptr;
}
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4051/files#diff-078826ec098dc190fe9aef215a8130a7f5d268dc36f9537b2bb9e6f255d589d6R991-R997'><strong>Precision Handling</strong></a>

Ensure that the precision handling in `getPersistentBufferSizeOfTensor` is correct and does not introduce any unintended behavior, especially when dealing with different data types.</summary>

```c++
if (auto upcast_input = getUpCastInputOf(buffer)) {
  dtype_size = dataTypeSize(
      upcast_input->getDataType().value(), runtime_info.getIndexType());
} else {
  dtype_size = dataTypeSize(
      buffer->getDataType().value(), runtime_info.getIndexType());
}
```

</details>

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

author_association: NONE
user: login: liqiangxl
created_at: 2025-03-09T02:33:40Z
body: !test
author_association: COLLABORATOR
user: login: liqiangxl
created_at: 2025-03-10T14:39:11Z
body: !test
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-03-12T21:07:30Z
body: Can you explain why the input projection isn't enough? Also, why does using a common cast op or separate unique ops make any change?
author_association: COLLABORATOR
user: login: liqiangxl
created_at: 2025-03-12T21:41:09Z
body: > Can you explain why the input projection isn't enough? 

(1) For example, in the [newly added test](https://github.com/NVIDIA/Fuser/blob/303f051966f23d043460446187dac323fc962bab/tests/cpp/test_persistent_buffer.cpp#L1564), the tensor with data type `bool` is not an input. So input projection can't solve this issue.

> Also, why does using a common cast op or separate unique ops make any change?

(2) `persistentBuffers()` only checks immediate `consumer-producer` pair. For example, in fusion
```
  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = set(tv0);
  auto tv2 = gt(tv1, IrBuilder::create<Val>(0.5));
  auto tv3 = castOp(DataType::Float, tv2);
  auto tv4 = sum(tv3, {1});
  auto tv5 = broadcast(tv4, {false, true});
  auto tv6 = add(tv5, tv3);
  fusion.addOutput(tv6);
```
it only checks the ID mapping between the following pairs, `tv0-tv1`, `tv1-tv2`, `tv2-tv3`, `tv3-tv4`, `tv3-tv6`, `tv4-tv5`, and `tv5-tv6`. The detected persistent buffer is `tv3` from pairs  `tv3-tv4` and `tv3-tv6`.

If we do sperate cast `auto tv6 = add(tv5, tv3);` becomes:
```
  auto tv7 = castOp(DataType::Float, tv2);
  auto tv6 = add(tv5, tv7);

```
`persistentBuffers()` checks `tv2-tv3` and `tv2-tv7`, tv2 is the buffer.

author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-03-12T22:15:18Z
body: It seems to me that the real problem is the algorithm to choose which tensor to save for a given normalization pattern. For example, which tensor would be cached in the following fusion? Would `tv2` be chosen?

```
  auto tv0 = makeSymbolicTensor(2);
  fusion.addInput(tv0);
  auto tv1 = set(tv0);
  auto tv2 = gt(tv1, IrBuilder::create<Val>(0.5));
  auto tv3 = castOp(DataType::Float, tv2);
  auto tv3_1 = set(tv3);
  auto tv4 = sum(tv3_1, {1});
  auto tv5 = broadcast(tv4, {false, true});
  auto tv6 = add(tv5, tv3_1);
  fusion.addOutput(tv6);
```


author_association: COLLABORATOR
user: login: liqiangxl
created_at: 2025-03-12T23:39:31Z
body: > tv3_1



> It seems to me that the real problem is the algorithm to choose which tensor to save for a given normalization pattern. For example, which tensor would be cached in the following fusion? Would `tv2` be chosen?
> 
> ```
>   auto tv0 = makeSymbolicTensor(2);
>   fusion.addInput(tv0);
>   auto tv1 = set(tv0);
>   auto tv2 = gt(tv1, IrBuilder::create<Val>(0.5));
>   auto tv3 = castOp(DataType::Float, tv2);
>   auto tv3_1 = set(tv3);
>   auto tv4 = sum(tv3_1, {1});
>   auto tv5 = broadcast(tv4, {false, true});
>   auto tv6 = add(tv5, tv3_1);
>   fusion.addOutput(tv6);
> ```

No, in this case, `tv3_1` is the buffer. If we want to keep `tv2` as the persistent buffer, we need to extend this PR to allow nested unary ops, e.g. unaryOp(unaryOp(...)). Since we don't know the cost of these unary ops, so I only implemented to allow castOp.
Or even more aggressive, try to find a procuder which has a smaller data type.
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-03-13T00:17:04Z
body: This is exactly what I was concerned about. The cast insertion and reversion effectively exposed the limitation of the persistent buffer analysis. Not reverting avoids the scheduling failure, but it's also important to acknowledge the limitation and think about what we could do.

> Since we don't know the cost of these unary ops

Not sure why it's important. We already project to inputs, maybe except when there's an expensive op. All we would need to do seems to be finding the smallest buffer along the path from the original persistent buffer to the fusion inputs. 

Again, I'm not totally against this PR. Seems like this is just enough for the immediate problem, and more general solutions may be desirable but probably not urgent.

Am I missing something? Does this make sense?

author_association: COLLABORATOR
user: login: liqiangxl
created_at: 2025-03-13T00:31:07Z
body: > This is exactly what I was concerned about. The cast insertion and reversion effectively exposed the limitation of the persistent buffer analysis. Not reverting avoids the scheduling failure, but it's also important to acknowledge the limitation and think about what we could do.
> 
> > Since we don't know the cost of these unary ops
> 
> Not sure why it's important. We already project to inputs, maybe except when there's an expensive op. All we would need to do seems to be finding the smallest buffer along the path from the original persistent buffer to the fusion inputs.
> 
> Again, I'm not totally against this PR. Seems like this is just enough for the immediate problem, and more general solutions may be desirable but probably not urgent.
> 
> Am I missing something? Does this make sense?

Your summary captures the key points well. This PR provides a straightforward solution to the immediate problem, which is effective for now. While a more general solution would be ideal in the long run, it isn’t an immediate priority.
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-03-13T01:17:30Z
body: > > This is exactly what I was concerned about. The cast insertion and reversion effectively exposed the limitation of the persistent buffer analysis. Not reverting avoids the scheduling failure, but it's also important to acknowledge the limitation and think about what we could do.
> > > Since we don't know the cost of these unary ops
> > 
> > 
> > Not sure why it's important. We already project to inputs, maybe except when there's an expensive op. All we would need to do seems to be finding the smallest buffer along the path from the original persistent buffer to the fusion inputs.
> > Again, I'm not totally against this PR. Seems like this is just enough for the immediate problem, and more general solutions may be desirable but probably not urgent.
> > Am I missing something? Does this make sense?
> 
> Your summary captures the key points well. This PR provides a straightforward solution to the immediate problem, which is effective for now. While a more general solution would be ideal in the long run, it isn’t an immediate priority.

This is an important context. Without this information, it would have been totally unclear what this PR addressed and why. Please expand code comments.
author_association: COLLABORATOR
user: login: liqiangxl
created_at: 2025-03-14T12:03:30Z
body: !test
author_association: COLLABORATOR
user: login: liqiangxl
created_at: 2025-03-14T15:43:25Z
body: !test
author_association: COLLABORATOR
user: login: liqiangxl
created_at: 2025-03-19T02:36:34Z
body: !test
author_association: COLLABORATOR