FFI 设备#
import pytest
import pickle
from tvm.ffi import Device
from tvm import ffi as tvm_ffi
tvm.ffi.Device#
device = tvm_ffi.Device("cuda", 0)
assert device.device_type == tvm_ffi.Device.kDLCUDA
assert device.device_id == 0
assert str(device) == "cuda:0"
assert device.__repr__() == "device(type='cuda', index=0)"
device = tvm_ffi.device("ext_dev:0")
assert device.device_type == tvm_ffi.Device.kDLExtDev
assert device.device_id == 0
assert str(device) == "ext_dev:0"
assert device.__repr__() == "device(type='ext_dev', index=0)"
dev = tvm_ffi.device(dev_type="cpu", dev_id=0)
dev
device(type='cpu', index=0)
device = tvm_ffi.device("cuda", 0)
device_pickled = pickle.loads(pickle.dumps(device))
assert device_pickled.device_type == device.device_type
assert device_pickled.device_id == device.device_id