/*
 * 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.
 */

/*!
 * \file tvm/target/virtual_device.h
 * \brief A compile time representation for where data is to be stored at runtime, and how to
 * compile code to compute it.
 */

#ifndef TVM_TARGET_VIRTUAL_DEVICE_H_
#define TVM_TARGET_VIRTUAL_DEVICE_H_

#include <tvm/ir/transform.h>
#include <tvm/target/target.h>

#include <string>
#include <unordered_set>
#include <utility>

namespace tvm {

/*!
 * \brief Abstract label for an area of memory.
 *
 * Currently uninterpreted and arbitrary. Likely to be replaced by a structured representation
 * of a memory pool in the future. Please try to use this alias instead of ffi::String to aid future
 * code migration.
 */
using MemoryScope = ffi::String;

// NOTE: cannot use enum as they are out of bound of the original enum
// and results in an undefined behavior
// A 'null' device type, does not correspond to any DLDeviceType enum.
// TODO(mbs): This is to help us as we transition away from representing the 'homogenous' case
// as a singleton target map indexed by the invalid DLDeviceType '0'.
constexpr int kNullDeviceType = 0;

// An 'invalid' device type, does not correspond to any DLDeviceType enum.
constexpr int kInvalidDeviceType = -1;

/*!
 * \brief Describes at compile time the constraints on where data is to be stored at runtime
 * down to the (virtual) device and memory scope level, and how to compile code to compute that
 * data. Used by the \p PlanDevices pass to collect and solve (virtual) device constraints for
 * the whole Relax program.
 *
 * Is a quadruple of:
 * - A \p device_type (\p DLDeviceType). May be \p kInvalidDeviceType if unconstrained.
 * - A \p virtual_device_id (\p int). This allows us to distinguish distinct devices
 *   with the same \p Target, for example in a multi-GPU system. May be -1 if unconstrained.
 *   See "Virtual Devices" below.
 * - A \p target (\p Target) describing how to compile code for the intended device. May be null
 *   if unconstrained.
 * - A \p memory_scope (\p MemoryScope, which is currently just \p String) describing which memory
 *   area is to be used to hold data. May be "" if unconstrained. See "Memory Scopes and Devices"
 *   below.
 *
 * Some or all of these fields may be unconstrained, signaling that device planning is free to
 * choose a value consistent with the whole program. However if a \p target is given then the \p
 * device_type must equal \p target->GetTargetDeviceType().
 *
 * Note that currently we assume if a function returns its result on a particular (virtual) device
 * then the function body is also executed on that device.
 *
 *
 * By 'execution' we include both (fused) primitive operators, and all the Relax expressions
 * surrounding them which coordinates data and control flow. Again, typically non-primitive
 * operators must be executed on a 'CPU'-like device with good support for control flow.
 *
 * Since TVM targets such a wide range of systems it is not possible for \p VirtualDevice to impose
 * much semantics on these fields, particularly for \p virtual_device_id and \p memory_scope.
 * Instead we assume downstream passes and codegen will interpret an validate these fields
 * appropriately.
 *
 * Targets vs Devices
 * ------------------
 * Generally \p Targets (a compile-time only datastructue) describe compiler options for a specific
 * microarchitecture and toolchain, while \p Devices (a runtime datastructure also available at
 * compile time) describe a physical device on the target system. Obviously the target must agree
 * with the device's microarchitecture, but we otherwise don't impose any constraints between them:
 *  - It's ok to use different \p Targets for the same \p Device, eg to squeeze some extra perf
 *    out of a particular primitive using particular compiler flags.
 *  - It's ok to use the same \p Target for multiple \p Devices, eg if we have multiple CPUs.
 *
 * Traditionally TVM assumes at most one \p Target per \p DLDeviceType. We are moving away from that
 * assumption.
 *
 * Virtual vs Physical Devices
 * ---------------------------
 * The \p virtual_device_id may be used by downstream passes or the runtime to help decide which
 * \p device_id to use for a particular physical runtime \p Device. For example:
 *  - Some runtimes may support passing in an array of actual `device` specifications, and the
 *    \p virtual_device_id can be used at runtime as an index into that array.
 *  - Some runtimes may support dynamically allocating computations to physical devices. On these
 *    systems a large space of \p virtual_device_ids could be used at compile time, even though
 *    at runtime only a few physical devices will be present.
 *
 * The \p virtual_device_id may also be left unconstrained if not needed.
 *
 * Memory Scopes and Devices
 * -------------------------
 * Multi-device systems can have complex memory hierarchies. For example
 * \code
 * (kDLCPU, 0, "llvm", "global")
 * \endcode
 * and
 * \code
 * (kDLCPU, 1, "llvm", "global")
 * \endcode
 * could denote:
 * - The same memory area accessible from two separate CPUs without any CPU affinity;
 * - Distinct memory areas in a NUMA architecture for which cross-device access is handled
 *   by the memory system;
 * - Outright distinct memory areas, where one device cannot directly address the memory of
 *   another.
 *
 * Similarly:
 * \code
 * (kDLCPU, 0, "llvm", "global")
 * \endcode
 * and
 * \code
 * (kDLCUDA, 0, "cuda", "host")
 * \endcode
 * could denote the same memory area, but with very different access costs.
 *
 * Furthermore, not all memory scopes are accessible to all devices, and it is possible for
 * a memory scope to only be accessible to a device when code is compiled with particular
 * \p Target options.
 *
 * \p VirtualDevices themselves have no system-level understanding. Currently the \p PlanDevices
 * pass will simply insert "device_copy" operators wherever \p VirtualDevices are not exactly
 * pointwise equal. We may revisit this in the future as the work on memory pools matures.
 *
 * Joining and Defaulting
 * ----------------------
 * It is possible to 'join' two \p VirtualDevices to yield the most constrained \p VirtualDevice
 * which agrees with both join arguments. Eg:
 * \code
 * Join((kDLCPU, -1, "llvm", ""), (kInvalidDeviceType, 3, null, "global))
 *   => (kDLCPU, 3, "llvm", "global")
 * Join((kDLCPU, -1, "llvm", ""), (kInvalidDeviceType, 3, null, "local))
 *   => null (no join possible)
 * \endcode
 *
 * Related to 'join' is 'default', which only takes constrained fields from the rhs when the
 * lhs is unconstrained:
 * \code
 * Default(kDLCPU, -1, "llvm", "local"), (kDLCPU, 3, null, "global"))
 *   => (kDLCPU, 3, "llvm", "local")
 * \endcode
 *
 * These operations are needed during device planning.
 */

class VirtualDeviceNode : public AttrsNode {
 private:
  /*!
   * \brief The \p DLDeviceType (represented as an int) of the virtual device. If \p target is
   * known then this will be equal to \p target->GetTargetDeviceType(). If \p target is null then
   * the target is to be determined later.
   *
   * This is needed to support the legacy "on_device" and "device_copy" calls which only allow
   * a \p DLDeviceTypes (as an integer) to be given.
   *
   * kInvalidDeviceType denotes unconstrained. An int since the DLDeviceType enum representation
   * is not fixed. Private to discourage further int vs DLDeviceType confusion.
   */
  int /* actually DLDeviceType */ device_type_int;

