Backport some changes from master. (pytorch#193)

Author: ZolotukhinM

torch/csrc/jit/tensorexpr/codegen.cpp

@@ -29,7 +29,7 @@ RegisterCodeGenList::StmtFactoryMethod RegisterCodeGenList::
 
 void RegisterCodeGenList::AddStmtFactoryMethod(
     const std::string& name,
-    StmtFactoryMethod stmt_factory_method) {
+    const StmtFactoryMethod& stmt_factory_method) {
   auto insert_ret =
       stmt_factory_methods_.insert(std::make_pair(name, stmt_factory_method));
   if (!insert_ret.second) {

torch/csrc/jit/tensorexpr/codegen.h

@@ -138,7 +138,7 @@ class RegisterCodeGenList {
   RegisterCodeGenList() {}
   TORCH_API void AddStmtFactoryMethod(
       const std::string& name,
-      StmtFactoryMethod stmt_factory_method);
+      const StmtFactoryMethod& stmt_factory_method);
   RegisterCodeGenList(const RegisterCodeGenList&) = delete;
   RegisterCodeGenList& operator=(const RegisterCodeGenList&) = delete;
 

torch/csrc/jit/tensorexpr/kernel.cpp

@@ -1026,7 +1026,7 @@ void TensorExprKernel::bindInput(const torch::jit::Value* input) {
 }
 
 TensorExprKernel::TensorExprKernel(const Graph& subgraph) {
-  KernelScope kernel_scope(kernel_arena_);
+  KernelScope kernel_scope(&kernel_arena_);
 
   // Bind inputs to buffers.
   n_inputs_ = subgraph.inputs().size();
@@ -1056,7 +1056,7 @@ TensorExprKernel::TensorExprKernel(const Graph& subgraph) {
 }
 
 void TensorExprKernel::run(Stack& stack) {
-  KernelScope kernel_scope(kernel_arena_);
+  KernelScope kernel_scope(&kernel_arena_);
   // Set up arguments (inputs, then outputs) for kernel call.
   auto inputs = last(stack, n_inputs_);
   PickAndCheckBackendType(inputs);

torch/csrc/jit/tensorexpr/mem_arena.cpp

@@ -1,55 +1,54 @@
-#include <stdexcept>
 #include "torch/csrc/jit/tensorexpr/mem_arena.h"
 
 namespace torch {
 namespace jit {
 namespace tensorexpr {
 
+namespace {
+// Define in an anonymous namespace to hide this symbol from other compilation
+// units
+thread_local KernelArena* current_arena = nullptr;
+}
+
 KernelArena::~KernelArena() {
   for (KernelScopedObject* p : kernel_objects_) {
     delete p;
   }
 }
 
 KernelScopedObject::KernelScopedObject() {
-  KernelArena& kernel = KernelArena::GetCurrentKernelArena();
-  kernel.kernel_objects_.push_back(this);
+  KernelArena* kernel = KernelArena::GetCurrentKernelArena();
+  kernel->kernel_objects_.push_back(this);
 }
 
 static std::vector<KernelArena*>& GetKernelArenaStack() {
   thread_local std::vector<KernelArena*> kernel_arena_stack;
   return kernel_arena_stack;
 }
 
-KernelArena& KernelArena::GetCurrentKernelArena() {
-  std::vector<KernelArena*>& kernel_arena_stack = GetKernelArenaStack();
-  if (kernel_arena_stack.empty()) {
-    throw std::runtime_error(
-        "A KernelScope must be bound before creating KernelScopedObject");
-  }
-  return *kernel_arena_stack.back();
+void KernelArena::SetCurrentKernelArena(KernelArena *new_kernel_arena) {
+  current_arena = new_kernel_arena;
 }
 
-KernelScope::KernelScope() : owning_kernel_arena_(true) {
-  kernel_arena_ = new KernelArena;
-  GetKernelArenaStack().push_back(kernel_arena_);
+KernelArena* KernelArena::GetCurrentKernelArena() {
+  return current_arena;
 }
 
-KernelScope::KernelScope(KernelArena& kernel_arena)
-    : owning_kernel_arena_(false) {
-  kernel_arena_ = &kernel_arena;
-  GetKernelArenaStack().push_back(&kernel_arena);
+KernelScope::KernelScope() : owning_(true) {
+  old_kernel_arena_ = KernelArena::GetCurrentKernelArena();
+  KernelArena::SetCurrentKernelArena(new KernelArena);
 }
 
-KernelScope::~KernelScope() noexcept(false) {
-  std::vector<KernelArena*>& kernel_arena_stack = GetKernelArenaStack();
-  if (kernel_arena_ != kernel_arena_stack.back()) {
-    throw std::runtime_error("Mismatch KernelScope and kernel");
-  }
-  if (owning_kernel_arena_) {
-    delete kernel_arena_;
+KernelScope::KernelScope(KernelArena* arena_) : owning_(false) {
+  old_kernel_arena_ = KernelArena::GetCurrentKernelArena();
+  KernelArena::SetCurrentKernelArena(arena_);
+}
+
+KernelScope::~KernelScope() {
+  if (owning_) {
+    delete KernelArena::GetCurrentKernelArena();
   }
-  kernel_arena_stack.pop_back();
+  KernelArena::SetCurrentKernelArena(old_kernel_arena_);
 }
 
 } // namespace tensorexpr

torch/csrc/jit/tensorexpr/mem_arena.h

@@ -11,7 +11,8 @@ class KernelScopedObject;
 // An arena that manages all the underlying kernel-scoped objects.
 class KernelArena {
  public:
-  static KernelArena& GetCurrentKernelArena();
+  static KernelArena* GetCurrentKernelArena();
+  static void SetCurrentKernelArena(KernelArena* new_arena);
   TORCH_API KernelArena() {}
   TORCH_API ~KernelArena();
 
@@ -23,20 +24,23 @@ class KernelArena {
 };
 
 // A RAII convenience wrapper on top of a kernel.
-// It either creates a Kernel, or take another existing Kernel, and sets it as
-// the current Kernel, as long as this KernelScope object is alive.
+// It either creates or takes an existing Kernel and sets it as the current
+// Kernel. When this object is destroyed, the previous Kernel is set as current,
+// and the created kernel is freed. If the kernel was passed, it stays alive.
 class KernelScope {
  public:
   TORCH_API KernelScope();
-  TORCH_API explicit KernelScope(KernelArena& kernel_arena);
-  TORCH_API ~KernelScope() noexcept(false);
+  TORCH_API explicit KernelScope(KernelArena* arena_);
+  TORCH_API ~KernelScope();
 
  private:
   KernelScope(const KernelScope&) = delete;
   KernelScope& operator=(const KernelScope&) = delete;
-  bool owning_kernel_arena_ = false;
-  KernelArena* kernel_arena_ =
-      nullptr; // possibly owned, if owning_kernel_arena_ == true
+  KernelArena* kernel_arena_ = nullptr; // arena to be used in this scope
+  KernelArena* old_kernel_arena_ =
+      nullptr; // previous arena, will be restored in destructor
+  bool owning_ = false; // determines whether the arena will be freed along with
+                        // the scope object
 };
 
 // The base object managed by the Kernel.
@@ -55,4 +59,3 @@ class TORCH_API KernelScopedObject {
 } // namespace tensorexpr
 } // namespace jit
 } // namespace torch
-

torch/csrc/jit/tensorexpr/unique_name_manager.cpp

@@ -23,7 +23,7 @@ const std::string& UniqueNameManager::get_unique_name(const Var* v) {
     name_hint = "v" + name_hint;
   }
   int& count = unique_name_count_[name_hint];
-  while (1) {
+  while (true) {
     // Even if with a new count, this name might already be used. For example
     // ("x", 1) could collidewith ("x_1", 0)
     int count_v = count++;