Small refactoring

Naoya Maruyama · Naoya Maruyama · commit 89521cc29e1f · 2020-08-03T14:01:47.000-07:00
diff --git a/torch/csrc/jit/codegen/cuda/ir_iostream.cpp b/torch/csrc/jit/codegen/cuda/ir_iostream.cpp
@@ -698,7 +698,7 @@ void IRPrinter::handle(const kir::GridReduction* gr) {
   indent();
   // Since block-level reduction is already done, those dimensions
   // with tidx/y/z being true do not participate in the grid reduction.
-  os << kir::getPredicateFlagName(out->view()) << " = "
+  os << kir::GridReduction::getPredicateFlagName(out->view()) << " = "
      << "reduction::gridReduce< " << (bidx ? "true" : "false") << ", "
      << (bidy ? "true" : "false") << ", " << (bidz ? "true" : "false") << ", "
      << (!tidx ? "true" : "false") << ", " << (!tidy ? "true" : "false") << ", "
diff --git a/torch/csrc/jit/codegen/cuda/kernel_ir.cpp b/torch/csrc/jit/codegen/cuda/kernel_ir.cpp
@@ -361,7 +361,7 @@ GridReduction::GridReduction(const GridReduction* src, IrCloner* ir_cloner)
       reduction_buffer_(ir_cloner->clone(src->reduction_buffer_)),
       sync_buffer_(ir_cloner->clone(src->sync_buffer_)) {}
 
-std::string getPredicateFlagName(const TensorView* val) {
+std::string GridReduction::getPredicateFlagName(const TensorView* val) {
   std::stringstream ss;
   ss << "T" << val->name() << "pred";
   return ss.str();
diff --git a/torch/csrc/jit/codegen/cuda/kernel_ir.h b/torch/csrc/jit/codegen/cuda/kernel_ir.h
@@ -529,6 +529,8 @@ class TORCH_CUDA_API GridReduction : public Expr {
     return sync_buffer_;
   }
 
+  static std::string getPredicateFlagName(const TensorView* val);
+
  private:
   ReductionOp* reduction_op_ = nullptr;
   Allocate* reduction_buffer_ = nullptr;
diff --git a/torch/csrc/jit/codegen/cuda/lower_index.cpp b/torch/csrc/jit/codegen/cuda/lower_index.cpp
@@ -138,7 +138,7 @@ void IndexLowering::handle(TernaryOp* top) {
 namespace {
 
 kir::Allocate* allocateGridReductionFlag(TensorView* out_tv) {
-  auto flag_name = kir::getPredicateFlagName(out_tv);
+  auto flag_name = kir::GridReduction::getPredicateFlagName(out_tv);
   auto flag_var = new kir::NamedScalar(flag_name, DataType::Bool);
   return new kir::Allocate(flag_var, MemoryType::Local, new kir::Int(1));
 }
diff --git a/torch/csrc/jit/codegen/cuda/lower_thread_predicate.cpp b/torch/csrc/jit/codegen/cuda/lower_thread_predicate.cpp
@@ -18,7 +18,7 @@ Val* getPredicatePerParallelType(
     TORCH_INTERNAL_ASSERT(!sources.empty(), "No predicate source found");
     TORCH_INTERNAL_ASSERT(sources.size() == 1, "Multiple sources detected");
     auto src = *sources.begin();
-    auto flag_name = kir::getPredicateFlagName(src);
+    auto flag_name = kir::GridReduction::getPredicateFlagName(src);
     return new kir::NamedScalar(flag_name, DataType::Bool);
   } else {
     return kir::eqExpr(kir::NamedScalar::getParallelIndex(pt), new kir::Int(0));

Original file line number	Diff line number	Diff line change
`@@ -529,6 +529,8 @@ class TORCH_CUDA_API GridReduction : public Expr {`
`529`	`529`	`return sync_buffer_;`
`530`	`530`	`}`
`531`	`531`
	`532`	`+ static std::string getPredicateFlagName(const TensorView* val);`
	`533`	`+`
`532`	`534`	`private:`
`533`	`535`	`ReductionOp* reduction_op_ = nullptr;`
`534`	`536`	`Allocate* reduction_buffer_ = nullptr;`
Original file line number	Diff line number	Diff line change
`@@ -138,7 +138,7 @@ void IndexLowering::handle(TernaryOp* top) {`
`138`	`138`	`namespace {`
`139`	`139`
`140`	`140`	`kir::Allocate* allocateGridReductionFlag(TensorView* out_tv) {`
`141`		`- auto flag_name = kir::getPredicateFlagName(out_tv);`
	`141`	`+ auto flag_name = kir::GridReduction::getPredicateFlagName(out_tv);`
`142`	`142`	`auto flag_var = new kir::NamedScalar(flag_name, DataType::Bool);`
`143`	`143`	`return new kir::Allocate(flag_var, MemoryType::Local, new kir::Int(1));`
`144`	`144`	`}`