| 22 | import org.apache.tvm.rpc.RPC; |
| 23 | |
| 24 | public class Device { |
| 25 | /** |
| 26 | * Provides the same information as the C++ enums DLDeviceType and |
| 27 | * TVMDeviceExtType. |
| 28 | */ |
| 29 | static final int kDLCPU = 1; |
| 30 | static final int kDLCUDA = 2; |
| 31 | static final int kDLCUDAHost = 3; |
| 32 | static final int kDLOpenCL = 4; |
| 33 | static final int kDLVulkan = 7; |
| 34 | static final int kDLMetal = 8; |
| 35 | static final int kDLVPI = 9; |
| 36 | static final int kDLROCM = 10; |
| 37 | static final int kDLROCMHost = 11; |
| 38 | static final int kDLExtDev = 12; |
| 39 | static final int kDLCUDAManaged = 13; |
| 40 | static final int kDLOneAPI = 14; |
| 41 | static final int kDLWebGPU = 15; |
| 42 | static final int kDLHexagon = 16; |
| 43 | |
| 44 | private static final Map<Integer, String> DEVICE_TYPE_TO_NAME = new HashMap<Integer, String>(); |
| 45 | private static final Map<String, Integer> DEVICE_NAME_TO_TYPE = new HashMap<String, Integer>(); |
| 46 | |
| 47 | static { |
| 48 | DEVICE_TYPE_TO_NAME.put(kDLCPU, "cpu"); |
| 49 | DEVICE_TYPE_TO_NAME.put(kDLCUDA, "cuda"); |
| 50 | DEVICE_TYPE_TO_NAME.put(kDLOpenCL, "opencl"); |
| 51 | DEVICE_TYPE_TO_NAME.put(kDLVulkan, "vulkan"); |
| 52 | DEVICE_TYPE_TO_NAME.put(kDLMetal, "metal"); |
| 53 | DEVICE_TYPE_TO_NAME.put(kDLVPI, "vpi"); |
| 54 | DEVICE_TYPE_TO_NAME.put(kDLHexagon, "hexagon"); |
| 55 | |
| 56 | DEVICE_NAME_TO_TYPE.put("cpu", kDLCPU); |
| 57 | DEVICE_NAME_TO_TYPE.put("cuda", kDLCUDA); |
| 58 | DEVICE_NAME_TO_TYPE.put("cl", kDLOpenCL); |
| 59 | DEVICE_NAME_TO_TYPE.put("opencl", kDLOpenCL); |
| 60 | DEVICE_NAME_TO_TYPE.put("vulkan", kDLVulkan); |
| 61 | DEVICE_NAME_TO_TYPE.put("metal", kDLMetal); |
| 62 | DEVICE_NAME_TO_TYPE.put("vpi", kDLVPI); |
| 63 | DEVICE_NAME_TO_TYPE.put("hexagon", kDLHexagon); |
| 64 | } |
| 65 | |
| 66 | /** |
| 67 | * Construct a CPU device. |
| 68 | * @param devId The device id |
| 69 | * @return The created device |
| 70 | */ |
| 71 | public static Device cpu(int devId) { |
| 72 | return new Device(kDLCPU, devId); |
| 73 | } |
| 74 | |
| 75 | public static Device cpu() { |
| 76 | return cpu(0); |
| 77 | } |
| 78 | |
| 79 | /** |
| 80 | * Construct a CUDA GPU device. |
| 81 | * @param devId The device id |
nothing calls this directly
no test coverage detected
searching dependent graphs…