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 AttrsNode<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;
  }

  TVM_DECLARE_ATTRS(VirtualDeviceNode, "VirtualDevice") {
    TVM_ATTR_FIELD(device_type_int)
        .describe("The type of the virtual device.")
        .set_default(kInvalidDeviceType);
    TVM_ATTR_FIELD(virtual_device_id)
        .describe("The device id of the virtual device.")
        .set_default(-1);
    TVM_ATTR_FIELD(target)
        .describe("The target describing how to compile for the virtual device.")
        .set_default(Target());
    TVM_ATTR_FIELD(memory_scope)
        .describe("The area of memory w.r.t. the virtual device where data is stored.")
        .set_default("");
  }

  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_