init devices and allocators.
(self, dev: Device, memory_cfg: str | dict[Device, str])
| 88 | self._setup_device(device, memory_cfg) |
| 89 | |
| 90 | def _setup_device(self, dev: Device, memory_cfg: str | dict[Device, str]) -> None: |
| 91 | """init devices and allocators.""" |
| 92 | devs = dev |
| 93 | if not isinstance(dev, list | tuple): |
| 94 | if not isinstance(dev, tvm.runtime.Device): |
| 95 | raise TypeError("dev is expected to be Device or List[Device]") |
| 96 | devs = [dev] |
| 97 | |
| 98 | # CPU is required for executing shape functions |
| 99 | if devs[-1].dlpack_device_type() % RPC_SESS_MASK != tvm.cpu().dlpack_device_type(): |
| 100 | devs.append(tvm.cpu()) |
| 101 | |
| 102 | default_alloc_type = VirtualMachine.POOLED_ALLOCATOR |
| 103 | if memory_cfg is None: |
| 104 | memory_cfg = {} |
| 105 | elif isinstance(memory_cfg, str): |
| 106 | assert memory_cfg in ["naive", "pooled"] |
| 107 | if memory_cfg == "naive": |
| 108 | default_alloc_type = VirtualMachine.NAIVE_ALLOCATOR |
| 109 | memory_cfg = {} |
| 110 | elif not isinstance(memory_cfg, dict): |
| 111 | raise TypeError( |
| 112 | "memory_cfg is expected be string or dictionary, " |
| 113 | + f"but received {type(memory_cfg)}" |
| 114 | ) |
| 115 | init_args = [] |
| 116 | for device in devs: |
| 117 | init_args.append(device.dlpack_device_type() % RPC_SESS_MASK) |
| 118 | init_args.append(device.index) |
| 119 | alloc_type = memory_cfg[device] if device in memory_cfg else default_alloc_type |
| 120 | init_args.append(alloc_type) |
| 121 | self.module["vm_initialization"](*init_args) |
| 122 | |
| 123 | def __getitem__(self, key: str) -> Function: |
| 124 | return self.module[key] |