 public:
  DLDeviceType device_type() const { return static_cast<DLDeviceType>(device_type_int); }

  /*!
   * \brief The device identifier for the virtual device. This must be resolved to a physical
   * device identifier either during compilation or at runtime.
   *
   * -1 denotes unconstrained.
   */
  int virtual_device_id;

  /*!
   * \brief The \p Target describing how to compile for the virtual device.
   *
   * Null denotes unconstrained. Note that if a target later becomes known for this \p VirtualDevice
   * then it must be consistent with the \p device_type if already known. This is enforced by the
   * Join and Default methods.
   */
  Target target;

  /*!
   * \brief The scope of memory w.r.t. the virtual device which holds data.
   *
   * Empty denotes unconstrained.
   */
  MemoryScope memory_scope;

  /*!
   * \brief Returns true if virtual device is 'fully unconstrained', ie no target/device type,
   * device id or memory scope is specified.
   */
  bool IsFullyUnconstrained() const {
    return !target.defined() && device_type() == kInvalidDeviceType && virtual_device_id == -1 &&
           memory_scope.empty();
  }

  /*!
   * \brief Returns true if virtual device is 'fully constrained', ie target, device id and memory
   * scope are all specified.
   */
  bool IsFullyConstrained() const {
    return target.defined() && virtual_device_id != -1 && !memory_scope.empty();
  }

