Add function to get temp allocator (pytorch#4034)

Summary:
Pull Request resolved: pytorch#4034

Expose the temp allocator in the `BackendExecutionContext`. For the custom allocator provided by the backend, it's possible that they can support more than allocate function and we want to provide the flexibility if they have the custom allocator.

Reviewed By: tarun292, cmt0

Differential Revision: D58881648

fbshipit-source-id: d124f529991ef3cbb9e3b3db4c6d447b809bdbee

runtime/backend/backend_execution_context.h

@@ -45,6 +45,13 @@ class BackendExecutionContext final {
     return temp_allocator_->allocate(size, alignment);
   }
 
+  /**
+   * Returns the temp allocator. This allocator will be reset every instruction.
+   */
+  MemoryAllocator* get_temp_allocator() {
+    return temp_allocator_;
+  }
+
  private:
   EventTracer* event_tracer_ = nullptr;
   MemoryAllocator* temp_allocator_ = nullptr;