feature.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_RELAY_FEATURE_H_
#define TVM_RELAY_FEATURE_H_

#include <tvm/ir/module.h>
#include <tvm/relay/expr.h>

#include <bitset>
#include <string>

namespace tvm {
namespace relay {

enum Feature : int {
  fVar = 0,
  fGlobalVar = 1,
  fConstant = 2,
  fTuple = 3,
  fTupleGetItem = 4,
  fFunction = 5,
  fOp = 6,
  fCall = 7,
  fLet = 8,
  fIf = 9,
  fRefCreate = 10,
  fRefRead = 11,
  fRefWrite = 12,
  fConstructor = 13,
  fMatch = 14,
  fGraph = 15,
  fLetRec = 16
};

constexpr size_t feature_count = 17;

class FeatureSet {
 public:
  FeatureSet(const FeatureSet&) = default;
  explicit FeatureSet(Feature ft) { bs_.set(static_cast<size_t>(ft)); }
  explicit FeatureSet(const tvm::Array<tvm::Integer>& ft) {
    for (Integer i : ft) {
      *this += Feature(i.IntValue());
    }
  }
  explicit operator Array<Integer>() const {
    Array<Integer> ret;
    for (size_t i = 0; i < feature_count; ++i) {
      if (bs_[i]) {
        ret.push_back(Integer(i));
      }
    }
    return ret;
  }
  static FeatureSet All() {
    FeatureSet fs;
    fs.bs_.flip();
    return fs;
  }
  static FeatureSet No() {
    FeatureSet fs;
    return fs;
  }
  template <typename T>
  FeatureSet& operator+=(const T& rhs) {
    bs_ |= FeatureSet(rhs).bs_;
    return *this;
  }
  template <typename T>
  FeatureSet operator+(const T& rhs) const {
    FeatureSet fs(*this);
    fs += rhs;
    return fs;
  }
  template <typename T>
  FeatureSet& operator-=(const T& rhs) {
    bs_ &= ~(FeatureSet(rhs)).bs_;
    return *this;
  }
  template <typename T>
  FeatureSet operator-(const T& rhs) const {
    FeatureSet fs(*this);
    fs -= rhs;
    return fs;
  }
  bool is_subset_of(const FeatureSet& rhs) const { return ((*this) - rhs).bs_.none(); }

  std::string ToString() const;

 private:
  std::bitset<feature_count> bs_;
  FeatureSet() = default;
  explicit FeatureSet(const std::bitset<feature_count>& bs) : bs_(bs) {}
};

FeatureSet DetectFeature(const RelayExpr& expr);

FeatureSet DetectFeature(const IRModule& mod);

inline FeatureSet DetectFeature(const Expr& expr, const IRModule& mod) {
  return DetectFeature(expr) + DetectFeature(mod);
}

void CheckFeature(const RelayExpr& expr, const FeatureSet& fs);

void CheckFeature(const IRModule& mod, const FeatureSet& fs);

inline void CheckFeature(const RelayExpr& expr, const IRModule& mod, const FeatureSet& fs) {
  CheckFeature(expr, fs);
  CheckFeature(mod, fs);
}

}  // namespace relay
}  // namespace tvm

#endif  // TVM_RELAY_FEATURE_H_