wehub-resource-sync
40 lines
1.5 KiB
Python
40 lines
1.5 KiB
Python
# 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.
|
|
# ruff: noqa: F821
|
|
"""Python bindings for creating VirtualDevices."""
|
|
|
|
import tvm_ffi
|
|
|
|
from . import _ffi_api
|
|
|
|
|
|
@tvm_ffi.register_object("target.VirtualDevice")
|
|
class VirtualDevice(tvm_ffi.core.Object):
|
|
"""A compile time representation for where data is to be stored at runtime,
|
|
and how to compile code to compute it."""
|
|
|
|
def __init__(self, device=None, target=None, memory_scope="") -> None:
|
|
if device is None:
|
|
# The 'unconstrained' device has device type -1 and device id -1.
|
|
device = tvm.device(-1, -1)
|
|
self.__init_handle_by_constructor__(
|
|
_ffi_api.VirtualDevice_ForDeviceTargetAndMemoryScope, device, target, memory_scope
|
|
)
|
|
|
|
def dlpack_device_type(self) -> int:
|
|
return self.device_type_int
|