Dynamic Resize of IterDomains

csrc/dynamic_transform.cpp

@@ -30,6 +30,9 @@ class DynamicTransformInfoBuilder : public IterVisitor {
   // Analyze a dynamic reshape and generate AnalyzeViewResult
   void handle(ViewOp* op) override;
 
+  // We handle IterDomain "Resize" ops at TensorView level
+  void handle(TensorView* tv) override;
+
   const auto& getInfo() const {
     return info_;
   }
@@ -77,20 +80,35 @@ bool DynamicTransformConcretizationInfo::operator==(
     }
   }
 
+  for (const auto i : c10::irange(resize_transforms_.size())) {
+    const auto& transform = resize_transforms_.at(i);
+    const auto& other_transform = other.resize_transforms_.at(i);
+    if (transform != other_transform) {
+      return false;
+    }
+  }
+
   return true;
 }
 
 DynamicTransformConcretizationInfo DynamicTransformConcretizationInfo::clone(
     IrCloner& ir_cloner) const {
   DynamicTransformConcretizationInfo cloned_info(
       (Fusion*)ir_cloner.container());
-  for (auto& pair : reshape_transforms_) {
+  for (const auto& [tv, analyze_result] : reshape_transforms_) {
     cloned_info.reshape_transforms_.emplace_back(
-        ir_cloner.clone(pair.first),
+        ir_cloner.clone(tv),
         // reshape_transforms_ holds pairs of TensorView* and AnalyzeViewResult
         // AnalyzeViewResult can be copied directly as it holds no references to
         // Statements that would need cloning, only integer indices of axes.
-        pair.second);
+        analyze_result);
+  }
+  for (const auto& [id, iter_type] : resize_transforms_) {
+    cloned_info.resize_transforms_.emplace_back(
+        ir_cloner.clone(id),
+        // Similar to reshape_transforms_, we only clone the IterDomains in
+        // resize_transforms_
+        iter_type);
   }
   return cloned_info;
 }
@@ -104,9 +122,56 @@ std::string DynamicTransformConcretizationInfo::toString() const {
     ss << indent << indent << kv.first->toString() << ", "
        << kv.second.toString() << "\n";
   }
+  ss << indent << "Resize:\n";
+  for (const auto& [id, iter_type] : resize_transforms_) {
+    ss << indent << indent << id->toString() << ", " << iter_type << "\n";
+  }
   return ss.str();
 }
 
