nvfuser-cu121-torch25 0.2.25.dev20250201__cp312-cp312-manylinux_2_28_x86_64.whl

nvfuser/include/nvfuser/id_model/to_string.h

@@ -0,0 +1,87 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <ir/all_nodes.h>
+#include <val_graph.h>
+
+#include <string>
+#include <vector>
+
+namespace nvfuser {
+
+std::string toString(const std::vector<Val*>& val_group, int indent_size = 0);
+
+std::string toString(
+    const std::vector<IterDomain*>& id_group,
+    int indent_size = 0);
+
+std::string toString(
+    const ValGroup& id_group,
+    int indent_size = 0,
+    bool with_ptr = false);
+
+std::string toString(
+    const std::vector<ValGroup>& id_groups,
+    int indent_size = 0,
+    bool with_ptr = false);
+
+std::string toString(
+    const ValGroups& id_groups,
+    int indent_size = 0,
+    bool with_ptr = false);
+
+std::string toInlineString(const std::vector<ValGroup>& id_groups);
+std::string toInlineString(const ValGroups& id_groups);
+
+std::string toString(const std::vector<Expr*>& expr_group, int indent_size = 0);
+std::string toString(
+    const ExprGroup& expr_group,
+    int indent_size = 0,
+    bool with_ptr = false);
+
+std::string toString(
+    const ValGraph& id_graph,
+    const std::vector<Expr*>& expr_group,
+    int indent_size = 0,
+    bool with_ptr = false);
+std::string toString(
+    const ValGraph& id_graph,
+    const ExprGroup& expr_groups,
+    int indent_size = 0,
+    bool with_ptr = false);
+
+std::string toString(
+    const ValGraph& id_graph,
+    const std::vector<ExprGroup>& expr_groups,
+    int indent_size = 0,
+    bool with_ptr = false);
+std::string toString(
+    const ValGraph& id_graph,
+    const ExprGroups& expr_groups,
+    int indent_size = 0,
+    bool with_ptr = false);
+
+std::string idGroupsString(
+    const ValGraph& id_graph,
+    int indent_size = 0,
+    bool with_ptr = false);
+std::string exprGroupsString(
+    const ValGraph& id_graph,
+    int indent_size = 0,
+    bool with_ptr = false);
+std::string definitionsString(
+    const ValGraph& id_graph,
+    int indent_size = 0,
+    bool with_ptr = false);
+std::string usesString(
+    const ValGraph& id_graph,
+    int indent_size = 0,
+    bool with_ptr = false);
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/transform_replay.h

@@ -0,0 +1,58 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <c10/macros/Export.h>
+
+#include <ir/all_nodes.h>
+
+#include <unordered_map>
+#include <vector>
+
+namespace nvfuser {
+
+// TODO: Consider merging this class with the existing replay
+// classes. The use cases are not exactly the same, so it isn't
+// immediately clear if they could be trivially merge.
+class ReplayTransform : OptInConstDispatch {
+ public:
+  // Replays expression_to_match with the provided ordered_inputs. Inputs should
+  // be ordered as they would be used in provided expression. Returns new
+  // replayed expression.
+  static Expr* replayAs(
+      const std::vector<IterDomain*>& ordered_inputs,
+      const Expr* expression_to_match);
+
+ private:
+  ReplayTransform(
+      const std::vector<IterDomain*>& ordered_inputs,
+      const Expr* expression_to_match);
+
+  using OptInConstDispatch::handle;
+
+  // We're going to replay this split operation on the corresponding ID
+  void handle(const Split* split) final;
+
+  // We're going to replay this merge operation on the corresponding IDs
+  void handle(const Merge* merge) final;
+
+  // We're going to replay this swizzle operation on the corresponding IDs
+  //  if replaying swizzle is enabled.
+  void handle(const Swizzle2D* swizzle_2d) final;
+
+  void handle(const Swizzle* swizzle) final;
+
+  // We're going to replay this resize operation on the corresponding IDs
+  //  if replaying resize is enabled.
+  void handle(const Resize* resize) final;
+
+  Expr* replayed_expr_ = nullptr;
+  const std::vector<IterDomain*>& input_ids_;
+};
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/utils.h

@@ -0,0 +1,176 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <expr_simplifier.h>
+#include <id_model/id_model.h>
+#include <id_model/to_string.h>
+#include <ir/utils.h>
+#include <options.h>
+#include <utils.h>
+
+#include <functional>
+#include <iostream>
+#include <sstream>
+
+namespace nvfuser {
+
+// Options to enable the IdModel-based tensor indexer selectively
+enum class IdModelEnableOption {
+  ConsumerIndex,
+  ProducerIndex,
+  InlinePredicate,
+  UnswitchPredicate,
+  // Uses the loop promotion to generate loops. Indexing and
+  // predication need to be enabled as well.
+  Loop,
+};
+
+inline std::unordered_set<IdModelEnableOption> getIdModelEnabledOptions() {
+  std::unordered_set<IdModelEnableOption> opts;
+
+  if (hasEnableOptionArgument(EnableOption::IdModel, "consumer_index") ||
+      hasEnableOptionArgument(EnableOption::IdModel, "index") ||
+      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
+    opts.insert(IdModelEnableOption::ConsumerIndex);
+  }
+
+  if (hasEnableOptionArgument(EnableOption::IdModel, "producer_index") ||
+      hasEnableOptionArgument(EnableOption::IdModel, "index") ||
+      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
+    opts.insert(IdModelEnableOption::ProducerIndex);
+  }
+
+  if (hasEnableOptionArgument(EnableOption::IdModel, "inline_predicate") ||
+      hasEnableOptionArgument(EnableOption::IdModel, "predicate") ||
+      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
+    opts.insert(IdModelEnableOption::InlinePredicate);
+  }
+
+  if (hasEnableOptionArgument(EnableOption::IdModel, "unswitch_predicate") ||
+      hasEnableOptionArgument(EnableOption::IdModel, "predicate") ||
+      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
+    opts.insert(IdModelEnableOption::UnswitchPredicate);
+  }
+
+  if (hasEnableOptionArgument(EnableOption::IdModel, "loop") ||
+      hasEnableOptionArgument(EnableOption::IdModel, "all")) {
+    opts.insert(IdModelEnableOption::Loop);
+  }
+
+  // Loop requires ConsumerIndex, ProducerIndex, InlinePredicate and
+  // UnswitchPredicate
+  if (opts.find(IdModelEnableOption::Loop) != opts.end()) {
+    NVF_ERROR(
+        opts.find(IdModelEnableOption::ConsumerIndex) != opts.end(),
+        "ConsumerIndex required for Loop");
+    NVF_ERROR(
+        opts.find(IdModelEnableOption::ProducerIndex) != opts.end(),
+        "ProducerIndex required for Loop");
+    NVF_ERROR(
+        opts.find(IdModelEnableOption::InlinePredicate) != opts.end(),
+        "InlinePredicate required for Loop");
+    NVF_ERROR(
+        opts.find(IdModelEnableOption::UnswitchPredicate) != opts.end(),
+        "UnswitchPredicate required for Loop");
+  }
+
+  return opts;
+}
+
+inline bool isIdModelOptionEnabled(IdModelEnableOption option) {
+  const auto opts = getIdModelEnabledOptions();
+  return opts.find(option) != opts.end();
+}
+
+// Get the promotion domain of a given loop domain.
+inline IterDomain* getLoopPromotion(
+    IterDomain* loop_id,
+    const IdModel& id_model) {
+  const auto& loop_graph = id_model.idGraph(IdMappingMode::LOOP);
+  const auto& loop_promotion_map = id_model.loopPromotionMap();
+  const auto& loop_group = loop_graph.toGroup(loop_id);
+
+  auto loop_promotion_map_it = loop_promotion_map.find(loop_group);
+  NVF_ERROR(
+      loop_promotion_map_it != loop_promotion_map.end(),
+      "No loop promotion found: ",
+      loop_id->toString(),
+      ". Loop group: ",
+      nvfuser::toString(loop_group));
+
+  return loop_promotion_map_it->second;
+}
+
+// Get the loop domains of a given expr. Currently, they're always
+// the loop domains of a consumer tensor, but in the future this
+// function may return the loop domains of a producer for
+// producer-based indexing.
+inline std::vector<IterDomain*> getLoopIds(
+    const Expr* expr,
+    const IdModel& id_model) {
+  // Assume consumer-based indexing. Needs to revisit for ops like
+  // scatter
+  NVF_ERROR(!expr->outputs().empty());
+  auto output_tv = ir_utils::getTvOutput(expr);
+  NVF_ERROR(output_tv != nullptr);
+  auto loop_ids = output_tv->getLoopDomain();
+
+  for (auto& loop_id : loop_ids) {
+    loop_id = getLoopPromotion(loop_id, id_model);
+  }
+
+  return loop_ids;
+}
+
+inline ParallelType getParallelType(const ValGroup& loop_group) {
+  ParallelType common_pt = ParallelType::Serial;
+  for (const auto val : *loop_group) {
+    auto pt = val->as<IterDomain>()->getParallelType();
+    if (common_pt == pt || pt == ParallelType::Serial) {
+      continue;
+    } else if (common_pt == ParallelType::Serial) {
+      common_pt = pt;
+    } else {
+      // Inconsistent parallelization
+      NVF_THROW(
+          "Inconsistent parallelization detected. ",
+          "Known type: ",
+          common_pt,
+          "New type: ",
+          pt);
+    }
+  }
+
+  return common_pt;
+}
+
+// Check if the loop index of a loop group should be always
+// just zero. For example, a loop group with an extent of one, i.e.,
+// a broadcast-only loop group, should just use zero.
+inline bool shouldUseZeroIndex(
+    const ValGroup& loop_group,
+    const IdModel& id_model) {
+  // Trivial loop
+  auto promotion_id =
+      getLoopPromotion(loop_group->front()->as<IterDomain>(), id_model);
+
+  // ExprSimplify should be disabled here as it would fail to
+  // recognize size-one IterDomain.
+  DisableOptionsGuard options_guard;
+  DisableOptionsGuard::getCurOptions().unset(DisableOption::ExprSimplify);
+
+  if (promotion_id->isBroadcast() ||
+      simplifyExpr(promotion_id->extent())->isOneInt()) {
+    return true;
+  }
+
+  return false;
+}
+
+} // namespace nvfuser

nvfuser/include/nvfuser/id_model/validation_utils.h

@@ -0,0 +1,55 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <compute_at_map.h>
+#include <id_model/id_model.h>
+#include <val_graph.h>
+
+namespace nvfuser {
+
+// Note that this class is a friend of ComputeAtMap as it needs to
+// have private access
+class IdModelValidator {
+ public:
+  IdModelValidator(Fusion* fusion, bool allow_self_mapping = false);
+
+  // Validate a given exact graph of IdModel by comparing it with
+  // ComputeAtMap. Their maps should
+  // be almost the same but there are some differences.
+  // - In ComputeAtMap, swizzles are just skipped no matter what swizzle
+  // type is used, so only swizzle outputs are mapped. In IdModel,
+  // only swizzle inputs are mapped, except for Loop swizzles where
+  // their inputs and outputs are mapped.
+  // - In ComputeAtMap, mappings are local. For example, if domain x0 is
+  // split to x1 and x2, and also domain y0 is split to y1 and
+  // y2. Suppose x0 and y0 are exactly mapped and the two splits are
+  // also considered exactly the same, IdModel maps x1 and y1, and x2
+  // and y2, respectively, whereas that doesn't happen with ComputeAtMap
+  //
+  // Accounting for the first difference doesn't seem trivial, so when
+  // swizzle is used we give up validating the exact graph. The second
+  // difference is whether mappings are propagated, which can be
+  // accounted for by updating the ComputeAtMap as is done in IdModel.
+  void checkExactGraphEquivalence(const ValGraph& exact_graph);
+
+  void checkAlmostExactGraphEquivalence(const ValGraph& almost_exact_graph);
+
+  void checkPermissiveGraphEquivalence(const ValGraph& permissive_graph);
+
+ private:
+  // Propagate mappings in a ComputeAtMap as is done in IdModel
+  static void fullyPropagateMappings(DisjointSets<IterDomain*>& id_sets);
+
+ private:
+  ComputeAtMap ca_map_;
+  // Validation is not enabled if swizzle is found. See the comment above
+  bool has_swizzle_ = false;
+};
+
+} // namespace nvfuser