This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 42bb5bee31 [REFACTOR][RPC] Represent remote devices and tensors
explicitly (#20667)
42bb5bee31 is described below
commit 42bb5bee31635f06469abc66a48a6ea60f058b84
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 11 18:50:13 2026 -0400
[REFACTOR][RPC] Represent remote devices and tensors explicitly (#20667)
RPC devices currently encode their session in the device type, requiring
core device dispatch and ordinary tensors to understand remote storage.
This change introduces RPC-owned remote device, tensor, and stream
objects carrying explicit session ownership and endpoint metadata.
Remote allocation, transfer, views, synchronization, and argument
conversion validate ownership through the RPC layer. Remote handles are
kept outside local tensor and DLPack storage, while server-side
execution uses ordinary devices and tensors. Remove masked-device
routing and migrate Python, JVM, tuning, and example consumers to
explicit remote APIs.
RemoteDevice dynamically resolves capability properties using the shared
Device query definitions and result conversions. Public remote
structures document descriptor, handle, storage, stream, and session
lifetimes.
---
apps/android_rpc/tests/android_rpc_test.py | 4 +-
apps/ios_rpc/tests/ios_rpc_test.py | 4 +-
docs/how_to/tutorials/cross_compilation_and_rpc.py | 16 +-
.../how_to/tutorials/export_and_load_executable.py | 2 +-
docs/reference/api/python/rpc.rst | 28 +++
include/tvm/runtime/device_api.h | 51 +---
jvm/core/src/main/java/org/apache/tvm/Device.java | 6 -
jvm/core/src/main/java/org/apache/tvm/TVMType.java | 2 +-
jvm/core/src/main/java/org/apache/tvm/rpc/RPC.java | 1 -
.../main/java/org/apache/tvm/rpc/RPCSession.java | 54 ++--
.../main/java/org/apache/tvm/rpc/RemoteTensor.java | 50 ++++
.../java/org/apache/tvm/rpc/TVMRemoteDevice.java | 51 +++-
jvm/native/src/main/native/jni_helper_func.h | 9 +-
python/tvm/contrib/coreml_runtime.py | 21 +-
python/tvm/rpc/__init__.py | 2 +
python/tvm/rpc/base.py | 3 -
python/tvm/rpc/client.py | 15 +-
python/tvm/rpc/remote.py | 270 ++++++++++++++++++++
python/tvm/runtime/device.py | 28 +--
python/tvm/runtime/module.py | 5 +-
python/tvm/runtime/vm.py | 6 +-
.../tvm/s_tir/meta_schedule/runner/rpc_runner.py | 2 +-
python/tvm/s_tir/meta_schedule/runner/utils.py | 7 +-
.../meta_schedule/testing/custom_builder_runner.py | 3 +-
.../meta_schedule/testing/validate_database.py | 2 +-
python/tvm/testing/runner.py | 14 +-
src/runtime/device_api.cc | 21 +-
src/runtime/module.cc | 2 +-
src/runtime/rpc/minrpc/minrpc_server.h | 2 +-
src/runtime/rpc/rpc_device_api.cc | 276 +++++++++++----------
src/runtime/rpc/rpc_endpoint.cc | 27 +-
src/runtime/rpc/rpc_local_session.h | 1 +
src/runtime/rpc/rpc_module.cc | 159 +++++-------
src/runtime/rpc/rpc_remote.h | 186 ++++++++++++++
src/runtime/rpc/rpc_session.cc | 45 ----
src/runtime/rpc/rpc_session.h | 112 +++++----
tests/nightly/python/test_nnapi/infrastructure.py | 4 +-
.../codegen/test_target_codegen_cross_llvm.py | 6 +-
tests/python/contrib/test_coreml_runtime.py | 10 +-
tests/python/contrib/test_random.py | 2 +-
tests/python/relax/backend/adreno/utils.py | 8 +-
tests/python/relax/test_vm_build.py | 32 ++-
tests/python/relax/texture/test_texture_nd.py | 8 +-
tests/python/runtime/test_runtime_rpc.py | 102 +++++---
web/tests/python/relax_rpc_test.py | 6 +-
web/tests/python/webgpu_rpc_test.py | 4 +-
46 files changed, 1100 insertions(+), 569 deletions(-)
diff --git a/apps/android_rpc/tests/android_rpc_test.py
b/apps/android_rpc/tests/android_rpc_test.py
index 6456881928..e9997a14e5 100644
--- a/apps/android_rpc/tests/android_rpc_test.py
+++ b/apps/android_rpc/tests/android_rpc_test.py
@@ -74,8 +74,8 @@ def test_rpc_module():
dev = remote.cl(0)
remote.upload(path_dso_cl)
f1 = remote.load_module("dev_lib_cl.so")
- a = tvm.runtime.tensor(a_np, dev)
- b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
+ a = dev.tensor(a_np)
+ b = dev.tensor(np.zeros(1024, dtype=A.dtype))
time_f = f1.time_evaluator(f1.entry_name, dev, number=10)
cost = time_f(a, b).mean
print(f"{cost:g} secs/op\n")
diff --git a/apps/ios_rpc/tests/ios_rpc_test.py
b/apps/ios_rpc/tests/ios_rpc_test.py
index 1d1c4f9b40..e5865b937d 100644
--- a/apps/ios_rpc/tests/ios_rpc_test.py
+++ b/apps/ios_rpc/tests/ios_rpc_test.py
@@ -70,8 +70,8 @@ def test_rpc_module(host, port, key, mode):
dev = remote.metal(0)
f1 = remote.load_module("dev_lib.dylib")
a_np = np.random.uniform(size=1024).astype(A.dtype)
- a = tvm.runtime.tensor(a_np, dev)
- b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
+ a = dev.tensor(a_np)
+ b = dev.tensor(np.zeros(1024, dtype=A.dtype))
time_f = f1.time_evaluator(f1.entry_name, dev, number=10)
cost = time_f(a, b).mean
print(f"Metal: {cost:g} secs/op")
diff --git a/docs/how_to/tutorials/cross_compilation_and_rpc.py
b/docs/how_to/tutorials/cross_compilation_and_rpc.py
index 76ba56d09e..9fc425c7a9 100644
--- a/docs/how_to/tutorials/cross_compilation_and_rpc.py
+++ b/docs/how_to/tutorials/cross_compilation_and_rpc.py
@@ -182,8 +182,8 @@ func = remote.load_module("lib.tar")
# create arrays on the remote device
dev = remote.cpu()
-a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(A.dtype), dev)
-b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
+a = dev.tensor(np.random.uniform(size=1024).astype(A.dtype))
+b = dev.tensor(np.zeros(1024, dtype=A.dtype))
# the function will run on the remote device
func(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
@@ -437,8 +437,8 @@ def run_opencl():
# run
dev = remote.cl()
- a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(A.dtype), dev)
- b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
+ a = dev.tensor(np.random.uniform(size=1024).astype(A.dtype))
+ b = dev.tensor(np.zeros(1024, dtype=A.dtype))
func(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
print("OpenCL test passed!")
@@ -606,7 +606,7 @@ def run_pytorch_model_via_rpc():
print(" remote.upload('model_params.npz')")
print("\n4. Load and run remotely:")
print(" lib = remote.load_module('model_deployed.so')")
- print(" vm = relax.VirtualMachine(lib, remote.cpu())")
+ print(" vm = rpc.virtual_machine(lib, remote.cpu())")
print(" result = vm['main'](input, *params)")
print("\nDevice examples:")
print(" - Raspberry Pi: 192.168.1.100")
@@ -640,12 +640,12 @@ def run_pytorch_model_via_rpc():
dev = remote.cpu()
# Create VM and load parameters
- vm = relax.VirtualMachine(lib, dev)
+ vm = rpc.virtual_machine(lib, dev)
# Load parameters from the uploaded file
# Note: In practice, you might load this from the remote filesystem
params_npz = np.load(params_path)
- remote_params = [tvm.runtime.tensor(params_npz[f"p_{i}"], dev) for i in
range(len(params_npz))]
+ remote_params = [dev.tensor(params_npz[f"p_{i}"]) for i in
range(len(params_npz))]
######################################################################
# Step 5: Run Inference on Remote Device
@@ -659,7 +659,7 @@ def run_pytorch_model_via_rpc():
# Prepare input data
input_data = np.random.randn(1, 1, 28, 28).astype("float32")
- remote_input = tvm.runtime.tensor(input_data, dev)
+ remote_input = dev.tensor(input_data)
# Run inference using set_input + invoke_stateful for RPC compatibility
vm.set_input("main", remote_input, *remote_params)
diff --git a/docs/how_to/tutorials/export_and_load_executable.py
b/docs/how_to/tutorials/export_and_load_executable.py
index d14e4ecd93..9972861c12 100644
--- a/docs/how_to/tutorials/export_and_load_executable.py
+++ b/docs/how_to/tutorials/export_and_load_executable.py
@@ -353,7 +353,7 @@ if RUN_EXAMPLE:
#
# # Step 4: Load and run on remote device
# lib = remote.load_module("mlp_arm.so")
-# vm = relax.VirtualMachine(lib, remote.cpu())
+# vm = rpc.virtual_machine(lib, remote.cpu())
# # ... prepare input and params, then run inference
#
# The key difference is using an ARM target triple during compilation and
diff --git a/docs/reference/api/python/rpc.rst
b/docs/reference/api/python/rpc.rst
index d5d1192018..1f93409522 100644
--- a/docs/reference/api/python/rpc.rst
+++ b/docs/reference/api/python/rpc.rst
@@ -17,6 +17,34 @@
tvm.rpc
-------
+RPC devices and tensors belong to a session. They are distinct from local
+``Device`` and ``Tensor`` values::
+
+ remote = tvm.rpc.connect(host, port)
+ device = remote.cpu()
+ x = device.tensor(numpy_data)
+ y = device.empty(x.shape, x.dtype)
+ module = remote.load_module("kernel.so")
+ module(x, y)
+ result = y.numpy()
+
+Use ``device.empty`` and ``device.tensor`` for remote allocations. Normal
+``tvm.runtime.empty`` and ``tvm.runtime.tensor`` allocate locally. Remote
+functions reject tensors, devices, modules, and objects from other sessions.
+``RemoteTensor.copyto(tvm.cpu())`` explicitly materializes a local tensor;
+remote tensors cannot export a local DLPack pointer. Views retain their remote
+backing storage, and returned tensors retain the remote tensor object.
+
+For a remote VM, use ``tvm.rpc.virtual_machine(module, device)``. This helper
+validates session ownership and passes ordinary endpoint device descriptors
+to VM initialization. A module's ``time_evaluator`` similarly validates the
+remote device through its RPC wrapper.
+
+Streams remain owned by the session until ``device.free_stream(stream)`` or
+session teardown. Selecting a temporary stream is safe, and aliases observe
+explicit release. ``remote.close()`` closes the session; existing remote
+values reject further operations while cleanup remains safe.
+
.. automodule:: tvm.rpc
:members:
:imported-members:
diff --git a/include/tvm/runtime/device_api.h b/include/tvm/runtime/device_api.h
index 70e40f9d48..e3b4726d32 100644
--- a/include/tvm/runtime/device_api.h
+++ b/include/tvm/runtime/device_api.h
@@ -366,60 +366,11 @@ inline const char* DLDeviceType2Str(int type) {
throw;
}
-/*! \brief The device type bigger than this is RPC device */
-constexpr int kRPCSessMask = 128;
-static_assert(kRPCSessMask >= TVMDeviceExtType_End);
-
-/*!
- * \brief Return true if a Device is owned by an RPC session.
- */
-inline bool IsRPCSessionDevice(Device dev) { return (dev.device_type /
kRPCSessMask) > 0; }
-
-/*!
- * \brief Return the RPCSessTable index of the RPC Session that owns this
device.
- * \return the table index.
- */
-inline int GetRPCSessionIndex(Device dev) {
- TVM_FFI_ICHECK(IsRPCSessionDevice(dev)) << "GetRPCSessionIndex: dev has no
RPC session";
- return dev.device_type / kRPCSessMask - 1;
-}
-
-/*!
- * \brief Remove the RPC session mask from a Device.
- * RPC clients typically do this when encoding a Device for transmission to an
RPC remote.
- * On the wire, RPCdevice are expected to be valid on the server without
interpretation.
- * \param dev A Device with non-zero RPC Session mask, valid on the RPC client.
- * \return A Device without any RPC Session mask, valid on the RPC server.
- */
-inline Device RemoveRPCSessionMask(Device dev) {
- dev.device_type = static_cast<DLDeviceType>(dev.device_type % kRPCSessMask);
- return dev;
-}
-
inline std::ostream& operator<<(std::ostream& os, DLDevice dev) { // NOLINT(*)
- if (tvm::runtime::IsRPCSessionDevice(dev)) {
- os << "remote[" << tvm::runtime::GetRPCSessionIndex(dev) << "]-";
- dev = tvm::runtime::RemoveRPCSessionMask(dev);
- }
- os << tvm::runtime::DLDeviceType2Str(static_cast<int>(dev.device_type)) <<
":" << dev.device_id;
+ os << DLDeviceType2Str(static_cast<int>(dev.device_type)) << ":" <<
dev.device_id;
return os;
}
-/*!
- * \brief Add a RPC session mask to a Device.
- * RPC clients typically do this when decoding a Device received from a RPC
remote.
- * \param dev A Device without any RPC Session mask, valid on the RPC server.
- * \param session_table_index Numeric index of the RPC session in the session
table.
- * \return A Device with RPC session mask added, valid on the RPC client.
- */
-inline Device AddRPCSessionMask(Device dev, int session_table_index) {
- TVM_FFI_ICHECK(!IsRPCSessionDevice(dev))
- << "AddRPCSessionMask: dev already non-zero RPCSessionIndex: " << dev;
- dev.device_type =
- static_cast<DLDeviceType>(dev.device_type | (kRPCSessMask *
(session_table_index + 1)));
- return dev;
-}
-
/*!
* \brief Check if runtime module is enabled for target.
* \param target The target module name.
diff --git a/jvm/core/src/main/java/org/apache/tvm/Device.java
b/jvm/core/src/main/java/org/apache/tvm/Device.java
index 5c7dbbe441..4661694ec1 100644
--- a/jvm/core/src/main/java/org/apache/tvm/Device.java
+++ b/jvm/core/src/main/java/org/apache/tvm/Device.java
@@ -19,7 +19,6 @@ package org.apache.tvm;
import java.util.HashMap;
import java.util.Map;
-import org.apache.tvm.rpc.RPC;
public class Device {
/**
@@ -228,11 +227,6 @@ public class Device {
@Override
public String toString() {
- if (deviceType >= RPC.RPC_SESS_MASK) {
- int tblId = deviceType / RPC.RPC_SESS_MASK - 1;
- int devType = deviceType % RPC.RPC_SESS_MASK;
- return String.format("remote[%d]:%s(%d)", tblId,
DEVICE_TYPE_TO_NAME.get(devType), deviceId);
- }
return String.format("%s(%d)", DEVICE_TYPE_TO_NAME.get(deviceType),
deviceId);
}
}
diff --git a/jvm/core/src/main/java/org/apache/tvm/TVMType.java
b/jvm/core/src/main/java/org/apache/tvm/TVMType.java
index 62693a93ac..76ee3f1ea6 100644
--- a/jvm/core/src/main/java/org/apache/tvm/TVMType.java
+++ b/jvm/core/src/main/java/org/apache/tvm/TVMType.java
@@ -99,7 +99,7 @@ public class TVMType {
}
String str = typeCodeStr + bits;
if (lanes != 1) {
- str += lanes;
+ str += "x" + lanes;
}
return str;
}
diff --git a/jvm/core/src/main/java/org/apache/tvm/rpc/RPC.java
b/jvm/core/src/main/java/org/apache/tvm/rpc/RPC.java
index f8ec822e5e..57be601cc5 100644
--- a/jvm/core/src/main/java/org/apache/tvm/rpc/RPC.java
+++ b/jvm/core/src/main/java/org/apache/tvm/rpc/RPC.java
@@ -25,7 +25,6 @@ public class RPC {
public static final int RPC_TRACKER_MAGIC = 0x2f271;
public static final int RPC_MAGIC = 0xff271;
public static final int RPC_CODE_MISMATCH = RPC_MAGIC + 2;
- public static final int RPC_SESS_MASK = 128;
public static final String TIMEOUT_ARG = "-timeout=";
diff --git a/jvm/core/src/main/java/org/apache/tvm/rpc/RPCSession.java
b/jvm/core/src/main/java/org/apache/tvm/rpc/RPCSession.java
index 07278f07b8..918988c9f0 100644
--- a/jvm/core/src/main/java/org/apache/tvm/rpc/RPCSession.java
+++ b/jvm/core/src/main/java/org/apache/tvm/rpc/RPCSession.java
@@ -33,12 +33,10 @@ import org.apache.tvm.Module;
*/
public class RPCSession {
private final Module session;
- private final int tblIndex;
private final Map<String, Function> remoteFuncs = new HashMap<String,
Function>();
RPCSession(Module sess) {
session = sess;
- tblIndex = (int)
RPC.getApi("SessTableIndex").pushArg(session).invoke().asLong();
}
/**
@@ -54,20 +52,19 @@ public class RPCSession {
* Construct a remote device.
* @param devType device type.
* @param devId device id.
- * @return The corresponding encoded remote device.
+ * @return The session-owned remote device.
*/
- public Device device(String devType, int devId) {
+ public TVMRemoteDevice device(String devType, int devId) {
Device dev = new Device(devType, devId);
- int encode = (tblIndex + 1) * RPC.RPC_SESS_MASK;
- return new TVMRemoteDevice(dev.deviceType + encode, devId, this);
+ return device(dev.deviceType, devId);
}
/**
* Construct a remote device.
* @param devType device type.
- * @return The corresponding encoded remote device.
+ * @return The session-owned remote device.
*/
- public Device device(String devType) {
+ public TVMRemoteDevice device(String devType) {
return device(devType, 0);
}
@@ -75,19 +72,18 @@ public class RPCSession {
* Construct a remote device.
* @param devType device type.
* @param devId device id.
- * @return The corresponding encoded remote device.
+ * @return The session-owned remote device.
*/
- public Device device(int devType, int devId) {
- int encode = (tblIndex + 1) * RPC.RPC_SESS_MASK;
- return new TVMRemoteDevice(devType + encode, devId, this);
+ public TVMRemoteDevice device(int devType, int devId) {
+ return (TVMRemoteDevice) RPC.getApi("RemoteDevice").call(session, new
Device(devType, devId));
}
/**
* Construct a remote device.
* @param devType device type.
- * @return The corresponding encoded remote device.
+ * @return The session-owned remote device.
*/
- public Device device(int devType) {
+ public TVMRemoteDevice device(int devType) {
return device(devType, 0);
}
@@ -96,15 +92,15 @@ public class RPCSession {
* @param devId device id.
* @return Remote CPU device.
*/
- public Device cpu(int devId) {
- return Device.cpu(devId);
+ public TVMRemoteDevice cpu(int devId) {
+ return device("cpu", devId);
}
/**
* Construct remote CPU device.
* @return Remote CPU device.
*/
- public Device cpu() {
+ public TVMRemoteDevice cpu() {
return cpu(0);
}
@@ -113,15 +109,15 @@ public class RPCSession {
* @param devId device id.
* @return Remote CUDA GPU device.
*/
- public Device cuda(int devId) {
- return Device.cuda(devId);
+ public TVMRemoteDevice cuda(int devId) {
+ return device("cuda", devId);
}
/**
* Construct remote CUDA GPU device.
* @return Remote CUDA GPU device.
*/
- public Device cuda() {
+ public TVMRemoteDevice cuda() {
return cuda(0);
}
@@ -130,15 +126,15 @@ public class RPCSession {
* @param devId device id.
* @return Remote OpenCL device.
*/
- public Device cl(int devId) {
- return Device.opencl(devId);
+ public TVMRemoteDevice cl(int devId) {
+ return device("opencl", devId);
}
/**
* Construct remote OpenCL device.
* @return Remote OpenCL device.
*/
- public Device cl() {
+ public TVMRemoteDevice cl() {
return cl(0);
}
@@ -147,15 +143,15 @@ public class RPCSession {
* @param devId device id.
* @return Remote OpenCL device.
*/
- public Device vulkan(int devId) {
- return Device.vulkan(devId);
+ public TVMRemoteDevice vulkan(int devId) {
+ return device("vulkan", devId);
}
/**
* Construct remote OpenCL device.
* @return Remote OpenCL device.
*/
- public Device vulkan() {
+ public TVMRemoteDevice vulkan() {
return vulkan(0);
}
@@ -164,15 +160,15 @@ public class RPCSession {
* @param devId device id.
* @return Remote metal device.
*/
- public Device metal(int devId) {
- return Device.metal(devId);
+ public TVMRemoteDevice metal(int devId) {
+ return device("metal", devId);
}
/**
* Construct remote Metal device.
* @return Remote metal device.
*/
- public Device metal() {
+ public TVMRemoteDevice metal() {
return metal(0);
}
diff --git a/jvm/core/src/main/java/org/apache/tvm/rpc/RemoteTensor.java
b/jvm/core/src/main/java/org/apache/tvm/rpc/RemoteTensor.java
new file mode 100644
index 0000000000..2020d9a4fe
--- /dev/null
+++ b/jvm/core/src/main/java/org/apache/tvm/rpc/RemoteTensor.java
@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.tvm.rpc;
+
+import org.apache.tvm.TVMObject;
+import org.apache.tvm.TVMValueNull;
+import org.apache.tvm.TensorBase;
+
+/**
+ * Remote storage whose opaque address cannot be used as a local Tensor.
+ * The native object retains the allocation or returned tensor handle and its
+ * session. Releasing this reference releases one owner; shared views retain
the
+ * backing allocation. Transfers use explicit local or same-session
destinations.
+ */
+public class RemoteTensor extends TVMObject {
+ public RemoteTensor(long handle, int typeIndex) {
+ super(handle, typeIndex);
+ }
+
+ public RemoteTensor copyFrom(TensorBase source) {
+ RPC.getApi("CopyFromLocal").call(source, this);
+ return this;
+ }
+
+ public TensorBase copyTo(TensorBase destination) {
+ RPC.getApi("CopyToLocal").call(this, destination);
+ return destination;
+ }
+
+ /** Copy within the owning session. Synchronize the device before observing
results. */
+ public RemoteTensor copyTo(RemoteTensor destination) {
+ RPC.getApi("CopyFromTo").call(this, destination, new TVMValueNull());
+ return destination;
+ }
+}
diff --git a/jvm/core/src/main/java/org/apache/tvm/rpc/TVMRemoteDevice.java
b/jvm/core/src/main/java/org/apache/tvm/rpc/TVMRemoteDevice.java
index 6f70fa9a29..0e71570170 100644
--- a/jvm/core/src/main/java/org/apache/tvm/rpc/TVMRemoteDevice.java
+++ b/jvm/core/src/main/java/org/apache/tvm/rpc/TVMRemoteDevice.java
@@ -17,14 +17,51 @@
package org.apache.tvm.rpc;
-import org.apache.tvm.Device;
+import org.apache.tvm.Function;
+import org.apache.tvm.TVMObject;
+import org.apache.tvm.TVMType;
+import org.apache.tvm.TVMValue;
+import org.apache.tvm.TVMValueNull;
-// always related to RPCSession. Cannot construct by users.
-public class TVMRemoteDevice extends Device {
- public final RPCSession rpcSession;
+/**
+ * A session-owned remote device, distinct from a local Device.
+ * Its native object retains the session and an ordinary endpoint device type
and
+ * index. The endpoint descriptor has no session identity by itself.
Allocation,
+ * queries and synchronization use the owning session and fail after its
closure.
+ */
+public class TVMRemoteDevice extends TVMObject {
+ public TVMRemoteDevice(long handle, int typeIndex) {
+ super(handle, typeIndex);
+ }
+
+ public boolean exist() {
+ return RPC.getApi("DeviceGetAttr").call(this, 0).asLong() != 0;
+ }
+
+ public long maxThreadsPerBlock() {
+ return RPC.getApi("DeviceGetAttr").call(this, 1).asLong();
+ }
+
+ public long warpSize() {
+ return RPC.getApi("DeviceGetAttr").call(this, 2).asLong();
+ }
+
+ public void sync() {
+ RPC.getApi("DeviceSync").call(this, new TVMValueNull());
+ }
- TVMRemoteDevice(int deviceType, int deviceId, RPCSession rpcSession) {
- super(deviceType, deviceId);
- this.rpcSession = rpcSession;
+ /** Allocate remote storage; shape metadata is constructed locally. */
+ public RemoteTensor empty(long[] dimensions, TVMType dtype) {
+ Function shapeConstructor = Function.getFunction("ffi.Shape");
+ for (long dim : dimensions) {
+ shapeConstructor.pushArg(dim);
+ }
+ TVMValue shape = shapeConstructor.invoke();
+ try {
+ return (RemoteTensor) RPC.getApi("Empty").call(
+ shape, dtype.toString(), this, new TVMValueNull());
+ } finally {
+ shape.release();
+ }
}
}
diff --git a/jvm/native/src/main/native/jni_helper_func.h
b/jvm/native/src/main/native/jni_helper_func.h
index 659c6e4f29..b2c7ac3616 100644
--- a/jvm/native/src/main/native/jni_helper_func.h
+++ b/jvm/native/src/main/native/jni_helper_func.h
@@ -168,7 +168,14 @@ jobject newTVMNull(JNIEnv* env) {
}
jobject newTVMObject(JNIEnv* env, jlong handle, jint type_index) {
- jclass cls = env->FindClass("org/apache/tvm/TVMObject");
+ const char* class_name = "org/apache/tvm/TVMObject";
+ std::string type_key =
reinterpret_cast<tvm::ffi::Object*>(handle)->GetTypeKey();
+ if (type_key == "rpc.RemoteDevice") {
+ class_name = "org/apache/tvm/rpc/TVMRemoteDevice";
+ } else if (type_key == "rpc.RemoteTensor") {
+ class_name = "org/apache/tvm/rpc/RemoteTensor";
+ }
+ jclass cls = env->FindClass(class_name);
jmethodID constructor = env->GetMethodID(cls, "<init>", "(JI)V");
jobject object = env->NewObject(cls, constructor, handle, type_index);
env->DeleteLocalRef(cls);
diff --git a/python/tvm/contrib/coreml_runtime.py
b/python/tvm/contrib/coreml_runtime.py
index a2c72df66c..30f6866e34 100644
--- a/python/tvm/contrib/coreml_runtime.py
+++ b/python/tvm/contrib/coreml_runtime.py
@@ -18,10 +18,8 @@
import tvm_ffi
-from ..rpc import base as rpc_base
-
-def create(symbol, compiled_model_path, device):
+def create(symbol, compiled_model_path, device, session=None):
"""Create a runtime executor module given a coreml model and context.
Parameters
@@ -31,21 +29,20 @@ def create(symbol, compiled_model_path, device):
compiled_model_path : str
The path of the compiled model to be deployed.
device : Device
- The device to deploy the module. It can be local or remote when there
- is only one Device.
+ The device to deploy the module.
+ session : RPCSession, optional
+ Explicit remote session when deploying on an RPC server.
Returns
-------
coreml_runtime : CoreMLModule
Runtime coreml module that can be used to execute the coreml model.
"""
- device_type = device.dlpack_device_type()
runtime_func = "tvm.coreml_runtime.create"
-
- if device_type >= rpc_base.RPC_SESS_MASK:
- fcreate = device._rpc_sess.get_function(runtime_func)
- else:
- fcreate = tvm_ffi.get_global_func(runtime_func)
- assert fcreate, "Cannot find `tvm.coreml_runtime.create` function."
+ fcreate = (
+ session.get_function(runtime_func)
+ if session is not None
+ else tvm_ffi.get_global_func(runtime_func)
+ )
return CoreMLModule(fcreate(symbol, compiled_model_path))
diff --git a/python/tvm/rpc/__init__.py b/python/tvm/rpc/__init__.py
index 38ee1884b5..dba7522678 100644
--- a/python/tvm/rpc/__init__.py
+++ b/python/tvm/rpc/__init__.py
@@ -35,3 +35,5 @@ from .server import Server
from .client import connect, connect_tracker
from .client import RPCSession, LocalSession, PopenSession, TrackerSession
from .minrpc import with_minrpc
+
+from .remote import RemoteDevice, RemoteTensor, RemoteStream, time_evaluator,
virtual_machine
diff --git a/python/tvm/rpc/base.py b/python/tvm/rpc/base.py
index dd1a33fc9e..271745e4d8 100644
--- a/python/tvm/rpc/base.py
+++ b/python/tvm/rpc/base.py
@@ -53,9 +53,6 @@ class TrackerCode:
GET_PENDING_MATCHKEYS = 7
-RPC_SESS_MASK = 128
-
-
# Use "127.0.0.1" or "::1" if there is a need to force ip4 or ip6
# connection for "localhost".
def get_addr_family(addr):
diff --git a/python/tvm/rpc/client.py b/python/tvm/rpc/client.py
index 7dde53cc2f..dc9a3cc16a 100644
--- a/python/tvm/rpc/client.py
+++ b/python/tvm/rpc/client.py
@@ -42,7 +42,6 @@ class RPCSession:
# pylint: disable=invalid-name
def __init__(self, sess):
self._sess = sess
- self._tbl_index = _ffi_api.SessTableIndex(sess)
self._remote_funcs = {}
def system_lib(self):
@@ -85,14 +84,14 @@ class RPCSession:
Returns
-------
- dev: Device
- The corresponding encoded remote device.
+ dev: RemoteDevice
+ The session-owned device.
"""
- dev = tvm.runtime.device(dev_type, dev_id)
- encode = (self._tbl_index + 1) * base.RPC_SESS_MASK
- dev = tvm.runtime.device(dev.dlpack_device_type() + encode, dev.index)
- dev._rpc_sess = self
- return dev
+ return _ffi_api.RemoteDevice(self._sess, tvm.runtime.device(dev_type,
dev_id))
+
+ def close(self):
+ """Close this session; existing remote values reject further
operations."""
+ self._sess.get_function("CloseRPCConnection")()
def upload(self, data, target=None):
"""Upload file to remote runtime temp folder
diff --git a/python/tvm/rpc/remote.py b/python/tvm/rpc/remote.py
new file mode 100644
index 0000000000..70db085d25
--- /dev/null
+++ b/python/tvm/rpc/remote.py
@@ -0,0 +1,270 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Explicit RPC-owned devices, tensors, and streams.
+
+A remote tensor is not a local tensor or a DLPack provider. Use ``numpy`` or
+``copyto(tvm.cpu())`` to materialize its contents in the client process.
+"""
+
+import numpy as np
+import tvm_ffi
+
+import tvm.runtime
+
+from . import _ffi_api
+
+
+def _register_remote_object(type_key):
+ # Keep RPC imports available for feature detection in builds without RPC.
+ if tvm_ffi.get_global_func("rpc.RemoteDevice", allow_missing=True) is None:
+ return lambda cls: cls
+ return tvm_ffi.register_object(type_key)
+
+
+class _RemoteDeviceAttributeQuery:
+ """Bind Device's shared property definitions to an RPC-owned query hook.
+
+ This is only a property receiver, never a local Device or tensor factory
+ argument. Keeping conversions in Device avoids a second attribute catalog.
+ """
+
+ def __init__(self, remote):
+ self.remote = remote
+ self.index = remote.index
+
+ def dlpack_device_type(self):
+ return self.remote.endpoint_device.dlpack_device_type()
+
+ def get_attr(self, attr):
+ return self.remote.get_attr(attr)
+
+ def _GetDeviceAttr(self, device_type, device_id, attr):
+ return self.get_attr(attr)
+
+
+@_register_remote_object("rpc.RemoteDevice")
+class RemoteDevice(tvm_ffi.Object):
+ """A device descriptor interpreted in one RPC session's execution endpoint.
+
+ ``endpoint_device`` is a read-only ordinary Device containing only the
device
+ type and index. The native object retains its session; the descriptor alone
+ does not identify that session and must not be used for local allocation.
+ Device capability attributes use the same definitions and conversions as
+ :class:`tvm.runtime.Device`, with queries dispatched by the owning session.
+ Explicit session closure invalidates allocation, queries and stream calls.
+ """
+
+ @property
+ def index(self):
+ return self.endpoint_device.index
+
+ @property
+ def type(self):
+ return self.endpoint_device.type
+
+ def __repr__(self):
+ return f"RemoteDevice({self.endpoint_device})"
+
+ def get_attr(self, attr):
+ """Query the shared DeviceAttrKind API through the owning session.
+
+ Results and backend errors are those of DeviceAPI.GetAttr; named Python
+ capabilities additionally use Device's existing result conversions.
+ """
+ return _ffi_api.DeviceGetAttr(self, attr)
+
+ def __getattr__(self, name):
+ # FFI reflection owns fields such as endpoint_device. Only missing
fields
+ # fall through to the ordinary Device's shared attribute definitions.
+ try:
+ return super().__getattr__(name)
+ except AttributeError:
+ attribute = getattr(tvm.runtime.Device, name, None)
+ if not isinstance(attribute, property):
+ raise AttributeError(f"RemoteDevice has no attribute
{name!r}") from None
+ return attribute.__get__(_RemoteDeviceAttributeQuery(self))
+
+ def texture_spatial_limit(self):
+ """Query the shared texture limit convenience through this session."""
+ return
tvm.runtime.Device.texture_spatial_limit(_RemoteDeviceAttributeQuery(self))
+
+ def empty(self, shape, dtype="float32", mem_scope=None):
+ """Allocate a remote tensor through this device's session."""
+ return _ffi_api.Empty(tvm_ffi.Shape(shape), tvm_ffi.dtype(dtype),
self, mem_scope)
+
+ def tensor(self, data, dtype=None):
+ """Allocate a remote tensor and copy data into it synchronously."""
+ if isinstance(data, RemoteTensor):
+ if dtype is not None and tvm_ffi.dtype(dtype) != data.dtype:
+ raise ValueError("Remote tensor copy cannot convert dtype;
materialize it first")
+ result = self.empty(data.shape, data.dtype)
+ return result.copyfrom(data)
+ local = (
+ data
+ if isinstance(data, tvm_ffi.Tensor) and dtype is None
+ else tvm.runtime.tensor(np.asarray(data, dtype=dtype))
+ )
+ return self.empty(local.shape, local.dtype).copyfrom(local)
+
+ def set_device(self):
+ _ffi_api.DeviceSetDevice(self)
+
+ def create_stream(self):
+ return _ffi_api.DeviceCreateStream(self)
+
+ def free_stream(self, stream):
+ _ffi_api.DeviceFreeStream(self, stream)
+
+ def current_stream(self):
+ return _ffi_api.DeviceCurrentStream(self)
+
+ def set_stream(self, stream=None):
+ _ffi_api.DeviceSetStream(self, stream)
+
+ def sync(self, stream=None):
+ _ffi_api.DeviceSync(self, stream)
+
+
+@_register_remote_object("rpc.RemoteStream")
+class RemoteStream(tvm_ffi.Object):
+ """An opaque stream tied to one remote device and its session.
+
+ The session retains created streams until ``device.free_stream`` or session
+ teardown, including streams selected using a temporary Python wrapper.
+ Current-stream aliases share closure state; releasing one invalidates them
+ all. Stream operations reject a different device or session before use.
+ """
+
+
+@_register_remote_object("rpc.RemoteTensor")
+class RemoteTensor(tvm_ffi.Object):
+ """Contiguous remote storage with explicit metadata and transfer
operations.
+
+ Read-only ``device`` identifies the owning RemoteDevice; ``shape`` and
+ ``dtype`` describe the logical tensor independently of remote metadata.
+ Views retain shared backing storage, and returned tensors retain the remote
+ object handle until the last local owner disappears. The backing owner also
+ retains the session. No data pointer is exposed for local use or DLPack;
+ ``numpy`` and ``copyto`` explicitly materialize data in client memory.
+ """
+
+ @property
+ def ndim(self):
+ return len(self.shape)
+
+ def __repr__(self):
+ return f"RemoteTensor(shape={tuple(self.shape)}, dtype={self.dtype},
device={self.device})"
+
+ def __dlpack__(self, *args, **kwargs):
+ raise TypeError("RemoteTensor cannot export local DLPack;
copyto(tvm.cpu()) first")
+
+ def __dlpack_device__(self):
+ raise TypeError("RemoteTensor has no local DLPack device;
copyto(tvm.cpu()) first")
+
+ def copyfrom(self, source):
+ """Copy a local value or a tensor from the same remote session."""
+ if isinstance(source, RemoteTensor):
+ _ffi_api.CopyFromTo(source, self, None)
+ source.device.sync()
+ self.device.sync()
+ else:
+ if not isinstance(source, tvm_ffi.Tensor):
+ source = tvm.runtime.tensor(np.asarray(source,
dtype=str(self.dtype)))
+ if tuple(source.shape) != tuple(self.shape):
+ raise ValueError("Tensor copy shape mismatch")
+ if str(source.dtype) != str(self.dtype):
+ raise ValueError("Tensor copy dtype mismatch")
+ _ffi_api.CopyFromLocal(source, self)
+ return self
+
+ def copyto(self, target):
+ """Synchronously copy into storage or allocate on an explicit
device."""
+ if isinstance(target, RemoteDevice):
+ target = target.empty(self.shape, self.dtype)
+ if isinstance(target, RemoteTensor):
+ return target.copyfrom(self)
+ if isinstance(target, tvm.runtime.Device):
+ is_cpu = target.dlpack_device_type() ==
tvm.runtime.cpu().dlpack_device_type()
+ local = tvm.runtime.empty(
+ self.shape, self.dtype, target if is_cpu else tvm.runtime.cpu()
+ )
+ _ffi_api.CopyToLocal(self, local)
+ if is_cpu:
+ return local
+ return local.copyto(target)
+ _ffi_api.CopyToLocal(self, target)
+ return target
+
+ def numpy(self):
+ """Materialize in client memory and convert to NumPy."""
+ return self.copyto(tvm.runtime.cpu()).numpy()
+
+ def _create_view(self, shape, dtype=None, relative_byte_offset=0):
+ return _ffi_api.CreateView(
+ self, tvm_ffi.Shape(shape), tvm_ffi.dtype(dtype or self.dtype),
relative_byte_offset
+ )
+
+ def __setitem__(self, key, value):
+ if (
+ not isinstance(key, slice)
+ or key.start is not None
+ or key.stop is not None
+ or key.step not in (None, 1)
+ ):
+ raise ValueError("RemoteTensor only supports full-slice
assignment")
+ self.copyfrom(value)
+
+
+def time_evaluator(
+ module,
+ func_name,
+ device,
+ number=10,
+ repeat=1,
+ min_repeat_ms=0,
+ limit_zero_time_iterations=100,
+ cooldown_interval_ms=0,
+ repeats_to_cooldown=1,
+ cache_flush_bytes=0,
+ f_preproc="",
+):
+ """Create an evaluator after validating the module and remote device
session."""
+ return module.time_evaluator(
+ func_name,
+ device,
+ number,
+ repeat,
+ min_repeat_ms,
+ limit_zero_time_iterations,
+ cooldown_interval_ms,
+ repeats_to_cooldown,
+ cache_flush_bytes,
+ f_preproc,
+ )
+
+
+def virtual_machine(module, device, memory_cfg=None):
+ """Initialize a remote VM with session-checked endpoint device
descriptors."""
+ from tvm.runtime.vm import VirtualMachine
+
+ devices = list(device) if isinstance(device, list | tuple) else [device]
+ local = [_ffi_api.ValidateDevice(module, dev) for dev in devices]
+ if isinstance(memory_cfg, dict):
+ memory_cfg = {
+ _ffi_api.ValidateDevice(module, dev): value for dev, value in
memory_cfg.items()
+ }
+ return VirtualMachine(module, local, memory_cfg)
diff --git a/python/tvm/runtime/device.py b/python/tvm/runtime/device.py
index 7439751c7d..7195e6dc6d 100644
--- a/python/tvm/runtime/device.py
+++ b/python/tvm/runtime/device.py
@@ -23,8 +23,6 @@ import tvm_ffi
from . import _ffi_api
-RPC_SESS_MASK = 128
-
class Device(tvm_ffi.core.Device):
"""TVM device strucure."""
@@ -55,8 +53,7 @@ class Device(tvm_ffi.core.Device):
"""Maximum number of threads on each block.
Returns device value for CUDA, Metal, ROCm, OpenCL, and Vulkan
- devices. Returns remote device value for RPC devices.
- Returns None for all other devices.
+ devices. Returns None for all other devices.
Returns
-------
@@ -72,7 +69,7 @@ class Device(tvm_ffi.core.Device):
Returns device value for CUDA, ROCm, and Vulkan. Returns
1 for Metal and OpenCL devices, regardless of the physical
- device. Returns remote device value for RPC devices. Returns
+ device. Returns
None for all other devices.
Returns
@@ -88,7 +85,7 @@ class Device(tvm_ffi.core.Device):
"""Total amount of shared memory per block in bytes.
Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
- Returns remote device value for RPC devices. Returns None for
+ Returns None for
all other devices.
Returns
@@ -107,7 +104,7 @@ class Device(tvm_ffi.core.Device):
supported by the device.
Returns device value for CUDA, ROCm, OpenCL, and
- Vulkan. Returns remote device value for RPC devices. Returns
+ Vulkan. Returns
None for all other devices.
Returns
@@ -123,7 +120,7 @@ class Device(tvm_ffi.core.Device):
"""Return the vendor-specific name of device.
Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
- Returns remote device value for RPC devices. Returns None for
+ Returns None for
all other devices.
Returns
@@ -139,7 +136,7 @@ class Device(tvm_ffi.core.Device):
"""Return the max clock frequency of device (kHz).
Returns device value for CUDA, ROCm, and OpenCL. Returns
- remote device value for RPC devices. Returns None for all
+ None for all
other devices.
Returns
@@ -155,7 +152,7 @@ class Device(tvm_ffi.core.Device):
"""Return the number of compute units in the device.
Returns device value for CUDA, ROCm, and OpenCL. Returns
- remote device value for RPC devices. Returns None for all
+ None for all
other devices.
Returns
@@ -171,7 +168,7 @@ class Device(tvm_ffi.core.Device):
"""Return the maximum size of each thread axis
Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
- Returns remote device value for RPC devices. Returns None for
+ Returns None for
all other devices.
Returns
@@ -190,7 +187,7 @@ class Device(tvm_ffi.core.Device):
Vulkan.
Returns device value for CUDA, ROCm, OpenCL, and Vulkan.
- Returns remote device value for RPC devices. Returns None for
+ Returns None for
all other devices.
Returns
@@ -208,8 +205,7 @@ class Device(tvm_ffi.core.Device):
Returns driver vendor's internal version number.
(e.g. "450.408.256" for nvidia-driver-450)
- Returns device value for opencl and vulkan. Returns remote
- device value for RPC devices. Returns None for all other
+ Returns device value for opencl and vulkan. Returns None for all other
devices.
Returns
@@ -320,10 +316,6 @@ class Device(tvm_ffi.core.Device):
_ffi_api.Device_StreamSync(self, stream or 0)
def __device_type_name__(self):
- if self.dlpack_device_type() >= RPC_SESS_MASK:
- tbl_id = self.dlpack_device_type() / RPC_SESS_MASK - 1
- dev_type = self.dlpack_device_type() % RPC_SESS_MASK
- return f"remote[{tbl_id}]:{Device._DEVICE_TYPE_TO_NAME[dev_type]}"
return Device._DEVICE_TYPE_TO_NAME[self.dlpack_device_type()]
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 8fde8cc443..82b725b72f 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -393,11 +393,10 @@ class Module(_Module):
The ProfileResult reports `repeat` time costs in seconds.
"""
try:
- feval = _ffi_api.RPCTimeEvaluator(
+ feval = _ffi_api.ModuleTimeEvaluator(
self,
func_name,
- dev.dlpack_device_type(),
- dev.index,
+ dev,
number,
repeat,
min_repeat_ms,
diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm.py
index 0adb446dbf..0e7eaee280 100644
--- a/python/tvm/runtime/vm.py
+++ b/python/tvm/runtime/vm.py
@@ -29,8 +29,6 @@ from tvm_ffi import Function, register_global_func
import tvm
from tvm.runtime import Device, Object
-from ..rpc.base import RPC_SESS_MASK
-
class VMInstrumentReturnKind(IntEnum):
NO_OP = 0
@@ -96,7 +94,7 @@ class VirtualMachine:
devs = [dev]
# CPU is required for executing shape functions
- if devs[-1].dlpack_device_type() % RPC_SESS_MASK !=
tvm.cpu().dlpack_device_type():
+ if devs[-1].dlpack_device_type() != tvm.cpu().dlpack_device_type():
devs.append(tvm.cpu())
default_alloc_type = VirtualMachine.POOLED_ALLOCATOR
@@ -114,7 +112,7 @@ class VirtualMachine:
)
init_args = []
for device in devs:
- init_args.append(device.dlpack_device_type() % RPC_SESS_MASK)
+ init_args.append(device.dlpack_device_type())
init_args.append(device.index)
alloc_type = memory_cfg[device] if device in memory_cfg else
default_alloc_type
init_args.append(alloc_type)
diff --git a/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
b/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
index 27ab71e669..898e2ffe7c 100644
--- a/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
+++ b/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
@@ -483,7 +483,7 @@ def default_alloc_argument(
"Please make sure 'USE_RANDOM' is turned ON in the config.cmake on the
RPC server.",
)
- return alloc_argument_common(f_random_fill, device, args_info,
alloc_repeat)
+ return alloc_argument_common(f_random_fill, device, args_info,
alloc_repeat, device.empty)
def default_run_evaluator(
diff --git a/python/tvm/s_tir/meta_schedule/runner/utils.py
b/python/tvm/s_tir/meta_schedule/runner/utils.py
index 1a18d80328..06716a75b7 100644
--- a/python/tvm/s_tir/meta_schedule/runner/utils.py
+++ b/python/tvm/s_tir/meta_schedule/runner/utils.py
@@ -36,6 +36,7 @@ def alloc_argument_common(
device: Device,
args_info: T_ARG_INFO_JSON_OBJ_LIST,
alloc_repeat: int,
+ f_alloc: Callable | None = None,
) -> list[T_ARGUMENT_LIST]:
"""Common function to allocate the arguments
@@ -57,7 +58,11 @@ def alloc_argument_common(
"""
def alloc_tensor(_, dtype, shape) -> tvm.runtime.Tensor:
- arg = tvm.runtime.empty(shape=shape, dtype=dtype, device=device)
+ arg = (
+ f_alloc(shape, dtype)
+ if f_alloc is not None
+ else tvm.runtime.empty(shape=shape, dtype=dtype, device=device)
+ )
f_random_fill(arg)
return arg
diff --git a/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py
b/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py
index 74c7a6c307..0f41b112d0 100644
--- a/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py
+++ b/python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py
@@ -37,7 +37,6 @@ def run_module_via_rpc(
import os
import tempfile
- from tvm.runtime import ndarray
from tvm.support.tar import tar
# pylint: enable=import-outside-toplevel
@@ -50,5 +49,5 @@ def run_module_via_rpc(
_, filename = os.path.split(filename)
rt_mod = session.load_module(filename)
dev = session.device(dev_type, 0)
- nd_args = {k: ndarray.array(v, dev) for k, v in args.items()}
+ nd_args = {k: dev.tensor(v) for k, v in args.items()}
return continuation(rt_mod, dev, nd_args)
diff --git a/python/tvm/s_tir/meta_schedule/testing/validate_database.py
b/python/tvm/s_tir/meta_schedule/testing/validate_database.py
index 27f57b4b04..410779178a 100644
--- a/python/tvm/s_tir/meta_schedule/testing/validate_database.py
+++ b/python/tvm/s_tir/meta_schedule/testing/validate_database.py
@@ -519,7 +519,7 @@ def make_alloc_arg_and_check(
args_info: ms.runner.rpc_runner.T_ARG_INFO_JSON_OBJ_LIST,
alloc_repeat: int,
) -> list[ms.runner.rpc_runner.T_ARGUMENT_LIST]:
- return f_with_args_alloc_argument_common(device, args_info,
alloc_repeat)
+ return [[device.tensor(arg) for arg in inputs] for _ in
range(alloc_repeat)]
def f_with_args_run_evaluator_rpc(
rpc_session: ms.runner.rpc_runner.RPCSession, # pylint:
disable=unused-argument
diff --git a/python/tvm/testing/runner.py b/python/tvm/testing/runner.py
index de14d717e1..3eb2e57539 100644
--- a/python/tvm/testing/runner.py
+++ b/python/tvm/testing/runner.py
@@ -29,7 +29,7 @@ if TYPE_CHECKING:
# pylint: disable=import-outside-toplevel,protected-access
-def _args_to_device(args, device):
+def _args_to_device(args, device, f_alloc=None):
import numpy as np
from tvm.runtime import Tensor, empty
@@ -37,7 +37,12 @@ def _args_to_device(args, device):
uploaded_args = []
for arg in args:
if isinstance(arg, np.ndarray | Tensor):
- uploaded_args.append(empty(arg.shape, dtype=arg.dtype,
device=device).copyfrom(arg))
+ allocated = (
+ f_alloc(arg.shape, arg.dtype)
+ if f_alloc
+ else empty(arg.shape, dtype=arg.dtype, device=device)
+ )
+ uploaded_args.append(allocated.copyfrom(arg))
elif isinstance(arg, int | float):
uploaded_args.append(arg)
else:
@@ -46,11 +51,12 @@ def _args_to_device(args, device):
def _args_to_numpy(args):
+ from tvm.rpc import RemoteTensor
from tvm.runtime import Tensor
downloaded_args = []
for arg in args:
- if isinstance(arg, Tensor):
+ if isinstance(arg, Tensor | RemoteTensor):
downloaded_args.append(arg.numpy())
else:
downloaded_args.append(arg)
@@ -210,7 +216,7 @@ def rpc_run( # pylint:
disable=too-many-arguments,too-many-locals
export_func(mod, artifact_path)
try:
session.upload(artifact_path, remote_path)
- args = _args_to_device(args, device)
+ args = _args_to_device(args, device, device.empty)
remote_mod = session.load_module(remote_path)
profile_result = remote_mod.time_evaluator(
func_name=remote_mod.entry_name,
diff --git a/src/runtime/device_api.cc b/src/runtime/device_api.cc
index c46d2ca9af..b19c03813f 100644
--- a/src/runtime/device_api.cc
+++ b/src/runtime/device_api.cc
@@ -56,7 +56,6 @@ class DeviceAPIManager {
private:
std::array<DeviceAPI*, kMaxDeviceAPI> api_;
- DeviceAPI* rpc_api_{nullptr};
std::mutex mutex_;
// constructor
DeviceAPIManager() { std::fill(api_.begin(), api_.end(), nullptr); }
@@ -67,20 +66,14 @@ class DeviceAPIManager {
}
// Get or initialize API.
DeviceAPI* GetAPI(int type, bool allow_missing) {
- if (type < kRPCSessMask) {
- if (api_[type] != nullptr) return api_[type];
- std::lock_guard<std::mutex> lock(mutex_);
- if (api_[type] != nullptr) return api_[type];
- api_[type] = GetAPI(DLDeviceType2Str(type), allow_missing);
- return api_[type];
- } else {
- if (rpc_api_ != nullptr) return rpc_api_;
- std::lock_guard<std::mutex> lock(mutex_);
- if (rpc_api_ != nullptr) return rpc_api_;
- rpc_api_ = GetAPI("rpc", allow_missing);
- return rpc_api_;
- }
+ TVM_FFI_CHECK(type >= 0 && type < kMaxDeviceAPI, ValueError) << "Invalid
device type " << type;
+ if (api_[type] != nullptr) return api_[type];
+ std::lock_guard<std::mutex> lock(mutex_);
+ if (api_[type] != nullptr) return api_[type];
+ api_[type] = GetAPI(DLDeviceType2Str(type), allow_missing);
+ return api_[type];
}
+
DeviceAPI* GetAPI(const std::string name, bool allow_missing) {
std::string factory = "device_api." + name;
const auto f = tvm::ffi::Function::GetGlobal(factory);
diff --git a/src/runtime/module.cc b/src/runtime/module.cc
index 427766a158..1bced00a39 100644
--- a/src/runtime/module.cc
+++ b/src/runtime/module.cc
@@ -51,7 +51,7 @@ bool RuntimeEnabled(const ffi::String& target_str) {
} else if (target == "vulkan") {
f_name = "device_api.vulkan";
} else if (target == "rpc") {
- f_name = "device_api.rpc";
+ f_name = "rpc.LocalSession";
} else if (target == "hexagon") {
f_name = "device_api.hexagon";
} else if (target.length() >= 5 && target.substr(0, 5) == "nvptx") {
diff --git a/src/runtime/rpc/minrpc/minrpc_server.h
b/src/runtime/rpc/minrpc/minrpc_server.h
index 059c2fc2e4..95d31fb7ec 100644
--- a/src/runtime/rpc/minrpc/minrpc_server.h
+++ b/src/runtime/rpc/minrpc/minrpc_server.h
@@ -37,7 +37,7 @@
#include <memory>
#include <utility>
-#include "../../support/arena.h"
+#include "../../../support/arena.h"
#include "rpc_reference.h"
namespace tvm {
diff --git a/src/runtime/rpc/rpc_device_api.cc
b/src/runtime/rpc/rpc_device_api.cc
index e828f752d9..09efe7d28b 100644
--- a/src/runtime/rpc/rpc_device_api.cc
+++ b/src/runtime/rpc/rpc_device_api.cc
@@ -17,147 +17,165 @@
* under the License.
*/
-/*!
- * \file rpc_device_api.cc
+/*! \file rpc_device_api.cc
+ * \brief Explicit device and tensor operations owned by RPC.
*/
-#include <tvm/ffi/error.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/device_api.h>
+#include <limits>
-#include <utility>
-
-#include "rpc_session.h"
+#include "rpc_remote.h"
namespace tvm {
namespace runtime {
-class RPCDeviceAPI final : public DeviceAPI {
- public:
- void SetDevice(Device dev) final {
- auto remote_dev = RemoveRPCSessionMask(dev);
- GetSess(dev)->GetDeviceAPI(remote_dev)->SetDevice(remote_dev);
- }
-
- void GetAttr(Device dev, DeviceAttrKind kind, ffi::Any* rv) final {
- auto remote_dev = RemoveRPCSessionMask(dev);
- GetSess(dev)->GetDeviceAPI(remote_dev)->GetAttr(remote_dev, kind, rv);
- }
-
- void* AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType
dtype,
- ffi::Optional<ffi::String> mem_scope) final {
- auto sess = GetSess(dev);
- auto remote_dev = RemoveRPCSessionMask(dev);
- void* data =
- sess->GetDeviceAPI(remote_dev)->AllocDataSpace(remote_dev, ndim,
shape, dtype, mem_scope);
- RemoteSpace* space = new RemoteSpace();
- space->data = data;
- space->sess = std::move(sess);
- return space;
- }
-
- void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType
type_hint) final {
- auto sess = GetSess(dev);
- auto remote_dev = RemoveRPCSessionMask(dev);
- void* data =
- sess->GetDeviceAPI(remote_dev)->AllocDataSpace(remote_dev, nbytes,
alignment, type_hint);
-
- RemoteSpace* space = new RemoteSpace();
- space->data = data;
- space->sess = std::move(sess);
- return space;
- }
- void FreeDataSpace(Device dev, void* ptr) final {
- RemoteSpace* space = static_cast<RemoteSpace*>(ptr);
- auto remote_dev = RemoveRPCSessionMask(dev);
- try {
- GetSess(dev)->GetDeviceAPI(remote_dev)->FreeDataSpace(remote_dev,
space->data);
- } catch (const ffi::Error& e) {
- // fault tolerance to remote close.
- }
- delete space;
- }
-
- void CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream)
final {
- DLDevice dev_from = from->device;
- DLDevice dev_to = to->device;
- if (IsRPCSessionDevice(dev_from) && IsRPCSessionDevice(dev_to)) {
- TVM_FFI_ICHECK(dev_from.device_type == dev_to.device_type)
- << "Cannot copy across two different remote session";
- DLTensor from_tensor = *from;
- from_tensor.device = RemoveRPCSessionMask(dev_from);
- from_tensor.data = static_cast<const RemoteSpace*>(from->data)->data;
- DLTensor to_tensor = *to;
- to_tensor.device = RemoveRPCSessionMask(dev_to);
- to_tensor.data = static_cast<const RemoteSpace*>(to->data)->data;
- auto remote_dev = from_tensor.device;
- if (remote_dev.device_type == kDLCPU) remote_dev = to_tensor.device;
-
GetSess(dev_from)->GetDeviceAPI(remote_dev)->CopyDataFromTo(&from_tensor,
&to_tensor, stream);
- } else if (IsRPCSessionDevice(dev_from) && dev_to.device_type == kDLCPU) {
- DLTensor from_tensor = *from;
- from_tensor.device = RemoveRPCSessionMask(dev_from);
- from_tensor.data = static_cast<const RemoteSpace*>(from->data)->data;
- void* to_bytes = static_cast<char*>(to->data) + to->byte_offset;
- size_t nbytes = ffi::GetDataSize(*to);
- GetSess(dev_from)->CopyFromRemote(&from_tensor, to_bytes, nbytes);
- } else if (dev_from.device_type == kDLCPU && IsRPCSessionDevice(dev_to)) {
- DLTensor to_tensor = *to;
- to_tensor.device = RemoveRPCSessionMask(dev_to);
- to_tensor.data = static_cast<const RemoteSpace*>(to->data)->data;
- void* from_bytes = static_cast<char*>(from->data) + from->byte_offset;
- size_t nbytes = ffi::GetDataSize(*from);
- GetSess(dev_to)->CopyToRemote(from_bytes, &to_tensor, nbytes);
- } else {
- TVM_FFI_THROW(InternalError) << "expect copy from/to remote or between
remote";
- }
- }
-
- TVMStreamHandle CreateStream(Device dev) {
- auto remote_dev = RemoveRPCSessionMask(dev);
- return GetSess(dev)->GetDeviceAPI(remote_dev)->CreateStream(remote_dev);
- }
-
- void FreeStream(Device dev, TVMStreamHandle stream) {
- auto remote_dev = RemoveRPCSessionMask(dev);
- GetSess(dev)->GetDeviceAPI(remote_dev)->FreeStream(remote_dev, stream);
- }
-
- void StreamSync(Device dev, TVMStreamHandle stream) final {
- auto remote_dev = RemoveRPCSessionMask(dev);
- GetSess(dev)->GetDeviceAPI(remote_dev)->StreamSync(remote_dev, stream);
- }
-
- void SetStream(Device dev, TVMStreamHandle stream) final {
- auto remote_dev = RemoveRPCSessionMask(dev);
- GetSess(dev)->GetDeviceAPI(remote_dev)->SetStream(remote_dev, stream);
- }
-
- TVMStreamHandle GetCurrentStream(Device dev) final {
- auto remote_dev = RemoveRPCSessionMask(dev);
- return
GetSess(dev)->GetDeviceAPI(remote_dev)->GetCurrentStream(remote_dev);
+// Check before allocation or view arithmetic; ffi::GetDataSize assumes valid
dimensions.
+static uint64_t RemoteTensorBytes(const ffi::Shape& shape, DLDataType dtype) {
+ TVM_FFI_CHECK(dtype.bits > 0 && dtype.lanes > 0, ValueError) << "Invalid
tensor dtype";
+ uint64_t size = 1;
+ for (int64_t dim : shape) {
+ TVM_FFI_CHECK_GE(dim, 0, ValueError) << "Negative tensor dimension";
+ TVM_FFI_CHECK(dim == 0 || size <= std::numeric_limits<uint64_t>::max() /
dim, ValueError)
+ << "Tensor size overflow";
+ size *= dim;
}
+ uint64_t element_bits = static_cast<uint64_t>(dtype.bits) * dtype.lanes;
+ TVM_FFI_CHECK_LE(size, (std::numeric_limits<uint64_t>::max() - 7) /
element_bits, ValueError)
+ << "Tensor size overflow";
+ return (size * element_bits + 7) / 8;
+}
- protected:
- void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t
to_offset,
- size_t num_bytes, Device dev_from, Device dev_to,
DLDataType type_hint,
- TVMStreamHandle stream) final {
- TVM_FFI_THROW(InternalError) << "Not implemented.";
- }
+static RemoteTensor RemoteEmpty(ffi::Shape shape, DLDataType dtype,
RemoteDevice dev,
+ ffi::Optional<ffi::String> scope) {
+ RemoteTensorBytes(shape, dtype);
+ auto n = ffi::make_object<RemoteTensorObj>();
+ n->device = dev;
+ n->shape = shape;
+ n->dtype = dtype;
+ auto storage = std::make_shared<RemoteTensorStorage>();
+ storage->sess = dev->sess;
+ storage->device = dev->endpoint_device;
+ storage->owns_allocation = true;
+ storage->data =
+ dev->api()->AllocDataSpace(dev->endpoint_device, shape.size(),
shape.data(), dtype, scope);
+ n->storage = std::move(storage);
+ return RemoteTensor(n);
+}
- private:
- std::shared_ptr<RPCSession> GetSess(Device dev) {
- int tbl_index = GetRPCSessionIndex(dev);
- return RPCSession::Get(tbl_index);
- }
-};
+static void CopyRemote(RemoteTensor from, RemoteTensor to,
ffi::Optional<RemoteStream> stream) {
+ TVM_FFI_CHECK(from->device->sess == to->device->sess, ValueError)
+ << "Cannot copy across different remote sessions";
+ TVM_FFI_CHECK_EQ(RemoteTensorBytes(from->shape, from->dtype),
+ RemoteTensorBytes(to->shape, to->dtype), ValueError)
+ << "Tensor copy size mismatch";
+ auto src = from->descriptor();
+ auto dst = to->descriptor();
+ auto dev = src.device.device_type == kDLCPU ? to->device : from->device;
+ auto handle = CheckRemoteStream(dev, stream);
+ dev->api()->CopyDataFromTo(&src, &dst, handle);
+}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def_packed("device_api.rpc", [](ffi::PackedArgs args,
ffi::Any* rv) {
- static RPCDeviceAPI inst;
- DeviceAPI* ptr = &inst;
- *rv = static_cast<void*>(ptr);
- });
+ refl::ObjectDef<RemoteDeviceObj>().def_ro("endpoint_device",
&RemoteDeviceObj::endpoint_device);
+ refl::ObjectDef<RemoteTensorObj>()
+ .def_ro("device", &RemoteTensorObj::device)
+ .def_ro("shape", &RemoteTensorObj::shape)
+ .def_ro("dtype", &RemoteTensorObj::dtype);
+ refl::ObjectDef<RemoteStreamObj>();
+ refl::GlobalDef()
+ .def("rpc.RemoteDevice",
+ [](ffi::Module mod, Device dev) {
+ return MakeRemoteDevice(RPCModuleGetSession(mod), dev);
+ })
+ .def("rpc.DeviceGetAttr",
+ [](RemoteDevice dev, int attr) {
+ dev->sess->EnsureOpen();
+ auto kind = static_cast<DeviceAttrKind>(attr);
+ auto* api = dev->sess->GetDeviceAPI(dev->endpoint_device, kind ==
kExist);
+ ffi::Any rv;
+ if (api != nullptr) {
+ api->GetAttr(dev->endpoint_device, kind, &rv);
+ } else {
+ rv = 0;
+ }
+ return rv;
+ })
+ .def("rpc.DeviceSetDevice",
+ [](RemoteDevice dev) { dev->api()->SetDevice(dev->endpoint_device);
})
+ .def("rpc.DeviceCreateStream",
+ [](RemoteDevice dev) {
+ auto n = ffi::make_object<RemoteStreamObj>();
+ n->device = dev;
+ n->state = dev->sess->TrackStream(
+ dev->endpoint_device,
dev->api()->CreateStream(dev->endpoint_device), true);
+ return RemoteStream(n);
+ })
+ .def("rpc.DeviceCurrentStream",
+ [](RemoteDevice dev) {
+ auto n = ffi::make_object<RemoteStreamObj>();
+ n->device = dev;
+ n->state = dev->sess->TrackStream(
+ dev->endpoint_device,
dev->api()->GetCurrentStream(dev->endpoint_device), false);
+ return RemoteStream(n);
+ })
+ .def("rpc.DeviceFreeStream",
+ [](RemoteDevice dev, RemoteStream stream) {
+ CheckRemoteStream(dev, stream);
+ dev->sess->FreeStream(stream->state);
+ })
+ .def("rpc.DeviceSetStream",
+ [](RemoteDevice dev, ffi::Optional<RemoteStream> stream) {
+ auto handle = CheckRemoteStream(dev, stream);
+ dev->api()->SetStream(dev->endpoint_device, handle);
+ })
+ .def("rpc.DeviceSync",
+ [](RemoteDevice dev, ffi::Optional<RemoteStream> stream) {
+ auto handle = CheckRemoteStream(dev, stream);
+ dev->api()->StreamSync(dev->endpoint_device, handle);
+ })
+ .def("rpc.Empty", RemoteEmpty)
+ .def("rpc.CopyFromTo", CopyRemote)
+ .def("rpc.CopyFromLocal",
+ [](ffi::TensorView from, RemoteTensor to) {
+ TVM_FFI_CHECK(from.device().device_type == kDLCPU &&
from.IsContiguous(), ValueError)
+ << "RPC transfer requires a contiguous local CPU tensor";
+ TVM_FFI_CHECK_EQ(ffi::GetDataSize(from),
RemoteTensorBytes(to->shape, to->dtype),
+ ValueError)
+ << "Tensor copy size mismatch";
+ auto dst = to->descriptor();
+ void* bytes = static_cast<char*>(from.data_ptr()) +
from.byte_offset();
+ to->device->sess->CopyToRemote(bytes, &dst,
ffi::GetDataSize(from));
+ })
+ .def("rpc.CopyToLocal",
+ [](RemoteTensor from, ffi::TensorView to) {
+ TVM_FFI_CHECK(to.device().device_type == kDLCPU &&
to.IsContiguous(), ValueError)
+ << "RPC transfer requires a contiguous local CPU tensor";
+ TVM_FFI_CHECK_EQ(RemoteTensorBytes(from->shape, from->dtype),
ffi::GetDataSize(to),
+ ValueError)
+ << "Tensor copy size mismatch";
+ auto src = from->descriptor();
+ void* bytes = static_cast<char*>(to.data_ptr()) +
to.byte_offset();
+ from->device->sess->CopyFromRemote(&src, bytes,
ffi::GetDataSize(to));
+ })
+ .def("rpc.CreateView", [](RemoteTensor source, ffi::Shape shape,
DLDataType dtype,
+ int64_t offset) {
+ source->device->sess->EnsureOpen();
+ uint64_t source_bytes = RemoteTensorBytes(source->shape,
source->dtype);
+ uint64_t view_bytes = RemoteTensorBytes(shape, dtype);
+ TVM_FFI_CHECK(offset >= 0 && static_cast<uint64_t>(offset) <=
source_bytes &&
+ view_bytes <= source_bytes - offset,
+ ValueError)
+ << "Remote tensor view exceeds source bounds";
+ TVM_FFI_CHECK_LE(static_cast<uint64_t>(offset),
+ std::numeric_limits<uint64_t>::max() -
source->byte_offset, ValueError)
+ << "Remote tensor view offset overflow";
+ auto n = ffi::make_object<RemoteTensorObj>();
+ n->device = source->device;
+ n->shape = shape;
+ n->dtype = dtype;
+ n->byte_offset = source->byte_offset + offset;
+ n->storage = source->storage;
+ return RemoteTensor(n);
+ });
}
} // namespace runtime
} // namespace tvm
diff --git a/src/runtime/rpc/rpc_endpoint.cc b/src/runtime/rpc/rpc_endpoint.cc
index e0ee1d691a..efaf0972b8 100644
--- a/src/runtime/rpc/rpc_endpoint.cc
+++ b/src/runtime/rpc/rpc_endpoint.cc
@@ -33,6 +33,7 @@
#include <array>
#include <chrono>
#include <cmath>
+#include <limits>
#include <memory>
#include <string>
#include <type_traits>
@@ -185,10 +186,6 @@ class RPCEndpoint::EventHandler : public support::Stream {
<< "Cannot pass argument " << i << ", type " << obj->GetTypeKey()
<< " (type_index = " << obj->type_index() << ")";
}
- } else if (auto opt_device = args[i].as<DLDevice>()) {
- DLDevice dev = opt_device.value();
- TVM_FFI_CHECK(!IsRPCSessionDevice(dev), InternalError)
- << "cannot pass RPC device in the channel";
}
}
}
@@ -921,10 +918,13 @@ void RPCEndpoint::CopyToRemote(void* from_bytes,
DLTensor* to, uint64_t nbytes)
RPCCode code = RPCCode::kCopyToRemote;
uint64_t tensor_total_size_bytes =
static_cast<uint64_t>(ffi::GetDataSize(*to));
- TVM_FFI_ICHECK_LE(to->byte_offset + nbytes, tensor_total_size_bytes)
+ // byte_offset is relative to backing storage, not to the logical view.
+ TVM_FFI_ICHECK_LE(nbytes, tensor_total_size_bytes)
<< "CopyToRemote: overflow in tensor size: (byte_offset=" <<
to->byte_offset
<< ", nbytes=" << nbytes << ", tensor_total_size=" <<
tensor_total_size_bytes << ")";
+ TVM_FFI_ICHECK_LE(nbytes, std::numeric_limits<uint64_t>::max() -
to->byte_offset);
+
uint64_t overhead = RemoteCopyCalculatePacketOverheadSize(to, code, nbytes);
uint64_t packet_nbytes = overhead + nbytes;
@@ -941,10 +941,13 @@ void RPCEndpoint::CopyFromRemote(DLTensor* from, void*
to_bytes, uint64_t nbytes
RPCCode code = RPCCode::kCopyFromRemote;
uint64_t tensor_total_size_bytes =
static_cast<uint64_t>(ffi::GetDataSize(*from));
- TVM_FFI_ICHECK_LE(from->byte_offset + nbytes, tensor_total_size_bytes)
+ // byte_offset is relative to backing storage, not to the logical view.
+ TVM_FFI_ICHECK_LE(nbytes, tensor_total_size_bytes)
<< "CopyFromRemote: overflow in tensor size: (byte_offset=" <<
from->byte_offset
<< ", nbytes=" << nbytes << ", tensor_total_size=" <<
tensor_total_size_bytes << ")";
+ TVM_FFI_ICHECK_LE(nbytes, std::numeric_limits<uint64_t>::max() -
from->byte_offset);
+
uint64_t overhead = RemoteCopyCalculatePacketOverheadSize(from, code,
nbytes);
uint64_t packet_nbytes = overhead;
@@ -1116,6 +1119,8 @@ class RPCClientSession : public RPCSession, public
DeviceAPI {
*/
explicit RPCClientSession(std::shared_ptr<RPCEndpoint> endpoint) :
endpoint_(endpoint) {}
+ ~RPCClientSession() override { ReleaseStreams(); }
+
// function overrides
PackedFuncHandle GetFunction(const std::string& name) final {
return endpoint_->SysCallRemote(RPCCode::kGetGlobalFunc,
name).cast<void*>();
@@ -1133,11 +1138,12 @@ class RPCClientSession : public RPCSession, public
DeviceAPI {
TVM_FFI_ICHECK_GT(rpc_max_size, overhead) << "CopyToRemote: Invalid block
size!";
const uint64_t block_size = rpc_max_size - overhead;
uint64_t block_count = 0;
+ const uint64_t initial_offset = remote_to->byte_offset;
const uint64_t num_blocks = nbytes / block_size;
void* from_bytes;
for (block_count = 0; block_count < num_blocks; block_count++) {
- remote_to->byte_offset = block_count * block_size;
+ remote_to->byte_offset = initial_offset + block_count * block_size;
from_bytes = reinterpret_cast<void*>(
(reinterpret_cast<uint8_t*>(local_from_bytes) + block_count *
block_size));
endpoint_->CopyToRemote(from_bytes, remote_to, block_size);
@@ -1145,7 +1151,7 @@ class RPCClientSession : public RPCSession, public
DeviceAPI {
const uint64_t remainder_bytes = nbytes % block_size;
if (remainder_bytes != 0) {
- remote_to->byte_offset = block_count * block_size;
+ remote_to->byte_offset = initial_offset + block_count * block_size;
from_bytes = reinterpret_cast<void*>(
(reinterpret_cast<uint8_t*>(local_from_bytes) + block_count *
block_size));
endpoint_->CopyToRemote(from_bytes, remote_to, remainder_bytes);
@@ -1159,11 +1165,12 @@ class RPCClientSession : public RPCSession, public
DeviceAPI {
TVM_FFI_ICHECK_GT(rpc_max_size, overhead) << "CopyFromRemote: Invalid
block size!";
const uint64_t block_size = rpc_max_size - overhead;
uint64_t block_count = 0;
+ const uint64_t initial_offset = remote_from->byte_offset;
const uint64_t num_blocks = nbytes / block_size;
void* to_bytes;
for (block_count = 0; block_count < num_blocks; block_count++) {
- remote_from->byte_offset = block_count * block_size;
+ remote_from->byte_offset = initial_offset + block_count * block_size;
to_bytes = reinterpret_cast<void*>(
(reinterpret_cast<uint8_t*>(local_to_bytes) + block_count *
block_size));
endpoint_->CopyFromRemote(remote_from, to_bytes, block_size);
@@ -1171,7 +1178,7 @@ class RPCClientSession : public RPCSession, public
DeviceAPI {
const uint64_t remainder_bytes = nbytes % block_size;
if (remainder_bytes != 0) {
- remote_from->byte_offset = block_count * block_size;
+ remote_from->byte_offset = initial_offset + block_count * block_size;
to_bytes = reinterpret_cast<void*>(
(reinterpret_cast<uint8_t*>(local_to_bytes) + block_count *
block_size));
endpoint_->CopyFromRemote(remote_from, to_bytes, remainder_bytes);
diff --git a/src/runtime/rpc/rpc_local_session.h
b/src/runtime/rpc/rpc_local_session.h
index 9035b486c9..93d7ae2565 100644
--- a/src/runtime/rpc/rpc_local_session.h
+++ b/src/runtime/rpc/rpc_local_session.h
@@ -42,6 +42,7 @@ namespace runtime {
*/
class LocalSession : public RPCSession {
public:
+ ~LocalSession() override { ReleaseStreams(); }
// function overrides
PackedFuncHandle GetFunction(const std::string& name) override;
diff --git a/src/runtime/rpc/rpc_module.cc b/src/runtime/rpc/rpc_module.cc
index bb84dc79c7..be1e45572f 100644
--- a/src/runtime/rpc/rpc_module.cc
+++ b/src/runtime/rpc/rpc_module.cc
@@ -38,53 +38,10 @@
#endif
#include "rpc_endpoint.h"
-#include "rpc_session.h"
+#include "rpc_remote.h"
namespace tvm {
namespace runtime {
-/*!
- * \brief Build a local Tensor with remote backing storage.
- * \param sess the RPCSession which owns the given handle.
- * \param handle A pointer valid on the remote end which should form the
`data` field of the
- * underlying DLTensor.
- * \param template_tensor An empty DLTensor whose shape and dtype fields are
used to fill the newly
- * created array. Needed because it's difficult to pass a shape vector as
a ffi::Function arg.
- * \param dev Remote device used with this tensor. Must have non-zero
RPCSessMask.
- * \param remote_tensor_handle The handle returned by RPC server to identify
the Tensor.
- */
-Tensor TensorFromRemoteOpaqueHandle(std::shared_ptr<RPCSession> sess, void*
handle,
- DLTensor* template_tensor, Device dev,
- void* remote_tensor_handle) {
- TVM_FFI_ICHECK_EQ(sess->table_index(), GetRPCSessionIndex(dev))
- << "The Device given does not belong to the given session";
- class RemoteSpaceAlloc {
- public:
- explicit RemoteSpaceAlloc(RemoteSpace space) : space_(space) {}
- void AllocData(DLTensor* tensor) {
- // the pointer to the remote space is passed in as the data pointer
- tensor->data = &(space_);
- }
- void FreeData(DLTensor* tensor) {
- if (space_.object_handle != nullptr) {
- try {
- space_.sess->FreeHandle(space_.object_handle);
- } catch (const ffi::Error& e) {
- // fault tolerance to remote close
- }
- }
- }
-
- private:
- RemoteSpace space_;
- };
- RemoteSpace space;
- space.sess = sess;
- space.data = handle;
- space.object_handle = remote_tensor_handle;
- ffi::Shape shape(template_tensor->shape, template_tensor->shape +
template_tensor->ndim);
- return Tensor::FromNDAlloc(RemoteSpaceAlloc(space), shape,
template_tensor->dtype, dev);
-}
-
/*!
* \brief A wrapped remote function as a ffi::Function.
*/
@@ -93,6 +50,7 @@ class RPCWrappedFunc : public ffi::Object {
RPCWrappedFunc(void* handle, std::shared_ptr<RPCSession> sess) :
handle_(handle), sess_(sess) {}
void operator()(ffi::PackedArgs args, ffi::Any* rv) const {
+ sess_->EnsureOpen();
std::vector<AnyView> packed_args(args.size());
std::vector<std::unique_ptr<DLTensor>> temp_dltensors;
@@ -112,31 +70,30 @@ class RPCWrappedFunc : public ffi::Object {
continue;
}
packed_args[i] = args[i];
- // run a remote translation to translate RPC related objects to
- // their remote counterparts.
+ if (auto tensor = args[i].as<RemoteTensor>()) {
+ TVM_FFI_CHECK((*tensor)->device->sess == sess_, ValueError)
+ << "Tensor belongs to a different remote session";
+ auto dptr = std::make_unique<DLTensor>((*tensor)->descriptor());
+ packed_args[i] = dptr.get();
+ temp_dltensors.emplace_back(std::move(dptr));
+ continue;
+ }
+ if (auto dev = args[i].as<RemoteDevice>()) {
+ TVM_FFI_CHECK((*dev)->sess == sess_, ValueError)
+ << "Device belongs to a different remote session";
+ packed_args[i] = (*dev)->endpoint_device;
+ continue;
+ }
+ if (auto obj = args[i].as<RPCObjectRef>()) {
+ TVM_FFI_CHECK((*obj)->sess() == sess_, ValueError)
+ << "Object belongs to a different remote session";
+ continue;
+ }
switch (args[i].type_index()) {
- case ffi::TypeIndex::kTVMFFITensor: {
- // Pass Tensor as DLTensor
- auto dptr =
std::make_unique<DLTensor>(*args[i].cast<Tensor>().operator->());
- dptr->device = RemoveSessMask(dptr->device);
- dptr->data = static_cast<RemoteSpace*>(dptr->data)->data;
- packed_args[i] = dptr.get();
- temp_dltensors.emplace_back(std::move(dptr));
- break;
- }
- case ffi::TypeIndex::kTVMFFIDLTensorPtr: {
- // translate to a remote view of DLTensor
- auto dptr = std::make_unique<DLTensor>(*args[i].cast<DLTensor*>());
- dptr->device = RemoveSessMask(dptr->device);
- dptr->data = static_cast<RemoteSpace*>(dptr->data)->data;
- packed_args[i] = dptr.get();
- temp_dltensors.emplace_back(std::move(dptr));
- break;
- }
- case ffi::TypeIndex::kTVMFFIDevice: {
- packed_args[i] = RemoveSessMask(args[i].cast<DLDevice>());
+ case ffi::TypeIndex::kTVMFFITensor:
+ case ffi::TypeIndex::kTVMFFIDLTensorPtr:
+ TVM_FFI_THROW(ValueError) << "Pass an RPC RemoteTensor, not a local
tensor";
break;
- }
case ffi::TypeIndex::kTVMFFIFunction:
case ffi::TypeIndex::kTVMFFIModule: {
packed_args[i] = UnwrapRemoteValueToHandle(args[i]);
@@ -146,6 +103,11 @@ class RPCWrappedFunc : public ffi::Object {
packed_args[i] = AnyView::CopyFromTVMFFIAny(temp);
break;
}
+ default:
+ TVM_FFI_CHECK(args[i].type_index() <
ffi::TypeIndex::kTVMFFIStaticObjectBegin ||
+ args[i].type_index() ==
ffi::TypeIndex::kTVMFFIBytes,
+ ValueError)
+ << "Cannot pass a local object to a remote function: " <<
args[i].GetTypeKey();
}
}
auto set_return = [this, rv](ffi::PackedArgs args) {
this->WrapRemoteReturnToValue(args, rv); };
@@ -170,14 +132,6 @@ class RPCWrappedFunc : public ffi::Object {
void* UnwrapRemoteValueToHandle(const ffi::AnyView& arg) const;
// wrap a remote return via Set
void WrapRemoteReturnToValue(ffi::PackedArgs args, ffi::Any* rv) const;
-
- // remove a remote session mask
- Device RemoveSessMask(Device dev) const {
- TVM_FFI_ICHECK(IsRPCSessionDevice(dev)) << "Can not pass in local device";
- TVM_FFI_ICHECK_EQ(GetRPCSessionIndex(dev), sess_->table_index())
- << "Can not pass in device with a different remote session";
- return RemoveRPCSessionMask(dev);
- }
};
// RPC that represents a remote module session.
@@ -203,8 +157,9 @@ class RPCModuleNode final : public ffi::ModuleObj {
ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) final {
if (name == "CloseRPCConnection") {
- return ffi::Function([this](ffi::PackedArgs, ffi::Any*) {
sess_->Shutdown(); });
+ return ffi::Function([sess = sess_](ffi::PackedArgs, ffi::Any*) {
sess->Close(); });
}
+ sess_->EnsureOpen();
if (module_handle_ == nullptr) {
return WrapRemoteFunc(sess_->GetFunction(name));
@@ -219,10 +174,6 @@ class RPCModuleNode final : public ffi::ModuleObj {
int cooldown_interval_ms, int
repeats_to_cooldown,
int cache_flush_bytes, const std::string&
f_preproc_name) {
InitRemoteFunc(&remote_get_time_evaluator_, "runtime.RPCTimeEvaluator");
- // Remove session mask because we pass dev by parts.
- TVM_FFI_CHECK_EQ(GetRPCSessionIndex(dev), sess_->table_index(), ValueError)
- << "Need to pass the matched remote device to
RPCModule.GetTimeEvaluator";
- dev = RemoveRPCSessionMask(dev);
if (module_handle_ != nullptr) {
return remote_get_time_evaluator_(
@@ -305,6 +256,8 @@ void
RPCWrappedFunc::WrapRemoteReturnToValue(ffi::PackedArgs args, ffi::Any* rv)
if (type_index == ffi::TypeIndex::kTVMFFINone) {
*rv = nullptr;
return;
+ } else if (type_index == ffi::TypeIndex::kTVMFFIDevice) {
+ *rv = MakeRemoteDevice(sess_, args[1].cast<DLDevice>());
} else if (type_index == ffi::TypeIndex::kTVMFFIFunction) {
TVM_FFI_ICHECK_EQ(args.size(), 2);
void* handle = args[1].cast<void*>();
@@ -321,9 +274,7 @@ void
RPCWrappedFunc::WrapRemoteReturnToValue(ffi::PackedArgs args, ffi::Any* rv)
TVM_FFI_ICHECK_EQ(args.size(), 3);
auto tensor = args[1].cast<DLTensor*>();
void* nd_handle = args[2].cast<void*>();
- *rv = TensorFromRemoteOpaqueHandle(sess_, tensor->data, tensor,
- AddRPCSessionMask(tensor->device,
sess_->table_index()),
- nd_handle);
+ *rv = RemoteTensorFromHandle(sess_, *tensor, nd_handle);
} else if (type_index == ffi::TypeIndex::kTVMFFIBytes ||
type_index == ffi::TypeIndex::kTVMFFIStr ||
type_index == ffi::TypeIndex::kTVMFFISmallStr ||
@@ -343,7 +294,6 @@ void
RPCWrappedFunc::WrapRemoteReturnToValue(ffi::PackedArgs args, ffi::Any* rv)
ffi::Module CreateRPCSessionModule(std::shared_ptr<RPCSession> sess) {
auto n = ffi::make_object<RPCModuleNode>(nullptr, sess);
- RPCSession::InsertToSessionTable(sess);
return ffi::Module(n);
}
@@ -486,18 +436,35 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_FFI_ICHECK_EQ(tkey, "rpc");
static_cast<RPCModuleNode*>(parent.operator->())->ImportModule(child);
})
- .def_packed("rpc.SessTableIndex",
- [](ffi::PackedArgs args, ffi::Any* rv) {
- ffi::Module m = args[0].cast<ffi::Module>();
- std::string tkey = m->kind();
- TVM_FFI_ICHECK_EQ(tkey, "rpc");
- *rv =
static_cast<RPCModuleNode*>(m.operator->())->sess()->table_index();
- })
- .def("tvm.rpc.TensorFromRemoteOpaqueHandle",
- [](ffi::Module mod, void* remote_array, DLTensor* template_tensor,
Device dev,
- void* tensor_handle) -> Tensor {
- return TensorFromRemoteOpaqueHandle(RPCModuleGetSession(mod),
remote_array,
- template_tensor, dev,
tensor_handle);
+ .def("rpc.ValidateDevice",
+ [](ffi::Module mod, RemoteDevice dev) {
+ auto sess = RPCModuleGetSession(mod);
+ sess->EnsureOpen();
+ TVM_FFI_CHECK(sess == dev->sess, ValueError)
+ << "Device belongs to a different remote session";
+ return dev->endpoint_device;
+ })
+ .def("runtime.ModuleTimeEvaluator",
+ [](ffi::Module mod, std::string name, ffi::AnyView device, int
number, int repeat,
+ int min_repeat_ms, int limit_zero_time_iterations, int
cooldown_interval_ms,
+ int repeats_to_cooldown, int cache_flush_bytes, std::string
f_preproc_name) {
+ if (std::string(mod->kind()) != "rpc") {
+ auto dev = device.cast<DLDevice>();
+ return
ffi::Function::GetGlobalRequired("runtime.RPCTimeEvaluator")(
+ mod, name, static_cast<int>(dev.device_type),
dev.device_id, number,
+ repeat, min_repeat_ms, limit_zero_time_iterations,
cooldown_interval_ms,
+ repeats_to_cooldown, cache_flush_bytes,
f_preproc_name)
+ .cast<ffi::Function>();
+ }
+ auto dev = device.cast<RemoteDevice>();
+ auto sess = RPCModuleGetSession(mod);
+ sess->EnsureOpen();
+ TVM_FFI_CHECK(sess == dev->sess, ValueError)
+ << "Device belongs to a different remote session";
+ return static_cast<RPCModuleNode*>(mod.operator->())
+ ->GetTimeEvaluator(name, dev->endpoint_device, number,
repeat, min_repeat_ms,
+ limit_zero_time_iterations,
cooldown_interval_ms,
+ repeats_to_cooldown, cache_flush_bytes,
f_preproc_name);
});
}
diff --git a/src/runtime/rpc/rpc_remote.h b/src/runtime/rpc/rpc_remote.h
new file mode 100644
index 0000000000..a94cb8da56
--- /dev/null
+++ b/src/runtime/rpc/rpc_remote.h
@@ -0,0 +1,186 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*! \file rpc_remote.h
+ * \brief Session-owned RPC values. Remote addresses never enter local tensor
storage.
+ */
+#ifndef TVM_RUNTIME_RPC_RPC_REMOTE_H_
+#define TVM_RUNTIME_RPC_RPC_REMOTE_H_
+
+#include <tvm/ffi/container/shape.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include <memory>
+#include <utility>
+
+#include "rpc_session.h"
+
+namespace tvm {
+namespace runtime {
+
+/*! \brief Session-owned device descriptor for an RPC execution endpoint.
+ *
+ * The endpoint descriptor never carries session identity. Only this object
+ * binds it to a session, and all device operations must use that session.
+ */
+class RemoteDeviceObj : public ffi::Object {
+ public:
+ /*! \brief Keeps the transport alive; explicit session closure still
invalidates operations. */
+ std::shared_ptr<RPCSession> sess;
+ /*! \brief Ordinary device type/index in the remote endpoint, not the client
process. */
+ Device endpoint_device;
+ DeviceAPI* api() const {
+ sess->EnsureOpen();
+ return sess->GetDeviceAPI(endpoint_device);
+ }
+ static constexpr bool _type_mutable = true;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("rpc.RemoteDevice", RemoteDeviceObj,
ffi::Object);
+};
+/*! \brief Owning reference to a session-bound remote device. */
+class RemoteDevice : public ffi::ObjectRef {
+ public:
+ explicit RemoteDevice(ffi::ObjectPtr<RemoteDeviceObj> data) :
ffi::ObjectRef(std::move(data)) {}
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RemoteDevice, ffi::ObjectRef,
RemoteDeviceObj);
+};
+inline RemoteDevice MakeRemoteDevice(std::shared_ptr<RPCSession> sess, Device
dev) {
+ sess->EnsureOpen();
+ auto n = ffi::make_object<RemoteDeviceObj>();
+ n->sess = std::move(sess);
+ n->endpoint_device = dev;
+ return RemoteDevice(n);
+}
+
+/*! \brief Shared ownership of a remote allocation or returned tensor object.
*/
+struct RemoteTensorStorage {
+ /*! \brief Session that allocated the storage or returned the owning tensor
handle. */
+ std::shared_ptr<RPCSession> sess;
+ /*! \brief Endpoint device needed to release a raw allocation through its
DeviceAPI. */
+ Device device;
+ /*! \brief Opaque endpoint data address, usable only in temporary RPC wire
descriptors. */
+ void* data{nullptr};
+ /*! \brief Non-null returned tensor handle; FreeHandle releases its remote
owner. */
+ void* object_handle{nullptr};
+ /*! \brief True for raw allocations released with FreeDataSpace, not
returned objects. */
+ bool owns_allocation{false};
+ ~RemoteTensorStorage() {
+ try {
+ if (object_handle != nullptr) {
+ sess->FreeHandle(object_handle);
+ } else if (owns_allocation && data != nullptr) {
+ sess->GetDeviceAPI(device)->FreeDataSpace(device, data);
+ }
+ } catch (const ffi::Error&) {
+ // A closed transport has already released its remote resources.
+ }
+ }
+};
+
+/*! \brief Logical contiguous tensor metadata over shared remote storage.
+ *
+ * Views copy metadata and retain the same storage owner. Neither this object
+ * nor its descriptor may be exported as a locally usable ffi::Tensor/DLPack.
+ */
+class RemoteTensorObj : public ffi::Object {
+ public:
+ /*! \brief Session-bound endpoint device, always from the storage owner's
session. */
+ RemoteDevice device{ffi::UnsafeInit{}};
+ /*! \brief Independently owned logical shape, never borrowed from a remote
return frame. */
+ ffi::Shape shape;
+ /*! \brief Element type for shape/byte-size validation and wire descriptors.
*/
+ DLDataType dtype;
+ /*! \brief Byte displacement from storage->data; bounded views accumulate
this offset. */
+ uint64_t byte_offset{0};
+ /*! \brief Shared allocation/object owner; the last view releases the remote
resource. */
+ std::shared_ptr<RemoteTensorStorage> storage;
+ /*! \brief Borrowed descriptor valid while this tensor remains alive during
an RPC call.
+ * Its data is an endpoint address and shape is local metadata; neither may
+ * escape into a locally usable tensor or a DLPack export.
+ */
+ DLTensor descriptor() const {
+ device->sess->EnsureOpen();
+ return {storage->data,
+ device->endpoint_device,
+ static_cast<int>(shape.size()),
+ dtype,
+ const_cast<int64_t*>(shape.data()),
+ nullptr,
+ byte_offset};
+ }
+ static constexpr bool _type_mutable = true;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("rpc.RemoteTensor", RemoteTensorObj,
ffi::Object);
+};
+/*! \brief Owning reference to remote metadata and its shared storage
lifetime. */
+class RemoteTensor : public ffi::ObjectRef {
+ public:
+ explicit RemoteTensor(ffi::ObjectPtr<RemoteTensorObj> data) :
ffi::ObjectRef(std::move(data)) {}
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RemoteTensor, ffi::ObjectRef,
RemoteTensorObj);
+};
+
+inline RemoteTensor RemoteTensorFromHandle(std::shared_ptr<RPCSession> sess,
const DLTensor& tensor,
+ void* object_handle) {
+ auto n = ffi::make_object<RemoteTensorObj>();
+ n->device = MakeRemoteDevice(sess, tensor.device);
+ n->shape = ffi::Shape(tensor.shape, tensor.shape + tensor.ndim);
+ n->dtype = tensor.dtype;
+ n->byte_offset = tensor.byte_offset;
+ n->storage = std::make_shared<RemoteTensorStorage>();
+ n->storage->sess = std::move(sess);
+ n->storage->device = tensor.device;
+ n->storage->data = tensor.data;
+ n->storage->object_handle = object_handle;
+ TVM_FFI_CHECK(ffi::IsContiguous(tensor), ValueError)
+ << "RPC requires a contiguous returned tensor";
+ return RemoteTensor(n);
+}
+
+/*! \brief Session-checked stream wrapper; aliases share one release state. */
+class RemoteStreamObj : public ffi::Object {
+ public:
+ /*! \brief Owning session and exact endpoint device accepted by stream
operations. */
+ RemoteDevice device{ffi::UnsafeInit{}};
+ /*! \brief Shared handle/closure state also retained by the session until
release. */
+ std::shared_ptr<RPCStreamState> state;
+ static constexpr bool _type_mutable = true;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("rpc.RemoteStream", RemoteStreamObj,
ffi::Object);
+};
+/*! \brief Owning reference to a session-bound stream alias. */
+class RemoteStream : public ffi::ObjectRef {
+ public:
+ explicit RemoteStream(ffi::ObjectPtr<RemoteStreamObj> data) :
ffi::ObjectRef(std::move(data)) {}
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RemoteStream, ffi::ObjectRef,
RemoteStreamObj);
+};
+
+inline TVMStreamHandle CheckRemoteStream(const RemoteDevice& dev,
+ const ffi::Optional<RemoteStream>&
stream) {
+ dev->sess->EnsureOpen();
+ if (!stream.has_value()) return nullptr;
+ const auto& s = stream.value();
+ TVM_FFI_CHECK(s->device->sess == dev->sess &&
+ s->device->endpoint_device.device_type ==
dev->endpoint_device.device_type &&
+ s->device->endpoint_device.device_id ==
dev->endpoint_device.device_id,
+ ValueError)
+ << "Stream belongs to a different remote device or session";
+ TVM_FFI_CHECK(!s->state->closed, ValueError) << "Remote stream is closed";
+ return s->state->handle;
+}
+
+} // namespace runtime
+} // namespace tvm
+#endif // TVM_RUNTIME_RPC_RPC_REMOTE_H_
diff --git a/src/runtime/rpc/rpc_session.cc b/src/runtime/rpc/rpc_session.cc
index 349568062e..c81930a8c6 100644
--- a/src/runtime/rpc/rpc_session.cc
+++ b/src/runtime/rpc/rpc_session.cc
@@ -27,9 +27,6 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
-#include <array>
-#include <mutex>
-
namespace tvm {
namespace runtime {
@@ -86,48 +83,6 @@ void RPCSession::AsyncStreamWait(Device dev, TVMStreamHandle
stream,
}
}
-class RPCSessTable {
- public:
- static constexpr int kMaxRPCSession = 32;
- // Get global singleton
- static RPCSessTable* Global() {
- static RPCSessTable inst;
- return &inst;
- }
- // Get session from table
- std::shared_ptr<RPCSession> Get(int index) {
- TVM_FFI_ICHECK(index >= 0 && index < kMaxRPCSession);
- return tbl_[index].lock();
- }
- // Insert session into table.
- int Insert(std::shared_ptr<RPCSession> ptr) {
- std::lock_guard<std::mutex> lock(mutex_);
- for (int i = 0; i < kMaxRPCSession; ++i) {
- if (tbl_[i].lock() == nullptr) {
- tbl_[i] = ptr;
- return i;
- }
- }
- TVM_FFI_THROW(InternalError) << "maximum number of RPC session reached";
- }
-
- private:
- // The mutex
- std::mutex mutex_;
- // Use weak_ptr intentionally
- // If the RPCSession get released, the pointer session will be released
- std::array<std::weak_ptr<RPCSession>, kMaxRPCSession> tbl_;
-};
-
-std::shared_ptr<RPCSession> RPCSession::Get(int table_index) {
- return RPCSessTable::Global()->Get(table_index);
-}
-
-void RPCSession::InsertToSessionTable(std::shared_ptr<RPCSession> sess) {
- TVM_FFI_ICHECK_EQ(sess->table_index_, 0);
- sess->table_index_ = RPCSessTable::Global()->Insert(sess);
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
tvm::ffi::reflection::ObjectDef<RPCObjectRefObj>(); }
} // namespace runtime
diff --git a/src/runtime/rpc/rpc_session.h b/src/runtime/rpc/rpc_session.h
index 1276d8e267..41a3ab9ef4 100644
--- a/src/runtime/rpc/rpc_session.h
+++ b/src/runtime/rpc/rpc_session.h
@@ -28,9 +28,13 @@
#include <tvm/ffi/function.h>
#include <tvm/runtime/device_api.h>
+#include <algorithm>
+#include <atomic>
#include <functional>
#include <memory>
+#include <mutex>
#include <string>
+#include <vector>
#include "minrpc/rpc_reference.h"
@@ -47,6 +51,21 @@ namespace runtime {
*/
constexpr int32_t kRuntimeRPCObjectRef =
TVMFFITypeIndex::kTVMFFIDynObjectBegin - 13;
+/*! \brief Stream ownership and alias state, retained until release or session
teardown. */
+struct RPCStreamState {
+ /*! \brief Ordinary endpoint device that created or selected this stream. */
+ Device device;
+ /*! \brief Opaque handle interpreted only by that endpoint's DeviceAPI. */
+ TVMStreamHandle handle;
+ /*! \brief Created by this session and freed on release; borrowed current
streams are not freed.
+ */
+ bool owned;
+ /*! \brief Shared invalidation flag observed by every wrapper alias before
handle use. */
+ std::atomic<bool> closed{false};
+ RPCStreamState(Device device, TVMStreamHandle handle, bool owned)
+ : device(device), handle(handle), owned(owned) {}
+};
+
/*!
* \brief The interface of all remote RPC sessions.
*
@@ -246,21 +265,54 @@ class RPCSession {
*/
virtual void AsyncStreamWait(Device dev, TVMStreamHandle stream,
FAsyncCallback on_compelte);
- /*!
- * \return The session table index of the session.
- */
- int table_index() const { return table_index_; }
-
- /*!
- * \brief Try get session from the global session table by table index.
- * \param table_index The table index of the session.
- * \return The shared_ptr to the session, can be nullptr.
- */
- static std::shared_ptr<RPCSession> Get(int table_index);
-
- /*!
- * \brief Shutdown RPC connection.
- */
+ /*! \brief Reject operations after explicit session closure. */
+ void EnsureOpen() const { TVM_FFI_CHECK(!closed_, RuntimeError) << "RPC
session is closed"; }
+ void Close() {
+ if (!closed_.exchange(true)) {
+ ReleaseStreams();
+ Shutdown();
+ }
+ }
+ std::shared_ptr<RPCStreamState> TrackStream(Device dev, TVMStreamHandle
handle, bool owned) {
+ std::lock_guard<std::mutex> lock(stream_mutex_);
+ EnsureOpen();
+ for (const auto& stream : streams_) {
+ if (!stream->closed && stream->handle == handle &&
+ stream->device.device_type == dev.device_type &&
+ stream->device.device_id == dev.device_id) {
+ stream->owned |= owned;
+ return stream;
+ }
+ }
+ auto stream = std::make_shared<RPCStreamState>(dev, handle, owned);
+ streams_.push_back(stream);
+ return stream;
+ }
+ void FreeStream(const std::shared_ptr<RPCStreamState>& stream) {
+ std::lock_guard<std::mutex> lock(stream_mutex_);
+ if (stream->closed) return;
+ stream->closed = true;
+ streams_.erase(std::remove(streams_.begin(), streams_.end(), stream),
streams_.end());
+ auto* api = GetDeviceAPI(stream->device);
+ if (api->GetCurrentStream(stream->device) == stream->handle) {
+ api->SetStream(stream->device, nullptr);
+ }
+ if (stream->owned) api->FreeStream(stream->device, stream->handle);
+ }
+ // Derived destructors call this while their DeviceAPI/transport is still
alive.
+ void ReleaseStreams() {
+ std::vector<std::shared_ptr<RPCStreamState>> streams;
+ {
+ std::lock_guard<std::mutex> lock(stream_mutex_);
+ streams.swap(streams_);
+ }
+ for (const auto& stream : streams) {
+ try {
+ FreeStream(stream);
+ } catch (const ffi::Error&) {
+ }
+ }
+ }
virtual void Shutdown() {}
protected:
@@ -271,33 +323,9 @@ class RPCSession {
void SendException(FAsyncCallback callback, const char* msg);
private:
- /*! \brief index of this session in RPC session table */
- int table_index_{0};
- /*! \brief Insert the current session to the session table.*/
- static void InsertToSessionTable(std::shared_ptr<RPCSession> sess);
- // friend declaration
- friend ffi::Module CreateRPCSessionModule(std::shared_ptr<RPCSession> sess);
-};
-
-/*!
- * \brief Remote space handle cell used by the RPC runtime API.
- *
- * When we allocate space using a rpc device, the data pointer
- * points to an allocated RemoteSpace.
- */
-struct RemoteSpace {
- /*! \brief The remote data handle. */
- void* data;
- /*! \brief Reference to the underlying RPC session. */
- std::shared_ptr<RPCSession> sess;
- /*!
- * \brief The remote Tensor object handle, if this RemoteSpace wraps a
returned Tensor.
- *
- * Returned RPC Tensors carry both the backing data pointer and a Tensor
object handle. The
- * object handle must be released with FreeHandle so the remote side can
correctly decrement the
- * Tensor refcount and free the backing storage when it is no longer shared.
- */
- void* object_handle{nullptr};
+ std::atomic<bool> closed_{false};
+ std::vector<std::shared_ptr<RPCStreamState>> streams_;
+ std::mutex stream_mutex_;
};
/*!
diff --git a/tests/nightly/python/test_nnapi/infrastructure.py
b/tests/nightly/python/test_nnapi/infrastructure.py
index 803925d5b4..58d4f52460 100644
--- a/tests/nightly/python/test_nnapi/infrastructure.py
+++ b/tests/nightly/python/test_nnapi/infrastructure.py
@@ -100,9 +100,9 @@ def _run(remote, tracker, ex, inputs):
try:
# Execute the model on the remote.
remote_ex = remote.load_module(so_name)
- vm = tvm.relax.VirtualMachine(remote_ex, device=dev)
+ vm = tvm.rpc.virtual_machine(remote_ex, dev)
- inputs = [x.copyto(dev) for x in inputs]
+ inputs = [dev.tensor(x) for x in inputs]
vm.set_input("main", *inputs)
vm.invoke_stateful("main")
diff --git a/tests/python/codegen/test_target_codegen_cross_llvm.py
b/tests/python/codegen/test_target_codegen_cross_llvm.py
index f58b6aac34..0b64b24832 100644
--- a/tests/python/codegen/test_target_codegen_cross_llvm.py
+++ b/tests/python/codegen/test_target_codegen_cross_llvm.py
@@ -85,9 +85,9 @@ def test_llvm_add_pipeline():
farm = remote.load_module("myadd.o")
dev = remote.cpu(0)
n = nn
- a =
tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
- b =
tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
- c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
+ a = dev.tensor(np.random.uniform(size=n).astype("float32"))
+ b = dev.tensor(np.random.uniform(size=n).astype("float32"))
+ c = dev.tensor(np.zeros(n, dtype="float32"))
farm(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
print("Verification finish on remote..")
diff --git a/tests/python/contrib/test_coreml_runtime.py
b/tests/python/contrib/test_coreml_runtime.py
index 514286100c..eb3866a230 100644
--- a/tests/python/contrib/test_coreml_runtime.py
+++ b/tests/python/contrib/test_coreml_runtime.py
@@ -57,7 +57,7 @@ def test_coreml_runtime():
)
return coremltools.models.MLModel(builder.spec)
- def verify(coreml_model, model_path, dev):
+ def verify(coreml_model, model_path, dev, session=None):
coreml_model = create_coreml_model()
out_spec = coreml_model.output_description._fd_spec
@@ -73,9 +73,11 @@ def test_coreml_runtime():
coreml_outputs = [coreml_model.predict(inputs)[name] for name in
out_names]
# inference via tvm coreml runtime
- runtime = coreml_runtime.create("main", model_path, dev)
+ runtime = coreml_runtime.create("main", model_path, dev,
session=session)
for name in inputs:
- runtime.set_input(name, tvm.runtime.tensor(inputs[name], dev))
+ runtime.set_input(
+ name, dev.tensor(inputs[name]) if session else
tvm.runtime.tensor(inputs[name], dev)
+ )
runtime.invoke()
tvm_outputs = [runtime.get_output(i).numpy() for i in
range(runtime.get_num_outputs())]
@@ -91,7 +93,7 @@ def test_coreml_runtime():
compiled_model = os.path.basename(compiled_model)
remote = rpc.connect(proxy_host, proxy_port, key=key)
dev = remote.cpu(0)
- verify(coreml_model, compiled_model, dev)
+ verify(coreml_model, compiled_model, dev, remote)
def check_local(coreml_model):
temp = utils.tempdir()
diff --git a/tests/python/contrib/test_random.py
b/tests/python/contrib/test_random.py
index a8c278c5d2..94e6913eae 100644
--- a/tests/python/contrib/test_random.py
+++ b/tests/python/contrib/test_random.py
@@ -130,7 +130,7 @@ def test_random_fill():
def check_remote(server):
remote = rpc.connect(server.host, server.port)
- value = tvm.runtime.empty((512, 512), dtype, remote.cpu())
+ value = remote.cpu().empty((512, 512), dtype)
random_fill = remote.get_function("tvm.contrib.random.random_fill")
random_fill(value)
diff --git a/tests/python/relax/backend/adreno/utils.py
b/tests/python/relax/backend/adreno/utils.py
index 1766409ec8..64fed4089a 100644
--- a/tests/python/relax/backend/adreno/utils.py
+++ b/tests/python/relax/backend/adreno/utils.py
@@ -180,8 +180,12 @@ def build_and_run(mod, inputs, tgt):
device_arr = [dev for _ in
range(len(mod.global_infos["vdevice"]))]
else:
device_arr = [dev]
- vm = relax.VirtualMachine(rexec, device_arr)
- device_inputs = [tvm.runtime.tensor(ip, dev) for ip in inputs]
+ if SessionManager.is_target_rpc():
+ vm = tvm.rpc.virtual_machine(rexec, device_arr)
+ device_inputs = [dev.tensor(ip) for ip in inputs]
+ else:
+ vm = relax.VirtualMachine(rexec, device_arr)
+ device_inputs = [tvm.runtime.tensor(ip, dev) for ip in inputs]
vm.set_input("main", *device_inputs)
vm.invoke_stateful("main")
tvm_output = vm.get_outputs("main")
diff --git a/tests/python/relax/test_vm_build.py
b/tests/python/relax/test_vm_build.py
index 19469fa331..9a4359c70b 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -1027,9 +1027,15 @@ def test_multi_systemlib():
worker.send(popen_check)
+def _trial_tensor(data, device):
+ if isinstance(device, rpc.RemoteDevice):
+ return device.tensor(data)
+ return tvm.runtime.tensor(data, device)
+
+
def set_input_trial(vm: relax.VirtualMachine, device: tvm.runtime.Device) ->
None:
- a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
- b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
+ a = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
+ b = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
vm.set_input("main", a, b)
vm.invoke_stateful("main")
res0 = vm.get_outputs("main")
@@ -1043,7 +1049,7 @@ def set_input_trial(vm: relax.VirtualMachine, device:
tvm.runtime.Device) -> Non
# bug! If you don't bind the Tensor to a var, the memory will get
corrupted.
# Possibly due to object lifecycles and other FFI issues
- a = tvm.runtime.tensor(np.array(2).astype("int32"), device)
+ a = _trial_tensor(np.array(2).astype("int32"), device)
vm.set_input("test_vm_tuple", a)
vm.invoke_stateful("test_vm_tuple")
res2 = vm.get_outputs("test_vm_tuple")
@@ -1051,7 +1057,7 @@ def set_input_trial(vm: relax.VirtualMachine, device:
tvm.runtime.Device) -> Non
# so we have to get the scalar out of the Tensor
assert tuple(map(lambda a: int(a.numpy()), res2)) == (2, 2)
- b = tvm.runtime.tensor(np.array(1).astype("int32"), device)
+ b = _trial_tensor(np.array(1).astype("int32"), device)
vm.set_input("test_vm_nested_tuple", b)
vm.invoke_stateful("test_vm_nested_tuple")
res3 = vm.get_outputs("test_vm_nested_tuple")
@@ -1062,8 +1068,8 @@ def set_input_trial(vm: relax.VirtualMachine, device:
tvm.runtime.Device) -> Non
def set_input_attempt_stateless(vm: relax.VirtualMachine, device:
tvm.runtime.Device) -> None:
# this should fail: once you set inputs, you cannot run statelessly
- a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
- b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
+ a = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
+ b = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
vm.set_input("main", a, b)
# must use invoke stateful!
vm["main"]()
@@ -1076,8 +1082,8 @@ def set_input_attempt_invoke(vm: relax.VirtualMachine,
device: tvm.runtime.Devic
def set_input_attempt_get(vm: relax.VirtualMachine, device:
tvm.runtime.Device) -> None:
# this should fail: you can't get outputs without invoking the function
first
- a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
- b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
+ a = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
+ b = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
vm.set_input("main", a, b)
_ = vm.get_outputs("main")
@@ -1120,7 +1126,7 @@ def run_on_rpc(
device = remote.cpu()
# Build a VM out of the executable and context.
- vm = relax.VirtualMachine(rexec, device=device)
+ vm = rpc.virtual_machine(rexec, device)
trial_func(vm, device)
check_remote(rpc.Server("127.0.0.1"))
@@ -1150,8 +1156,8 @@ def test_set_input_tuple():
def save_function_kwargs_trial(vm: relax.VirtualMachine, device:
tvm.runtime.Device) -> None:
# just checking that we can use kwargs for the args when saving a function
- a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
- b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
+ a = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
+ b = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
vm.save_function("main", "saved_main", x=a, w=b)
res0 = vm["saved_main"]()
tvm.testing.assert_allclose(res0.numpy(), a.numpy() * b.numpy(),
rtol=1e-7, atol=1e-7)
@@ -1171,8 +1177,8 @@ def save_function_time_evaluator_trial(
vm: relax.VirtualMachine, device: tvm.runtime.Device
) -> None:
# just checking that the saved function can be called in the time evaluator
- a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
- b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
+ a = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
+ b = _trial_tensor(np.random.rand(32, 32).astype("float32"), device)
vm.save_function("main", "saved_main", a, b)
vm.time_evaluator("saved_main", device)()
diff --git a/tests/python/relax/texture/test_texture_nd.py
b/tests/python/relax/texture/test_texture_nd.py
index dfe6233fbe..a627c0cb1f 100644
--- a/tests/python/relax/texture/test_texture_nd.py
+++ b/tests/python/relax/texture/test_texture_nd.py
@@ -184,8 +184,12 @@ def test_texture_copy(backend, dtype, channel_size,
read_width):
device_arr = [dev for _ in range(len(mod.global_infos["vdevice"]))]
else:
device_arr = [dev]
- vm = relax.VirtualMachine(rexec, device_arr)
- inps = [tvm.runtime.tensor(inp, dev) for inp in inputs]
+ if remote_session is not None:
+ vm = tvm.rpc.virtual_machine(rexec, device_arr)
+ inps = [dev.tensor(inp) for inp in inputs]
+ else:
+ vm = relax.VirtualMachine(rexec, device_arr)
+ inps = [tvm.runtime.tensor(inp, dev) for inp in inputs]
vm["main"](*inps)
np.testing.assert_equal(inps[-1].numpy(), inps[0].numpy())
diff --git a/tests/python/runtime/test_runtime_rpc.py
b/tests/python/runtime/test_runtime_rpc.py
index b42615fb9b..965af262ac 100644
--- a/tests/python/runtime/test_runtime_rpc.py
+++ b/tests/python/runtime/test_runtime_rpc.py
@@ -67,7 +67,9 @@ pytestmark = pytest.mark.skipif(
# to ensure all the remote resources destructs before the server terminates
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_bigendian_rpc():
"""Test big endian rpc when there is a PowerPC RPC server available"""
host = os.environ.get("TVM_POWERPC_TEST_HOST", None)
@@ -81,8 +83,8 @@ def test_bigendian_rpc():
f = tvm.compile(te.create_function([A, B]), target=target)
dev = remote.cpu(0)
- a = tvm.runtime.tensor(np.random.randint(0, 256,
size=shape).astype(A.dtype), device=dev)
- b = tvm.runtime.tensor(np.zeros(shape).astype(A.dtype), device=dev)
+ a = dev.tensor(np.random.randint(0, 256,
size=shape).astype(str(A.dtype)))
+ b = dev.tensor(np.zeros(shape).astype(str(A.dtype)))
temp = utils.tempdir()
path_dso = temp.relpath("dev_lib.o")
f.write_to_file(path_dso)
@@ -98,7 +100,9 @@ def test_bigendian_rpc():
verify_rpc(remote, target, (10,), dtype)
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_simple():
server = rpc.Server(key="x1")
client = rpc.connect("127.0.0.1", server.port, key="x1")
@@ -117,7 +121,9 @@ def test_rpc_simple():
check_remote()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_runtime_string():
server = rpc.Server(key="x1")
client = rpc.connect("127.0.0.1", server.port, key="x1")
@@ -131,15 +137,17 @@ def test_rpc_runtime_string():
check_remote()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_array():
server = rpc.Server()
remote = rpc.connect("127.0.0.1", server.port)
def check_remote():
x = np.ones((3, 4))
- r_cpu = tvm.runtime.tensor(x, remote.cpu(0))
- assert str(r_cpu.device).startswith("remote")
+ r_cpu = remote.cpu(0).tensor(x)
+ assert isinstance(r_cpu.device, rpc.RemoteDevice)
np.testing.assert_equal(r_cpu.numpy(), x)
fremote = remote.get_function("rpc.test.remote_tensor_func")
fremote(r_cpu)
@@ -147,7 +155,9 @@ def test_rpc_array():
check_remote()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_large_array():
# testcase of large array creation
server = rpc.Server()
@@ -157,8 +167,8 @@ def test_rpc_large_array():
dev = remote.cpu(0)
a_np = np.ones((5041, 720)).astype("float32")
b_np = np.ones((720, 192)).astype("float32")
- a = tvm.runtime.tensor(a_np, dev)
- b = tvm.runtime.tensor(b_np, dev)
+ a = dev.tensor(a_np)
+ b = dev.tensor(b_np)
np.testing.assert_equal(a.numpy(), a_np)
np.testing.assert_equal(b.numpy(), b_np)
@@ -166,7 +176,9 @@ def test_rpc_large_array():
@tvm.testing.skip_if_32bit(reason="skipping test for i386.")
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_echo():
def check(remote, local_session):
fecho = remote.get_function("testing.echo")
@@ -215,7 +227,9 @@ def test_rpc_echo():
# check_minrpc()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_file_exchange():
server = rpc.Server()
remote = rpc.connect("127.0.0.1", server.port)
@@ -229,7 +243,9 @@ def test_rpc_file_exchange():
check_remote()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_rpc_remote_module():
# graph
@@ -256,9 +272,9 @@ def test_rpc_remote_module():
f.export_library(path_dso)
remote.upload(path_dso)
f1 = remote.load_module("dev_lib.so")
- a = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype),
dev)
- b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), dev)
- time_f = f1.time_evaluator(f1.entry_name, remote.cpu(0), number=10)
+ a = dev.tensor(np.random.uniform(size=102).astype(str(A.dtype)))
+ b = dev.tensor(np.zeros(102, dtype=str(A.dtype)))
+ time_f = rpc.time_evaluator(f1, f1.entry_name, remote.cpu(0),
number=10)
cost = time_f(a, b).mean
print(f"{cost:g} secs/op")
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
@@ -271,8 +287,8 @@ def test_rpc_remote_module():
with open(local_download_path, "wb") as fo:
fo.write(remote.download_linked_module("dev_lib.tar"))
fupdated = tvm.runtime.load_module(local_download_path)
- a = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype),
tvm.cpu(0))
- b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), tvm.cpu(0))
+ a =
tvm.runtime.tensor(np.random.uniform(size=102).astype(str(A.dtype)), tvm.cpu(0))
+ b = tvm.runtime.tensor(np.zeros(102, dtype=str(A.dtype)), tvm.cpu(0))
fupdated(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
@@ -294,9 +310,9 @@ def test_rpc_remote_module():
dev = remote.cpu(0)
f1 = remote.system_lib()
- a = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype),
dev)
- b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), dev)
- time_f = f1.time_evaluator("myadd", remote.cpu(0), number=1)
+ a = dev.tensor(np.random.uniform(size=102).astype(str(A.dtype)))
+ b = dev.tensor(np.zeros(102, dtype=str(A.dtype)))
+ time_f = rpc.time_evaluator(f1, "myadd", remote.cpu(0), number=1)
cost = time_f(a, b).mean
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
@@ -330,8 +346,8 @@ def test_rpc_remote_module():
f.export_library(path_tar)
remote.upload(path_tar)
fhost = remote.load_module("myadd.tar")
- a = tvm.runtime.tensor(np.random.uniform(size=102).astype(A.dtype),
dev)
- b = tvm.runtime.tensor(np.zeros(102, dtype=A.dtype), dev)
+ a = dev.tensor(np.random.uniform(size=102).astype(str(A.dtype)))
+ b = dev.tensor(np.zeros(102, dtype=str(A.dtype)))
fhost(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
@@ -340,7 +356,9 @@ def test_rpc_remote_module():
check_minrpc()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_return_func():
server = rpc.Server(key="x1")
client = rpc.connect("127.0.0.1", server.port, key="x1")
@@ -353,7 +371,9 @@ def test_rpc_return_func():
check_remote()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_session_constructor_args():
# start server
server0 = rpc.Server(key="x0")
@@ -374,7 +394,7 @@ def test_rpc_session_constructor_args():
assert fecho("xyz") == "xyz"
assert bytes(fecho(bytearray(b"123"))) == b"123"
- nd = tvm.runtime.tensor([1, 2, 3], device=client.cpu(0))
+ nd = client.cpu(0).tensor([1, 2, 3])
assert nd.numpy()[1] == 2
def check_error_handling():
@@ -390,7 +410,9 @@ def test_rpc_session_constructor_args():
check_error_handling()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_return_tensor():
def run_arr_test():
server = rpc.Server(key="x1")
@@ -411,7 +433,9 @@ def test_rpc_return_tensor():
run_arr_test()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_rpc_return_remote_object():
def check(client, is_local):
make_shape = client.get_function("ffi.Shape")
@@ -457,7 +481,9 @@ def test_rpc_return_remote_object():
check_minrpc()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
def test_local_func():
client = rpc.LocalSession()
@@ -474,7 +500,9 @@ def test_local_func():
check_remote()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
@pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"])
def test_rpc_tracker_register(device_key):
# test registration
@@ -547,7 +575,9 @@ def _target(host, port, device_key, timeout):
remote.cpu()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
@pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"])
def test_rpc_tracker_request(device_key):
# test concurrent request
@@ -588,7 +618,9 @@ def test_rpc_tracker_request(device_key):
tracker.terminate()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
@pytest.mark.parametrize("device_key", ["test_device", "127.0.0.1:5555"])
def test_rpc_tracker_via_proxy(device_key):
"""
@@ -630,7 +662,9 @@ def test_rpc_tracker_via_proxy(device_key):
tracker_server.terminate()
[email protected](not env.build_flag_enabled("USE_RPC"), reason="need rpc")
[email protected](
+ tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None,
reason="need rpc"
+)
@pytest.mark.parametrize("with_proxy", (True, False))
def test_rpc_session_timeout_error(with_proxy):
port = 9000
diff --git a/web/tests/python/relax_rpc_test.py
b/web/tests/python/relax_rpc_test.py
index 579ed014ce..13f0a4cbdb 100644
--- a/web/tests/python/relax_rpc_test.py
+++ b/web/tests/python/relax_rpc_test.py
@@ -75,11 +75,11 @@ def test_rpc():
def check(remote):
dev = remote.webgpu(0)
# invoke the function
- vm = relax.VirtualMachine(remote.system_lib(), device=dev)
+ vm = rpc.virtual_machine(remote.system_lib(), dev)
adata = np.random.uniform(size=n).astype(dtype)
bdata = np.random.uniform(size=n).astype(dtype)
- a = tvm.runtime.tensor(adata, dev)
- b = tvm.runtime.tensor(bdata, dev)
+ a = dev.tensor(adata)
+ b = dev.tensor(bdata)
vm.set_input("main", a, b)
vm.invoke_stateful("main")
c = vm.get_outputs("main")
diff --git a/web/tests/python/webgpu_rpc_test.py
b/web/tests/python/webgpu_rpc_test.py
index 999c4db50c..a216ac1e8f 100644
--- a/web/tests/python/webgpu_rpc_test.py
+++ b/web/tests/python/webgpu_rpc_test.py
@@ -67,8 +67,8 @@ def test_rpc():
# basic function checks.
dev = remote.webgpu(0)
adata = np.random.uniform(size=size).astype(A.dtype)
- a = tvm.runtime.tensor(adata, dev)
- b = tvm.runtime.tensor(np.zeros(size, dtype=A.dtype), dev)
+ a = dev.tensor(adata)
+ b = dev.tensor(np.zeros(size, dtype=A.dtype))
np.testing.assert_equal(a.numpy(), adata)
f1 = remote.system_lib()