+void DynamicTransformInfoBuilder::handle(TensorView* tv) {
+  const auto& rfd = tv->getMaybeRFactorDomain();
+  for (auto id : rfd) {
+    if (!id->definition()) {
+      continue;
+    }
+    if (auto op = dynamic_cast<Resize*>(id->definition());
+        id->getIterType() == IterType::Symbolic && op != nullptr) {
+      auto out_extent_val = expr_eval_->evaluate(id->extent());
+      TORCH_INTERNAL_ASSERT(
+          out_extent_val.has_value(),
+          "Cannot evaluate the extent of a resized IterDomain: ",
+          id->toString());
+
+      auto in_id = op->in()->as<IterDomain>();
+      auto in_extent_val = expr_eval_->evaluate(in_id->extent());
+      TORCH_INTERNAL_ASSERT(
+          in_extent_val.has_value(),
+          "Cannot evaluate the extent of input to an IterDomain resize: ",
+          in_id->toString());
+
+      auto left = op->leftExpand()->as<Int>();
+      auto left_val = expr_eval_->evaluate(left);
+      TORCH_INTERNAL_ASSERT(
+          left_val.has_value(),
+          "Cannot evaluate the left expansion of an IterDomain resize: ",
+          left->toString());
+
+      auto right = op->rightExpand()->as<Int>();
+      auto right_val = expr_eval_->evaluate(right);
+      TORCH_INTERNAL_ASSERT(
+          right_val.has_value(),
+          "Cannot evaluate the right expansion of an IterDomain resize: ",
+          right->toString());
+
+      auto out_itertype = out_extent_val->as<int64_t>() == 1
+          ? IterType::Broadcast
+          : IterType::Iteration;
+      info_.resize_transforms_.emplace_back(id, out_itertype);
+    }
+  }
+}
+
 void DynamicTransformInfoBuilder::handle(ViewOp* op) {
   auto inp_tv = op->in()->as<TensorView>();
   auto out_tv = op->out()->as<TensorView>();
@@ -204,6 +269,8 @@ class DynamicTransformConcretizer : public OptOutMutator {
 
   void concretizeReshape();
 
+  void concretizeResize();
+
   using OptOutMutator::mutate;
 
   void mutate(TensorView* tv) final;
@@ -216,15 +283,17 @@ class DynamicTransformConcretizer : public OptOutMutator {
 
  private:
   const DynamicTransformConcretizationInfo& info_;
-  std::unordered_map<IterDomain*, IterDomain*> update_map_;
 };
 
 void DynamicTransformConcretizer::concretize() {
   // First, concretize all dynamic reshape ops
   concretizeReshape();
 
-  // Second, propagate concretized domains
-  auto all_stmts = StmtSort::getStmts(info_.fusion(), false);
+  // Set output IterTypes for dynamic resize ops
+  concretizeResize();
+
+  // Finally, propagate concretized domains
+  auto all_stmts = StmtSort::getStmts(info_.fusion(), true);
   for (auto stmt : all_stmts) {
     if (stmt->isA<Val>()) {
       mutate(stmt);
@@ -257,6 +326,24 @@ void DynamicTransformConcretizer::concretizeReshape() {
   }
 }
 
+void DynamicTransformConcretizer::concretizeResize() {
+  // Concretize each resize op.
+  for (const auto& [id, iter_type] : info_.getResizeTransforms()) {
+    TORCH_CHECK(
+        id->definition() && id->definition()->isA<Resize>(),
+        "Resized IterDomain must have a Resize definition");
+    auto def = id->definition()->as<Resize>();
+    auto new_id = IterDomain::resize(
+        def->in(),
+        def->leftExpand(),
+        def->rightExpand(),
+        id->isRFactorProduct(),
+        iter_type);
+
+    registerMutation(id, new_id);
+  }
+}
+
 // Concretizes inherited symbolic domains. Note that when this is
 // called, it is assumed that all dynamic ops themselves are
 // concretized. Since symbolic IDs may be propagated down to
@@ -268,15 +355,12 @@ void DynamicTransformConcretizer::mutate(TensorView* tv) {
 
   // First, try to concretize the root domain as there may be symbolic
   // axes inherited from the producers
-  auto propagated = propagateFromProducerToConsumer(tv);
-
-  // If no root domain is altered, nothing to do further
-  if (!propagated) {
-    return;
-  }
+  propagateFromProducerToConsumer(tv);
 
-  // Root IDs are altered. Need to propagate the changes to rfactor
-  // domain
+  // If no root domain is altered by producer, we don't need to propagate back
+  // up to rfactor. We could return early, but instead we go ahead and check the
+  // root to rfactor transforms to be sure we have concretized any intermediate
+  // IterDomains.
 
   // At this point, there should be no expr beyond rfactor root
   TORCH_INTERNAL_ASSERT(
@@ -308,20 +392,21 @@ void DynamicTransformConcretizer::mutate(TensorView* tv) {
             ". IterDomain was expected.");
       }
 
-      // If none of the output IDs is symbolic, nothing to concretize
-      if (std::all_of(
-              expr->outputs().begin(), expr->outputs().end(), [](Val* output) {
-                return output->as<IterDomain>()->getIterType() !=
-                    IterType::Symbolic;
-              })) {
-        continue;
-      }
-      // If any of output IDs is symbolic, all outputs should be symbolic
-      TORCH_INTERNAL_ASSERT(std::all_of(
-          expr->outputs().begin(), expr->outputs().end(), [](Val* output) {
-            return output->as<IterDomain>()->getIterType() ==
-                IterType::Symbolic;
-          }));
+      // NOTE: We do not return early if all outputs are concrete as there may
+      // still be concrete inputs. For example, a Symbolic IterDomain might be
+      // padded with constant pad widths (1, 1), in which case although we do
+      // not know the exact extent of the output, we know it is at least as
+      // large as the sum of the pad widths, 2. In such cases, the output
+      // IterDomain is concrete at definition, since if the extent is >1 we know
+      // the IterType is Iteration. In these cases, we must continue to
+      // concretize intermediate expressions between the root and R-factor
+      // domain. See test DynamicTransform5_CUDA which demonstrates this
+      // behavior.
+      // NOTE: We also do not assume that if one output ID is symbolic, that
+      // they all must be. See test FusionSliceForNanoGPT3_CUDA for an example
+      // that does a static split by a factor of 16 of a symbolic input domain.
+      // The static split in that case results in a concrete IterDomain with
+      // extent 16 along with a symbolic one (extent ceilDiv(n / 16)).
 
       // Determine the output IterType
       IterType iter_type = IterType::Symbolic;
@@ -336,13 +421,13 @@ void DynamicTransformConcretizer::mutate(TensorView* tv) {
 
       // Update the IterType of each output
       for (auto out_id : ir_utils::filterByType<IterDomain>(expr->outputs())) {
-        auto concreteized_out_id =
+        auto concretized_out_id =
             IterDomainBuilder(out_id).iter_type(iter_type).build();
-        registerMutation(out_id, concreteized_out_id);
+        registerMutation(out_id, concretized_out_id);
       }
 
-      // Outputs are mutated. The expr itself needs to be mutated as
-      // well, which can be done by the mutate method
+      // The expr itself needs to be mutated as well in case the outputs are
+      // mutated, which can be done by the mutate method
       OptOutMutator::mutate(expr);
     }
   }
@@ -457,7 +542,14 @@ bool DynamicTransformConcretizer::propagateFromProducerToConsumer(
     }
 
     TORCH_INTERNAL_ASSERT(
-        id_type.has_value() && id_type != IterType::Symbolic,
+        id_type.has_value(),
+        "Did not find id_type for consumer root domain ",
+        root_id->toString(),
+        ". Perhaps consumer def has no inputs. Consumer definition = ",
+        def->toString());
+
+    TORCH_INTERNAL_ASSERT(
+        id_type != IterType::Symbolic,
         "Failed to concretize ",
         root_id->toString(),
         " of ",

csrc/dynamic_transform.h

@@ -27,11 +27,16 @@ class DynamicTransformInfoBuilder;
 //! of the fusion inputs
 class TORCH_CUDA_CU_API DynamicTransformConcretizationInfo {
  public:
-  const std::vector<std::pair<TensorView*, AnalyzeViewResult>>
+  const std::vector<std::pair<TensorView*, AnalyzeViewResult>>&
   getReshapeTransforms() const {
     return reshape_transforms_;
   }
 
+  const std::vector<std::pair<IterDomain*, IterType>>& getResizeTransforms()
+      const {
+    return resize_transforms_;
+  }
+
   bool operator==(const DynamicTransformConcretizationInfo& other) const;
 
   bool operator!=(const DynamicTransformConcretizationInfo& other) const {
@@ -53,8 +58,15 @@ class TORCH_CUDA_CU_API DynamicTransformConcretizationInfo {
 
  private:
   Fusion* fusion_ = nullptr;
+
+  // Holds, for each dynamic reshape, the output TensorView, and the result of
+  // analyzeView
   std::vector<std::pair<TensorView*, AnalyzeViewResult>> reshape_transforms_;
 
+  // Holds the resized IterDomain (output of the Resize op) along with the
+  // TensorView where it appears, and its concretized IterType
+  std::vector<std::pair<IterDomain*, IterType>> resize_transforms_;
+
   friend class DynamicTransformInfoBuilder;
 };
 

csrc/ir_internal_base_nodes.h

@@ -152,11 +152,25 @@ class TORCH_CUDA_CU_API IterDomain : public Val {
   //! is marked as an rfactor domain. For example, expressions such as
   //! PadOp and SliceOp resize IterDomains and generate rfactor
   //! resized domains.
+  //!
+  //! Usually, the IterType of the output IterDomain will be Symbolic. This is
+  //! because unless the left and right expansions are known at Fusion
+  //! definition we cannot be sure that the output will have an extent != 1. In
+  //! case the output extent is in fact 1, we will set the IterType to
+  //! Broadcast. If the left and right expansions are constant, and sum to at
+  //! least two, then even an empty input will result in an Iteration IterType.
+  //! In these cases, we will set the output IterType to Iteration at
+  //! definition. Otherwise, it will be set to Symbolic and will be resolved
+  //! when concretization is performed by FusionExecutorCache.
+  //!
+  //! The optional iter_type argument can be used to force the output IterType,
+  //! but for safety its use should typically be confined to concretization.
   static IterDomain* resize(
       IterDomain* in,
       Val* left_expansion,
       Val* right_expansion,
-      bool mark_as_rfactor = false);
+      bool mark_as_rfactor = false,
+      std::optional<IterType> iter_type = std::nullopt);
 
   bool isReduction() const {
     return getIterType() == IterType::Reduction;

csrc/ir_nodes.cpp

@@ -7,6 +7,7 @@
 // clang-format on
 #include <device_lower/lower2device.h>
 #include <disjoint_set.h>
+#include <dynamic_transform.h>
 #include <ir_cloner.h>
 #include <ir_interface_nodes.h>
 #include <ir_iostream.h>
@@ -2101,9 +2102,12 @@ IterDomain::IterDomain(
       is_padded_dimension_(is_padded_dimension),
       padded_to_size_(padded_to_size),
       is_mma_swizzled_(is_mma_swizzled) {
-  TORCH_CHECK(
-      !(isRFactorProduct() && isBroadcast()),
-      "IterDomain cannot be both a broadcast and rfactor domain.");
+  // NOTE: We previously asserted !(isRFactorProduct() && isBroadcast()), i.e.
+  // that an IterDomain could not be both a broadcast and an rfactor domain.
+  // However, since the introduction of the resize op, we now have a legitimate
+  // case where this may be true; namely, whenever we resize an IterDomain to
+  // size 1, we will mark it as Broadcast, but the resize must lie between root
+  // and rfactor.
 
   TORCH_INTERNAL_ASSERT(
       extent->isIntegralScalar(),
@@ -2459,7 +2463,8 @@ IterDomain* IterDomain::resize(
     IterDomain* in,
     Val* left_expansion,
     Val* right_expansion,
-    bool mark_as_rfactor) {
+    bool mark_as_rfactor,
+    std::optional<IterType> iter_type_opt) {
   TORCH_CHECK(
       left_expansion->isIntegralScalar(),
       "Expansion factor must be an integer scalar: ",
@@ -2502,10 +2507,28 @@ IterDomain* IterDomain::resize(
         right_expansion);
   }
 
+  // If output IterType is provided, use it. Otherwise, if we can prove the
+  // resized extent is 1, set to Broadcast, if we can prove it is >1 set to
+  // Iteration, and otherwise fall back to Symbolic.
+  IterType iter_type = IterType::Symbolic;
+  if (iter_type_opt.has_value()) {
+    iter_type = iter_type_opt.value();
+  } else if (left_expansion->isConstInt() && right_expansion->isConstInt()) {
+    if (resized_id_size->isConstInt()) {
+      // Means input extent is also known
+      auto out_extent = resized_id_size->evaluateInt();
+      iter_type = out_extent == 1 ? IterType::Broadcast : IterType::Iteration;
+    } else if (
+        left_expansion->evaluateInt() + right_expansion->evaluateInt() > 1) {
+      // Input extent is non-negative, so we know out_extent > 1
+      iter_type = IterType::Iteration;
+    }
+  }
+
   auto resized_id =
       IterDomainBuilder(in->container()->zeroVal(), resized_id_size->as<Int>())
           .is_rfactor_domain(mark_as_rfactor)
-          .iter_type(in->getIterType())
+          .iter_type(iter_type)
           .build();
 
   IrBuilder::create<Resize>(