ForLoop::isGrouped()

zasdfgbnm · zasdfgbnm · commit 52c7ded87a81 · 2022-12-02T04:59:34.000-08:00
diff --git a/torch/csrc/jit/codegen/cuda/codegen.cpp b/torch/csrc/jit/codegen/cuda/codegen.cpp
@@ -115,38 +115,6 @@ std::string genCall(
   return ss.str();
 }
 
-//! A utility class to check if an expression of a particular type exists
-class ExprFinder : kir::ConstIrVisitor {
- public:
-  //! True if expr or any of its nested expressions is included in
-  //! expr_types
-  static bool exists(
-      const Expr* expr,
-      const std::unordered_set<std::type_index>& expr_types) {
-    ExprFinder finder(expr_types);
-    finder.handle(std::vector<const Expr*>{expr});
-    return finder.is_found_;
-  }
-
- private:
-  ExprFinder(const std::unordered_set<std::type_index>& expr_types)
-      : expr_types_(expr_types) {}
-
-  using kir::ConstIrVisitor::handle;
-
-  void handle(const Expr* expr) final {
-    if (expr_types_.find(typeid(*expr)) != expr_types_.end()) {
-      is_found_ = true;
-      return;
-    }
-    kir::ConstIrVisitor::handle(expr);
-  }
-
- private:
-  const std::unordered_set<std::type_index>& expr_types_;
-  bool is_found_ = false;
-};
-
 class CudaKernelGenerator : private OptOutConstDispatch {
   static constexpr const char* kTab = "  ";
 
@@ -2482,19 +2450,6 @@ class CudaKernelGenerator : private OptOutConstDispatch {
         " which is handled by its own handler");
   }
 
-  //! True if loop is grouped. The IterDomain of the loop must have
-  //! ParallelType::Group, but it isn't sufficient as the loop may be
-  //! for an initialization expression, for which the loop shold not
-  //! be grouped. Make sure a GroupedGridReduction is found.
-  bool isGroupedLoop(const kir::ForLoop* loop) {
-    if (loop->iter_domain()->getParallelType() != ParallelType::Group) {
-      return false;
-    }
-    return ExprFinder::exists(
-        loop,
-        {typeid(kir::GroupedGridReduction), typeid(kir::GroupedGridWelford)});
-  }
-
   void handle(const kir::ForLoop* loop) final {
     if (loop->isTrivial()) {
       handleTrivialLoop(loop);
@@ -2503,7 +2458,7 @@ class CudaKernelGenerator : private OptOutConstDispatch {
 
     // If a loop is grouped, no loop is created, but it isn't
     // considered trivial as the loop trip count is not one.
-    if (isGroupedLoop(loop)) {
+    if (loop->isGrouped()) {
       grouped_loops_.push_back(loop);
       handleScope(loop->body());
       grouped_loops_.pop_back();
diff --git a/torch/csrc/jit/codegen/cuda/kernel_ir.cpp b/torch/csrc/jit/codegen/cuda/kernel_ir.cpp
@@ -475,6 +475,54 @@ bool ForLoop::isTrivial() const {
   return false;
 }
 
+namespace {
+
+//! A utility class to check if an expression of a particular type exists
+class ExprFinder : kir::ConstIrVisitor {
+ public:
+  //! True if expr or any of its nested expressions is included in
+  //! expr_types
+  static bool exists(
+      const Expr* expr,
+      const std::unordered_set<std::type_index>& expr_types) {
+    ExprFinder finder(expr_types);
+    finder.handle(std::vector<const Expr*>{expr});
+    return finder.is_found_;
+  }
+
+ private:
+  ExprFinder(const std::unordered_set<std::type_index>& expr_types)
+      : expr_types_(expr_types) {}
+
+  using kir::ConstIrVisitor::handle;
+
+  void handle(const Expr* expr) final {
+    if (expr_types_.find(typeid(*expr)) != expr_types_.end()) {
+      is_found_ = true;
+      return;
+    }
+    kir::ConstIrVisitor::handle(expr);
+  }
+
+ private:
+  const std::unordered_set<std::type_index>& expr_types_;
+  bool is_found_ = false;
+};
+
+} // namespace
+
+bool ForLoop::isGrouped() const {
+  //! The IterDomain of the loop must have ParallelType::Group, but it isn't
+  //! sufficient as the loop may be for an initialization expression, for which
+  //! the loop shold not be grouped. Make sure a GroupedGridReduction is found.
+  if (iter_domain()->getParallelType() != ParallelType::Group) {
+    return false;
+  }
+  return ExprFinder::exists(
+      this,
+      {typeid(kir::GroupedGridReduction), typeid(kir::GroupedGridWelford)});
+}
+
 IfThenElse::IfThenElse(IrBuilderPasskey passkey, Predicate* cond)
     : Expr(passkey) {
   setPredicate(cond);
diff --git a/torch/csrc/jit/codegen/cuda/kernel_ir.h b/torch/csrc/jit/codegen/cuda/kernel_ir.h
@@ -502,6 +502,9 @@ class TORCH_CUDA_CU_API ForLoop final : public Expr {
   //! True if no actual for-loop is materialized
   bool isTrivial() const;
 
+  //! True if loop is grouped.
+  bool isGrouped() const;
+
   //! Returns the stage of a double buffered iterdomain
   //!  that this for loop materializes.
   auto doubleBufferLoopStage() const {
