apache--tvm/python/tvm/target/target.py

# 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