expr.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_IR_EXPR_H_
#define TVM_IR_EXPR_H_
 
#include <tvm/ir/source_map.h>
#include <tvm/ir/type.h>
#include <tvm/node/node.h>
#include <tvm/runtime/container/string.h>
#include <tvm/runtime/object.h>
 
#include <algorithm>
#include <functional>
#include <limits>
#include <string>
#include <type_traits>
 
namespace tvm {
 
using tvm::runtime::String;
 
// Forward-declare VirtualDevice to avoid circular imports.
class VirtualDevice;
 
class BaseExprNode : public Object {
 public:
  mutable Span span;
 
  static constexpr const char* _type_key = "BaseExpr";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  static constexpr const uint32_t _type_child_slots = 64;
  TVM_DECLARE_BASE_OBJECT_INFO(BaseExprNode, Object);
};
 
class BaseExpr : public ObjectRef {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(BaseExpr, ObjectRef, BaseExprNode);
};
 
class PrimExprNode : public BaseExprNode {
 public:
  DataType dtype;
 
  TVM_OBJECT_ENABLE_SCRIPT_PRINTER();
 
  static constexpr const char* _type_key = "PrimExpr";
  static constexpr const uint32_t _type_child_slots = 40;
  TVM_DECLARE_BASE_OBJECT_INFO(PrimExprNode, BaseExprNode);
};
 
class PrimExpr : public BaseExpr {
 public:
  TVM_DLL PrimExpr(int32_t value);  // NOLINT(*)
  TVM_DLL PrimExpr(float value);  // NOLINT(*)
 
  DataType dtype() const { return static_cast<const PrimExprNode*>(get())->dtype; }
 
  TVM_DEFINE_OBJECT_REF_METHODS(PrimExpr, BaseExpr, PrimExprNode);
 
