AliasInfo Class — pytorch Architecture

Source Code

aten/src/ATen/core/alias_info.h lines 21–92

class AliasInfo {
 public:
  AliasInfo() = default;
  AliasInfo(bool is_write, const std::set<std::string>& before_qual_strings, const std::set<std::string>& after_qual_strings) : isWrite_(is_write) {
    for (const auto& s: before_qual_strings) {
      beforeSets_.insert(Symbol::fromQualString(s));
    }
    for (const auto& s : after_qual_strings) {
      afterSets_.insert(Symbol::fromQualString(s));
    }
  }
  // Symbol for the set that can alias anything
  static Symbol wildcardSet() {
    static const Symbol wc = Symbol::fromQualString("alias::*");
    return wc;
  }

  void setIsWrite(bool isWrite) {
    isWrite_ = isWrite;
  }

  bool isWrite() const {
    return isWrite_;
  }

  void addBeforeSet(Symbol aliasSet) {
    beforeSets_.insert(aliasSet);
  }

  void addAfterSet(Symbol aliasSet) {
    afterSets_.insert(aliasSet);
  }

  const std::unordered_set<Symbol>& beforeSets() const {
    return beforeSets_;
  }

  const std::unordered_set<Symbol>& afterSets() const {
    return afterSets_;
  }

  Symbol beforeSet() const {
    AT_ASSERT(beforeSets_.size() == 1);
    return *beforeSets_.begin();
  }

  bool isWildcardBefore() const {
    return beforeSets_.count(wildcardSet()) != 0;
  }

  bool isWildcardAfter() const {
    return afterSets_.count(wildcardSet()) != 0;
  }

  // the alias info for the contained types of the type
  // e.g. if this is an annotation on List[T], `sets` refers to
  // the alias sets that the list may be in
  // while containedTypes()[0] refers to the sets that members of the list
  // may be in
  void addContainedType(AliasInfo aliasInfo) {
    containedTypes_.push_back(std::move(aliasInfo));
  }
  const std::vector<AliasInfo>& containedTypes() const {
    return containedTypes_;
  }

 private:
  std::unordered_set<Symbol> beforeSets_;
  std::unordered_set<Symbol> afterSets_;
  std::vector<AliasInfo> containedTypes_;
  bool isWrite_ = false;
};