virtual_device.h

Go to the documentation of this file.

/*
 * 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.
 */
 
#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 {
 
using MemoryScope = 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;
 
class VirtualDeviceNode : public AttrsNodeReflAdapter<VirtualDeviceNode> {
 private:
  int /* actually DLDeviceType */ device_type_int;
 
 public:
  DLDeviceType device_type() const { return static_cast<DLDeviceType>(device_type_int); }
 
  int virtual_device_id;
 
  Target target;
 
  MemoryScope memory_scope;
 
  bool IsFullyUnconstrained() const {
    return !target.defined() && device_type() == kInvalidDeviceType && virtual_device_id == -1 &&
           memory_scope.empty();
  }
 
  bool IsFullyConstrained() const {
    return target.defined() && virtual_device_id != -1 && !memory_scope.empty();
  }
 
  Device ToDevice() const {
    ICHECK(device_type_int != kInvalidDeviceType);
    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(""));
  }
 
  static constexpr const char* _type_key = "target.VirtualDevice";
  TVM_FFI_DECLARE_FINAL_OBJECT_INFO(VirtualDeviceNode, BaseAttrsNode);
 
  friend class VirtualDevice;
};
 
class VirtualDevice : public ObjectRef {
 public:
  explicit VirtualDevice(int device_type_int = kInvalidDeviceType, int virtual_device_id = -1,
                         Target target = {}, MemoryScope memory_scope = {});
 
  static VirtualDevice FullyUnconstrained();
 
  static VirtualDevice ForDeviceType(DLDeviceType device_type, int virtual_device_id = -1) {
    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(const Integer& device_type, int virtual_device_id = -1) {
    return ForDeviceType(static_cast<int>(device_type->value), virtual_device_id);
  }
 
  static VirtualDevice ForDevice(const Device& device) {
    return ForDeviceType(device.device_type, device.device_id);
  }
 
  static VirtualDevice ForDeviceAndTarget(const Device& device, Target target) {
    return VirtualDevice(device.device_type, device.device_id, std::move(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));
  }
 
  static VirtualDevice ForMemoryScope(MemoryScope memory_scope) {
    return VirtualDevice(kInvalidDeviceType, -1, {}, std::move(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));
  }
 
  static Optional<VirtualDevice> Join(const VirtualDevice& lhs, const VirtualDevice& rhs);
 
  static VirtualDevice Default(const VirtualDevice& lhs, const VirtualDevice& rhs);
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(VirtualDevice, ObjectRef, VirtualDeviceNode);
 
  friend class VirtualDeviceCache;  // Private implementation helper.
};
 
class VirtualDeviceCache {
 public:
  VirtualDevice Make(int device_type = kInvalidDeviceType, int virtual_device_id = -1,
                     Target target = {}, MemoryScope memory_scope = {});
 
  VirtualDevice Unique(const VirtualDevice& virtual_device);
 
 private:
  std::unordered_set<VirtualDevice, StructuralHash, StructuralEqual> cache_;
};
 
constexpr const char* kVirtualDevice = "virtual_device";
 
}  // namespace tvm
 
#endif  //  TVM_TARGET_VIRTUAL_DEVICE_H_