Repository: NVIDIA/Fuser
Pull request #4158: Unroll for-loops whose indices are used for indexing register tensors
Author: login: naoyam
Created: 2025-03-31T17:12:29Z | Merged: 2025-03-31T20:37:45Z

This PR adds a `TensorIndexer` version of [ensureStaticIndexing](https://github.com/NVIDIA/Fuser/blob/6ff60e2a320733a2f49de57007d6bb45000107cd/csrc/index_compute.cpp#L1194).



Diff:
diff --git a/csrc/id_model/indexing.cpp b/csrc/id_model/indexing.cpp
index 4b700d8bd6f..bf2608a14ea 100644
--- a/csrc/id_model/indexing.cpp
+++ b/csrc/id_model/indexing.cpp
@@ -143,7 +143,7 @@ std::vector<Val*> TensorIndexer::getIndexFor(
     const std::vector<ForLoop*>& for_loops) const {
   auto info = computeIndex(expr, index_ids, for_loops);
   const auto& replacement_map = getIndexReplacementMap(
-      expr, as_consumer, info.loop_domains, for_loops, info.index_map);
+      expr, as_consumer, info.loop_ids, for_loops, info.index_map);
 
   // Note that IDs of index_ids may be mapped as the traversal graph
   // is the AlmostExact graph.
@@ -187,7 +187,7 @@ Val* TensorIndexer::getLinearIndex(
   const auto& alloc_info = getIndexAllocationInfo(tv);
 
   const auto [contig_indices, contig_strides] = getContigIndexFor(
-      expr, as_consumer, alloc_info, for_loops, override_index);
+      tv, expr, as_consumer, alloc_info, for_loops, override_index);
 
   // Linearize the indices with strides.
   Val* linear_index = tv->fusion()->zeroVal();
@@ -228,10 +228,10 @@ IndexingInfo TensorIndexer::computeIndex(
     const Expr* expr,
     const std::vector<IterDomain*>& index_ids,
     const std::vector<ForLoop*>& for_loops) const {
-  const auto loop_domains = getLoopIds(expr, id_model_);
+  const auto loop_ids = getLoopIds(expr, id_model_);
   const ExprPath<ExprGroup> traversal_path = getIndexingPath(expr, index_ids);
   const std::unordered_map<ValGroup, Val*> initial_index_map =
-      getInitialIndexMap(loop_domains, for_loops);
+      getInitialIndexMap(loop_ids, for_loops);
 
   IdGraphIndexCompute index_compute(traversalGraph(), initial_index_map);
 
@@ -242,7 +242,7 @@ IndexingInfo TensorIndexer::computeIndex(
   std::unordered_map<ValGroup, ValGroups> loop_group_dependencies;
 
   // Initialize the loop dependency mappings
-  for (const auto& loop_domain : loop_domains) {
+  for (const auto& loop_domain : loop_ids) {
     const auto& traversal_graph_group = traversalGraph().toGroup(loop_domain);
     const auto& loop_graph_group =
         id_model_.idGraph(IdMappingMode::LOOP).toGroup(loop_domain);
@@ -277,7 +277,7 @@ IndexingInfo TensorIndexer::computeIndex(
   }
 
   IndexingInfo info{
-      loop_domains, traversal_path, index_map, loop_group_dependencies};
+      loop_ids, index_ids, traversal_path, index_map, loop_group_dependencies};
   return info;
 }
 
@@ -597,8 +597,33 @@ std::pair<std::vector<ValGroup>, std::vector<Val*>> TensorIndexer::
       {contig_strides.begin(), contig_strides.end()}};
 }
 
+ValGroups TensorIndexer::getUsedLoopGroups(
+    const IndexingInfo& index_info) const {
+  ValGroups used_loop_groups;
+  for (const auto& index_id : index_info.index_ids) {
+    const ValGroups& loop_groups = index_info.loop_group_dependencies.at(
+        traversalGraph().toGroup(index_id));
+    used_loop_groups.pushBack(loop_groups);
+  }
+  return used_loop_groups;
+}
+
+void TensorIndexer::ensureStaticIndexing(
+    const std::vector<ForLoop*>& for_loops,
+    const IndexingInfo& index_info) const {
+  const ValGroups used_loop_groups = getUsedLoopGroups(index_info);
+
+  for (auto for_loop : for_loops) {
+    if (used_loop_groups.has(id_model_.idGraph(IdMappingMode::LOOP)
+                                 .toGroup(for_loop->iter_domain()))) {
+      for_loop->requireUnroll();
+    }
+  }
+}
+
 std::pair<std::vector<Val*>, std::vector<Val*>> TensorIndexer::
     getContigIndexFor(
+        TensorView* tv,
         const Expr* expr,
         bool as_consumer,
         const AllocationDomainInfo& alloc_info,
@@ -617,7 +642,7 @@ std::pair<std::vector<Val*>, std::vector<Val*>> TensorIndexer::
   }
   const auto& index_map = index_info.index_map;
   const auto& replacement_map = getIndexReplacementMap(
-      expr, as_consumer, index_info.loop_domains, for_loops, index_map);
+      expr, as_consumer, index_info.loop_ids, for_loops, index_map);
 
   std::vector<ValGroup> contig_alloc_groups;
   std::vector<Val*> contig_strides;
