Move unwrap logic from c10 to caffe2 (#21620)

Will Feng · facebook-github-bot · commit 04f09d4235e5 · 2019-06-14T22:02:43.000-07:00
Summary: After #17072, we are allowed to pass Variables into ATen ops, thus there is no need to unwrap input variables in the c10 call path. Note that since Caffe2 still expects inputs to be pure Tensors, we moved the unwrapping logic to the Caffe2 wrapper. Pull Request resolved: #21620 Differential Revision: D15763560 Pulled By: yf225 fbshipit-source-id: 5375f0e51eb320f380ae599ebf98e6b259f0bff8
diff --git a/caffe2/core/export_caffe2_op_to_c10.h b/caffe2/core/export_caffe2_op_to_c10.h
@@ -27,6 +27,44 @@ inline c10::ListPtr<at::Tensor> _call_caffe2_op(
   return std::move(op).move_newstyle_outputs();
 }
 
+inline at::Tensor unwrap_tensor(at::Tensor&& tensor) {
+  if (tensor.is_variable()) {
+    auto tensor_impl = tensor.unsafeGetTensorImpl();
+    auto tensor_impl_copy = tensor_impl->shallow_copy_and_detach(
+    /*version_counter=*/tensor_impl->version_counter(),
+    /*allow_tensor_metadata_change=*/tensor_impl->allow_tensor_metadata_change());
+    return at::Tensor(tensor_impl_copy);
+  } else {
+    return std::move(tensor);
+  }
+}
+
+inline IValue unwrap(IValue&& ivalue) {
+  // TODO Remove the .defined() check once we don't have undefined tensors on the stack anymore (@wanchaol is working on this)
+  if (ivalue.isTensor() && ivalue.toTensor().defined()) {
+    return unwrap_tensor(std::move(ivalue).toTensor());
+  } else if (ivalue.isTensorList()) {
+    c10::ListPtr<at::Tensor> list = std::move(ivalue).toTensorList();
+    for (size_t i = 0; i < list.size(); ++i) {
+      list[i] = unwrap_tensor(list.extract(i));
+    }
+    return std::move(list);
+  } else if (ivalue.isGenericList()) {
+    c10::impl::GenericListPtr list = std::move(ivalue).toGenericList();
+    for (size_t i = 0; i < list.size(); ++i) {
+      list[i] = unwrap(list.extract(i));
+    }
+    return std::move(list);
+  } else if (ivalue.isGenericDict()) {
+    for (auto& item : ivalue.toGenericDict()) {
+      item.setValue(unwrap(item.value()));
+    }
+    return std::move(ivalue);
+  } else {
+    return std::move(ivalue);
+  }
+}
+
 // This function is inline in the hope that compilers optimizing for speed will
 // inline it into call_caffe2_op_from_c10, allowing call_op to be inlined and
 // avoiding the function pointer indirection, while compilers optimizing for
@@ -65,6 +103,11 @@ inline void _call_caffe2_op_from_c10(
     outputs = std::move(preallocated_outputs).toTensorList();
   }
 
+  // unwrap tensor inputs from variable
+  for (auto iter = stack->end() - num_inputs; iter != stack->end(); ++iter) {
+    *iter = unwrap(std::move(*iter));
+  }
+
   // TODO Avoid vector allocation. One idea would be to keep the std::vector
   // instances in the cache.
   std::vector<IValue> inputs = torch::jit::pop(*stack, num_inputs);
diff --git a/torch/csrc/jit/register_c10_ops.cpp b/torch/csrc/jit/register_c10_ops.cpp
@@ -7,40 +7,6 @@ namespace torch {
 namespace jit {
 namespace {
 
-at::Tensor unwrap_tensor(at::Tensor&& tensor) {
-  if (tensor.is_variable()) {
-    return torch::autograd::Variable(std::move(tensor)).tensor_data();
-  } else {
-    return std::move(tensor);
-  }
-}
-
-IValue unwrap(IValue&& ivalue) {
-  // TODO Remove the .defined() check once we don't have undefined tensors on the stack anymore (@wanchaol is working on this)
-  if (ivalue.isTensor() && ivalue.toTensor().defined()) {
-    return unwrap_tensor(std::move(ivalue).toTensor());
-  } else if (ivalue.isTensorList()) {
-    c10::ListPtr<at::Tensor> list = std::move(ivalue).toTensorList();
-    for (size_t i = 0; i < list.size(); ++i) {
-      list[i] = unwrap_tensor(list.extract(i));
-    }
-    return std::move(list);
-  } else if (ivalue.isGenericList()) {
-    c10::impl::GenericListPtr list = std::move(ivalue).toGenericList();
-    for (size_t i = 0; i < list.size(); ++i) {
-      list[i] = unwrap(list.extract(i));
-    }
-    return std::move(list);
-  } else if (ivalue.isGenericDict()) {
-    for (auto& item : ivalue.toGenericDict()) {
-      item.setValue(unwrap(item.value()));
-    }
-    return std::move(ivalue);
-  } else {
-    return std::move(ivalue);
-  }
-}
-
 at::Tensor wrap_tensor(at::Tensor&& tensor) {
   if (tensor.is_variable()) {
     return std::move(tensor);
@@ -170,11 +136,6 @@ Operator createOperatorFromC10(const c10::OperatorHandle& op) {
         graph->insertNode(node);
       }
 
-      // unwrap tensor inputs from variable
-      for (auto iter = stack.end() - input_size; iter != stack.end(); ++iter) {
-        *iter = unwrap(std::move(*iter));
-      }
-
       c10::Dispatcher::singleton().lookup(op, &stack).call(&stack);
 
       // wrap tensor outputs as variable
