VariableCountingVector.h

Go to the documentation of this file.

 
#pragma once
 
#include <tvm/api.h>
#include <tvm/defs.h>
 
#include <tvm/Space.h>
#include <tvm/VariableVector.h>
#include <tvm/internal/RangeCounting.h>
#include <tvm/utils/internal/map.h>
 
namespace tvm::internal
{
class TVM_DLLAPI VariableCountingVector
{
public:
  VariableCountingVector(bool split = false) : split_(split) {}
  bool add(VariablePtr v);
  void add(const VariableVector & v);
  bool remove(const Variable & v);
  void remove(const VariableVector & v);
 
  void clear();
 
  void set(const VectorConstRef & val);
 
  const VariableVector & variables() const;
 
  const std::vector<uint8_t> simple() const;
 
  bool split() const { return split_; }
 
  bool isDisjointUnion();
 
private:
  struct SpaceRangeCounting
  {
    tvm::internal::RangeCounting mSize_;
    tvm::internal::RangeCounting rSize_;
    tvm::internal::RangeCounting tSize_;
 
    bool add(const Space & start, const Space & dim);
    bool remove(const Space & start, const Space & dim);
    bool empty() const { return mSize_.empty(); }
    int maxCount() const { return mSize_.maxCount(); }
  };
 
  void update() const;
 
  tvm::utils::internal::map<Variable *, std::pair<SpaceRangeCounting, size_t>> count_;
  bool split_;
  mutable bool upToDate_ = false;
  mutable VariableVector variables_;
  mutable std::vector<uint8_t> simple_;
};
} // namespace tvm::internal