Repository: NVIDIA/Fuser
Pull request #4108: Enable the new indexer for gather-like indirect indexing (Part 1)
Author: login: naoyam
Created: 2025-03-19T19:37:06Z | Merged: 2025-03-20T21:16:15Z

This PR enables indexing of ops like `index_select` and `take_along_axis` with `TensorIndexer`. This first PR adds minimal logic with the following restrictions:

- No support for `gather`. Both `gather` and `take_along_axis` are represented as `GatherOp` with the only difference of the [exactness](https://github.com/NVIDIA/Fuser/blob/main/csrc/ir/internal_nodes.h#L136-L138). Only the exact case is supported.
- No support for indexing intermediate tensors. Some more refactoring of allocation domains is needed.
- No support for scatter. I actually wonder if we should just remove the current implementation as there's probably not much we could reuse.


Diff:
diff --git a/csrc/id_model/indexing.cpp b/csrc/id_model/indexing.cpp
index 75585db7f72..3411bbfba18 100644
--- a/csrc/id_model/indexing.cpp
+++ b/csrc/id_model/indexing.cpp
@@ -876,7 +876,8 @@ std::vector<Val*> TensorIndexer::getIndexFor(
 Val* TensorIndexer::getLinearIndex(
     TensorView* tv,
     const Expr* expr,
-    const std::vector<ForLoop*>& for_loops) const {
+    const std::vector<ForLoop*>& for_loops,
+    const std::unordered_map<IterDomain*, Val*>& override_index) const {
   NVF_ERROR(tv != nullptr);
   NVF_ERROR(expr != nullptr);
   NVF_ERROR(
@@ -895,8 +896,8 @@ Val* TensorIndexer::getLinearIndex(
 
   const auto alloc_info = getIndexingAllocationInfo(tv);
 
-  const auto [contig_indices, contig_strides] =
-      getContigIndexFor(expr, as_consumer, alloc_info, for_loops);
+  const auto [contig_indices, contig_strides] = getContigIndexFor(
+      expr, as_consumer, alloc_info, for_loops, override_index);
 
   // Linearize the indices with strides.
   Val* linear_index = tv->fusion()->zeroVal();
@@ -1317,8 +1318,19 @@ std::pair<std::vector<Val*>, std::vector<Val*>> TensorIndexer::
         const Expr* expr,
         bool as_consumer,
         const IndexingAllocationInfo& alloc_info,
-        const std::vector<ForLoop*>& for_loops) const {
-  auto index_info = computeIndex(expr, alloc_info.domains, for_loops);
+        const std::vector<ForLoop*>& for_loops,
+        const std::unordered_map<IterDomain*, Val*>& override_index) const {
+  std::vector<IterDomain*> indexed_ids;
+  indexed_ids.reserve(alloc_info.domains.size());
+  for (const auto& id : alloc_info.domains) {
+    if (!override_index.count(id)) {
+      indexed_ids.push_back(id);
+    }
+  }
+  auto index_info = computeIndex(expr, indexed_ids, for_loops);
+  for (const auto& [indexed_id, index] : override_index) {
+    index_info.index_map.emplace(traversalGraph().toGroup(indexed_id), index);
+  }
   const auto& index_map = index_info.index_map;
   const auto& replacement_map = getIndexReplacementMap(
       expr, as_consumer, index_info.loop_domains, for_loops, index_map);
diff --git a/csrc/id_model/indexing.h b/csrc/id_model/indexing.h
index 21d9a931a96..56802f6e01f 100644
--- a/csrc/id_model/indexing.h
+++ b/csrc/id_model/indexing.h
@@ -63,11 +63,18 @@ class TensorIndexer {
   // as a consumer or a producer. The predicate indexing will have a
   // separate interface.
   //
-  // The actual for-loops are required for handling circular buffering
+  // The actual for-loops are required for handling circular
+  // buffering.
+  //
+  // The override_index parameter is used to enable indirect indexing
+  // such as index_select. When the parameter is given, the index of a
+  // ValGroup corresponding to a given iter domain is given by its
+  // mapped index Val.
   Val* getLinearIndex(
       TensorView* tv,
       const Expr* expr,
-      const std::vector<ForLoop*>& loops) const;
+      const std::vector<ForLoop*>& loops,
+      const std::unordered_map<IterDomain*, Val*>& override_index = {}) const;
 
   // Get the index of a loop domain.
   Val* getLoopIndex(IterDomain* loop_id, const std::vector<ForLoop*>& for_loops)
@@ -85,7 +92,8 @@ class TensorIndexer {
       const Expr* expr,
       bool as_consumer,
       const IndexingAllocationInfo& alloc_info,
-      const std::vector<ForLoop*>& loops) const;
+      const std::vector<ForLoop*>& loops,
+      const std::unordered_map<IterDomain*, Val*>& override_index) const;
 
   // The AlmostExact graph is used since size-1 splits and merges
   // should not affect actual index exprs.
diff --git a/csrc/index_compute.cpp b/csrc/index_compute.cpp
index eb0e7eb03d6..897cda002d6 100644
--- a/csrc/index_compute.cpp
+++ b/csrc/index_compute.cpp
@@ -2160,7 +2160,7 @@ kir::TensorIndex* Index::getProducerIndex(
       GpuLower::current()->idModelOptions().producerIndex() ||
       GpuLower::current()->tmemInfo().hasTMemTensor()) {
     index = GpuLower::current()->tensorIndexer().getLinearIndex(
-        producer, consumer->definition(), loops);
+        producer, consumer->definition(), loops, override_index);
     if (generate_pointer) {
       auto address_offset = index;
       if (producer->getMemoryType() == MemoryType::Shared) {
diff --git a/tests/cpp/test_gather.cpp b/tests/cpp/test_gather.cpp
index b16886c2d27..12da7e79925 100644
--- a/tests/cpp/test_gather.cpp
+++ b/tests/cpp/test_gather.cpp
@@ -87,6 +87,13 @@ TEST_F(GatherTest, GatherAllRankAllSelectedDim) {
                                            : gather(tv1, dim, tv_idx);
         fusion.addOutput(tv_out);
 
+        if (is_take_along) {
+          EnableOptionsGuard::getCurOptions().set(
+              EnableOption::IdModel, {"all"});
+        } else {
+          EnableOptionsGuard::getCurOptions().unset(EnableOption::IdModel);
+        }
+
         auto input_dims = randomVector(2, max_dim_size, rank);
         auto index_dims =
             randomIndexVector(input_dims, 1, rank, is_take_along, dim);
@@ -123,6 +130,13 @@ TEST_F(GatherTest, GatherAddMul) {
         auto tv_out = mul(tv_gather, tv_add);
         fusion.addOutput(tv_out);
 
+        if (is_take_along) {
+          EnableOptionsGuard::getCurOptions().set(
+              EnableOption::IdModel, {"all"});
+        } else {
+          EnableOptionsGuard::getCurOptions().unset(EnableOption::IdModel);
+        }
+
         auto input_dims = randomVector(2, max_dim_size, rank);
         auto index_dims =
             randomIndexVector(input_dims, 1, rank, is_take_along, dim);
@@ -163,6 +177,13 @@ TEST_F(GatherTest, AddGatherSumAdd) {
 
         fusion.addOutput(tv_out);
 
+        if (is_take_along) {
+          EnableOptionsGuard::getCurOptions().set(
+              EnableOption::IdModel, {"all"});
+        } else {
+          EnableOptionsGuard::getCurOptions().unset(EnableOption::IdModel);
+        }
+
         auto input_dims = randomVector(2, max_dim_size, rank);
         auto index_dims =
             randomIndexVector(input_dims, 1, rank, is_take_along, dim);
@@ -213,6 +234,13 @@ TEST_F(GatherTest, GatherSumAdd) {
 
         fusion.addOutput(tv_out);
 
+        if (is_take_along) {
+          EnableOptionsGuard::getCurOptions().set(
+              EnableOption::IdModel, {"all"});
+        } else {
+          EnableOptionsGuard::getCurOptions().unset(EnableOption::IdModel);
+        }
+
         auto input_dims = randomVector(2, max_dim_size, rank);
         auto index_dims =
             randomIndexVector(input_dims, 1, rank, is_take_along, dim);
@@ -255,6 +283,13 @@ TEST_F(GatherTest, GatherAddMulHugeSize) {
         auto tv_out = mul(tv_gather, tv_add);
         fusion.addOutput(tv_out);
 
+        if (is_take_along) {
+          EnableOptionsGuard::getCurOptions().set(
+              EnableOption::IdModel, {"all"});
+        } else {
+          EnableOptionsGuard::getCurOptions().unset(EnableOption::IdModel);
+        }
+
         auto input_dims = randomVector(2, max_dim_size, rank);
         auto index_dims =
             randomIndexVector(input_dims, 1, rank, is_take_along, dim);
@@ -357,6 +392,8 @@ TEST_F(GatherTest, TakeAlongBroadcastIndex) {
     auto tv5 = add(tv4, tv2);
     fusion.addOutput(tv5);
 
+    EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"});
+
     std::vector<int64_t> input_dims{10, 11, 12};
     std::vector<int64_t> index_dims{index_dim};
     std::vector<int64_t> out_dims = input_dims;
@@ -417,6 +454,13 @@ TEST_F(GatherTest, GatherBroadcastInput) {
         auto tv5 = add(tv4, tv2);
         fusion.addOutput(tv5);
 
+        if (is_take_along) {
+          EnableOptionsGuard::getCurOptions().set(
+              EnableOption::IdModel, {"all"});
+        } else {
+          EnableOptionsGuard::getCurOptions().unset(EnableOption::IdModel);
+        }
+
         auto options =
             at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
         auto options_i =
@@ -726,6 +770,8 @@ TEST_F(GatherTest, TakeAlongAxisIntermediateTensorNormalization1) {
   auto tv6 = takeAlongAxis(tv4, tv5, 1);
   fusion.addOutput(tv6);
 
+  at::manual_seed(0);
+
   auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
   auto options_i = at::TensorOptions().dtype(at::kLong).device(at::kCUDA, 0);
   auto t0 = at::randn(shape, options);
diff --git a/tests/cpp/test_index_select.cpp b/tests/cpp/test_index_select.cpp
index 012903eb29d..a84747ac4de 100644
--- a/tests/cpp/test_index_select.cpp
+++ b/tests/cpp/test_index_select.cpp
@@ -71,7 +71,13 @@ void checkIndexSelectVectorization(
 
 } // namespace
 
-using IndexSelectTest = NVFuserTest;
+class IndexSelectTest : public NVFuserTest {
+ protected:
+  void SetUp() override {
+    EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"});
+    NVFuserTest::SetUp();
+  }
+};
 
 TEST_F(IndexSelectTest, Simple1) {
   for (int i = 0; i < 5; ++i) {
diff --git a/tests/cpp/test_select.cpp b/tests/cpp/test_select.cpp
index 7cbe0ff1c6b..3d8e45b4c10 100644
--- a/tests/cpp/test_select.cpp
+++ b/tests/cpp/test_select.cpp
@@ -15,7 +15,13 @@
 
 namespace nvfuser {
 
-using SelectTest = NVFuserTest;
+class SelectTest : public NVFuserTest {
+ protected:
+  void SetUp() override {
+    EnableOptionsGuard::getCurOptions().set(EnableOption::IdModel, {"all"});
+    NVFuserTest::SetUp();
+  }
+};
 
 TEST_F(SelectTest, Pointwise) {
   auto fusion_ptr = std::make_unique<Fusion>();


Reviews:
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T23:20:11Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T23:20:36Z
user: login: zasdfgbnm
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-20T05:29:58Z

Review comments:
diff_hunk: @@ -726,6 +770,8 @@ TEST_F(GatherTest, TakeAlongAxisIntermediateTensorNormalization1) {
   auto tv6 = takeAlongAxis(tv4, tv5, 1);
   fusion.addOutput(tv6);
 
+  at::manual_seed(0);
path: tests/cpp/test_gather.cpp
user: login: naoyam
body: Not directly related, but this test exhibited non-deterministic rounding errors. This should make it reproducible.
created_at: 2025-03-19T23:20:11Z
author_association: COLLABORATOR
path: tests/cpp/test_gather.cpp
user: login: naoyam
body: Selectively enabled TensorIndexer when supported.
created_at: 2025-03-19T23:20:36Z
author_association: COLLABORATOR

PR comments:
user: login: naoyam
created_at: 2025-03-19T19:38:29Z
body: !test --diff
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-03-19T19:39:20Z
body: <!-- llm action combined persistent comment, dont modify this line -->

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


### **Description**
- Added support for indirect indexing with TensorIndexer.

- Enabled new indexer for gather-like operations.

- Introduced `override_index` parameter for indirect indexing.

- Updated tests to enable IdModel for gather and take_along_axis.


___



### **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>Add override_index parameter for indirect indexing</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/id_model/indexing.cpp

<li>Added <code>override_index</code> parameter to <code>getLinearIndex</code> and <br><code>getContigIndexFor</code>.<br> <li> Modified <code>getContigIndexFor</code> to handle <code>override_index</code> by filtering and <br>updating <code>index_info</code>.


</details>


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

</tr>

<tr>
  <td>
    <details>
      <summary><strong>index_compute.cpp</strong><dd><code>Pass override_index to getLinearIndex</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/index_compute.cpp

- Passed `override_index` to `getLinearIndex` in `getProducerIndex`.


</details>


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

</tr>

<tr>
  <td>
    <details>
      <summary><strong>indexing.h</strong><dd><code>Update function signatures with override_index</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

csrc/id_model/indexing.h

<li>Updated <code>getLinearIndex</code> and <code>getContigIndexFor</code> signatures to include <br><code>override_index</code>.


</details>


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

</tr>
</table></td></tr><tr><td><strong>Tests</strong></td><td><table>
<tr>
  <td>
    <details>
      <summary><strong>test_gather.cpp</strong><dd><code>Enable IdModel for gather and take_along tests</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></summary>
<hr>

tests/cpp/test_gather.cpp

- Enabled IdModel for `take_along` and `gather` tests.


</details>


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

</tr>

<tr>
  <td>
    <details>
      <summary><strong>test_index_select.cpp</strong><dd><code>Enable IdModel in IndexSelectTest</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_index_select.cpp

- Enabled IdModel in `IndexSelectTest` setup.


</details>


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

</tr>

<tr>
  <td>
    <details>
      <summary><strong>test_select.cpp</strong><dd><code>Enable IdModel in SelectTest</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; </dd></summary>
<hr>

tests/cpp/test_select.cpp

- Enabled IdModel in `SelectTest` setup.


</details>


  </td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/4108/files#diff-5e8825ea42e19aafef7ded04bb10e9f9c49d8b8c515d197f3ed9f805ec13ce41">+7/-1</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>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/4108/files#diff-0c47d4d562dba268498573ced9dd7fc7d9638b5882aca4c4cc20e70a7ccec891R876-R880'><strong>Performance Impact</strong></a>

Ensure that the new parameters `for_loops` and `override_index` do not introduce performance regressions. Validate the performance impact of enabling the new indexer for gather-like indirect indexing.</summary>

```c++
Val* TensorIndexer::getLinearIndex(
    TensorView* tv,
    const Expr* expr,
    const std::vector<ForLoop*>& for_loops,
    const std::unordered_map<IterDomain*, Val*>& override_index) const {
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4108/files#diff-3ca379373c5ed245cdaa8252faa22e55337129e4115424213d7ac7b4e2dd3095R84-R133'><strong>Test Coverage</strong></a>

Verify that the added tests cover all cases, including edge cases, for both `gather` and `take_along_axis` operations. Ensure that the tests validate the correctness of the new indexer implementation.</summary>

```c++
        fusion.addInput(tv1);
        fusion.addInput(tv_idx);
        TensorView* tv_out = is_take_along ? takeAlongAxis(tv1, tv_idx, dim)
                                           : gather(tv1, dim, tv_idx);
        fusion.addOutput(tv_out);

        if (is_take_along) {
          EnableOptionsGuard::getCurOptions().set(
              EnableOption::IdModel, {"all"});
        } else {
          EnableOptionsGuard::getCurOptions().unset(EnableOption::IdModel);
        }

        auto input_dims = randomVector(2, max_dim_size, rank);
        auto index_dims =
            randomIndexVector(input_dims, 1, rank, is_take_along, dim);
        at::Tensor t0 = at::randn(input_dims, options);
        at::Tensor idx = at::randint(0, input_dims[dim], index_dims, options_i);
        at::Tensor output = at::zeros(index_dims, options);

        FusionExecutorCache executor_cache(std::move(fusion_ptr));
        auto cg_outputs = executor_cache.runFusionWithInputs({t0, idx});
        testValidate(&fusion, cg_outputs, {t0, idx}, __LINE__, __FILE__);
      }
    }
  }
}
// Test the fusion support of gather operator(producer) and elemetwise(consumer)
TEST_F(GatherTest, GatherAddMul) {
  const int max_dim_size = 64;
  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto options_i = at::TensorOptions().dtype(at::kLong).device(at::kCUDA, 0);
  for (const auto is_take_along : {false, true}) {
    for (int rank = 1; rank <= 4; ++rank) {
      for (int dim = 0; dim < rank; ++dim) {
        auto fusion_ptr = std::make_unique<Fusion>();
        Fusion& fusion = *fusion_ptr.get();
        FusionGuard fg(&fusion);

        TensorView* tv1 = makeContigTensor(rank);
        TensorView* tv_idx = makeContigTensor(rank, DataType::Int);
        fusion.addInput(tv1);
        fusion.addInput(tv_idx);
        auto tv_gather = is_take_along ? takeAlongAxis(tv1, tv_idx, dim)
                                       : gather(tv1, dim, tv_idx);
        auto tv_add = add(tv_gather, tv_gather);
        auto tv_out = mul(tv_gather, tv_add);
        fusion.addOutput(tv_out);

        if (is_take_along) {
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/4108/files#diff-9de5ff259388291ea4c43075b334585933623e6632ffdd9103acf391c97e40c5R60-R97'><strong>Documentation</strong></a>

Update the documentation to reflect the new parameters `for_loops` and `override_index` in the `getLinearIndex` and `getContigIndexFor` methods. Provide clear explanations of the purpose and usage of these parameters.</summary>

```c
}

// Get a linear index of a given tensor appearing in a given expr, either
// as a consumer or a producer. The predicate indexing will have a
// separate interface.
//
// The actual for-loops are required for handling circular
// buffering.
//
// The override_index parameter is used to enable indirect indexing
// such as index_select. When the parameter is given, the index of a
// ValGroup corresponding to a given iter domain is given by its
// mapped index Val.
Val* getLinearIndex(
    TensorView* tv,
    const Expr* expr,
    const std::vector<ForLoop*>& loops,
    const std::unordered_map<IterDomain*, Val*>& override_index = {}) const;

// Get the index of a loop domain.
Val* getLoopIndex(IterDomain* loop_id, const std::vector<ForLoop*>& for_loops)
    const;

// Get the index of the given ID groups
std::vector<Val*> getIndexFor(
    const Expr* expr,
    bool as_consumer,
    const std::vector<IterDomain*>& index_ids,
    const std::vector<ForLoop*>& loops) const;

// Get the contig indices of the given ID groups with their strides
std::pair<std::vector<Val*>, std::vector<Val*>> getContigIndexFor(
    const Expr* expr,
    bool as_consumer,
    const IndexingAllocationInfo& alloc_info,
    const std::vector<ForLoop*>& loops,
    const std::unordered_map<IterDomain*, Val*>& override_index) const;

```

</details>

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

author_association: NONE
user: login: naoyam
created_at: 2025-03-19T21:21:32Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-03-19T23:21:25Z
body: The errors with the codegen diff tests are likely unrelated to this PR. @jacobhinkle is looking into it.
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-03-19T23:45:35Z
body: !test --diff
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-03-20T15:05:07Z
body: !test --diff
author_association: COLLABORATOR