# 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. """Target data structure.""" import tvm_ffi from tvm_ffi import Map from tvm_ffi.core import String from tvm.runtime import Device, Object, convert from . import _ffi_api @tvm_ffi.register_object("target.TargetKind") class TargetKind(Object): """Kind of a compilation target""" @property def options(self): """Returns the dict of available option names and types""" return dict(_ffi_api.ListTargetKindOptions(self)) @staticmethod def options_from_name(kind_name: str): """Returns the dict of available option names and types from a name of TargetKind""" return dict(_ffi_api.ListTargetKindOptionsFromName(kind_name)) class TargetFeatures: def __init__(self, target): self.target = target def __getattr__(self, name: str): return _ffi_api.TargetGetFeature(self.target, name) @tvm_ffi.register_object("target.Target") class Target(Object): """Target device information, use through TVM API. Targets can be constructed from: - A JSON config dictionary: ``Target({"kind": "cuda", "arch": "sm_80"})`` - A tag name: ``Target("nvidia/nvidia-a100")`` - A tag with overrides: ``Target({"tag": "nvidia/nvidia-a100", "l2_cache_size_bytes": 12345})`` - A kind name: ``Target("cuda")`` Use ``target.attrs["key"]`` to access target attributes. Examples -------- .. code-block:: python # From a tag target = Target("nvidia/nvidia-a100") # From a tag with attribute overrides target = Target({"tag": "qcom/hexagon-v68", "vtcm-capacity": 70000}) # From a config dictionary target = Target({"kind": "cuda", "arch": "sm_80"}) """ def __init__(self, target, host=None): """Construct a TVM target object from 1) Raw target string 2) Target config dict 3) Target tag string 4) Tag with overrides dict Parameters ---------- target : Union[str, Dict[str, Any]] Can be one of a literal target string, a json string describing a configuration, or a dictionary of configuration options. When using a dictionary or json string to configure target, the possible values are: tag : str (optional) A registered tag name (e.g. ``"nvidia/nvidia-a100"``). When ``tag`` is present, the tag's base config is loaded and any additional fields in the dict override the base values. The ``kind`` field is not needed when ``tag`` is specified. kind : str (required unless tag is specified) Which codegen path to use, for example 'llvm' or 'cuda'. keys : List of str (optional) A set of strategies that can be dispatched to. When using "kind=opencl" for example, one could set keys to ["mali", "opencl", "gpu"]. device : str (optional) A single key that corresponds to the actual device being run on. This will be effectively appended to the keys. libs : List of str (optional) The set of external libraries to use. For example ['cblas', 'mkl']. system-lib : bool (optional) If True, build a module that contains self registered functions. Useful for environments where dynamic loading like dlopen is banned. mcpu : str (optional) The specific cpu being run on. Serves only as an annotation. model : str (optional) An annotation indicating what model a workload came from. runtime : str (optional) An annotation indicating which runtime to use with a workload. mtriple : str (optional) The llvm triplet describing the target, for example "arm64-linux-android". mattr : List of str (optional) The llvm features to compile with, for example ["+avx512f", "+mmx"]. mfloat-abi : str (optional) An llvm setting that is one of 'hard' or 'soft' indicating whether to use hardware or software floating-point operations. mabi : str (optional) An llvm setting. Generate code for the specified ABI, for example "lp64d". host : Union[str, Dict[str, Any]] (optional) Description for target host. Can be recursive. Similar to target. host : Optional[Union[str, Dict[str, Any]]] Similar to target but for target host. Can be one of a literal target host string, a json string describing a configuration, or a dictionary of configuration options. When using a dictionary or json string to configure target, the possible values are same as target. """ if isinstance(target, dict | str): target = convert(target) if isinstance(host, dict | str): host = convert(host) if target is None or not isinstance(target, Map | String | Target | str): raise ValueError(f"target has to be a string or dictionary. instead get {type(target)}") if host is not None: if not isinstance(host, Map | String | Target | str): raise ValueError("target host has to be a string or dictionary.") self.__init_handle_by_constructor__(_ffi_api.Target, Target(target), Target(host)) else: self.__init_handle_by_constructor__(_ffi_api.Target, target) def __enter__(self): _ffi_api.TargetEnterScope(self) return self def __exit__(self, ptype, value, trace): _ffi_api.TargetExitScope(self) def export(self): return _ffi_api.TargetExport(self) def with_host(self, host=None): return _ffi_api.WithHost(self, Target(host)) @staticmethod def from_device(device: str | Device) -> "Target": """Detects Target associated with the given device. If the device does not exist, there will be an Error. Parameters ---------- dev : Union[str, Device] The device to detect the target for. Supported device types: ["cuda", "metal", "rocm", "vulkan", "opencl", "cpu"] Returns ------- target : Target The detected target. """ from .detect_target import ( # pylint: disable=import-outside-toplevel detect_target_from_device, ) return detect_target_from_device(device) @staticmethod def current(allow_none=True): """Returns the current target. Parameters ---------- allow_none : bool Whether allow the current target to be none Raises ------ ValueError if current target is not set. """ return _ffi_api.TargetCurrent(allow_none) @property def features(self): return TargetFeatures(self) def __getattr__(self, name: str): """Backward-compatible attribute access for target attrs. Historically, code accessed target options via attribute syntax (e.g. ``target.arch``). Newer APIs prefer ``target.attrs["arch"]``. """ attrs = self.attrs if name in attrs: value = attrs[name] return str(value) if isinstance(value, String) else value raise AttributeError(f"'Target' object has no attribute '{name}'") def get_kind_attr(self, attr_name): """Get additional attribute about the target kind. Parameters ---------- attr_name : str The attribute name. Returns ------- value : object The attribute value """ return _ffi_api.TargetKindGetAttr(self.kind, attr_name) def get_target_device_type(self): """Returns the device_type for this target.""" return _ffi_api.TargetGetDeviceType(self) @staticmethod def list_kinds(): """Returns the list of available target names.""" return list(_ffi_api.ListTargetKinds()) @staticmethod def target_or_current(target): """Returns target, or the current target in the environment if target is None""" if target is None: target = Target.current() if target is None: raise ValueError("Target is not set in env or passed as argument.") return target