@@ -653,6 +678,13 @@ std::pair<std::vector<Val*>, std::vector<Val*>> TensorIndexer::
     result.push_back(replaced_idx);
   }
 
+  // It's a bit confusing for the function named as "getContigIndexFor"
+  // to change the property of ForLoops, but the local variable of
+  // index_info is needed.
+  if (tv->getMemoryType() == MemoryType::Local) {
+    ensureStaticIndexing(for_loops, index_info);
+  }
+
   return {result, contig_strides};
 }
 
diff --git a/csrc/id_model/indexing.h b/csrc/id_model/indexing.h
index b832297b85a..e5eadb2b53a 100644
--- a/csrc/id_model/indexing.h
+++ b/csrc/id_model/indexing.h
@@ -24,7 +24,8 @@
 namespace nvfuser {
 
 struct IndexingInfo {
-  std::vector<IterDomain*> loop_domains;
+  std::vector<IterDomain*> loop_ids;
+  std::vector<IterDomain*> index_ids;
   // Indexing traversal path from loop domains
   ExprPath<ExprGroup> traversal_path;
   // Index mappings of ID groups along the traversal path
@@ -85,12 +86,23 @@ class TensorIndexer {
 
   // Get the contig indices of the given ID groups with their strides
   std::pair<std::vector<Val*>, std::vector<Val*>> getContigIndexFor(
+      TensorView* tv,
       const Expr* expr,
       bool as_consumer,
       const AllocationDomainInfo& alloc_info,
       const std::vector<ForLoop*>& loops,
       const std::unordered_map<IterDomain*, Val*>& override_index) const;
 
+  // Get all ValGroups of the Loop graph representing the used
+  // for-loops for the given indexing
+  ValGroups getUsedLoopGroups(const IndexingInfo& index_info) const;
+
+  // Add "pragma unroll" to for-loops whose loop indices are used for
+  // the given indexing. This is meant to be used for register tensors.
+  void ensureStaticIndexing(
+      const std::vector<ForLoop*>& loops,
+      const IndexingInfo& index_info) const;
+
   // The AlmostExact graph is used since size-1 splits and merges
   // should not affect actual index exprs.
   // Returns non-const reference because indexing may create new domains and
diff --git a/tests/cpp/test_indexing.cpp b/tests/cpp/test_indexing.cpp
index 11f0bff3582..d536888ec08 100644
--- a/tests/cpp/test_indexing.cpp
+++ b/tests/cpp/test_indexing.cpp
@@ -5932,4 +5932,48 @@ TEST_F(IndexingTest, Rng) {
   testValidate(&fusion, outputs, {1}, {randn_sample}, __LINE__, __FILE__);
 }
 
+// Loops should be annotated with "pragma unroll" when their indices
+// are used for indexing of register tensors. This is one example a
+// loop may not be unrolled.
+TEST_F(IndexingTest, StaticIndexing) {
+  EnableOptionsGuard enable_options_guard;
+  EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"});
+
+  auto fusion_ptr = std::make_unique<Fusion>();
+  auto& fusion = *fusion_ptr;
+  FusionGuard fg(fusion_ptr.get());
+
+  auto tv0 = makeSymbolicTensor(1);
+  fusion.addInput(tv0);
+  auto tv1 = set(tv0);
+  auto tv2 = set(tv1);
+  fusion.addOutput(tv2);
+
+  tv1->split(0, 4);
+  tv2->split(0, 4);
+
+  tv1->inlineAt(1);
+  // Unswitched loops are not unrolled by default. This should be
+  // overridden because tv1 is a register tensor.
+  tv1->axis(1)->parallelize(ParallelType::Unswitch);
+
+  // Check if tv1's innermost loop is required to be unrolled
+  class Validator : public kir::IrVisitor {
+   public:
+    using kir::IrVisitor::handle;
+
+    void handle(LoadStoreOp* ldst) override {
+      if (ir_utils::getTvOutput(ldst)->name() == 1) {
+        ASSERT_FALSE(for_loops_.empty());
+        EXPECT_TRUE(for_loops_.back()->isUnrollRequired());
+      }
+    }
+  };
+
+  GpuLower lower(&fusion);
+  kir::Kernel* kernel = lower.run();
+  Validator validator;
+  validator.handle(kernel->topLevelExprs());
+}
+
 } // namespace nvfuser


Reviews:
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-31T17:43:33Z
user: login: jacobhinkle
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-31T20:00:30Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-31T20:23:56Z

Review comments:
diff_hunk: @@ -24,7 +24,8 @@
 namespace nvfuser {
 
 struct IndexingInfo {
-  std::vector<IterDomain*> loop_domains;
+  std::vector<IterDomain*> loop_ids;
path: csrc/id_model/indexing.h
user: login: naoyam
body: Renamed `domains` to `ids`. This is just my (recent) personal preference because "id" is a bit more explicit that it refers to an IterDomain.
created_at: 2025-03-31T17:43:33Z
author_association: COLLABORATOR
diff_hunk: @@ -597,8 +597,37 @@ std::pair<std::vector<ValGroup>, std::vector<Val*>> TensorIndexer::
       {contig_strides.begin(), contig_strides.end()}};
 }
 
+ValGroups TensorIndexer::getUsedLoopGroups(
+    const IndexingInfo& index_info) const {
+  ValGroups used_loop_groups;
+  for (const auto& index_id : index_info.index_ids) {
+    const ValGroups& loop_groups = index_info.loop_group_dependencies.at(
+        traversalGraph().toGroup(index_id));
+    used_loop_groups.pushBack(loop_groups);
+  }
+  return used_loop_groups;
+}
+
+void TensorIndexer::ensureStaticIndexing(
+    const std::vector<ForLoop*>& for_loops,
+    const IndexingInfo& index_info) const {
+  const ValGroups used_loop_groups = getUsedLoopGroups(index_info);
+
+  for (auto for_loop : for_loops) {
+    if (std::any_of(
+            used_loop_groups.begin(),
+            used_loop_groups.end(),
+            [&](const ValGroup& used_loop_group) {
+              return used_loop_group->has(for_loop->iter_domain());
+            })) {
path: csrc/id_model/indexing.cpp
user: login: jacobhinkle
body: Nit: would it be clearer to use `used_loop_groups.has(traversalGraph().toGroup(for_loop->iter_domain())`?
created_at: 2025-03-31T20:00:08Z
author_association: COLLABORATOR
diff_hunk: @@ -597,8 +597,37 @@ std::pair<std::vector<ValGroup>, std::vector<Val*>> TensorIndexer::
       {contig_strides.begin(), contig_strides.end()}};
 }
 
+ValGroups TensorIndexer::getUsedLoopGroups(
+    const IndexingInfo& index_info) const {
+  ValGroups used_loop_groups;
+  for (const auto& index_id : index_info.index_ids) {
+    const ValGroups& loop_groups = index_info.loop_group_dependencies.at(
+        traversalGraph().toGroup(index_id));
+    used_loop_groups.pushBack(loop_groups);
+  }
+  return used_loop_groups;
+}
+
+void TensorIndexer::ensureStaticIndexing(
+    const std::vector<ForLoop*>& for_loops,
+    const IndexingInfo& index_info) const {
+  const ValGroups used_loop_groups = getUsedLoopGroups(index_info);
+
+  for (auto for_loop : for_loops) {
+    if (std::any_of(
+            used_loop_groups.begin(),
+            used_loop_groups.end(),
+            [&](const ValGroup& used_loop_group) {
+              return used_loop_group->has(for_loop->iter_domain());
+            })) {
path: csrc/id_model/indexing.cpp
user: login: naoyam
body: Thanks. Updated with a slight change. `used_loop_groups` is a vector of LoopGraph val groups, whereas `traversalGraph()` returns the AlmostExact graph.
created_at: 2025-03-31T20:23:56Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-31T17:13:18Z
body: <!-- llm action combined persistent comment, dont modify this line -->

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


### **Description**
- Added `ensureStaticIndexing` to unroll loops for register tensors

- Updated `TensorIndexer` to use `loop_ids` and `index_ids`

- Added unit test for static indexing


___



### **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>indexing.cpp</strong><dd><code>Enhance TensorIndexer with static indexing and loop unrolling</code></dd></summary>
<hr>

csrc/id_model/indexing.cpp

<li>Updated <code>computeIndex</code> to use <code>loop_ids</code> and <code>index_ids</code><br> <li> Added <code>ensureStaticIndexing</code> to unroll loops<br> <li> Updated <code>getContigIndexFor</code> to include <code>tv</code> parameter<br> <li> Added <code>getUsedLoopGroups</code> to determine used loop groups


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4158/files#diff-0c47d4d562dba268498573ced9dd7fc7d9638b5882aca4c4cc20e70a7ccec891">+39/-7</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td>
    <details>
      <summary><strong>indexing.h</strong><dd><code>Update IndexingInfo and TensorIndexer declarations</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/id_model/indexing.h

<li>Updated <code>IndexingInfo</code> to include <code>loop_ids</code> and <code>index_ids</code><br> <li> Added <code>getUsedLoopGroups</code> and <code>ensureStaticIndexing</code> declarations


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4158/files#diff-9de5ff259388291ea4c43075b334585933623e6632ffdd9103acf391c97e40c5">+13/-1</a>&nbsp; &nbsp; </td>

</tr>
</table></td></tr><tr><td><strong>Tests</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>test_indexing.cpp</strong><dd><code>Add unit test for static indexing</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; </dd></summary>
<hr>

tests/cpp/test_indexing.cpp

- Added unit test `StaticIndexing` to validate loop unrolling


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4158/files#diff-4ce155b5873f7974a9bb52817ba0d401172eb8520092000370c0938daf9eb79c">+44/-0</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/4158/files#diff-0c47d4d562dba268498573ced9dd7fc7d9638b5882aca4c4cc20e70a7ccec891R146-R146'><strong>Naming Consistency</strong></a>

The variable names `loop_ids` and `loop_domains` are used interchangeably in different parts of the code. Ensure consistency in naming to avoid confusion.</summary>

```c++
expr, as_consumer, info.loop_ids, for_loops, info.index_map);
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4158/files#diff-0c47d4d562dba268498573ced9dd7fc7d9638b5882aca4c4cc20e70a7ccec891R685-R688'><strong>Function Side Effects</strong></a>

The function `getContigIndexFor` modifies the `ForLoop` properties by calling `ensureStaticIndexing`. This side effect is not immediately clear from the function name. Consider refactoring to separate concerns or renaming the function to reflect its behavior.</summary>

```c++
  ensureStaticIndexing(for_loops, index_info);
}

return {result, contig_strides};
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4158/files#diff-9de5ff259388291ea4c43075b334585933623e6632ffdd9103acf391c97e40c5R27-R28'><strong>Struct Member Order</strong></a>

The order of members in `IndexingInfo` struct has changed. Ensure that the order is logical and does not affect performance or readability.</summary>

```c
std::vector<IterDomain*> loop_ids;
std::vector<IterDomain*> index_ids;
```

</details>

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

author_association: NONE
user: login: naoyam
created_at: 2025-03-31T17:45:54Z
body: !test
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-03-31T20:24:31Z
body: !build
author_association: COLLABORATOR