 private:
  // Internal function for conversion.
  friend struct runtime::PackedFuncValueConverter<PrimExpr>;
  TVM_DLL static PrimExpr FromObject_(ObjectRef ref);
};
 
TVM_DLL PrimExpr operator+(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator-(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator-(PrimExpr a);
 
TVM_DLL PrimExpr operator*(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator/(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator<<(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator>>(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator>(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator>=(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator<(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator<=(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator==(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator!=(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator&&(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator||(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator!(PrimExpr a);
 
TVM_DLL PrimExpr operator&(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator|(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator^(PrimExpr a, PrimExpr b);
 
TVM_DLL PrimExpr operator~(PrimExpr a);
 
class RelaxExprNode : public BaseExprNode {
 public:
  mutable Type checked_type_ = Type(nullptr);
 
  mutable Optional<ObjectRef> struct_info_ = Optional<ObjectRef>();
 
  inline const Type& checked_type() const;
  template <typename TTypeNode>
  inline const TTypeNode* type_as() const;
 
  static constexpr const char* _type_key = "RelaxExpr";
  static constexpr const uint32_t _type_child_slots = 22;
  TVM_DECLARE_BASE_OBJECT_INFO(RelaxExprNode, BaseExprNode);
};
 
class RelaxExpr : public BaseExpr {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(RelaxExpr, BaseExpr, RelaxExprNode);
};
 
class GlobalVar;
class GlobalVarNode : public RelaxExprNode {
 public:
  String name_hint;
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("name_hint", &name_hint);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
    v->Visit("struct_info_", &struct_info_);
  }
 
  bool SEqualReduce(const GlobalVarNode* other, SEqualReducer equal) const {
    // name matters for global var.
    return equal(name_hint, other->name_hint) && equal.FreeVarEqualImpl(this, other);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(name_hint);
    hash_reduce.FreeVarHashImpl(this);
  }
 
  static constexpr const char* _type_key = "GlobalVar";
  TVM_DECLARE_FINAL_OBJECT_INFO(GlobalVarNode, RelaxExprNode);
};
 
class GlobalVar : public RelaxExpr {
 public:
  TVM_DLL explicit GlobalVar(String name_hint, Type type = {}, Span span = {});
 
  TVM_DEFINE_OBJECT_REF_METHODS(GlobalVar, RelaxExpr, GlobalVarNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(GlobalVarNode);
};
 
// PrimExprs that are useful as runtime containers.
//
class IntImmNode : public PrimExprNode {
 public:
  int64_t value;
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("value", &value);
    v->Visit("span", &span);
  }
 
  bool SEqualReduce(const IntImmNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(value, other->value);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(value);
  }
 
  static constexpr const char* _type_key = "IntImm";
  TVM_DECLARE_FINAL_OBJECT_INFO(IntImmNode, PrimExprNode);
};
 
class IntImm : public PrimExpr {
 public:
  TVM_DLL IntImm(DataType dtype, int64_t value, Span span = Span());
 
  TVM_DEFINE_OBJECT_REF_METHODS(IntImm, PrimExpr, IntImmNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IntImmNode);
};
 
class FloatImmNode : public PrimExprNode {
 public:
  double value;
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("value", &value);
    v->Visit("span", &span);
  }
 
  bool SEqualReduce(const FloatImmNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(value, other->value);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(value);
  }
 
  static constexpr const char* _type_key = "FloatImm";
  TVM_DECLARE_FINAL_OBJECT_INFO(FloatImmNode, PrimExprNode);
};
 
class FloatImm : public PrimExpr {
 public:
  TVM_DLL FloatImm(DataType dtype, double value, Span span = Span());
 
  TVM_DEFINE_OBJECT_REF_METHODS(FloatImm, PrimExpr, FloatImmNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(FloatImmNode);
};
 
class Bool : public IntImm {
 public:
  explicit Bool(bool value, Span span = Span()) : IntImm(DataType::Bool(), value, span) {}
  Bool operator!() const { return Bool((*this)->value == 0); }
  operator bool() const { return (*this)->value != 0; }
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Bool, IntImm, IntImmNode);
};
 
// Overload operators to make sure we have the most fine grained types.
inline Bool operator||(const Bool& a, bool b) { return Bool(a.operator bool() || b); }
inline Bool operator||(bool a, const Bool& b) { return Bool(a || b.operator bool()); }
inline Bool operator||(const Bool& a, const Bool& b) {
  return Bool(a.operator bool() || b.operator bool());
}
inline Bool operator&&(const Bool& a, bool b) { return Bool(a.operator bool() && b); }
inline Bool operator&&(bool a, const Bool& b) { return Bool(a && b.operator bool()); }
inline Bool operator&&(const Bool& a, const Bool& b) {
  return Bool(a.operator bool() && b.operator bool());
}
 
inline bool operator==(const Bool& a, bool b) { return a.operator bool() == b; }
inline bool operator==(bool a, const Bool& b) { return a == b.operator bool(); }
inline bool operator==(const Bool& a, const Bool& b) {
  return a.operator bool() == b.operator bool();
}
 
class Integer : public IntImm {
 public:
  Integer() {}
  explicit Integer(ObjectPtr<Object> node) : IntImm(node) {}
  Integer(int value, Span span = Span()) : IntImm(DataType::Int(32), value, span) {}  // NOLINT(*)
  Integer(IntImm other) : IntImm(std::move(other)) {}  // NOLINT(*)
  template <typename Enum, typename = typename std::enable_if<std::is_enum<Enum>::value>::type>
  explicit Integer(Enum value) : Integer(static_cast<int>(value)) {
    static_assert(std::is_same<int, typename std::underlying_type<Enum>::type>::value,
                  "declare enum to be enum int to use visitor");
  }
  Integer& operator=(const IntImm& other) {
    data_ = ObjectRef::GetDataPtr<Object>(other);
    return *this;
  }
  int64_t IntValue() const {
    ICHECK(data_ != nullptr) << " Trying to reference a null Integer";
    return (*this)->value;
  }
  // comparators
  Bool operator==(int other) const {
    if (data_ == nullptr) return Bool(false);
    return Bool((*this)->value == other);
  }
  Bool operator!=(int other) const { return !(*this == other); }
  template <typename Enum, typename = typename std::enable_if<std::is_enum<Enum>::value>::type>
  Bool operator==(Enum other) const {
    return *this == static_cast<int>(other);
  }
  template <typename Enum, typename = typename std::enable_if<std::is_enum<Enum>::value>::type>
  Bool operator!=(Enum other) const {
    return *this != static_cast<int>(other);
  }
};
 
class RangeNode : public Object {
 public:
  PrimExpr min;
  PrimExpr extent;
  mutable Span span;
  RangeNode() {}
  RangeNode(PrimExpr min, PrimExpr extent, Span span = Span())
      : min(min), extent(extent), span(span) {}
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("min", &min);
    v->Visit("extent", &extent);
    v->Visit("span", &span);
  }
 
  bool SEqualReduce(const RangeNode* other, SEqualReducer equal) const {
    return equal(min, other->min) && equal(extent, other->extent);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(min);
    hash_reduce(extent);
  }
 
  static constexpr const char* _type_key = "Range";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  TVM_DECLARE_FINAL_OBJECT_INFO(RangeNode, Object);
};
 
class Range : public ObjectRef {
 public:
  TVM_DLL Range(PrimExpr begin, PrimExpr end, Span span = Span());
  static Range FromMinExtent(PrimExpr min, PrimExpr extent, Span span = Span());
  // declare range.
  TVM_DEFINE_OBJECT_REF_METHODS(Range, ObjectRef, RangeNode);
};
 
// implementations
inline const Type& RelaxExprNode::checked_type() const {
  ICHECK(checked_type_.defined()) << "internal error: the type checker has "
                                  << "not populated the checked_type "
                                  << "field for " << GetRef<RelaxExpr>(this);
  return this->checked_type_;
}
 
template <typename TTypeNode>
inline const TTypeNode* RelaxExprNode::type_as() const {
  static_assert(std::is_base_of<TypeNode, TTypeNode>::value,
                "TType must be a special case of type");
  ICHECK(checked_type_.defined())
      << "Type inference for this Expr has not completed. Try to call infer_type pass.";
  const TTypeNode* node = checked_type_.as<TTypeNode>();
  ICHECK(node != nullptr) << "Expected type to be " << TTypeNode::_type_key << ", but get "
                          << checked_type_->GetTypeKey();
  return node;
}
 
}  // namespace tvm
 
namespace tvm {
namespace runtime {
 
// Automatic conversion into IntImm, Integer, and Bool, when called
// through the FFI.  Automatic conversions into PrimExpr are
// registered in "tvm/tir/expr.h", as it includes conversions to the
// TIR-only StringImm.
//
// While the FFI only requires the From() method, these
// implementations also define a TryFrom() method to avoid duplicate
// logic in the PrimExpr conversion.
 
template <>
struct PackedFuncValueConverter<tvm::IntImm> {
  template <typename PODSubclass>
  static Optional<tvm::IntImm> TryFrom(const PODSubclass& val) {
    if (auto opt = val.TryAsInt()) {
      int64_t value = opt.value();
      auto dtype =
          (value > std::numeric_limits<int>::max() || value < std::numeric_limits<int>::min())
              ? DataType::Int(64)
              : DataType::Int(32);
      return IntImm(dtype, value);
    } else if (auto opt = val.TryAsBool()) {
      return IntImm(DataType::Int(32), opt.value());
    } else {
      return NullOpt;
    }
  }
 
  template <typename PODSubclass>
  static tvm::IntImm From(const PODSubclass& val) {
    if (auto opt = TryFrom(val)) {
      return opt.value();
    } else {
      return val.template AsObjectRef<tvm::IntImm>();
    }
  }
};
 
template <>
struct PackedFuncValueConverter<tvm::Integer> {
  template <typename PODSubclass>
  static tvm::Integer From(const PODSubclass& val) {
    if (auto opt = PackedFuncValueConverter<tvm::IntImm>::TryFrom(val)) {
      return Integer(opt.value());
    } else {
      return val.template AsObjectRef<tvm::Integer>();
    }
  }
};
 
template <>
struct PackedFuncValueConverter<tvm::Bool> {
  template <typename PODSubclass>
  static Optional<tvm::Bool> TryFrom(const PODSubclass& val) {
    if (auto opt = val.TryAsBool()) {
      return tvm::Bool(opt.value());
    } else if (auto opt = val.TryAsInt()) {
      int value = opt.value();
      ICHECK(value == 0 || value == 1)
          << "ValueError: boolean value can only be 0 or 1, but get " << value;
      return tvm::Bool(static_cast<bool>(value));
    } else {
      return NullOpt;
    }
  }
 
  template <typename PODSubclass>
  static tvm::Bool From(const PODSubclass& val) {
    if (auto opt = TryFrom(val)) {
      return opt.value();
    } else {
      return val.template AsObjectRef<tvm::Bool>();
    }
  }
};
 
template <>
struct PackedFuncValueConverter<tvm::FloatImm> {
  static Optional<tvm::FloatImm> TryFrom(const TVMPODValue_& val) {
    if (auto opt = val.TryAsFloat()) {
      return FloatImm(runtime::DataType::Float(32), opt.value());
    } else {
      return NullOpt;
    }
  }
 
  template <typename PODSubclass>
  static tvm::FloatImm From(const PODSubclass& val) {
    if (auto opt = TryFrom(val)) {
      return opt.value();
    } else {
      return val.template AsObjectRef<tvm::FloatImm>();
    }
  }
};
 
/* \brief Backwards compatibility wrapper for IntImm arguments
 *
 * In previous versions of TVM, IntImm was the default FFI type for
 * integer arguments, instead of runtime::Int.  For backwards
 * compatibility where the callee has been updated to expected a
 * runtime::Int, the caller has not been updated to provide a
 * runtime::Int, and the auto-unboxing of
 * runtime::Int does not apply (e.g. making an `Array<runtime::Int>`),
 * allow the IntImm to be generated.
 */
template <>
struct PackedFuncValueConverter<runtime::Int> {
  template <typename PODSubclass>
  static runtime::Int From(const PODSubclass& val) {
    if (val.template IsObjectRef<tvm::IntImm>()) {
      return runtime::Int(val.template AsObjectRef<tvm::IntImm>()->value);
    } else {
      return val.template AsObjectRef<runtime::Int>();
    }
  }
};
 
}  // namespace runtime
}  // namespace tvm
 
/* \brief Allow tvm.GLobalVar as key in STL tables
 *
 * For most IR expressions, it would be ambiguous whether the
 * expression should follow reference equality or structural equality.
 * This is not the case for variables, which do not contain nested
 * internal structure, and are frequently used as keys in lookup
 * tables.
 *
 * Providing `std::hash` and `std::equal_to` specializations for
 * `tvm::GlobalVar` allows it to be used as a key in STL tables.  For
 * other IR expressions, the user must specify the type of equality
 * used (e.g. `std::unordered_set<T, StructuralHash, StructuralEqual>`
 * or `std::unordered_set<T, ObjectPtrHash, ObjectPtrEqual>`).
 */
template <>
struct std::hash<tvm::GlobalVar> {
  std::size_t operator()(const tvm::GlobalVar& var) const {
    return tvm::runtime::ObjectPtrHash()(var);
  }
};
 
template <>
struct std::equal_to<tvm::GlobalVar> {
  bool operator()(const tvm::GlobalVar& var_a, const tvm::GlobalVar& var_b) const {
    return tvm::runtime::ObjectPtrEqual()(var_a, var_b);
  }
};
#endif  // TVM_IR_EXPR_H_