Remove getScriptRemoteCallType (#57854)

lw · facebook-github-bot · commit 20d02cb7ddc5 · 2021-05-21T13:15:07.000-07:00
Summary: Pull Request resolved: #57854 Because OwnerRRefs used to be created before their value was computed, we had to figure out their type ahead of time. After the previous diff, we inverted the order of operations, and we can now first compute the result and then create the OwnerRRef. Which means we can just inspect the value to get its type. Much simpler, and much less likely to get it wrong. ghstack-source-id: 129567060 Test Plan: CI Reviewed By: mrshenli Differential Revision: D28253843 fbshipit-source-id: f13c9b294f477ae66fcbdbc85c642fdc69b2740f
diff --git a/torch/csrc/distributed/rpc/request_callback_impl.cpp b/torch/csrc/distributed/rpc/request_callback_impl.cpp
@@ -195,23 +195,6 @@ void RequestCallbackImpl::processPythonCall(
   }
 }
 
-TypePtr RequestCallbackImpl::getScriptRemoteCallType(
-    ScriptRemoteCall& scriptRemoteCall) const {
-  TypePtr returnType;
-  if (scriptRemoteCall.hasOp()) {
-    returnType = scriptRemoteCall.op()->schema().returns()[0].type();
-  } else {
-    returnType = PythonRpcHandler::getInstance()
-                     .jitCompilationUnit()
-                     ->get_function(scriptRemoteCall.qualifiedName())
-                     .getSchema()
-                     .returns()
-                     .at(0)
-                     .type();
-  }
-  return returnType;
-}
-
 c10::intrusive_ptr<JitFuture> RequestCallbackImpl::processScriptRemoteCall(
     ScriptRemoteCall& scriptRemoteCall,
     std::vector<at::IValue>& stack) const {
diff --git a/torch/csrc/distributed/rpc/request_callback_impl.h b/torch/csrc/distributed/rpc/request_callback_impl.h
@@ -24,9 +24,6 @@ class TORCH_API RequestCallbackImpl : public RequestCallbackNoPython {
       const std::function<void(Message)>& markComplete,
       const c10::intrusive_ptr<JitFuture>& responseFuture) const override;
 
-  TypePtr getScriptRemoteCallType(
-      ScriptRemoteCall& scriptRemoteCall) const override;
-
   c10::intrusive_ptr<JitFuture> processScriptRemoteCall(
       ScriptRemoteCall& scriptRemoteCall,
       std::vector<at::IValue>& stack) const override;
diff --git a/torch/csrc/distributed/rpc/request_callback_no_python.cpp b/torch/csrc/distributed/rpc/request_callback_no_python.cpp
@@ -158,14 +158,6 @@ void RequestCallbackNoPython::processScriptCallOp(
   });
 }
 
-TypePtr RequestCallbackNoPython::getScriptRemoteCallType(
-    ScriptRemoteCall& scriptRemoteCall) const {
-  TORCH_CHECK(
-      scriptRemoteCall.hasOp(),
-      "Only supports the case where ScriptCall has an op");
-  return scriptRemoteCall.op()->schema().returns()[0].type();
-}
-
 void RequestCallbackNoPython::processPythonCall(
     RpcCommandBase& rpc,
     const std::function<void(Message)>& markComplete,
@@ -216,10 +208,12 @@ void RequestCallbackNoPython::processBaseScriptRemoteCall(
         IValue(c10::make_intrusive<Message>(std::move(m))));
   };
 
+  auto& stack = scriptRemoteCall.stackRef();
+  auto jitFuture = processScriptRemoteCall(scriptRemoteCall, stack);
+
   // scriptRemoteCall is only alive within this block, use reference to
   // avoid copy. If the underlying code runs with a continuation, runAsync()
   // below will std::move the appropriate portion of the stack.
-  TypePtr returnType = getScriptRemoteCallType(scriptRemoteCall);
   c10::intrusive_ptr<OwnerRRef> ownerRRef;
   if (rrefId == forkId) {
     // Creating an owner RRef on self, should already exist in owners map
@@ -228,12 +222,9 @@ void RequestCallbackNoPython::processBaseScriptRemoteCall(
                               ->constValue()
                               .toRRef());
   } else {
-    ownerRRef = ctx.getOrCreateOwnerRRef(rrefId, returnType);
+    ownerRRef = ctx.getOrCreateOwnerRRef(rrefId, jitFuture->elementType());
   }
 
-  auto& stack = scriptRemoteCall.stackRef();
-  auto jitFuture = processScriptRemoteCall(scriptRemoteCall, stack);
-
   jitFuture->addCallback([postProcessing, ownerRRef](JitFuture& future) {
     if (future.hasError()) {
       ownerRRef->setError(future.exception_ptr());
diff --git a/torch/csrc/distributed/rpc/request_callback_no_python.h b/torch/csrc/distributed/rpc/request_callback_no_python.h
@@ -38,9 +38,6 @@ class TORCH_API RequestCallbackNoPython : public RequestCallback {
       const std::function<void(Message)>& markComplete,
       const c10::intrusive_ptr<JitFuture>& responseFuture) const;
 
-  virtual TypePtr getScriptRemoteCallType(
-      ScriptRemoteCall& scriptRemoteCall) const;
-
   virtual c10::intrusive_ptr<JitFuture> processScriptRemoteCall(
       ScriptRemoteCall& scriptRemoteCall,
       std::vector<at::IValue>& stack) const;
diff --git a/torch/csrc/distributed/rpc/rref_context.cpp b/torch/csrc/distributed/rpc/rref_context.cpp
@@ -348,7 +348,7 @@ c10::intrusive_ptr<OwnerRRef> RRefContext::getOrCreateOwnerRRef(
     // since Tensor can only get specialized with a previous run of local
     // JIT function, and we shouldn't preserve the specialized SubTensorType
     // information on other workers because it's only information only.
-    if (type == TensorType::get()) {
+    if (type->isSubtypeOf(TensorType::get())) {
       TORCH_INTERNAL_ASSERT(
           ownerRRef->type()->isSubtypeOf(TensorType::get()),
           "Expect OwnerRRef to be a sub-type of TensorType, but got ",
