[DO NOT REVIEW] see if it blows up

csrc/fusion_segmenter.cpp

@@ -1816,11 +1816,12 @@ void eraseInputDistinctRootDomains(Fusion* fusion) {
       if (tv->getLoopDomain() == tv->getAllocationDomain()) {
         new_loop = new_alloc;
       } else {
-        NVF_ERROR(
-            tv->getLoopDomain() == tv->getLogicalDomain(),
-            tv,
-            " has an unexpected loop domain:\n",
-            tv->domain()->toString(0, /*loop_only=*/false));
+        // we shouldn't assert on loopdomain of inputs, because they carry no meaning.
+        // NVF_ERROR(
+        //     tv->getLoopDomain() == tv->getLogicalDomain(),
+        //     tv,
+        //     " has an unexpected loop domain:\n",
+        //     tv->domain()->toString(0, /*loop_only=*/false));
 
         new_loop = new_logical_domain;
       }

csrc/id_model/indexing.cpp

@@ -1027,11 +1027,6 @@ bool TensorIndexer::isSupported(Fusion* fusion) {
         if (auto swizzle2d = dynamic_cast<Swizzle2D*>(id->definition())) {
           reason << "Swizzle2D not supported: " << swizzle2d->toString();
           break;
-        } else if (ir_utils::isIndexedConsumerID(tv, id)) {
-          reason << "Indirect indexing of consumer ID not supported: "
-                 << tv->toString() << ", " << id->toString() << ", "
-                 << tv->definition()->toString();
-          break;
         }
       }
     }

csrc/id_model/predicate_indexing.cpp

@@ -21,7 +21,9 @@ std::vector<IterDomain*> getPredicateDomains(
   // domains need to be predicated. Note that the non-divisible split
   // info does not seem to cover non-divisible reduction rfactor
   // splits.
-  std::vector<IterDomain*> predicate_domains = consumer_tv->hasReduction()
+
+  // NOTE: we should categories this as predicate on root, I think reduction is similar to how scatter is being represented
+  std::vector<IterDomain*> predicate_domains = (consumer_tv->hasReduction() || expr->isA<ScatterOp>())
       ? consumer_tv->getMaybeRootDomain()
       : consumer_tv->getLogicalDomain();
 
@@ -51,6 +53,16 @@ std::vector<IterDomain*> getPredicateDomains(
     }
   }
 
+  // NOTE: we don't need to predicate on the ScatterGather ID. we should probably remove it.
+  if (expr->isA<ScatterOp>()) {
+    predicate_domains.erase(
+        std::remove_if(
+            predicate_domains.begin(),
+            predicate_domains.end(),
+            [&](IterDomain* id) -> bool { return id == expr->as<ScatterOp>()->getConsumerLogicalID(); }),
+        predicate_domains.end());
+  }
+
   return predicate_domains;
 }
 

csrc/index_compute.cpp

