Add SpanningTreePrinter (csarofeen#1786)

shmsong · Jul 1, 2022 · 635ebfc · 635ebfc
1 parent 59f3c32
commit 635ebfc
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git a/torch/csrc/jit/codegen/cuda/maxinfo_propagator.cpp b/torch/csrc/jit/codegen/cuda/maxinfo_propagator.cpp
@@ -416,6 +416,24 @@ std::shared_ptr<MaxInfoSpanningTree::Information> MaxRootDomainInfoSpanningTree:
  return from_info;
 }
 
+void SpanningTreePrinter::propagateTvPasC(TensorView* from, TensorView* to) {
+ stream_ << "propagateTvPasC" << std::endl;
+ stream_ << " from: " << from->toString() << std::endl;
+ stream_ << " to: " << to->toString() << std::endl;
+}
+
+void SpanningTreePrinter::propagateTvCasP(TensorView* from, TensorView* to) {
+ stream_ << "propagateTvCasP" << std::endl;
+ stream_ << " from: " << from->toString() << std::endl;
+ stream_ << " to: " << to->toString() << std::endl;
+}
+
+void SpanningTreePrinter::propagateTvSibling(TensorView* from, TensorView* to) {
+ stream_ << "propagateTvSibling" << std::endl;
+ stream_ << " from: " << from->toString() << std::endl;
+ stream_ << " to: " << to->toString() << std::endl;
+}
+
 } // namespace cuda
 } // namespace fuser
 } // namespace jit

diff --git a/torch/csrc/jit/codegen/cuda/maxinfo_propagator.h b/torch/csrc/jit/codegen/cuda/maxinfo_propagator.h
@@ -232,6 +232,17 @@ class TORCH_CUDA_CU_API MaxRootDomainInfoSpanningTree
  selector) {}
 };
 
+class TORCH_CUDA_CU_API SpanningTreePrinter
+ : public MaxInfoSpanningTree::Propagator {
+ std::ostream& stream_;
+
+ public:
+ virtual void propagateTvPasC(TensorView* from, TensorView* to) override;
+ virtual void propagateTvCasP(TensorView* from, TensorView* to) override;
+ virtual void propagateTvSibling(TensorView* from, TensorView* to) override;
+ SpanningTreePrinter(std::ostream& stream = std::cout) : stream_(stream) {}
+};
+
 } // namespace cuda
 } // namespace fuser
 } // namespace jit