mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
9ae4d38a21
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/33189 Add RRefInterface to Aten/Core, which will later be used by IValue Switch all the rpc code base to use intrusive_ptr instead of shared_ptr, so that we could add it to IValue. Actual adding to IValue and JIT will be in next PR Test Plan: Imported from OSS Differential Revision: D19871241 Pulled By: wanchaol fbshipit-source-id: d7e1fd04b46320e0f26c18591b49c92ad30a4032
37 lines
1.4 KiB
C++
37 lines
1.4 KiB
C++
#pragma once
|
|
|
|
#include <ATen/core/ivalue.h>
|
|
#include <torch/csrc/distributed/autograd/utils.h>
|
|
#include <torch/csrc/distributed/rpc/rref_context.h>
|
|
#include <torch/csrc/distributed/rpc/script_remote_call.h>
|
|
|
|
namespace torch {
|
|
namespace distributed {
|
|
namespace rpc {
|
|
|
|
// This function sends an rpc call to run torchscript function, currently the
|
|
// torchscript function could only be a user defined python function with
|
|
// "@torch.jit.script" annotation. The torchscript function could not be
|
|
// a class constructor, class method, instance method or a script module.
|
|
// dst: destination worker name
|
|
// qualifiedName: torchscript function qualified name string like
|
|
// "moduleName::torchscriptFunctionName", e.g,
|
|
// "dist_autograd_test::my_py_add"
|
|
// stack: a bag of IValue args passed to torchscriptFunctionName
|
|
// It returns c10::intrusive_ptr<ivalue::Future>
|
|
c10::intrusive_ptr<c10::ivalue::Future> TORCH_API rpcTorchscript(
|
|
const std::string& dstWorkerName,
|
|
const c10::QualifiedName& qualifiedName,
|
|
const c10::FunctionSchema& functionSchema,
|
|
std::vector<c10::IValue>& stack);
|
|
|
|
c10::intrusive_ptr<UserRRef> TORCH_API remoteTorchscript(
|
|
const std::string& dstWorkerName,
|
|
const c10::QualifiedName& qualifiedName,
|
|
const c10::FunctionSchema& functionSchema,
|
|
std::vector<c10::IValue>& stack);
|
|
|
|
} // namespace rpc
|
|
} // namespace distributed
|
|
} // namespace torch
|