Fix contiguity analysis of predicates to match updated contiguity. (#1991)

Author: csarofeen

torch/csrc/jit/codegen/cuda/contiguity.cpp

@@ -390,16 +390,20 @@ ContigIDs::ContigIDs(
     const std::vector<IterDomain*>& ids,
     const std::vector<IterDomain*>& root_domain,
     const std::vector<bool>& root_contiguity,
-    std::unordered_map<IterDomain*, IterDomain*> concrete_to_ref,
+    const std::unordered_set<IterDomain*>& final_ids,
+    const std::unordered_map<IterDomain*, Val*>& index_map,
     const std::unordered_set<Split*>& divisible_splits,
     std::unordered_map<IterDomain*, IterDomain*> p2c_id_map,
-    bool ignore_indexability)
+    bool ignore_indexability,
+    bool ignore_consistent_ordering)
     : root_domain_(root_domain),
       root_contiguity_(root_contiguity),
-      concrete_to_ref_(std::move(concrete_to_ref)),
+      final_ids_(final_ids),
+      index_map_(index_map),
       divisible_splits_(divisible_splits),
       p2c_id_map_(std::move(p2c_id_map)),
       ignore_indexability_(ignore_indexability),
+      ignore_consistent_ordering_(ignore_consistent_ordering),
       non_divisible_id_info_(ids, root_domain_, divisible_splits_) {
   if (ids.size() > 0) {
     // This constructor doesn't provide the following information so it needs to
@@ -419,22 +423,26 @@ ContigIDs::ContigIDs(
     const std::vector<IterDomain*>& ids,
     const std::vector<IterDomain*>& root_domain,
     const std::vector<bool>& root_contiguity,
-    std::unordered_map<IterDomain*, IterDomain*> concrete_to_ref,
+    const std::unordered_set<IterDomain*>& final_ids,
+    const std::unordered_map<IterDomain*, Val*>& index_map,
     const std::unordered_set<Split*>& divisible_splits,
     std::shared_ptr<const ComputeAtMap> ca_map,
     std::shared_ptr<const HaloInfo> halo_info,
     std::shared_ptr<const ConcretizedBroadcastDomains> concrete_info,
     std::unordered_map<IterDomain*, IterDomain*> p2c_id_map,
-    bool ignore_indexability)
+    bool ignore_indexability,
+    bool ignore_consistent_ordering)
     : root_domain_(root_domain),
       root_contiguity_(root_contiguity),
-      concrete_to_ref_(std::move(concrete_to_ref)),
+      final_ids_(final_ids),
+      index_map_(index_map),
       divisible_splits_(divisible_splits),
       ca_map_(ca_map),
       halo_info_(halo_info),
       concrete_info_(concrete_info),
       p2c_id_map_(std::move(p2c_id_map)),
       ignore_indexability_(ignore_indexability),
+      ignore_consistent_ordering_(ignore_consistent_ordering),
       consistent_transform_info_(std::make_unique<const OrderedIdInformation>(
           ids,
           root_domain,
@@ -443,6 +451,10 @@ ContigIDs::ContigIDs(
   build(ids);
 }
 
+ContigIDs ContigIDs::getNonContigIDs() {
+  return ContigIDs({}, {}, {}, {}, {}, {});
+}
+
 void ContigIDs::build(const std::vector<IterDomain*>& ids) {
   if (ids.empty() || root_domain_.empty()) {
     return;
@@ -488,8 +500,12 @@ void ContigIDs::handle(Merge* merge) {
   // If output is not consistently ordered or doesn't solely consume all root
   // domains in its dependencies, then it can't be a contiguously indexable
   // iterdomain.
-  if (!(consistent_transform_info_->isConsistentlyOrdered(merge->out()) &&
-        consistent_transform_info_->exclusivelyConsumesRoots(merge->out()))) {
+  if (!(ignore_consistent_ordering_ ||
+        consistent_transform_info_->isConsistentlyOrdered(merge->out()))) {
+    return;
+  }
+
+  if (!consistent_transform_info_->exclusivelyConsumesRoots(merge->out())) {
     return;
   }
 
@@ -499,6 +515,11 @@ void ContigIDs::handle(Merge* merge) {
     return;
   }
 
+  // If inputs are marked as final, stop
+  if (final_ids_.count(merge->inner()) || final_ids_.count(merge->outer())) {
+    return;
+  }
+
   // Check root domains for contiguity
   auto root_ids_it =
       consistent_transform_info_->idToRootIds().find(merge->out());
@@ -512,17 +533,25 @@ void ContigIDs::handle(Merge* merge) {
 
   VectorOfUniqueEntries<IterDomain*> root_ids = root_ids_it->second;
 
+  bool is_indexing_pass = !ignore_consistent_ordering_;
+
   IterDomain* last_root = nullptr;
   for (auto root_id_i : c10::irange(root_domain_.size())) {
     auto root_id = root_domain_[root_id_i];
     if (root_ids.has(root_id)) {
       // ID found, remove it
       root_ids.erase(root_id);
-      // If the last id isn't contiguous that's fine, we can use the stride of
-      // the last iter domain to multiply the contig index.
-      if (!root_contiguity_[root_id_i] && !root_ids.empty()) {
-        // Otherwise this merge->out() isn't a contiguously indexable ID
-        return;
+      // If we're indexing:
+      // we could still potentially consider this ID linearly indexable, as we
+      // could multiple the index by the last root's stride.
+      //
+      // If we're computing predicates (ignore_consistent_ordering_==true),
+      // then we don't have this same constraint, we can just ignore
+      // contiguity of the roots all together.
+      if (!root_contiguity_[root_id_i] && is_indexing_pass) {
+        if (!root_ids.empty()) {
+          return;
+        }
       }
       last_root = root_id;
     }
@@ -581,7 +610,7 @@ bool ContigIDs::isIndexable(IterDomain* id) const {
   }
   auto c_id =
       ca_map_->getConcreteMappedID(getMappedId(id), IdMappingMode::EXACT);
-  return concrete_to_ref_.find(c_id) != concrete_to_ref_.end();
+  return index_map_.find(c_id) != index_map_.end();
 }
 
 } // namespace cuda

torch/csrc/jit/codegen/cuda/contiguity.h

@@ -161,22 +161,49 @@ class ContigIDs : public OptInDispatch {
       const std::vector<IterDomain*>& ids,
       const std::vector<IterDomain*>& root_domain,
       const std::vector<bool>& root_contiguity,
-      std::unordered_map<IterDomain*, IterDomain*> concrete_to_ref,
+      const std::unordered_set<IterDomain*>& final_ids,
+      const std::unordered_map<IterDomain*, Val*>& index_map,
       const std::unordered_set<Split*>& divisible_splits,
       std::unordered_map<IterDomain*, IterDomain*> p2c_id_map = {},
-      bool ignore_indexability = false);
-
+      bool ignore_indexability = false,
+      bool ignore_consistent_ordering = false);
+
+  //! \param ids IterDomains on the leaves of the domain we're looking for
+  //! contiguous indexing into.
+  //! \param root_domain the root domain of the domain we're looking for
+  //! contiguous indexing into.
+  //! \param root_contiguity the contiguity of the root_domain.
+  //! \param concrete_to_ref concrete ids of the exact map that the reference
+  //! index is using for indexing.
+  //! \param divisible_splits a set of all splits in the fusion that are
+  //! divisible.
+  //! \param ca_map compute at map of the fusion.
+  //! \param halo_info halo information of the fusion.
+  //! \param concrete_info concretized broadcast information of the fusion.
+  //! \param p2c_id_map map from producer to consumer ids used for indexing
+  //! producer tensors.
+  //! \param ignore_consistent_ordering true for actual indexing into tensors
+  //! but false for predicate analysis. Ordering of merges don't matter for
+  //! predicate generation as they don't map to a physical address.
+  //! \param ignore_indexability can only be true if providing a real
+  //! concrete_to_ref map. As what it's checking is if the index is actually
+  //! indexable based on the reference.
   ContigIDs(
       const std::vector<IterDomain*>& ids,
       const std::vector<IterDomain*>& root_domain,
       const std::vector<bool>& root_contiguity,
-      std::unordered_map<IterDomain*, IterDomain*> concrete_to_ref,
+      const std::unordered_set<IterDomain*>& final_ids,
+      const std::unordered_map<IterDomain*, Val*>& index_map,
       const std::unordered_set<Split*>& divisible_splits,
       std::shared_ptr<const ComputeAtMap> ca_map,
       std::shared_ptr<const HaloInfo> halo_info,
       std::shared_ptr<const ConcretizedBroadcastDomains> concrete_info,
       std::unordered_map<IterDomain*, IterDomain*> p2c_id_map = {},
-      bool ignore_indexability = false);
+      bool ignore_indexability = false,
+      bool ignore_consistent_ordering = false);
+
+  //! Return an empty ContigIDs with no contiguous ID
+  static ContigIDs getNonContigIDs();
 
   const std::unordered_set<IterDomain*>& contigIDs() const {
     return contig_ids_;
@@ -191,6 +218,14 @@ class ContigIDs : public OptInDispatch {
     return root_to_indexed_id_;
   }
 
+  VectorOfUniqueEntries<IterDomain*> indexedRootIDs(IterDomain* id) const {
+    auto root_ids_it = consistent_transform_info_->idToRootIds().find(id);
+    if (root_ids_it == consistent_transform_info_->idToRootIds().end()) {
+      return {};
+    }
+    return root_ids_it->second;
+  }
+
  private:
   using OptInDispatch::handle;
 
@@ -233,9 +268,12 @@ class ContigIDs : public OptInDispatch {
   const std::vector<IterDomain*>& root_domain_;
   //! Contiguity of root_domain_
   const std::vector<bool>& root_contiguity_;
-  //! Mapping of concrete to reference domains. If a concrete domain
-  //! is not mapped, it is not indexable as there's no mapped index.
-  const std::unordered_map<IterDomain*, IterDomain*> concrete_to_ref_;
+  //! Domains where indexing/predicates cannot be done with their
+  //! consumers domains
+  const std::unordered_set<IterDomain*>& final_ids_;
+  //! Mapping of concrete domains to indices. Just used to check if
+  //! there's an index for an IterDomain.
+  const std::unordered_map<IterDomain*, Val*> index_map_;
   // Divisible split information as we can still consider iter domains
   // contiguous through divisible splits.
   const std::unordered_set<Split*>& divisible_splits_;
@@ -247,7 +285,9 @@ class ContigIDs : public OptInDispatch {
   //! Producer-to-consumer index map in the case of analyzing replayed
   //! producer tensors
   const std::unordered_map<IterDomain*, IterDomain*> p2c_id_map_;
+
   const bool ignore_indexability_ = false;
+  const bool ignore_consistent_ordering_ = false;
 
   //! Mapping of root domain to bool indicating contiguity
   std::unordered_map<IterDomain*, bool> is_contig_root_;

torch/csrc/jit/codegen/cuda/index_compute.cpp

@@ -594,7 +594,7 @@ IndexCompute::IndexCompute(
           std::move(extent_map),
           std::move(zero_domains),
           std::move(zero_merged_in),
-          ContigIDs({}, {}, {}, {}, {}),
+          ContigIDs::getNonContigIDs(),
           std::move(preferred_paths),
           std::move(halo_extent_map)) {}
 
@@ -2326,103 +2326,71 @@ std::vector<PredicateDomainInfo> getPredicateContigIds(
 
   const auto& consumer_root_domain = consumer_tv->getRootDomain();
 
-  std::vector<IterDomain*> contiguous_ids = consumer_root_domain;
-
-  if (contiguous_ids.empty()) {
+  if (consumer_root_domain.empty()) {
     return std::vector<PredicateDomainInfo>();
   }
 
-  // If root IDs are partial, i.e., start is non-zero and stop is not
-  // equal to extent, predication can't be done with merged domains as
-  // start and stop information is only available with root
-  // domains. Similarly, merged domains don't have enough information
-  // about halo to do correct predication, so they must be excluded.
-  std::unordered_set<IterDomain*> excluded_ids;
+  std::unordered_map<IterDomain*, Val*> concrete_index_map;
+  for (auto entry : consumer_index_map) {
+    auto c_id = gpu_lower->caMap()->getConcreteMappedID(
+        entry.first, IdMappingMode::EXACT);
+    concrete_index_map[c_id] = entry.second;
+  }
 
-  for (auto consumer_root_id : consumer_root_domain) {
-    if (gpu_lower->haloInfo()->getRootAxisInfo(consumer_root_id).hasHalo()) {
-      excluded_ids.insert(consumer_root_id);
-      continue;
-    }
-    if (consumer_root_id->maybePartial()) {
-      excluded_ids.insert(consumer_root_id);
-      continue;
-    }
-    // When consumer_root_id is a broadcast domain, do not allow contig
-    // predication as the merged output is not mapped with the
-    // reference unless the concrete domain is also a broadcast
-    // domain.
-    if (consumer_root_id->isBroadcast() &&
-        !GpuLower::current()
-             ->caMap()
-             ->getConcreteMappedID(consumer_root_id, IdMappingMode::PERMISSIVE)
-             ->isBroadcast()) {
-      excluded_ids.insert(consumer_root_id);
+  std::vector<bool> predicate_contiguity(consumer_root_domain.size(), true);
+  std::unordered_set<IterDomain*> final_ids;
+  for (auto root_i : c10::irange(predicate_contiguity.size())) {
+    auto root_id = consumer_root_domain[root_i];
+    if (root_id->maybePartial()) {
+      final_ids.insert(root_id);
       continue;
     }
     // Shifted or gathered axes need to be predicated at the root domain
     auto shift_expr = dynamic_cast<ShiftOp*>(consumer_tv->definition());
     auto gather_expr = dynamic_cast<GatherOp*>(consumer_tv->definition());
-    if (shift_expr == nullptr && gather_expr == nullptr) {
-      continue;
-    }
-    auto consumer_root_pos = consumer_tv->domain()->rootPosOf(consumer_root_id);
-    if ((shift_expr && shift_expr->offset(consumer_root_pos) != 0) ||
-        (gather_expr && consumer_root_pos < gather_expr->windowShape().size() &&
-         gather_expr->windowShape().at(consumer_root_pos) != 1)) {
-      excluded_ids.insert(consumer_root_id);
+    if ((shift_expr && shift_expr->offset(root_i) != 0) ||
+        (gather_expr && root_i < gather_expr->windowShape().size() &&
+         gather_expr->windowShape().at(root_i) != 1)) {
+      final_ids.insert(root_id);
     }
   }
 
-  // Run through iteration domain history
-  auto exprs = StmtSort::getExprs(
-      consumer_tv->fusion(),
-      {consumer_tv->domain()->domain().begin(),
-       consumer_tv->domain()->domain().end()});
+  ContigIDs contig_finder(
+      consumer_tv->domain()->domain(),
+      consumer_root_domain,
+      predicate_contiguity,
+      final_ids,
+      concrete_index_map,
+      GpuLower::current()->divisbleSplitSet(),
+      GpuLower::current()->caMap(),
+      GpuLower::current()->haloInfo(),
+      GpuLower::current()->concretizedBroadcastDomains(),
+      {},
+      false,
+      true);
 
-  for (auto expr : exprs) {
-    // If not a merge, output is not contiguous
-    if (expr->isA<Merge>()) {
-      auto merge = expr->as<Merge>();
-      auto inner_contig_it = std::find(
-          contiguous_ids.begin(), contiguous_ids.end(), merge->inner());
-      auto outer_contig_it = std::find(
-          contiguous_ids.begin(), contiguous_ids.end(), merge->outer());
+  std::vector<PredicateDomainInfo> contig_id_infos;
+  std::unordered_set<IterDomain*> covered_roots;
 
-      if (excluded_ids.count(merge->inner()) > 0 ||
-          excluded_ids.count(merge->outer()) > 0) {
-        continue;
-      }
+  // Create entries and return them
+  for (auto root_id : consumer_root_domain) {
+    if (covered_roots.count(root_id) > 0) {
+      continue;
+    }
 
-      // Do not try to predicate the merge output domain if the output
-      // domain has not a predicate that is mapped from the reference.
-      // See FusionContigPredicate_CUDA for a concrete example.
-      if (consumer_index_map.find(merge->out()) == consumer_index_map.end()) {
-        continue;
-      }
+    auto contig_id_it = contig_finder.rootToIndexedID().find(root_id);
 
-      if (inner_contig_it != contiguous_ids.end() &&
-          outer_contig_it != contiguous_ids.end()) {
-        // If inner and outer are contiguous, out must be contiguous. Remove
-        // inner and outer, and add out.
-        contiguous_ids.erase(outer_contig_it);
-        contiguous_ids.erase(std::find(
-            contiguous_ids.begin(), contiguous_ids.end(), merge->inner()));
-        contiguous_ids.emplace_back(merge->out());
-      }
-    }
-  }
+    TORCH_INTERNAL_ASSERT(
+        contig_id_it != contig_finder.rootToIndexedID().end(),
+        "Error in predicate contiguity analysis, missing index for root ",
+        root_id->toString());
 
-  std::vector<PredicateDomainInfo> contig_id_infos;
+    auto contig_id = contig_id_it->second;
 
-  // Create entries and return them
-  for (auto contig_id : contiguous_ids) {
     // Pick inputs from the starting domains, i.e.,
     // reference_predicated_root_domain.
-    auto contig_root_vals = IterVisitor::getInputsTo(
-        {contig_id},
-        {consumer_root_domain.begin(), consumer_root_domain.end()});
-    auto contig_root_ids = ir_utils::filterByType<IterDomain>(contig_root_vals);
+    auto contig_root_ids = contig_finder.indexedRootIDs(contig_id);
+    covered_roots.insert(contig_root_ids.begin(), contig_root_ids.end());
     PredicateDomainInfo contig_id_info;
     contig_id_info.id = contig_id;
     contig_id_info.covered_ids = std::unordered_set<IterDomain*>(

torch/csrc/jit/codegen/cuda/iter_visitor.cpp

@@ -214,10 +214,10 @@ void IterVisitor::traverseBetween(
 
 void IterVisitor::traverseTo(
     Fusion* fusion,
-    const std::vector<Val*>& To,
+    const std::vector<Val*>& to,
     bool traverse_all_paths,
     bool traverse_into_members) {
-  traverseBetween(fusion, {}, To, traverse_all_paths, traverse_into_members);
+  traverseBetween(fusion, {}, to, traverse_all_paths, traverse_into_members);
 }
 
 void IterVisitor::traverseHelper(Fusion* fusion, bool traverse_all_paths) {