  /*!
   * \brief Returns the (virtual) \p Device implied by this \p VirtualDevice. Both the \p
   * device_type and \p virtual_device_must be constrained. The returned \p Device may not
   * correspond to any physical device available at compile time or even runtime: see "Virtual vs
   * Physical Devices" above.
   */
  Device ToDevice() const {
    TVM_FFI_ICHECK(device_type_int != kInvalidDeviceType);
    TVM_FFI_ICHECK(virtual_device_id != -1);
    Device device;
    device.device_type = device_type();
    device.device_id = virtual_device_id;
    return device;
  }

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<VirtualDeviceNode>()
        .def_ro("device_type_int", &VirtualDeviceNode::device_type_int,
                "The type of the virtual device.", refl::DefaultValue(kInvalidDeviceType))
        .def_ro("virtual_device_id", &VirtualDeviceNode::virtual_device_id,
                "The device id of the virtual device.", refl::DefaultValue(-1))
        .def_ro("target", &VirtualDeviceNode::target,
                "The target describing how to compile for the virtual device.",
                refl::DefaultValue(Target()))
        .def_ro("memory_scope", &VirtualDeviceNode::memory_scope,
                "The area of memory w.r.t. the virtual device where data is stored.",
                refl::DefaultValue(""));
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.VirtualDevice", VirtualDeviceNode, AttrsNode);

  friend class VirtualDevice;
};

/*!
 * \brief Managed reference class to \p VirtualDeviceNode.
 */
class VirtualDevice : public ffi::ObjectRef {
 public:
  /*!
   * \brief Construct a virtual device.
   * \param device_type_int The device type for the virtual device, or \p kInvalidDeviceType if
   * unconstrained.  If \p target is defined then must match its \p target->GetTargetDeviceType().
   * \param virtual_device_id The device id for the virtual device, or -1 if unconstrained.
   * \param target The target describing how to compile for the virtual device, or null if
   * unconstrained.
   * \param memory_scope The memory scope w.r.t. the virtual device which holds data, or "" if
   * unconstrained.
   */
  TVM_DLL explicit VirtualDevice(int device_type_int = kInvalidDeviceType,
                                 int virtual_device_id = -1, Target target = {},
                                 MemoryScope memory_scope = {});

  /*! \brief Returns the unique fully unconstrained \p VirtualDevice. */
  static VirtualDevice FullyUnconstrained();

  /*!
   * \brief Returns the \p VirtualDevice for \p device_type and (if not -1) \p virtual_device_id.
   * The target and memory scope will be unconstrained.
   */
  static VirtualDevice ForDeviceType(DLDeviceType device_type, int virtual_device_id = -1) {
    TVM_FFI_ICHECK_GT(device_type, 0);
    return VirtualDevice(device_type, virtual_device_id);
  }
  static VirtualDevice ForDeviceType(int device_type, int virtual_device_id = -1) {
    return ForDeviceType(static_cast<DLDeviceType>(device_type), virtual_device_id);
  }
  static VirtualDevice ForDeviceType(int64_t device_type, int virtual_device_id = -1) {
    return ForDeviceType(static_cast<int>(device_type), virtual_device_id);
  }