@@ -1739,13 +1739,18 @@ std::vector<Val*> Index::getStrides(TensorView* tv) {
 std::vector<Val*> Index::getConsumerAllocationIndices(
     const TensorView* tv,
     const std::vector<ForLoop*>& loops,
-    const IndexFromIdGraph& index_from_id_graph) {
+    const IndexFromIdGraph& index_from_id_graph,
+    const std::unordered_map<int, Val*>& override_index) {
   const auto& alloc_dom = tv->getMaybeAllocationDomain();
   auto indexing = index_from_id_graph.index;
 
   std::vector<Val*> alloc_inds(
       alloc_dom.size(), GpuLower::current()->kernel()->zeroVal());
-  for (const auto i : arange(alloc_dom.size())) {
+  for (const auto i : arange((int)alloc_dom.size())) {
+    if (override_index.count(i)) {
+      alloc_inds[i] = nullptr;
+      continue;
+    }
     // See a comment in indexing to allocation domains in
     // getGlobalProducerIndex.
     if (alloc_dom[i]->isReduction() || alloc_dom[i]->isBroadcast() ||
@@ -1898,7 +1903,7 @@ std::vector<Val*> Index::getGlobalConsumerStridedIndices(
   // if we need to override index, we need to generate the index from each
   // allocation axis firstly.
   auto alloc_inds =
-      getConsumerAllocationIndices(consumer_tv, loops, index_from_id_graph);
+      getConsumerAllocationIndices(consumer_tv, loops, index_from_id_graph, override_index);
 
   // Global striding
   auto vectorize_shift =
@@ -1908,8 +1913,10 @@ std::vector<Val*> Index::getGlobalConsumerStridedIndices(
   for (const auto i : arange(alloc_inds.size())) {
     auto override_it = override_index.find((int)i);
     if (override_it != override_index.end()) {
+      NVF_ERROR(alloc_inds[i] == nullptr);
       alloc_inds[i] = override_it->second;
     }
+    NVF_ERROR(alloc_inds[i] != nullptr);
     if (alloc_inds[i]->isZeroInt()) {
       continue;
     } else {
@@ -2326,7 +2333,7 @@ kir::TensorIndex* Index::getConsumerIndex(
     bool generate_pointer,
     DataType as_type) {
   Val* index = nullptr;
-  if (!ir_utils::hasRootToLoopLinearTransformations(consumer) ||
+  if (!ir_utils::hasRootToLoopLinearTransformations(consumer, override_index) ||
       ir_utils::isCpAsyncBulkLoad(consumer->definition()) ||
       GpuLower::current()->idModelOptions().consumerIndex() ||
       GpuLower::current()->tmemInfo().hasTMemTensor()) {

csrc/index_compute.h

@@ -450,7 +450,8 @@ class Index {
   static std::vector<Val*> getConsumerAllocationIndices(
       const TensorView* tv,
       const std::vector<ForLoop*>& loops,
-      const IndexFromIdGraph& index_from_id_graph);
+      const IndexFromIdGraph& index_from_id_graph,
+      const std::unordered_map<int, Val*>& override_index = {});
 
   // get the allocation indices of a producer tensor
   static std::vector<Val*> getProducerAllocationIndices(

csrc/ir/internal_base_nodes.h

@@ -443,7 +443,8 @@ class TensorDomain : public Val {
       std::vector<IterDomain*> allocation,
       std::vector<IterDomain*> loop_domain,
       std::vector<std::optional<bool>> contiguity = {},
-      std::vector<IterDomain*> additional_ids = {});
+      std::vector<IterDomain*> additional_ids = {},
+      std::vector<IterDomain*> no_loop_ids = {});
 
   TensorDomain(IrBuilderPasskey, const TensorDomain* src);
 
@@ -625,6 +626,13 @@ class TensorDomain : public Val {
     return additional_ids_;
   }
 
+  // Additional IDs that are not on the path from one of
+  // root/logical/allocation/loop domain to another. We need to keep track of
+  // these IDs to ensure that we can find all paths/IDs of interest.
+  const std::vector<IterDomain*>& noLoopIDs() const {
+    return no_loop_ids_;
+  }
+
   // Set the loop domain of this TensorDomain.
   void setLoopDomain(std::vector<IterDomain*> new_loop_domain);
 
@@ -739,6 +747,7 @@ class TensorDomain : public Val {
   // setLoopDomain
   std::vector<IterDomain*> initial_loop_domain_;
   std::vector<IterDomain*> additional_ids_;
+  std::vector<IterDomain*> no_loop_ids_;
 
   std::vector<IterDomain*> no_bcast_domain_;
   std::vector<IterDomain*> no_reduction_domain_;

csrc/ir/internal_nodes.h

@@ -275,6 +275,10 @@ class ScatterOp : public Expr {
 
   IterDomain* getIndexedID() const;
 
+  IterDomain* getConsumerLoopID() const;
+
+  IterDomain* getConsumerLogicalID() const;
+
   ScatterOpType getScatterOpType() const {
     return attribute<ScatterOpType>(1);
   }

csrc/ir/nodes.cpp

@@ -315,7 +315,15 @@ std::string ScatterOp::toInlineString(int indent_size) const {
 }
 
 IterDomain* ScatterOp::getIndexedID() const {
-  return ir_utils::getTvOutput(this)->getLogicalDomain().at(dim());
+  return ir_utils::getTvOutput(this)->getRootDomain().at(dim());
+}
+
+IterDomain* ScatterOp::getConsumerLoopID() const {
+  return ir_utils::getTvOutput(this)->getRootDomain().at(dim() + 1);
+}
+
+IterDomain* ScatterOp::getConsumerLogicalID() const {
+  return ir_utils::getTvOutput(this)->getRootDomain().at(dim());
 }
 
 std::vector<PolymorphicValue> ScatterOp::evaluate(
@@ -3134,7 +3142,8 @@ void validateContiguity(
 void validateLoopDomain(
     const std::vector<IterDomain*>& logical_domain,
     const std::vector<IterDomain*>& loop_domain,
-    const std::vector<IterDomain*>& additional_ids) {
+    const std::vector<IterDomain*>& additional_ids,
+    const std::vector<IterDomain*>& no_reference_ids) {
   // Skip if there's any symbolic ID
   if (std::any_of(
           logical_domain.begin(),
@@ -3159,8 +3168,16 @@ void validateLoopDomain(
   reference.insert(
       reference.end(), additional_ids.begin(), additional_ids.end());
 
+  std::vector<IterDomain*> covered_loop_domain;
+  covered_loop_domain.reserve(loop_domain.size() + no_reference_ids.size());
+  covered_loop_domain.insert(
+      covered_loop_domain.end(), loop_domain.begin(), loop_domain.end());
+  // no_reference_ids are also considered part of the loop domain
+  covered_loop_domain.insert(
+      covered_loop_domain.end(), no_reference_ids.begin(), no_reference_ids.end());
+
   auto [redundant_ids, _, unreachable_reference_ids] =
-      ir_utils::compareDomainWithReference(loop_domain, reference);
+      ir_utils::compareDomainWithReference(covered_loop_domain, reference);
 
   auto empty_or_broadcast = [](const auto& ids) {
     return std::all_of(ids.begin(), ids.end(), [](IterDomain* id) {
@@ -3177,6 +3194,8 @@ void validateLoopDomain(
       empty_or_broadcast(unreachable_reference_ids),
       "Not all logical IDs are covered by loop domain. Loop: ",
       toDelimitedString(loop_domain),
+      ". no loop logical IDs: ",
+      toDelimitedString(no_reference_ids),
       ". Unreachable logical IDs: ",
       toDelimitedString(unreachable_reference_ids));
 }
@@ -3251,7 +3270,7 @@ TensorDomain::TensorDomain(
   NVF_CHECK(
       loop_domain_.empty() == logical_domain_.empty(),
       "logical domain and loop domain can only be both empty or neither empty");
-  validateLoopDomain(logical_domain_, loop_domain_, additional_ids_);
+  validateLoopDomain(logical_domain_, loop_domain_, additional_ids_, no_loop_ids_);
 
   // resetDomains initializes other member variables, required by clang-tidy
   resetDomains();
@@ -3275,7 +3294,7 @@ TensorDomain::TensorDomain(
   NVF_CHECK(
       loop_domain_.empty() == logical_domain_.empty(),
       "logical domain and loop domain can only be both empty or neither empty");
-  validateLoopDomain(logical_domain_, loop_domain_, additional_ids_);
+  validateLoopDomain(logical_domain_, loop_domain_, additional_ids_, no_loop_ids_);
   if (!root_domain_.empty()) {
     ir_utils::validateDomainEquivalence(
         logical_domain_, root_domain_, additional_ids_);
@@ -3292,14 +3311,16 @@ TensorDomain::TensorDomain(
     std::vector<IterDomain*> allocation_domain,
     std::vector<IterDomain*> loop_domain,
     std::vector<std::optional<bool>> contiguity,
-    std::vector<IterDomain*> additional_ids)
+    std::vector<IterDomain*> additional_ids,
+    std::vector<IterDomain*> no_loop_ids)
     : Val(passkey, ValType::TensorDomain, DataType::Null),
       root_domain_(std::move(root_domain)),
       logical_domain_(std::move(logical_domain)),
       allocation_domain_(std::move(allocation_domain)),
       loop_domain_(std::move(loop_domain)),
       initial_loop_domain_(loop_domain_),
       additional_ids_(std::move(additional_ids)),
+      no_loop_ids_(std::move(no_loop_ids)),
       contiguity_(
           contiguity.empty() ? getContiguityFilledWith(maybeAllocation(), false)
                              : std::move(contiguity)) {
@@ -3308,11 +3329,11 @@ TensorDomain::TensorDomain(
   NVF_CHECK(
       loop_domain_.empty() == logical_domain_.empty(),
       "logical domain and loop domain can only be both empty or neither empty");
-  validateLoopDomain(logical_domain_, loop_domain_, additional_ids_);
-  if (!root_domain_.empty()) {
-    ir_utils::validateDomainEquivalence(
-        logical_domain_, root_domain_, additional_ids_);
-  }
+  validateLoopDomain(logical_domain_, loop_domain_, additional_ids_, no_loop_ids_);
+  // if (!root_domain_.empty()) {
+  //   ir_utils::validateDomainEquivalence(
+  //       logical_domain_, root_domain_, additional_ids_);
+  // }
   if (!allocation_domain_.empty()) {
     ir_utils::validateDomainEquivalence(
         logical_domain_, allocation_domain_, additional_ids_);
@@ -3330,6 +3351,7 @@ TensorDomain::TensorDomain(IrBuilderPasskey passkey, const TensorDomain* src)
       loop_domain_(src->loop_domain_),
       initial_loop_domain_(src->initial_loop_domain_),
       additional_ids_(src->additional_ids_),
+      no_loop_ids_(src->no_loop_ids_),
       no_bcast_domain_(src->no_bcast_domain_),
       no_reduction_domain_(src->no_reduction_domain_),
       contiguity_(src->contiguity_),
@@ -3343,6 +3365,7 @@ TensorDomain::TensorDomain(const TensorDomain* src, IrCloner* ir_cloner)
       loop_domain_(ir_cloner->clone(src->loop_domain_)),
       initial_loop_domain_(ir_cloner->clone(src->initial_loop_domain_)),
       additional_ids_(ir_cloner->clone(src->additional_ids_)),
+      no_loop_ids_(ir_cloner->clone(src->no_loop_ids_)),
       no_bcast_domain_(ir_cloner->clone(src->no_bcast_domain_)),
       no_reduction_domain_(ir_cloner->clone(src->no_reduction_domain_)),
       contiguity_(src->contiguity()),
@@ -3902,7 +3925,7 @@ std::pair<TensorDomain*, TensorDomain*> TensorDomain::rFactor(
 }
 
 void TensorDomain::setLoopDomain(std::vector<IterDomain*> new_loop_domain) {
-  validateLoopDomain(logical(), new_loop_domain, additionalIDs());
+  validateLoopDomain(logical(), new_loop_domain, additionalIDs(), noLoopIDs());
   loop_domain_ = std::move(new_loop_domain);
   initial_loop_domain_ = loop_domain_;
   resetDomains();

csrc/ir/utils.cpp

@@ -781,6 +781,29 @@ bool isIndexSelectLookupTv(const TensorView* tv) {
   return false;
 }
 
+bool isScatterSelfTv(const TensorView* tv) {
+  for (auto expr : tv->uses()) {
+    if (expr->isA<ScatterOp>()) {
+      auto idx_sel = expr->as<ScatterOp>();
+      if (idx_sel->selfTv() == tv) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+bool isScatterIndexTv(const TensorView* tv) {
+  for (auto expr : tv->uses()) {
+    if (expr->isA<ScatterOp>()) {
+      auto idx_sel = expr->as<ScatterOp>();
+      if (idx_sel->indexTv() == tv) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
 bool isIndexSelectIndicesTv(const TensorView* tv) {
   for (auto expr : tv->uses()) {
     if (expr->isA<IndexSelectOp>()) {
@@ -1343,7 +1366,7 @@ bool hasUniformSiblings(Expr* expr) {
   return !expr->isOneOf<SdpaFwdOp, SdpaBwdOp>();
 }
 
-bool hasRootToLoopLinearTransformations(const TensorView* tv) {
+bool hasRootToLoopLinearTransformations(const TensorView* tv, const std::unordered_map<int, Val*>& override_index) {
   auto root = tv->getMaybeRootDomain();
   auto loop = tv->getLoopDomain();
   std::vector<Val*> loop_val(loop.begin(), loop.end());
@@ -1352,11 +1375,16 @@ bool hasRootToLoopLinearTransformations(const TensorView* tv) {
   std::unordered_set<Val*> all_ids_set(all_ids_vec.begin(), all_ids_vec.end());
   auto alloc = tv->getMaybeAllocationDomain();
   auto logical = tv->getLogicalDomain();
-  bool all_alloc_id_on_path = std::all_of(
-      alloc.begin(), alloc.end(), [&](Val* v) { return all_ids_set.count(v); });
-  bool all_logical_id_on_path =
-      std::all_of(logical.begin(), logical.end(), [&](Val* v) {
-        return all_ids_set.count(v);
+
+  bool all_alloc_id_on_path = std::ranges::all_of(
+      nvfuser::views::enumerate_view(alloc), [&](const auto& enumerator) {
+        const auto& [index, v] = enumerator;
+        return override_index.count(index) || all_ids_set.count(v);
+      });
+  bool all_logical_id_on_path = std::ranges::all_of(
+      nvfuser::views::enumerate_view(logical), [&](const auto& enumerator) {
+        const auto& [index, v] = enumerator;
+        return override_index.count(index) || all_ids_set.count(v);
       });
   return all_alloc_id_on_path && all_logical_id_on_path;
 }

csrc/ir/utils.h

@@ -440,6 +440,9 @@ bool isIndexSelectLookupTv(const TensorView* tv);
 // Check if the given tv is third argment of indexSelect(lookup, dim, indices)
 bool isIndexSelectIndicesTv(const TensorView* tv);
 
+bool isScatterSelfTv(const TensorView* tv);
+bool isScatterIndexTv(const TensorView* tv);
+
 bool isGatherLookupTv(const Val* tv);
 
 std::string varName(const Val* val);
@@ -749,7 +752,7 @@ inline bool isMemorySharedAcross(
 //! Check if the given tv has a root domain -> loop domain linear
 //! transformation. This is a temporary check used to incrementally enable
 //! IdModel. Eventually, this should be removed.
-bool hasRootToLoopLinearTransformations(const TensorView* tv);
+bool hasRootToLoopLinearTransformations(const TensorView* tv, const std::unordered_map<int, Val*>& override_index = {});
 
 //! In addition to the above hasRootToLoopLinearTransformations, it
 //! also checks the loop domain has any extra domain