se_scope.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_SE_SCOPE_H_
#define TVM_TARGET_SE_SCOPE_H_

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

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

namespace tvm {

using MemoryScope = String;

class SEScopeNode : public AttrsNode<SEScopeNode> {
 public:
  int device_type_int;

  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() != kInvalidDeviceType);
    ICHECK(virtual_device_id != -1);
    Device device;
    device.device_type = device_type();
    device.device_id = virtual_device_id;
    return device;
  }

  TVM_DECLARE_ATTRS(SEScopeNode, "SEScope") {
    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 SEScope;
};

class SEScope : public ObjectRef {
 public:
  explicit SEScope(DLDeviceType device_type = kInvalidDeviceType, int virtual_device_id = -1,
                   Target target = {}, MemoryScope memory_scope = {});

  static SEScope FullyUnconstrained();

  static SEScope ForDeviceType(DLDeviceType device_type, int virtual_device_id = -1) {
    ICHECK_GT(device_type, 0);
    return SEScope(device_type, virtual_device_id);
  }
  static SEScope ForDeviceType(int device_type, int virtual_device_id = -1) {
    return ForDeviceType(static_cast<DLDeviceType>(device_type), virtual_device_id);
  }
  static SEScope ForDeviceType(const Integer& device_type, int virtual_device_id = -1) {
    return ForDeviceType(static_cast<int>(device_type->value), virtual_device_id);
  }

  static SEScope ForDevice(const Device& device) {
    return ForDeviceType(device.device_type, device.device_id);
  }

  static SEScope ForDeviceAndTarget(const Device& device, Target target) {
    return SEScope(device.device_type, device.device_id, std::move(target));
  }

  TVM_DLL static SEScope ForDeviceTargetAndMemoryScope(const Device& device, Target target,
                                                       MemoryScope memory_scope) {
    return SEScope(device.device_type, device.device_id, std::move(target),
                   std::move(memory_scope));
  }

  static Optional<SEScope> Join(const SEScope& lhs, const SEScope& rhs);

  static SEScope Default(const SEScope& lhs, const SEScope& rhs);

  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(SEScope, ObjectRef, SEScopeNode);

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

class SEScopeCache {
 public:
  SEScope Make(DLDeviceType device_type = kInvalidDeviceType, int virtual_device_id = -1,
               Target target = {}, MemoryScope memory_scope = {});

  SEScope Unique(const SEScope& scope);

 private:
  std::unordered_set<SEScope, StructuralHash, StructuralEqual> cache_;
};

}  // namespace tvm

#endif  //  TVM_TARGET_SE_SCOPE_H_