  /*! \brief Returns the \p VirtualDevice for \p device. */
  static VirtualDevice ForDevice(const Device& device) {
    return ForDeviceType(device.device_type, device.device_id);
  }

  /*! \brief Returns the \p VirtualDevice for \p device and \p target. */
  static VirtualDevice ForDeviceAndTarget(const Device& device, Target target) {
    return VirtualDevice(device.device_type, device.device_id, std::move(target));
  }

  /*! \brief Returns the \p VirtualDevice for \p target. */
  static VirtualDevice ForTarget(Target target) {
    DLDeviceType device_type = static_cast<DLDeviceType>(target->GetTargetDeviceType());
    return VirtualDevice(device_type, /*virtual_device_id=*/0, std::move(target));
  }

  /*! \brief Returns the \p VirtualDevice for \p memory_scope alone. */
  static VirtualDevice ForMemoryScope(MemoryScope memory_scope) {
    return VirtualDevice(kInvalidDeviceType, -1, {}, std::move(memory_scope));
  }

  /*! \brief Returns the \p VirtualDevice for \p device, \p target and \p memory_scope. */
  TVM_DLL static VirtualDevice ForDeviceTargetAndMemoryScope(const Device& device, Target target,
                                                             MemoryScope memory_scope) {
    return VirtualDevice(device.device_type, device.device_id, std::move(target),
                         std::move(memory_scope));
  }

  /*!
   * \brief Returns the 'join' of \p lhs and \p rhs. The result will agree pointwise with
   * \p lhs and \p rhs on all their constrained fields. Returns the null optional if no such
   * join exists, ie there's disagreement on at least one constrained field.
   */
  TVM_DLL static ffi::Optional<VirtualDevice> Join(const VirtualDevice& lhs,
                                                   const VirtualDevice& rhs);

  /*!
   * \brief Returns the 'default' of \p lhs and \p rhs. The result will be \p lhs, except any
   * unconstrained fields in \p lhs will take their value from \p rhs. Always well-defined.
   */
  TVM_DLL static VirtualDevice Default(const VirtualDevice& lhs, const VirtualDevice& rhs);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VirtualDevice, ffi::ObjectRef, VirtualDeviceNode);

  friend class VirtualDeviceCache;  // Private implementation helper.
};

/*!
 * \brief A cache of \p VirtualDevices. This can be used:
 *  - To avoid ending up with lots of identical instances, since the space of VirtualDevices for any
 *    one compilation is very small but the number of points they need to be constructed can
 *    be very large (eg during device planning).
 *  - So we can assume \p VirtualDevices are pointer equal if and only if they are structurally
 * equal. This simplifies the unification of 'device domains' which are built on \p VirtualDevices.
 */
class TVM_DLL VirtualDeviceCache {
 public:
  /*! \brief Returns the unique \p VirtualDevice representing given fields. */
  VirtualDevice Make(int device_type = kInvalidDeviceType, int virtual_device_id = -1,
                     Target target = {}, MemoryScope memory_scope = {});

  /*!
   * \brief Returns the unique \p VirtualDevice structurally equal to the given \p virtual_device.
   */
  VirtualDevice Unique(const VirtualDevice& virtual_device);

 private:
  /*! \brief Already constructed VirtualDevices. */
  std::unordered_set<VirtualDevice, ffi::StructuralHash, ffi::StructuralEqual> cache_;
};

/*! brief The attribute key for the virtual device. This key will be promoted to first class on
 * functions. For use in the parser and printer only.
 *
 * Type: VirtualDevice
 */
constexpr const char* kVirtualDevice = "virtual_device";

}  // namespace tvm

#endif  //  TVM_TARGET_VIRTUAL_DEVICE_H_