Ensure that deleter is called even for a no-data tensor.

Summary:

When using a custom deleter InefficientStdFunctionContext was using a
std::unique_ptr<> to store the pointer and call the deleter - but this failed to
call the deleter if the pointer was null. Since we have a separate holder class
anyway take out the std::unique_ptr<> and call the deleter directly.

Fixes #117273

Test Plan:

ghstack-source-id: b5f8bf8
Pull Request resolved: #117418

Author: aorenste

aten/src/ATen/ops/from_blob.h

@@ -83,7 +83,7 @@ class TORCH_API TensorMaker {
 
   std::size_t computeStorageSize() const noexcept;
 
-  DataPtr makeDataPtrFromDeleter() const;
+  DataPtr makeDataPtrFromDeleter() noexcept;
 
   DataPtr makeDataPtrFromContext() noexcept;
 

aten/src/ATen/templates/Functions.cpp

@@ -78,8 +78,8 @@ Tensor TensorMaker::make_tensor() {
    return storage_size;
  }
 
- inline DataPtr TensorMaker::makeDataPtrFromDeleter() const {
-   return InefficientStdFunctionContext::makeDataPtr(data_, deleter_, *device_);
+ inline DataPtr TensorMaker::makeDataPtrFromDeleter() noexcept {
+   return InefficientStdFunctionContext::makeDataPtr(data_, std::move(deleter_), *device_);
  }
 
  inline DataPtr TensorMaker::makeDataPtrFromContext() noexcept {

c10/core/Allocator.cpp

@@ -27,11 +27,11 @@ static void deleteInefficientStdFunctionContext(void* ptr) {
 
 at::DataPtr InefficientStdFunctionContext::makeDataPtr(
     void* ptr,
-    const std::function<void(void*)>& deleter,
+    std::function<void(void*)> deleter,
     Device device) {
   return {
       ptr,
-      new InefficientStdFunctionContext({ptr, deleter}),
+      new InefficientStdFunctionContext(ptr, std::move(deleter)),
       &deleteInefficientStdFunctionContext,
       device};
 }

c10/core/Allocator.h

@@ -223,13 +223,18 @@ struct C10_API Allocator {
 // allocation InefficientStdFunctionContext, on top of the dynamic
 // allocation which is implied by std::function itself.
 struct C10_API InefficientStdFunctionContext {
-  std::unique_ptr<void, std::function<void(void*)>> ptr_;
-  InefficientStdFunctionContext(
-      std::unique_ptr<void, std::function<void(void*)>>&& ptr)
-      : ptr_(std::move(ptr)) {}
+  void* ptr_;
+  std::function<void(void*)> deleter_;
+  InefficientStdFunctionContext(void* ptr, std::function<void(void*)> deleter)
+      : ptr_(ptr), deleter_(std::move(deleter)) {}
+  ~InefficientStdFunctionContext() {
+    if (deleter_) {
+      deleter_(ptr_);
+    }
+  }
   static DataPtr makeDataPtr(
       void* ptr,
-      const std::function<void(void*)>& deleter,
+      std::function<void(void*)> deleter,
       Device device);
 };