chore: import upstream snapshot with attribution

include/tvm/ir/attr_registry_map.h

@@ -0,0 +1,144 @@
+/*
+ * 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.
+ */
+/*!
+ * \file tvm/ir/attr_registry_map.h
+ * \brief Attribute map used in registry.
+ */
+#ifndef TVM_IR_ATTR_REGISTRY_MAP_H_
+#define TVM_IR_ATTR_REGISTRY_MAP_H_
+
+#include <tvm/ffi/string.h>
+
+#include <utility>
+#include <vector>
+
+namespace tvm {
+
+/*!
+ * \brief Generic attribute map.
+ * \tparam KeyType the type of the key.
+ */
+template <typename KeyType>
+class AttrRegistryMapContainerMap {
+ public:
+  /*!
+   * \brief Check if the map has key.
+   * \param key The key to the map
+   * \return 1 if key is contained in map, 0 otherwise.
+   */
+  int count(const KeyType& key) const {
+    if (key.defined()) {
+      const uint32_t idx = key->AttrRegistryIndex();
+      return idx < data_.size() ? (data_[idx].second != 0) : 0;
+    } else {
+      return 0;
+    }
+  }
+  /*!
+   * \brief get the corresponding value element at key.
+   * \param key The key to the map
+   * \return the const reference to the content value.
+   */
+  const ffi::Any& operator[](const KeyType& key) const {
+    TVM_FFI_ICHECK(key.defined());
+    const uint32_t idx = key->AttrRegistryIndex();
+    TVM_FFI_ICHECK(idx < data_.size() && data_[idx].second != 0)
+        << "Attribute " << attr_name_ << " has not been registered for " << key->name;
+    return data_[idx].first;
+  }
+  /*!
+   * \brief get the corresponding value element at key with default value.
+   * \param key The key to the map
+   * \param def_value The default value when the key does not exist.
+   * \return the const reference to the content value.
+   * \tparam ValueType The content value type.
+   */
+  template <typename ValueType>
+  ValueType get(const KeyType& key, ValueType def_value) const {
+    TVM_FFI_ICHECK(key.defined());
+    const uint32_t idx = key->AttrRegistryIndex();
+    if (idx < data_.size() && data_[idx].second != 0) {
+      if constexpr (std::is_same_v<ValueType, ffi::Any>) {
+        return data_[idx].first;
+      } else {
+        return data_[idx].first.template cast<ValueType>();
+      }
+    } else {
+      return def_value;
+    }
+  }
+
+ private:
+  /*! \brief The name of the attr field */
+  ffi::String attr_name_;
+  /*! \brief The internal data. */
+  std::vector<std::pair<ffi::Any, int>> data_;
+  /*! \brief The constructor */
+  AttrRegistryMapContainerMap() = default;
+  template <typename, typename>
+  friend class AttrRegistry;
+  friend class OpRegEntry;
+};
+
+/*!
+ * \brief ffi::Map<Key, ValueType> used to store meta-data.
+ * \tparam KeyType The type of the key
+ * \tparam ValueType The type of the value stored in map.
+ */
+template <typename KeyType, typename ValueType>
+class AttrRegistryMap {
+ public:
+  /*!
+   * \brief constructor
+   * \param map The internal map.
+   */
+  explicit AttrRegistryMap(const AttrRegistryMapContainerMap<KeyType>& map) : map_(map) {}
+  /*!
+   * \brief Check if the map has op as key.
+   * \param key The key to the map
+   * \return 1 if op is contained in map, 0 otherwise.
+   */
+  int count(const KeyType& key) const { return map_.count(key); }
+  /*!
+   * \brief get the corresponding value element at key.
+   * \param key The key to the map
+   * \return the const reference to the content value.
+   */
+  ValueType operator[](const KeyType& key) const {
+    if constexpr (std::is_same_v<ValueType, ffi::Any>) {
+      return map_[key];
+    } else {
+      return map_[key].template cast<ValueType>();
+    }
+  }
+  /*!
+   * \brief get the corresponding value element at key with default value.
+   * \param key The key to the map
+   * \param def_value The default value when the key does not exist.
+   * \return the const reference to the content value.
+   */
+  ValueType get(const KeyType& key, ValueType def_value) const { return map_.get(key, def_value); }
+
+ protected:
+  /*! \brief The internal map field */
+  const AttrRegistryMapContainerMap<KeyType>& map_;
+};
+
+}  // namespace tvm
+#endif  // TVM_IR_ATTR_REGISTRY_MAP_H_

include/tvm/ir/attrs.h

@@ -0,0 +1,325 @@
+/*
+ * 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.
+ */
+/*!
+ * \file tvm/ir/attrs.h
+ * \brief Helpers for attribute objects.
+ *
+ *  This module enables declaration of named attributes
+ *  which support default value setup and bound checking.
+ *
+ * \sa AttrsNode
+ */
+#ifndef TVM_IR_ATTRS_H_
+#define TVM_IR_ATTRS_H_
+
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ffi/extra/structural_hash.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/cow.h>
+
+#include <string>
+#include <type_traits>
+#include <unordered_map>
+#include <utility>
+
+namespace tvm {
+
+/*!
+ * \brief Base class of all attribute class
+ * \sa Attrs
+ */
+class AttrsNode : public ffi::Object {
+ public:
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+  TVM_FFI_DECLARE_OBJECT_INFO("ir.Attrs", AttrsNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference to AttrsNode.
+ * \sa AttrsNode
+ */
+class Attrs : public ffi::ObjectRef {
+ public:
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Attrs, ffi::ObjectRef, AttrsNode);
+};
+
+/*!
+ * \brief Specialized attribute type that is backed by a map.
+ *  The DictAttrsNode implements the Attrs behavior,
+ *  its fields are directly accessible via object.field_name
+ *  like other normal nodes.
+ */
+class DictAttrsNode : public AttrsNode {
+ public:
+  /*! \brief internal attrs map */
+  ffi::Map<ffi::String, ffi::Any> dict;
+
+  static void RegisterReflection() {
+    namespace rfl = ffi::reflection;
+    rfl::ObjectDef<DictAttrsNode>().def_ro("__dict__", &DictAttrsNode::dict);
+  }
+
+  // type info
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DictAttrs", DictAttrsNode, AttrsNode);
+};
+
+/*!
+ * \brief Managed reference to DictAttrsNode
+ * \sa DictAttrsNode.
+ *
+ * \note DictAttrs is NOTNULLABLE: every instance must hold a backing
+ *       DictAttrsNode. The class enforces this end-to-end by:
+ *       - the default constructor (no args) allocating an empty backing,
+ *       - the copy/move ctors and assignments leaving the moved-from
+ *         instance in a defined-but-empty state rather than null,
+ *       - the FFI type traits rejecting None at deserialization boundaries
+ *         (since `_type_is_nullable == false`), and
+ *       - the FFI lambda for ``ir.IRModule`` explicitly normalizing a
+ *         missing/None attrs argument to ``DictAttrs()`` before forwarding
+ *         to the C++ constructor.
+ *       Callers (including third-party code via templates like ``WithAttr``)
+ *       can therefore rely on ``attrs->dict`` being safe to dereference
+ *       without a ``.defined()`` guard.
+ */
+class DictAttrs : public Attrs {
+ public:
+  /*!
+   * \brief Construct a DictAttrs backed by DictAttrsNode.
+   *
+   * The no-argument form constructs an empty (but always defined) DictAttrs.
+   * \param dict The attributes.
+   */
+  explicit DictAttrs(ffi::Map<ffi::String, Any> dict = {}) {
+    ffi::ObjectPtr<DictAttrsNode> n = ffi::make_object<DictAttrsNode>();
+    n->dict = std::move(dict);
+    data_ = std::move(n);
+  }
+
+  /*!
+   * \brief Move constructor that leaves the source in a defined-but-empty
+   *        state rather than null, preserving the NOTNULLABLE invariant
+   *        even after `std::move`.
+   */
+  DictAttrs(DictAttrs&& other) noexcept : Attrs(ffi::UnsafeInit{}) {
+    data_ = std::move(other.data_);
+    other.data_ = ffi::make_object<DictAttrsNode>();
+  }
+
+  /*!
+   * \brief Move assignment that leaves the source in a defined-but-empty
+   *        state rather than null, preserving the NOTNULLABLE invariant
+   *        even after `std::move`.
+   */
+  DictAttrs& operator=(DictAttrs&& other) noexcept {
+    if (this != &other) {
+      data_ = std::move(other.data_);
+      other.data_ = ffi::make_object<DictAttrsNode>();
+    }
+    return *this;
+  }
+
+  // Explicit copy ctor/assign defaults. Declaring the move members above
+  // would otherwise suppress the implicit copy members.
+  DictAttrs(const DictAttrs& other) = default;
+  DictAttrs& operator=(const DictAttrs& other) = default;
+
+  // Utils for accessing attributes
+  /*!
+   * \brief Get a function attribute.
+   *
+   * \param attr_key The attribute key.
+   * \param default_value The default value if the key does not exist, defaults to nullptr.
+   *
+   * \return The result
+   *
+   * \tparam TOBjectRef the expected object type.
+   * \throw Error if the key exists but the value does not match TObjectRef
+   *
+   * \code
+   *
+   *  void GetAttrExample(const BaseFunc& f) {
+   *    auto value = f->attrs.GetAttr<int64_t>("AttrKey", 0);
+   *  }
+   *
+   * \endcode
+   */
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetAttr(
+      const std::string& attr_key,
+      ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
+    const DictAttrsNode* node = get();
+    auto it = node->dict.find(attr_key);
+    if (it != node->dict.end()) {
+      return (*it).second.cast<TObjectRef>();
+    } else {
+      return default_value;
+    }
+  }
+  // variant that uses TObjectRef to enable implicit conversion to default value.
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
+    return GetAttr<TObjectRef>(attr_key, ffi::Optional<TObjectRef>(default_value));
+  }
+  /*!
+   * \brief Check whether the function has an non-zero integer attr.
+   *
+   * This function can be used to check whether an optional
+   * attribute mark(e.g. inline) exists.
+   *
+   * \param attr_key The key to the attribute.
+   * \return The check result.
+   *
+   * \code
+   *
+   *  void HasNonzeroAttrExample(const BaseFunc& f) {
+   *    if (f->HasNonzeroAttr(attr::kInline)) {
+   *      // inline the function.
+   *    }
+   *  }
+   *
+   * \endcode
+   */
+  bool HasNonzeroAttr(const std::string& attr_key) const {
+    return GetAttr<int64_t>(attr_key, 0).value_or(0) != 0;
+  }
+
+  // Inline-expand TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE here, minus
+  // the default copy/move it normally injects (we define our own move members
+  // above so the moved-from instance stays defined-but-empty).
+  explicit DictAttrs(::tvm::ffi::UnsafeInit tag) : Attrs(tag) {}
+  using __PtrType =
+      std::conditional_t<DictAttrsNode::_type_mutable, DictAttrsNode*, const DictAttrsNode*>;
+  __PtrType operator->() const { return static_cast<__PtrType>(data_.get()); }
+  __PtrType get() const { return static_cast<__PtrType>(data_.get()); }
+  static constexpr bool _type_is_nullable = false;
+  using ContainerType = DictAttrsNode;
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(DictAttrsNode);
+};
+
+/*!
+ * \brief Copy the function or module, but overrides
+ *        the attribute value key with the value.
+ *
+ * \param input The thing to annotate (BaseFunc or IRModule)
+ * \param attr_key The attribute key.
+ * \param attr_value The value attribute value.
+ *
+ * \tparam TFunc The corresponding function or module type.
+ *
+ * \returns The new function or module with updated attributes.
+ *
+ * \note This function performs copy on write optimization for func and module.
+ *       If we move a uniquely referenced func or module into WithAttr,
+ *       then no additional copy will be performed.
+ *
+ *       This is also why we make it as a function instead of a member function
+ *       and why we pass by value in the first argument.
+ *
+ * \code
+ *
+ *  // Recommended way to trigger copy on write
+ *  func = WithAttr(std::move(func), "key1", value1);
+ *  func = WithAttr(std::move(func), "key2", value2);
+ *
+ * \endcode
+ */
+template <typename TFunc>
+inline TFunc WithAttr(TFunc input, const std::string& attr_key, Any attr_value) {
+  using TNode = typename TFunc::ContainerType;
+  static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
+  TNode* node = input.CopyOnWrite();
+  // node->attrs is NOTNULLABLE by contract, but defend against a caller
+  // that left a moved-from DictAttrs in place by re-initializing here.
+  if (!node->attrs.defined()) node->attrs = DictAttrs();
+  node->attrs.CopyOnWrite()->dict.Set(attr_key, std::move(attr_value));
+  return input;
+}
+
+/*!
+ * \brief Copy the function or module, but overrides the attributes with the entries from \p attrs.
+ *
+ * \param input The thing to annotate (BaseFunc or IRModule)
+ * \param attrs Key/values attributes to add to \p input.
+ *
+ * \tparam TFunc The corresponding function or module type.
+ *
+ * \returns The new function or module with updated attributes.
+ */
+template <typename TFunc>
+inline TFunc WithAttrs(TFunc input, ffi::Map<ffi::String, Any> attrs) {
+  using TNode = typename TFunc::ContainerType;
+  static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
+  if (attrs.empty()) return input;
+  TNode* node = input.CopyOnWrite();
+  // node->attrs is NOTNULLABLE by contract, but defend against a caller
+  // that left a moved-from DictAttrs in place by re-initializing here.
+  if (!node->attrs.defined()) node->attrs = DictAttrs();
+  auto* dict_node = node->attrs.CopyOnWrite();
+  for (const auto& [k, v] : attrs) {
+    dict_node->dict.Set(k, v);
+  }
+  return input;
+}
+
+/*!
+ * \brief Copy the function or module, but removes the specified
+ *        attribute.
+ *
+ * \param input The thing to annotate (BaseFunc or IRModule)
+ * \param attr_key The attribute key.
+ *
+ * \tparam TFunc The corresponding function or module type.
+ *
+ * \returns The new function or module with removed attribute.
+ *
+ * \note This function performs copy on write optimization for func and module.
+ *       If we move a uniquely referenced func or module into WithoutAttr,
+ *       then no additional copy will be performed.
+ *
+ *       This is also why we make it as a function instead of a member function
+ *       and why we pass by value in the first argument.
+ *
+ * \code
+ *
+ *  // Recommended way to trigger copy on write
+ *  func = WithoutAttr(std::move(func), "key1");
+ *  func = WithoutAttr(std::move(func), "key2");
+ *
+ * \endcode
+ */
+template <typename TFunc>
+inline TFunc WithoutAttr(TFunc input, const std::string& attr_key) {
+  using TNode = typename TFunc::ContainerType;
+  static_assert(TNode::_type_final, "Can only operate on the leaf nodes");
+  TNode* node = input.CopyOnWrite();
+  // node->attrs is NOTNULLABLE by contract, but defend against a caller
+  // that left a moved-from DictAttrs in place; nothing to erase from an
+  // empty dict.
+  if (!node->attrs.defined()) {
+    node->attrs = DictAttrs();
+    return input;
+  }
+  node->attrs.CopyOnWrite()->dict.erase(attr_key);
+  return input;
+}
+
+}  // namespace tvm
+#endif  // TVM_IR_ATTRS_H_

include/tvm/ir/base_expr.h

@@ -0,0 +1,521 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/base_expr.h
+ * \brief Base expression and primitive type nodes.
+ */
+#ifndef TVM_IR_BASE_EXPR_H_
+#define TVM_IR_BASE_EXPR_H_
+
+#include <tvm/ffi/cast.h>
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/string.h>
+#include <tvm/ir/source_map.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+#include <type_traits>
+
+namespace tvm {
+
+/*!
+ * \brief Type is the base type of all types.
+ *
+ * TVM's type system contains following subclasses:
+ *
+ * - PrimType: type of primitive type values used in the low-level IR.
+ * - FuncType: type of a function.
+ * - TensorType: type of certain Tensor values in the expression.
+ *
+ * There are also advanced types to support generic(polymorphic types).
+ * \sa Type
+ */
+class TypeNode : public ffi::Object {
+ public:
+  /*!
+   * \brief Span that points to the original source code.
+   *        Reserved debug information.
+   */
+  mutable Span span;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    // span do not participate in structural equal and hash.
+    refl::ObjectDef<TypeNode>().def_ro("span", &TypeNode::span, refl::DefaultValue(Span()),
+                                       refl::AttachFieldFlag::SEqHashIgnore());
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+
+  static constexpr const uint32_t _type_child_slots = 14;
+  TVM_FFI_DECLARE_OBJECT_INFO("ir.Type", TypeNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference to TypeNode.
+ * \sa TypeNode
+ */
+class Type : public ffi::ObjectRef {
+ public:
+  /*! \brief Sentinel for a type that has not been populated yet. */
+  TVM_DLL static Type Missing();
+
+  /*! \return whether this is the missing-type sentinel. */
+  TVM_DLL bool IsMissing() const;
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Type, ffi::ObjectRef, TypeNode);
+};
+
+/*!
+ * \brief Primitive data types used in the low-level IR.
+ *
+ * PrimType represents primitive POD values and the void sentinel.
+ *
+ * \sa PrimType
+ */
+class PrimTypeNode final : public TypeNode {
+ public:
+  /*!
+   * \brief The raw DLPack dtype represented by this primitive type.
+   */
+  DLDataType dtype;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<PrimTypeNode>().def_ro("dtype", &PrimTypeNode::dtype);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.PrimType", PrimTypeNode, TypeNode);
+};
+
+/*
+ * \brief Managed reference to PrimTypeNode.
+ * \sa PrimTypeNode
+ */
+class PrimType final : public Type {
+ public:
+  /*!
+   * \brief Construct from a raw DLPack dtype.
+   * \param dtype The corresponding DLPack dtype.
+   */
+  TVM_DLL explicit PrimType(DLDataType dtype);
+
+  /*!
+   * \brief Construct from DLPack dtype fields.
+   * \param code The DLPack dtype code.
+   * \param bits The scalar bit width.
+   * \param lanes The fixed lane count.
+   */
+  TVM_DLL PrimType(DLDataTypeCode code, int bits, int lanes = 1);
+
+  /*! \brief Construct a signed integer type with fixed lanes. */
+  TVM_DLL static PrimType Int(int bits, int lanes = 1);
+  /*! \brief Construct an unsigned integer type with fixed lanes. */
+  TVM_DLL static PrimType UInt(int bits, int lanes = 1);
+  /*! \brief Construct a floating-point type with fixed lanes. */
+  TVM_DLL static PrimType Float(int bits, int lanes = 1);
+  /*! \brief Construct a bfloat type with fixed lanes. */
+  TVM_DLL static PrimType BFloat(int bits, int lanes = 1);
+  /*! \brief Construct a boolean type with fixed lanes. */
+  TVM_DLL static PrimType Bool(int lanes = 1);
+  /*! \brief Construct the void sentinel type, encoded as handle(0, 0). */
+  TVM_DLL static PrimType Void();
+  /*!
+   * \brief Construct a scalable vector type.
+   * \param code The DLPack dtype code.
+   * \param bits The scalar bit width.
+   * \param lanes The positive vscale factor to encode in the DLPack lane field.
+   */
+  TVM_DLL static PrimType ScalableVector(DLDataTypeCode code, int bits, int lanes);
+
+  /*! \return The DLPack dtype code. */
+  TVM_FFI_INLINE DLDataTypeCode code() const {
+    return static_cast<DLDataTypeCode>(static_cast<int>(get()->dtype.code));
+  }
+
+  /*! \return The scalar bit width. */
+  TVM_FFI_INLINE int32_t bits() const { return get()->dtype.bits; }
+
+  /*!
+   * \return The fixed lane count.
+   * \note Throws on scalable vector types, where the encoded lane field stores a vscale factor.
+   */
+  TVM_FFI_INLINE int32_t lanes() const {
+    int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
+    if (TVM_FFI_PREDICT_FALSE(encoded_lanes < 0)) {
+      TVM_FFI_THROW(InternalError)
+          << "Can't fetch the lanes of a scalable vector at a compile time.";
+    }
+    return encoded_lanes;
+  }
+
+  /*!
+   * \brief Check the scalar element code and bit width.
+   * \note Lane count and scalable-vector encoding are intentionally ignored.
+   */
+  TVM_FFI_INLINE bool MatchesElementType(DLDataTypeCode code, int bits) const {
+    DLDataType dtype = get()->dtype;
+    return dtype.code == static_cast<uint8_t>(code) && dtype.bits == bits;
+  }
+
+  /*!
+   * \brief Check whether the dtype code matches any of the provided DLPack codes.
+   * \note Bit width and lanes are intentionally ignored.
+   */
+  template <typename... Codes>
+  TVM_FFI_INLINE bool MatchesCode(Codes... codes) const {
+    uint8_t dtype_code = get()->dtype.code;
+    return ((dtype_code == static_cast<uint8_t>(codes)) || ...);
+  }
+
+  /*! \brief Whether this type is a scalar, excluding fixed and scalable vectors. */
+  TVM_FFI_INLINE bool IsScalar() const {
+    int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
+    return encoded_lanes == 1;
+  }
+
+  /*! \brief Whether this type is the void sentinel `handle(0, 0)`. */
+  TVM_FFI_INLINE bool IsVoid() const {
+    DLDataType dtype = get()->dtype;
+    return dtype.code == static_cast<uint8_t>(DLDataTypeCode::kDLOpaqueHandle) && dtype.bits == 0 &&
+           static_cast<int16_t>(dtype.lanes) == 0;
+  }
+
+  /*! \brief Whether this type is a scalable vector. */
+  TVM_FFI_INLINE bool IsScalableVector() const {
+    return static_cast<int16_t>(get()->dtype.lanes) < -1;
+  }
+
+  /*! \brief Whether this type is a fixed-length vector. */
+  TVM_FFI_INLINE bool IsFixedLengthVector() const {
+    return static_cast<int16_t>(get()->dtype.lanes) > 1;
+  }
+
+  /*!
+   * \brief Return the number of bytes needed to store one value of this type.
+   *
+   * This uses the same packed sub-byte dtype sizing rule as runtime tensors.
+   * Scalable vector types have no compile-time storage size and are rejected.
+   */
+  TVM_FFI_INLINE size_t StorageBytes() const {
+    DLDataType dtype = get()->dtype;
+    int16_t encoded_lanes = static_cast<int16_t>(dtype.lanes);
+    if (TVM_FFI_PREDICT_FALSE(encoded_lanes < 0)) {
+      TVM_FFI_THROW(InternalError)
+          << "Cannot compute compile-time storage bytes for non-fixed vector type " << dtype;
+    }
+    return static_cast<size_t>(
+        (static_cast<uint64_t>(dtype.bits) * static_cast<uint64_t>(dtype.lanes) + 7) / 8);
+  }
+
+  /*! \brief Return the same type with a different dtype code, preserving bits and lanes. */
+  TVM_FFI_INLINE PrimType WithCode(DLDataTypeCode code) const {
+    DLDataType dtype = get()->dtype;
+    int16_t encoded_lanes = static_cast<int16_t>(dtype.lanes);
+    if (encoded_lanes < -1) {
+      return ScalableVector(code, dtype.bits, -encoded_lanes);
+    }
+    return PrimType(code, dtype.bits, encoded_lanes);
+  }
+
+  /*! \brief Return the same type with a different scalar bit width, preserving code and lanes. */
+  TVM_FFI_INLINE PrimType WithBits(int bits) const {
+    DLDataType dtype = get()->dtype;
+    int16_t encoded_lanes = static_cast<int16_t>(dtype.lanes);
+    if (encoded_lanes < -1) {
+      return ScalableVector(this->code(), bits, -encoded_lanes);
+    }
+    return PrimType(this->code(), bits, encoded_lanes);
+  }
+
+  /*! \brief Return the same scalar element type with a fixed lane count. */
+  TVM_FFI_INLINE PrimType WithLanes(int lanes) const {
+    return PrimType(this->code(), this->bits(), lanes);
+  }
+
+  /*! \return The vscale factor encoded in a scalable vector type. */
+  TVM_FFI_INLINE int32_t VScaleFactor() const {
+    int16_t encoded_lanes = static_cast<int16_t>(get()->dtype.lanes);
+    if (encoded_lanes >= -1) {
+      TVM_FFI_THROW(InternalError) << "A fixed length vector doesn't have a vscale factor.";
+    }
+    return -encoded_lanes;
+  }
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimType, Type, PrimTypeNode);
+};
+
+inline bool operator==(const PrimType& lhs, const PrimType& rhs) {
+  return lhs->dtype == rhs->dtype;
+}
+
+inline bool operator!=(const PrimType& lhs, const PrimType& rhs) { return !(lhs == rhs); }
+
+/*!
+ * \brief Base type of all the expressions.
+ * \sa Expr
+ */
+class ExprNode : public ffi::Object {
+ public:
+  /*!
+   * \brief Span that points to the original source code.
+   *        Reserved debug information.
+   */
+  mutable Span span;
+
+  /*!
+   * \brief The deduced or annotated type of the expression.
+   *
+   * Type::Missing() denotes type information that will be populated by
+   * later analysis passes instead of expression constructors.
+   */
+  mutable Type ty = Type::Missing();
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    // span does not participate in structural equal and hash.
+    refl::ObjectDef<ExprNode>()
+        .def_ro("span", &ExprNode::span, refl::DefaultValue(Span()),
+                refl::AttachFieldFlag::SEqHashIgnore())
+        .def_ro("ty", &ExprNode::ty, refl::DefaultValue(Type::Missing()));
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+
+  static constexpr const uint32_t _type_child_slots = 64;
+  TVM_FFI_DECLARE_OBJECT_INFO("ir.Expr", ExprNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference to ExprNode.
+ * \sa ExprNode
+ */
+class Expr : public ffi::ObjectRef {
+ public:
+  // Expressions do not implicitly compare by object identity or address.  Callers must name
+  // whether they intend object identity, structural equality, or primitive symbolic comparison.
+  bool operator==(const Expr& other) const = delete;
+  bool operator!=(const Expr& other) const = delete;
+  bool operator<(const Expr& other) const = delete;
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Expr, ffi::ObjectRef, ExprNode);
+};
+
+class Call;
+
+/*!
+ * \brief Typed reference/view over an expression whose result type is a
+ * specific Type subtype.
+ * \tparam ExpectedType The expected expression result type.
+ */
+template <typename ExpectedType>
+class TypedExpr : public Expr {
+ public:
+  /*! \return the typed result of this expression. */
+  ExpectedType ty() const {
+    const auto* node = get();
+    TVM_FFI_DCHECK(node != nullptr);
+    const auto* ty_node = node->ExprNode::ty.template as<typename ExpectedType::ContainerType>();
+    TVM_FFI_DCHECK(ty_node != nullptr);
+    return ffi::GetRef<ExpectedType>(ty_node);
+  }
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TypedExpr, Expr, ExprNode);
+  static constexpr bool _type_container_is_exact = false;
+};
+
+/*!
+ * \brief Typed reference/view over any Expr whose `ExprNode::ty` is PrimType.
+ *
+ * PrimExpr is a type category rather than a dedicated runtime node category.
+ * It can contain intrinsic primitive nodes such as IntImmNode and FloatImmNode,
+ * or a general ExprNode such as CallNode, when that expression's `ty` field is
+ * a PrimType. This keeps primitive-only APIs explicit while allowing shared
+ * Expr nodes for cross-dialect values with richer result types when needed.
+ */
+class PrimExpr : public TypedExpr<PrimType> {
+ public:
+  using TypedExpr<PrimType>::ty;
+
+  /*!
+   * \brief Construct from a call after checking that its result type is
+   * PrimType.
+   * \param call The call to view as a primitive expression.
+   */
+  TVM_DLL PrimExpr(Call call);  // NOLINT(*)
+
+  /*!
+   * \brief construct from integer.
+   * \param value The value to be constructed.
+   */
+  TVM_DLL PrimExpr(int32_t value);  // NOLINT(*)
+  /*!
+   * \brief construct from float.
+   * \param value The value to be constructed.
+   */
+  TVM_DLL PrimExpr(float value);  // NOLINT(*)
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimExpr, TypedExpr<PrimType>, ExprNode);
+  static constexpr bool _type_container_is_exact = false;
+
+  /*!
+   * \brief construct from string to form a StringImm.
+   * \param value The value to be constructed.
+   */
+  TVM_DLL static PrimExpr ConvertFallbackValue(ffi::String value);  // NOLINT(*)
+};
+
+/*!
+ * \brief Base class for other IR constructs that can be converted to PrimExpr.
+ * This is useful for the FFI to convert the expressions to PrimExpr.
+ * \sa PrimExpr
+ */
+class PrimExprConvertibleNode : public ffi::Object {
+ public:
+  virtual ~PrimExprConvertibleNode() {}
+  virtual PrimExpr ToPrimExpr() const = 0;
+  TVM_FFI_DECLARE_OBJECT_INFO("ir.PrimExprConvertible", PrimExprConvertibleNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference to PrimExprConvertibleNode.
+ * \sa PrimExprConvertibleNode
+ */
+class PrimExprConvertible : public ffi::ObjectRef {
+ public:
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimExprConvertible, ffi::ObjectRef,
+                                             PrimExprConvertibleNode);
+};
+
+namespace ffi {
+template <>
+inline constexpr bool use_default_type_traits_v<PrimType> = false;
+
+template <>
+struct TypeTraits<PrimType> : public ObjectRefWithFallbackTraitsBase<PrimType, DLDataType> {
+  TVM_FFI_INLINE static PrimType ConvertFallbackValue(DLDataType dtype) { return PrimType(dtype); }
+};
+
+template <typename ExpectedType>
+inline constexpr bool use_default_type_traits_v<TypedExpr<ExpectedType>> = false;
+
+template <typename ExpectedType>
+struct TypeTraits<TypedExpr<ExpectedType>>
+    : public ObjectRefTypeTraitsBase<TypedExpr<ExpectedType>> {
+  using Base = ObjectRefTypeTraitsBase<TypedExpr<ExpectedType>>;
+  using Base::CopyFromAnyViewAfterCheck;
+  using Base::CopyToAnyView;
+  using Base::GetMismatchTypeInfo;
+  using Base::MoveFromAnyAfterCheck;
+  using Base::MoveToAny;
+  using Base::TypeSchema;
+  using Base::TypeStr;
+
+  TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
+    if (src->type_index == TypeIndex::kTVMFFINone) {
+      return TypedExpr<ExpectedType>::_type_is_nullable;
+    }
+    if (src->type_index < TypeIndex::kTVMFFIStaticObjectBegin ||
+        !details::IsObjectInstance<ExprNode>(src->type_index)) {
+      return false;
+    }
+    const auto* expr = static_cast<const ExprNode*>(
+        details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(src->v_obj).get());
+    return details::AnyUnsafe::CheckAnyStrict<ExpectedType>(expr->ty);
+  }
+
+  TVM_FFI_INLINE static std::optional<TypedExpr<ExpectedType>> TryCastFromAnyView(
+      const TVMFFIAny* src) {
+    if (CheckAnyStrict(src)) {
+      if (src->type_index == TypeIndex::kTVMFFINone) {
+        return details::ObjectUnsafe::ObjectRefFromObjectPtr<TypedExpr<ExpectedType>>(nullptr);
+      }
+      return details::ObjectUnsafe::ObjectRefFromObjectPtr<TypedExpr<ExpectedType>>(
+          details::ObjectUnsafe::ObjectPtrFromUnowned<ExprNode>(src->v_obj));
+    }
+    return std::nullopt;
+  }
+};
+
+template <>
+inline constexpr bool use_default_type_traits_v<PrimExpr> = false;
+
+template <typename ObjectRefType, typename ExpectedType, typename... FallbackTypes>
+struct TypedExprWithFallbackTraitsBase
+    : public ObjectRefWithFallbackTraitsBase<ObjectRefType, FallbackTypes...> {
+  using Base = ObjectRefWithFallbackTraitsBase<ObjectRefType, FallbackTypes...>;
+
+  TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
+    return TypeTraits<TypedExpr<ExpectedType>>::CheckAnyStrict(src);
+  }
+
+  TVM_FFI_INLINE static std::optional<ObjectRefType> TryCastFromAnyView(const TVMFFIAny* src) {
+    if (TypeTraits<TypedExpr<ExpectedType>>::TryCastFromAnyView(src)) {
+      return details::ObjectUnsafe::ObjectRefFromObjectPtr<ObjectRefType>(
+          details::ObjectUnsafe::ObjectPtrFromUnowned<ExprNode>(src->v_obj));
+    }
+    return Base::template TryFallbackTypes<FallbackTypes...>(src);
+  }
+};
+
+// define automatic conversion from bool, int64_t, double, ffi::String to PrimExpr
+// These functions are declared early to avoid circular dependency
+template <>
+struct TypeTraits<PrimExpr>
+    : public TypedExprWithFallbackTraitsBase<PrimExpr, PrimType, StrictBool, int64_t, double,
+                                             ffi::String, PrimExprConvertible> {
+  using Base = TypedExprWithFallbackTraitsBase<PrimExpr, PrimType, StrictBool, int64_t, double,
+                                               ffi::String, PrimExprConvertible>;
+  using Base::CheckAnyStrict;
+  using Base::CopyFromAnyViewAfterCheck;
+  using Base::CopyToAnyView;
+  using Base::GetMismatchTypeInfo;
+  using Base::MoveFromAnyAfterCheck;
+  using Base::MoveToAny;
+  using Base::TryCastFromAnyView;
+  using Base::TypeSchema;
+  using Base::TypeStr;
+
+  TVM_DLL static PrimExpr ConvertFallbackValue(StrictBool value);
+  TVM_DLL static PrimExpr ConvertFallbackValue(int64_t value);
+  TVM_DLL static PrimExpr ConvertFallbackValue(double value);
+  TVM_FFI_INLINE static PrimExpr ConvertFallbackValue(ffi::String value) {
+    return PrimExpr::ConvertFallbackValue(value);
+  }
+  TVM_FFI_INLINE static PrimExpr ConvertFallbackValue(PrimExprConvertible value) {
+    return value->ToPrimExpr();
+  }
+};
+
+template <>
+inline constexpr bool use_default_type_traits_v<Expr> = false;
+
+// Allow generic Expr arguments to use the primitive-literal conversions
+// already defined by PrimExpr.
+template <>
+struct TypeTraits<Expr> : public ObjectRefWithFallbackTraitsBase<Expr, PrimExpr> {
+  TVM_FFI_INLINE static Expr ConvertFallbackValue(PrimExpr value) { return value; }
+};
+}  // namespace ffi
+
+}  // namespace tvm
+
+#endif  // TVM_IR_BASE_EXPR_H_

include/tvm/ir/config_schema.h

@@ -0,0 +1,242 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/config_schema.h
+ * \brief Minimal schema for dynamic config canonicalization and validation.
+ *
+ * This utility is intended for dynamic map-like configs (e.g. Target options),
+ * where we still want type checking, optional defaulting, and canonicalization.
+ */
+#ifndef TVM_IR_CONFIG_SCHEMA_H_
+#define TVM_IR_CONFIG_SCHEMA_H_
+
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include <sstream>
+#include <string>
+#include <type_traits>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace ir {
+
+/*!
+ * \brief Dynamic config schema for map-like options.
+ *
+ * The schema supports:
+ * - Option declaration (`def_option<T>`)
+ * - Optional canonicalizer (`set_canonicalizer`)
+ * - Resolution (`Resolve`) that performs validation/defaulting, unknown-key policy,
+ *   and canonicalization (last).
+ */
+class ConfigSchema {
+ public:
+  using ConfigMap = ffi::Map<ffi::String, ffi::Any>;
+  using Canonicalizer = ffi::TypedFunction<ConfigMap(ConfigMap)>;
+
+  /*! \brief Schema entry for one declared option. */
+  struct OptionEntry {
+    /*! \brief Option key. */
+    ffi::String key;
+    /*! \brief Type string for this option. */
+    ffi::String type_str;
+    /*! \brief Per-option validator/coercer (Any -> Any). */
+    ffi::TypedFunction<ffi::Any(ffi::Any)> validator;
+    /*! \brief Whether this option has a default value or factory. */
+    bool has_default = false;
+    /*! \brief Whether the default is produced by a factory function. */
+    bool default_from_factory = false;
+    /*! \brief Default value, or factory function () -> Any when default_from_factory is true. */
+    ffi::Any default_value;
+  };
+
+  /*!
+   * \brief Declare a typed option.
+   *
+   * Validation/coercion is implicitly generated from `T`.
+   * Additional optional traits may be supplied (e.g. `refl::DefaultValue`, `const char*` doc).
+   *
+   * \tparam T The canonical value type of this option.
+   * \tparam Traits Optional metadata/traits.
+   * \param key Option key.
+   * \param traits Optional traits.
+   * \return Reference to `*this` for chaining.
+   */
+  template <typename T, typename... Traits>
+  ConfigSchema& def_option(const ffi::String& key, Traits&&... traits) {
+    std::string skey(key);
+    if (key_to_index_.count(skey)) {
+      TVM_FFI_THROW(ValueError) << "Duplicate config option key: '" << key << "'";
+    }
+    key_to_index_[skey] = options_.size();
+    options_.push_back(MakeEntry<T>(key, std::forward<Traits>(traits)...));
+    return *this;
+  }
+
+  /*! \brief Set whole-object canonicalizer. */
+  void set_canonicalizer(Canonicalizer f) { canonicalizer_ = std::move(f); }
+
+  /*! \brief Trait to set a custom validator for a config option. */
+  struct AttrValidator {
+    ffi::TypedFunction<ffi::Any(ffi::Any)> func;
+    explicit AttrValidator(ffi::TypedFunction<ffi::Any(ffi::Any)> f) : func(std::move(f)) {}
+  };
+
+  /*! \brief Set whether unknown keys trigger an error. */
+  void set_error_on_unknown(bool value) { error_on_unknown_ = value; }
+
+  /*!
+   * \brief Default/validate, then canonicalize a config object.
+   *
+   * Resolve flow:
+   * 1) Validate/coerce declared options in declaration order.
+   * 2) Materialize defaults and enforce required options.
+   * 3) Apply unknown-key policy.
+   * 4) Run canonicalizer as final step.
+   *
+   * \param config Input config object.
+   * \return Canonical validated config object.
+   * \throws ValueError/TypeError with option context.
+   */
+  ConfigMap Resolve(ConfigMap config) const {
+    ConfigMap result;
+
+    // Step 1: validate/coerce and materialize options in declaration order
+    for (const auto& e : options_) {
+      auto it = config.find(e.key);
+      if (it != config.end()) {
+        result.Set(e.key, e.validator((*it).second));
+      } else if (e.has_default) {
+        if (e.default_from_factory) {
+          result.Set(e.key, e.default_value.cast<ffi::Function>()());
+        } else {
+          result.Set(e.key, e.default_value);
+        }
+      }
+      // else: missing non-required option, stays absent
+    }
+
+    // Step 2: unknown-key policy
+    if (error_on_unknown_) {
+      for (const auto& kv : config) {
+        if (!key_to_index_.count(std::string(kv.first))) {
+          std::ostringstream os;
+          os << "Unknown config option '" << kv.first << "'. Known options: ";
+          bool first = true;
+          for (const auto& e : options_) {
+            if (!first) os << ", ";
+            os << "'" << e.key << "'";
+            first = false;
+          }
+          TVM_FFI_THROW(ValueError) << os.str();
+        }
+      }
+    }
+
+    // Step 3: whole-object canonicalization (last)
+    if (canonicalizer_ != nullptr) {
+      result = canonicalizer_(result);
+    }
+
+    return result;
+  }
+
+  /*!
+   * \brief List declared options in declaration order.
+   * \return Const reference to the option entries vector.
+   */
+  const std::vector<OptionEntry>& ListOptions() const { return options_; }
+
+  /*! \brief Check if an option with the given key exists. */
+  bool HasOption(const ffi::String& key) const { return key_to_index_.count(std::string(key)) > 0; }
+
+ private:
+  template <typename T>
+  static ffi::TypedFunction<ffi::Any(ffi::Any)> MakeValidator(const ffi::String& key) {
+    return ffi::TypedFunction<ffi::Any(ffi::Any)>([key](ffi::Any val) -> ffi::Any {
+      auto opt = val.try_cast<T>();
+      if (!opt.has_value()) {
+        TVM_FFI_THROW(TypeError) << "Option '" << key << "': expected type '"
+                                 << ffi::TypeTraits<T>::TypeStr() << "' but got '"
+                                 << val.GetTypeKey() << "'";
+      }
+      return ffi::Any(opt.value());
+    });
+  }
+
+  template <typename Trait>
+  static void ApplyTrait(OptionEntry* entry, ffi::reflection::FieldInfoBuilder* info,
+                         Trait&& trait) {
+    using T = std::decay_t<Trait>;
+    if constexpr (std::is_same_v<T, AttrValidator>) {
+      entry->validator = std::move(trait.func);
+    } else if constexpr (std::is_base_of_v<ffi::reflection::InfoTrait, T>) {
+      trait.Apply(info);
+    } else if constexpr (std::is_same_v<T, const char*> || std::is_same_v<T, char*>) {
+      const char* doc = trait;
+      if (doc != nullptr && doc[0] != '\0') {
+        info->doc = TVMFFIByteArray{doc, std::char_traits<char>::length(doc)};
+      }
+    }
+  }
+
+  template <typename T, typename... Traits>
+  OptionEntry MakeEntry(const ffi::String& key, Traits&&... traits) {
+    OptionEntry e;
+    e.key = key;
+    e.type_str = ffi::String(ffi::TypeTraits<T>::TypeStr());
+    e.validator = MakeValidator<T>(key);
+    // Apply traits through a temporary FieldInfoBuilder so existing
+    // reflection traits (notably refl::DefaultValue) are reused unchanged.
+    ffi::reflection::FieldInfoBuilder info{};
+    info.flags = 0;
+    info.default_value_or_factory = ffi::AnyView(nullptr).CopyToTVMFFIAny();
+    info.doc = TVMFFIByteArray{nullptr, 0};
+    (ApplyTrait(&e, &info, std::forward<Traits>(traits)), ...);
+    if (info.flags & kTVMFFIFieldFlagBitMaskHasDefault) {
+      e.has_default = true;
+      e.default_from_factory = (info.flags & kTVMFFIFieldFlagBitMaskDefaultFromFactory) != 0;
+      e.default_value = ffi::AnyView::CopyFromTVMFFIAny(info.default_value_or_factory);
+      // Release the extra ref created by CopyToTVMFFIAny in Apply
+      if (info.default_value_or_factory.type_index >= TVMFFITypeIndex::kTVMFFIStaticObjectBegin) {
+        ffi::details::ObjectUnsafe::DecRefObjectHandle(info.default_value_or_factory.v_obj);
+      }
+    }
+    return e;
+  }
+
+  /*! \brief Declared options in declaration order. */
+  std::vector<OptionEntry> options_;
+  /*! \brief Map from key string to index in options_. */
+  std::unordered_map<std::string, size_t> key_to_index_;
+  /*! \brief Optional whole-config canonicalizer. */
+  Canonicalizer canonicalizer_{nullptr};
+  /*! \brief Whether unknown keys trigger an error. */
+  bool error_on_unknown_ = true;
+};
+
+}  // namespace ir
+}  // namespace tvm
+
+#endif  // TVM_IR_CONFIG_SCHEMA_H_

include/tvm/ir/cow.h

@@ -0,0 +1,70 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/cow.h
+ * \brief Copy-on-write helper macro for IR ffi::ObjectRef types.
+ */
+#ifndef TVM_IR_COW_H_
+#define TVM_IR_COW_H_
+
+#include <tvm/ffi/object.h>
+
+#include <utility>
+
+namespace tvm {
+
+/*!
+ * \brief Define CopyOnWrite function in an ffi::ObjectRef.
+ * \param ObjectName The Type of the Node.
+ *
+ *  CopyOnWrite will generate a unique copy of the internal node.
+ *  The node will be copied if it is referenced by multiple places.
+ *  The function returns the raw pointer to the node to allow modification
+ *  of the content.
+ *
+ * \code
+ *
+ *  MyCOWObjectRef ref, ref2;
+ *  ref2 = ref;
+ *  ref.CopyOnWrite()->value = new_value;
+ *  assert(ref2->value == old_value);
+ *  assert(ref->value == new_value);
+ *
+ * \endcode
+ */
+#ifndef TVM_DEFINE_OBJECT_REF_COW_METHOD
+#define TVM_DEFINE_OBJECT_REF_COW_METHOD(ObjectName)                       \
+  static_assert(ObjectName::_type_final,                                   \
+                "TVM's CopyOnWrite may only be used for "                  \
+                "Object types that are declared as final, "                \
+                "using the TVM_FFI_DECLARE_OBJECT_INFO_FINAL macro.");     \
+  ObjectName* CopyOnWrite() {                                              \
+    TVM_FFI_ICHECK(data_ != nullptr);                                      \
+    if (!data_.unique()) {                                                 \
+      auto n = ::tvm::ffi::make_object<ObjectName>(*(operator->()));       \
+      ::tvm::ffi::ObjectPtr<::tvm::ffi::Object>(std::move(n)).swap(data_); \
+    }                                                                      \
+    return static_cast<ObjectName*>(data_.get());                          \
+  }
+#endif  // TVM_DEFINE_OBJECT_REF_COW_METHOD
+
+}  // namespace tvm
+
+#endif  // TVM_IR_COW_H_

include/tvm/ir/env_func.h

@@ -0,0 +1,162 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/env_func.h
+ * \brief Serializable global function used in IR.
+ */
+#ifndef TVM_IR_ENV_FUNC_H_
+#define TVM_IR_ENV_FUNC_H_
+
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/base.h>
+
+#include <string>
+#include <utility>
+
+namespace tvm {
+
+/*!
+ * \brief A serializable function backed by TVM's global environment.
+ *
+ * This is a wrapper to enable serializable global ffi::Function.
+ * An EnvFunc is saved by its name in the global registry
+ * under the assumption that the same function is registered during load.
+ * \sa EnvFunc
+ */
+class EnvFuncNode : public ffi::Object {
+ public:
+  /*! \brief Unique name of the global function */
+  ffi::String name;
+  /*! \brief The internal packed function */
+  ffi::Function func;
+  /*! \brief constructor */
+  EnvFuncNode() {}
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    // func do not participate in structural equal and hash.
+    refl::ObjectDef<EnvFuncNode>()
+        .def_ro("name", &EnvFuncNode::name)
+        .def_ro("func", &EnvFuncNode::func, refl::AttachFieldFlag::SEqHashIgnore());
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.EnvFunc", EnvFuncNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference to EnvFuncNode.
+ * \sa EnvFuncNode
+ */
+class EnvFunc : public ffi::ObjectRef {
+ public:
+  EnvFunc() {}
+  explicit EnvFunc(ffi::ObjectPtr<ffi::Object> n) : ffi::ObjectRef(n) {}
+  /*!
+   * \brief constructor with UnsafeInit
+   */
+  explicit EnvFunc(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
+  /*! \return The internal global function pointer */
+  const EnvFuncNode* operator->() const { return static_cast<const EnvFuncNode*>(get()); }
+  /*!
+   * \brief Invoke the function.
+   * \param args The arguments
+   * \returns The return value.
+   */
+  template <typename... Args>
+  ffi::Any operator()(Args&&... args) const {
+    const EnvFuncNode* n = operator->();
+    TVM_FFI_ICHECK(n != nullptr);
+    return n->func(std::forward<Args>(args)...);
+  }
+  /*!
+   * \brief Get a global function based on the name.
+   * \param name The name of the global function.
+   * \return The created global function.
+   * \note The function can be unique
+   */
+  TVM_DLL static EnvFunc Get(const ffi::String& name);
+  /*! \brief specify container node */
+  using ContainerType = EnvFuncNode;
+};
+
+/*!
+ * \brief Please refer to \ref TypedEnvFuncAnchor "TypedEnvFunc<R(Args..)>"
+ */
+template <typename FType>
+class TypedEnvFunc;
+
+/*!
+ * \anchor TypedEnvFuncAnchor
+ * \brief A typed version of EnvFunc.
+ * It is backed by a GlobalFuncNode internally.
+ *
+ * \tparam R The return value of the function.
+ * \tparam Args The argument signature of the function.
+ * \sa EnvFunc
+ */
+template <typename R, typename... Args>
+class TypedEnvFunc<R(Args...)> : public ffi::ObjectRef {
+ public:
+  /*! \brief short hand for this function type */
+  using TSelf = TypedEnvFunc<R(Args...)>;
+  TypedEnvFunc() {}
+  explicit TypedEnvFunc(ffi::ObjectPtr<ffi::Object> n) : ffi::ObjectRef(n) {}
+  /*!
+   * \brief constructor with UnsafeInit
+   */
+  explicit TypedEnvFunc(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
+  /*!
+   * \brief Assign global function to a TypedEnvFunc
+   * \param other Another global function.
+   * \return reference to self.
+   */
+  TSelf& operator=(const EnvFunc& other) {
+    ffi::ObjectRef::operator=(other);
+    return *this;
+  }
+  /*! \return The internal global function pointer */
+  const EnvFuncNode* operator->() const { return static_cast<const EnvFuncNode*>(get()); }
+  /*!
+   * \brief Invoke the function.
+   * \param args The arguments
+   * \returns The return value.
+   */
+  R operator()(Args... args) const {
+    const EnvFuncNode* n = operator->();
+    TVM_FFI_ICHECK(n != nullptr);
+    if constexpr (std::is_same_v<R, void>) {
+      n->func(std::forward<Args>(args)...);
+    } else {
+      ffi::Any res = n->func(std::forward<Args>(args)...);
+      if constexpr (std::is_same_v<R, ffi::Any>) {
+        return res;
+      } else {
+        return std::move(res).cast<R>();
+      }
+    }
+  }
+  /*! \brief specify container node */
+  using ContainerType = EnvFuncNode;
+};
+
+}  // namespace tvm
+#endif  // TVM_IR_ENV_FUNC_H_

include/tvm/ir/expr.h

@@ -0,0 +1,568 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/expr.h
+ * \brief Base expr nodes in TVM.
+ */
+#ifndef TVM_IR_EXPR_H_
+#define TVM_IR_EXPR_H_
+
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/dataclass.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/string.h>
+#include <tvm/ir/attrs.h>
+#include <tvm/ir/base_expr.h>
+#include <tvm/ir/cow.h>
+#include <tvm/ir/source_map.h>
+
+#include <algorithm>
+#include <functional>
+#include <limits>
+#include <optional>
+#include <string>
+#include <type_traits>
+
+namespace tvm {
+
+// Forward-declare VirtualDevice to avoid circular imports.
+class VirtualDevice;
+
+/*!
+ * \brief add operator
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator+(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief subtraction operator
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator-(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief negation.
+ *
+ * \param a input.
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator-(PrimExpr a);
+
+/*!
+ * \brief multiplication operator
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator*(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief division operator
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator/(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief left shift operator
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator<<(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief right shift operator
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator>>(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief greater
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator>(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief greater_equal
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator>=(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief less
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator<(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief less_equal
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator<=(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief equal
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator==(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief not_equal
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator!=(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief and
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note This operator does eager constant folding.
+ */
+TVM_DLL PrimExpr operator&&(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief or
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note This operator does eager constant folding.
+ */
+TVM_DLL PrimExpr operator||(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief not
+ *
+ * \param a left operand
+ * \return The result expression.
+ * \note This operator does eager constant folding.
+ */
+TVM_DLL PrimExpr operator!(PrimExpr a);
+
+/*!
+ * \brief take bitwise and of two values
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator&(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief take bitwise or of two values
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator|(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief take bitwise xor of two values
+ *
+ * \param a left operand
+ * \param b right operand
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator^(PrimExpr a, PrimExpr b);
+
+/*!
+ * \brief take bitwise negation of two values
+ *
+ * \param a the input expression.
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr operator~(PrimExpr a);
+
+class GlobalVar;
+/*!
+ * \brief Global variable that lives in the top-level module.
+ *
+ * A GlobalVar only refers to function definitions.
+ * This is used to enable recursive calls between function.
+ *
+ * \sa GlobalVarNode
+ */
+class GlobalVarNode : public ExprNode {
+ public:
+  /*! \brief The name of the variable, this only acts as a hint. */
+  ffi::String name_hint;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<GlobalVarNode>().def_ro("name_hint", &GlobalVarNode::name_hint);
+    // A GlobalVar identifies a module-level symbol.  Its type is derived from the
+    // corresponding function definition and is not part of the symbol identity.
+    refl::TypeAttrDef<GlobalVarNode>()
+        .def("__s_equal__", &GlobalVarNode::SEqual)
+        .def("__s_hash__", &GlobalVarNode::SHash);
+  }
+
+  bool SEqual(const GlobalVarNode* other,
+              ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const {
+    return equal(name_hint, other->name_hint, false, "name_hint");
+  }
+
+  int64_t SHash(int64_t init_hash, ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const {
+    return hash(name_hint, init_hash, false);
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.GlobalVar", GlobalVarNode, ExprNode);
+};
+
+/*!
+ * \brief Managed reference to GlobalVarNode.
+ * \sa GlobalVarNode
+ */
+class GlobalVar : public Expr {
+ public:
+  TVM_DLL explicit GlobalVar(ffi::String name_hint, Span span = {});
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GlobalVar, Expr, GlobalVarNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(GlobalVarNode);
+};
+
+/*!
+ * \brief Call corresponds to callable invocation.
+ */
+class CallNode : public ExprNode {
+ public:
+  /*!
+   * \brief The operator/function being invoked.
+   *
+   * It can be an Op, a GlobalVar, a local function value, or another callable
+   * expression.
+   */
+  Expr op;
+
+  /*! \brief The arguments of the call. */
+  ffi::Array<Expr> args;
+
+  /*! \brief The additional attributes. */
+  Attrs attrs;
+
+  /*! \brief The type information arguments passed to the callee. */
+  ffi::Array<Type> ty_args;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<CallNode>()
+        .def_ro("op", &CallNode::op)
+        .def_ro("args", &CallNode::args)
+        .def_ro("attrs", &CallNode::attrs)
+        .def_ro("ty_args", &CallNode::ty_args);
+  }
+
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Call", CallNode, ExprNode);
+};
+
+/*!
+ * \brief Managed reference to CallNode.
+ */
+class Call : public Expr {
+ public:
+  TVM_DLL Call(Type ret_ty, Expr op, ffi::Array<Expr> args, Attrs attrs = Attrs(),
+               ffi::Array<Type> ty_args = ffi::Array<Type>(), Span span = Span());
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Call, Expr, CallNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
+};
+
+/*!
+ * \brief Constant integer literals in the program.
+ * \sa IntImm
+ */
+class IntImmNode : public ExprNode {
+ public:
+  /*! \brief the Internal value. */
+  int64_t value;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<IntImmNode>().def_ro("value", &IntImmNode::value);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IntImm", IntImmNode, ExprNode);
+};
+
+/*!
+ * \brief Managed reference class to IntImmNode.
+ *
+ * \sa IntImmNode
+ */
+class IntImm : public PrimExpr {
+ public:
+  /*!
+   * \brief Constructor.
+   * \param value_ty The primitive type of the value.
+   * \param value The internal value.
+   * \param span The location of this object in the source code.
+   */
+  TVM_DLL IntImm(PrimType value_ty, int64_t value, Span span = Span());
+
+  /*!
+   * \brief Construct a scalar boolean constant.
+   * \param value The boolean value.
+   * \param span The location of this object in the source code.
+   */
+  static IntImm Bool(bool value, Span span = Span()) {
+    return IntImm(PrimType::Bool(), value, span);
+  }
+
+  /*!
+   * \brief Construct a scalar int32 constant.
+   * \param value The integer value.
+   * \param span The location of this object in the source code.
+   */
+  static IntImm Int32(int64_t value, Span span = Span()) {
+    return IntImm(PrimType::Int(32), value, span);
+  }
+
+  /*!
+   * \brief Construct a scalar int64 constant.
+   * \param value The integer value.
+   * \param span The location of this object in the source code.
+   */
+  static IntImm Int64(int64_t value, Span span = Span()) {
+    return IntImm(PrimType::Int(64), value, span);
+  }
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntImm, PrimExpr, IntImmNode);
+  static constexpr bool _type_container_is_exact = true;
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(IntImmNode);
+};
+
+/*!
+ * \brief Constant floating point literals in the program.
+ * \sa FloatImm
+ */
+class FloatImmNode : public ExprNode {
+ public:
+  /*! \brief The constant value content. */
+  double value;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<FloatImmNode>().def_ro("value", &FloatImmNode::value);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.FloatImm", FloatImmNode, ExprNode);
+};
+
+/*!
+ * \brief Managed reference class to FloatImmNode.
+ *
+ * \sa FloatImmNode
+ */
+class FloatImm : public PrimExpr {
+ public:
+  /*!
+   * \brief Constructor.
+   * \param value_ty The primitive type of the value.
+   * \param value The internal value.
+   * \param span The location in the source code.
+   */
+  TVM_DLL FloatImm(PrimType value_ty, double value, Span span = Span());
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(FloatImm, PrimExpr, FloatImmNode);
+  static constexpr bool _type_container_is_exact = true;
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(FloatImmNode);
+};
+
+/*! \brief range over one dimension */
+class RangeNode : public ffi::Object {
+ public:
+  /*! \brief beginning of the node */
+  PrimExpr min;
+  /*! \brief the extend of range */
+  PrimExpr extent;
+  /*! \brief the location of this range in the source */
+  mutable Span span;
+  /*! \brief constructor */
+  RangeNode() {}
+  RangeNode(PrimExpr min, PrimExpr extent, Span span = Span())
+      : min(min), extent(extent), span(span) {}
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<RangeNode>()
+        .def_ro("min", &RangeNode::min)
+        .def_ro("extent", &RangeNode::extent)
+        .def_ro("span", &RangeNode::span, refl::AttachFieldFlag::SEqHashIgnore());
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Range", RangeNode, ffi::Object);
+};
+
+/*! \brief Range container  */
+class Range : public ffi::ObjectRef {
+ public:
+  /*!
+   * \brief constructor by begin and end
+   * \param begin The begin of the range.
+   * \param end The end of the range.
+   * \param span The location of the Range in the source.
+   */
+  TVM_DLL Range(PrimExpr begin, PrimExpr end, Span span = Span());
+  /*!
+   * \brief construct a new range with min and extent
+   *  The corresponding constructor is removed,
+   *  because that is counter convention of tradition meaning
+   *  of range(begin, end)
+   *
+   * \param min The minimum range.
+   * \param extent The extent of the range.
+   * \param span The location of the Range in the source.
+   */
+  TVM_DLL static Range FromMinExtent(PrimExpr min, PrimExpr extent, Span span = Span());
+  // declare range.
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Range, ffi::ObjectRef, RangeNode);
+};
+
+namespace ffi {
+template <>
+inline constexpr bool object_ref_contains_v<PrimExpr, IntImmNode> = true;
+template <>
+inline constexpr bool object_ref_contains_v<PrimExpr, FloatImmNode> = true;
+
+// Type traits to enable automatic conversion into IntImm, Integer, and Bool
+// when called through the FFI
+template <>
+inline constexpr bool use_default_type_traits_v<IntImm> = false;
+
+// specialize to enable implicit conversion from const char*
+template <>
+struct TypeTraits<IntImm> : public ObjectRefWithFallbackTraitsBase<IntImm, int64_t> {
+  TVM_FFI_INLINE static IntImm ConvertFallbackValue(int64_t value) {
+    auto value_ty =
+        (value > std::numeric_limits<int>::max() || value < std::numeric_limits<int>::min())
+            ? PrimType::Int(64)
+            : PrimType::Int(32);
+    return IntImm(value_ty, value);
+  }
+};
+
+template <>
+inline constexpr bool use_default_type_traits_v<FloatImm> = false;
+
+template <>
+struct TypeTraits<FloatImm> : public ObjectRefWithFallbackTraitsBase<FloatImm, double> {
+  TVM_FFI_INLINE static FloatImm ConvertFallbackValue(double value) {
+    return FloatImm(PrimType::Float(32), value);
+  }
+};
+}  // namespace ffi
+}  // 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, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>`).
+ */
+template <>
+struct std::hash<tvm::GlobalVar> {
+  std::size_t operator()(const tvm::GlobalVar& var) const { return tvm::ffi::ObjectPtrHash()(var); }
+};
+
+template <>
+struct std::equal_to<tvm::GlobalVar> {
+  bool operator()(const tvm::GlobalVar& var_a, const tvm::GlobalVar& var_b) const {
+    return tvm::ffi::ObjectPtrEqual()(var_a, var_b);
+  }
+};
+#endif  // TVM_IR_EXPR_H_

include/tvm/ir/function.h

@@ -0,0 +1,256 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/function.h
+ * \brief Function nodes.
+ */
+#ifndef TVM_IR_FUNCTION_H_
+#define TVM_IR_FUNCTION_H_
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/string.h>
+#include <tvm/ir/attrs.h>
+#include <tvm/ir/expr.h>
+
+#include <string>
+#include <type_traits>
+
+namespace tvm {
+
+/*!
+ * \brief Possible Calling conventions.
+ *
+ *  NOTE: The calling convention also implies
+ *  the way we implement the function during lowering.
+ */
+enum class CallingConv : int {
+  /*!
+   * \brief Default calling convention.
+   *
+   * - Uses the native calling convention of the target.
+   * - Implementation: specified by the native target.
+   */
+  kDefault = 0,
+  /*!
+   * \brief ffi::Function that exposes a Cffi::Function signature.
+   *
+   * - Calling by ffi::Function calling convention.
+   * - Implementation: Expose a function with the Cffi::Function signature.
+   */
+  kCPackedFunc = 1,
+  /*!
+   * \brief Device kernel launch
+   *
+   * - Call by ffi::Function calling convention.
+   * - Implementation: defined by device runtime(e.g. runtime/cuda)
+   */
+  kDeviceKernelLaunch = 2,
+};
+
+/*!
+ * \brief Supported linkage types.
+ */
+enum class LinkageType : int {
+  /*!
+   * \brief Internal linkage.
+   */
+  kInternal = 0,
+  /*!
+   * \brief External linkage.
+   - Function with external linkage should have a global symbol attached to it.
+   */
+  kExternal = 1
+};
+
+/*!
+ * \brief Generic attribute names that can be attached to any function.
+ *
+ * \sa tvm::tirx::attr, tvm::relax::attr
+ */
+namespace attr {
+/*!
+ * \brief Indicates the special calling convention.
+ *
+ * Type: IntImm
+ *
+ * \sa tvm::CallingConv
+ */
+constexpr const char* kCallingConv = "calling_conv";
+
+/*!
+ * \brief Compilation target of the function.
+ *
+ * Type: Target
+ *
+ * \sa tvm::Target
+ */
+constexpr const char* kTarget = "target";
+
+/*!
+ * \brief Global linker symbol of the function in generated code.
+ *
+ *  This option forces the code generator to name the
+ *  function with the given.
+ *
+ *  For example, we could set a global_symbol of a function
+ *  early to make sure that we can always refer to it by
+ *  the symbol name in the generated DLL.
+ *
+ *  We should not set the attribute for local functions,
+ *  so that the compiler can freely rename them.
+ *
+ *  A unique global symbol will be automatically assigned
+ *  to each function in the module before the target code
+ *  generation phase.
+ *
+ * Type: String
+ */
+constexpr const char* kGlobalSymbol = "global_symbol";
+
+/*!
+ * \brief The function uses s_tir (apache-derived TIR) semantics:
+ *        parser fills layout=None, ScriptComplete wraps body in a root SBlock,
+ *        and printer emits `s_tir=True` on the decorator.
+ *        Default (attr absent or False) is tirx semantics.
+ *
+ * Type: IntImm (bool dtype)
+ */
+constexpr const char* kSTir = "s_tir";
+
+/*!
+ * \brief Number of inputs of the Primfunc
+ *
+ * Type: Int
+ */
+constexpr const char* kNumInputs = "num_inputs";
+
+}  // namespace attr
+
+/*!
+ * \brief Base node of all functions.
+ *
+ * We support several variants of functions throughout the stack.
+ * All of the functions share the same type system
+ * to support cross variant calls.
+ *
+ * \sa BaseFunc
+ */
+class BaseFuncNode : public ExprNode {
+ public:
+  /*! \brief Additional attributes storing the meta-data */
+  DictAttrs attrs;
+
+  /*!
+   * \brief Get a function attribute.
+   *
+   * \param attr_key The attribute key.
+   * \param default_value The default value if the key does not exist, defaults to nullptr.
+   *
+   * \return The result
+   *
+   * \tparam TOBjectRef the expected object type.
+   * \throw Error if the key exists but the value does not match TObjectRef
+   *
+   * \code
+   *
+   *  void GetAttrExample(const BaseFunc& f) {
+   *    auto value = f->GetAttr<int64_t>("AttrKey", 0);
+   *  }
+   *
+   * \endcode
+   */
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key,
+                                    ffi::Optional<TObjectRef> default_value = std::nullopt) const {
+    return attrs.GetAttr(attr_key, default_value);
+  }
+  // variant that uses TObjectRef to enable implicit conversion to default value.
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
+    return GetAttr<TObjectRef>(attr_key, ffi::Optional<TObjectRef>(default_value));
+  }
+
+  /*!
+   * \brief Check whether the function has an non-zero integer attr.
+   *
+   * This function can be used to check whether an optional
+   * attribute mark(e.g. inline) exists.
+   *
+   * \param attr_key The key to the attribute.
+   * \return The check result.
+   *
+   * \code
+   *
+   *  void HasNonzeroAttrExample(const BaseFunc& f) {
+   *    if (f->HasNonzeroAttr(attr::kInline)) {
+   *      // inline the function.
+   *    }
+   *  }
+   *
+   * \endcode
+   */
+  bool HasNonzeroAttr(const std::string& attr_key) const { return attrs.HasNonzeroAttr(attr_key); }
+  /*!
+   * \brief Get the type of the linkage.
+   *
+   * Currently, we only consider external/internal linkage.
+   * This can be extended in the future when necessary.
+   *
+   * \return Linkage type.
+   *
+   * \code
+   *
+   *  void Example(const BaseFunc& f) {
+   *    if (f->GetLinkageType() == tvm::LinkageType::kExternal) {
+   *      // Do not remove a function with external linkage
+   *    }
+   *  }
+   *
+   * \endcode
+   */
+
+  LinkageType GetLinkageType() const {
+    if (GetAttr<ffi::String>(attr::kGlobalSymbol))
+      return LinkageType::kExternal;
+    else
+      return LinkageType::kInternal;
+  }
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<BaseFuncNode>().def_ro("attrs", &BaseFuncNode::attrs);
+  }
+
+  static constexpr const uint32_t _type_child_slots = 2;
+  TVM_FFI_DECLARE_OBJECT_INFO("ir.BaseFunc", BaseFuncNode, ExprNode);
+};
+
+/*!
+ * \brief Managed reference to BaseFuncNode.
+ * \sa BaseFuncNode
+ */
+class BaseFunc : public Expr {
+ public:
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BaseFunc, Expr, BaseFuncNode);
+};
+
+}  // namespace tvm
+#endif  // TVM_IR_FUNCTION_H_

include/tvm/ir/global_info.h

@@ -0,0 +1,117 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/global_info.h
+ * \brief GlobalInfo are globally static object that are referred by the IR itself.
+ */
+
+#ifndef TVM_IR_GLOBAL_INFO_H_
+#define TVM_IR_GLOBAL_INFO_H_
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/expr.h>
+#include <tvm/target/target.h>
+
+namespace tvm {
+
+/*!
+ * \brief Abstract label for an area of memory.
+ */
+using MemoryScope = ffi::String;
+
+/*!
+ * \brief GlobalInfo are globally static object that are referred by the IR itself.
+ *        Base node for all global info that can appear in the IR
+ */
+class GlobalInfoNode : public ffi::Object {
+ public:
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+
+  TVM_FFI_DECLARE_OBJECT_INFO("ir.GlobalInfo", GlobalInfoNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference to GlobalInfoNode.
+ * \sa GlobalInfoNode
+ */
+class GlobalInfo : public ffi::ObjectRef {
+ public:
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GlobalInfo, ffi::ObjectRef, GlobalInfoNode);
+};
+
+/*!
+ * \brief A global info subclass for virtual devices.
+ */
+class VDeviceNode : public GlobalInfoNode {
+ public:
+  /*! \brief The \p Target describing how to compile for the virtual device. */
+  Target target;
+  /*! \brief The device identifier for the virtual device. This enables us to
+   * differentiate between distinct devices with same Target, such as multiple GPUs.
+   */
+  int vdevice_id;
+  MemoryScope memory_scope;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<VDeviceNode>()
+        .def_ro("target", &VDeviceNode::target)
+        .def_ro("vdevice_id", &VDeviceNode::vdevice_id)
+        .def_ro("memory_scope", &VDeviceNode::memory_scope);
+  }
+
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.VDevice", VDeviceNode, GlobalInfoNode);
+};
+
+/*!
+ * \brief Managed reference to VDeviceNode.
+ * \sa VDeviceNode
+ */
+class VDevice : public GlobalInfo {
+ public:
+  TVM_DLL explicit VDevice(Target tgt, int dev_id, MemoryScope mem_scope);
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VDevice, GlobalInfo, VDeviceNode);
+};
+
+/*!
+ * \brief A dummy global info sub-class for testing purpose.
+ */
+class DummyGlobalInfoNode : public GlobalInfoNode {
+ public:
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<DummyGlobalInfoNode>();
+  }
+
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DummyGlobalInfo", DummyGlobalInfoNode, GlobalInfoNode);
+};
+
+/*!
+ * \brief Managed reference to DummyGlobalInfoNode.
+ * \sa DummyGlobalInfoNode
+ */
+class DummyGlobalInfo : public GlobalInfo {
+ public:
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DummyGlobalInfo, GlobalInfo, DummyGlobalInfoNode);
+};
+
+}  // namespace tvm
+
+#endif  // TVM_IR_GLOBAL_INFO_H_

include/tvm/ir/instrument.h

@@ -0,0 +1,158 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/instrument.h
+ *
+ * This file introduces a pass instrument infrastructure, inspired by LLVM and MLIR.
+ * It inserts instrumentation points around passes.
+ */
+#ifndef TVM_IR_INSTRUMENT_H_
+#define TVM_IR_INSTRUMENT_H_
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/string.h>
+
+#include <utility>
+#include <vector>
+
+namespace tvm {
+
+class IRModule;
+
+// Forward class for PassInstrumentNode methods
+namespace transform {
+class PassInfo;
+}  // namespace transform
+
+namespace instrument {
+
+/*!
+ * \brief PassInstrumentNode forms an instrument implementation.
+ * It provides API for users to register callbacks at different instrumentation points.
+ *
+ * Within a PassContext, call sequence of a PassInstrument implementation is like:
+ *
+ *   with PassContext(instruments=[pi]):  # pi = a PassInstrument implementation
+ *       pi.EnterPassContext()
+ *
+ *       if pi.ShouldRun(Pass1):
+ *           pi.RunBeforePass()
+ *           Pass1()
+ *           pi.RunAfterPass()
+ *
+ *       if pi.ShouldRun(Pass2):
+ *           pi.RunBeforePass()
+ *           Pass2()
+ *           pi.RunAfterPass()
+ *
+ *       pi.ExitPassContext()
+ *
+ * `EnterPassContext` and `ExitPassContext` are only called once when entering/exiting a
+ * PassContext. `ShouldRun`, `RunBeforePass` and `RunAfterPass` are called multiple times depending
+ * on how many passes.
+ *
+ * If there are multiple pass instrumentations provided, the instrument points are the same.
+ * PassInstrument implementations' callbacks are called in order:
+ *
+ *   with PassContext(instruments=[pi1, pi2]):  # pi1, pi2 = two distinct PassInstrument impls
+ *       pi.EnterPassContext() for pi in instruments
+ *
+ *       should_run = all([pi.ShoudRun(Pass1) for pi in instruments)])
+ *       if (should_run)
+ *           pi.RunBeforePass() for pi in instruments
+ *           Pass1()
+ *           pi.RunAfterPass()  for pi in instruments
+ *
+ *       should_run = all([pi.ShouldRun(Pass2) for pi in instruments)])
+ *       if (should_run)
+ *           pi.RunBeforePass() for pi in instruments
+ *           Pass2()
+ *           pi.RunAfterPass() for pi in instruments
+ *
+ *       pi.ExitPassContext() for pi in instruments
+ *
+ * Note:
+ *   1. Assume there is no dependency between PassInstrument implementations in `instruments` .
+ *   2. `EnterPassContext` and `ExitPassContext` have `with` behavior (see PassContext and its FFI):
+ *        If there is any exception raised in `ShouldRun()`, `RunBeforePass()`, `RunAfterPass()` and
+ *        `Pass()`, `ExitPassContext()` is still called.
+ *   3. In mutiple PassInstrument instances scenario, callbacks are called in order:
+ *        If one throws exceptions, remainings will not be called.
+ *
+ * \sa PassInstrument
+ * \sa src/ir/transform.cc
+ */
+class PassInstrumentNode : public ffi::Object {
+ public:
+  /*! \brief Name of this pass instrument object. */
+  ffi::String name;
+
+  virtual ~PassInstrumentNode() {}
+
+  /*! \brief Instrument when entering PassContext. Called once within a PassContext. */
+  virtual void EnterPassContext() const = 0;
+
+  /*! \brief Instrument when exiting PassContext. Called once within a PassContext. */
+  virtual void ExitPassContext() const = 0;
+
+  /*!
+   * \brief Determine whether to run the pass or not. Called multiple times depend on number of
+   *        passes.
+   * \param mod The module that an optimization pass runs on.
+   * \param info The pass information.
+   *
+   * \return true to run the pass; false to skip the pass.
+   */
+  virtual bool ShouldRun(const IRModule& mod, const transform::PassInfo& info) const = 0;
+
+  /*!
+   * \brief Instrument before pass run. Called multiple times depend on number of passes.
+   * \param mod The module that an optimization pass runs on.
+   * \param info The pass information.
+   */
+  virtual void RunBeforePass(const IRModule& mod, const transform::PassInfo& info) const = 0;
+
+  /*!
+   * \brief Instrument after pass run. Called multiple time depend on number of passes.
+   * \param mod The module that an optimization pass runs on.
+   * \param info The pass information.
+   */
+  virtual void RunAfterPass(const IRModule& mod, const transform::PassInfo& info) const = 0;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<PassInstrumentNode>().def_ro("name", &PassInstrumentNode::name);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO("instrument.PassInstrument", PassInstrumentNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference class for PassInstrumentNode
+ * \sa PassInstrumentNode
+ */
+class PassInstrument : public ffi::ObjectRef {
+ public:
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PassInstrument, ffi::ObjectRef, PassInstrumentNode);
+};
+
+}  // namespace instrument
+}  // namespace tvm
+
+#endif  // TVM_IR_INSTRUMENT_H_

include/tvm/ir/module.h

@@ -0,0 +1,366 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/module.h
+ * \brief IRModule that holds the functions and type definitions.
+ */
+#ifndef TVM_IR_MODULE_H_
+#define TVM_IR_MODULE_H_
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/string.h>
+#include <tvm/ir/cow.h>
+#include <tvm/ir/expr.h>
+#include <tvm/ir/function.h>
+#include <tvm/ir/global_info.h>
+#include <tvm/ir/source_map.h>
+#include <tvm/ir/type.h>
+
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+
+class IRModule;
+
+/*!
+ * \brief IRModule that holds functions and type definitions.
+ *
+ *  IRModule is the basic unit for all IR transformations across the stack.
+ *
+ *  Many operations require access to the global IRModule.
+ *  We pass the IRModule by value in a functional style as an explicit argument,
+ *  but we mutate the Module while optimizing programs.
+ * \sa IRModule
+ */
+class IRModuleNode : public ffi::Object {
+ public:
+  /*! \brief A map from ids to all global functions. */
+  ffi::Map<GlobalVar, BaseFunc> functions;
+  /*! \brief The source map for the module. */
+  SourceMap source_map;
+  /* \brief Additional attributes storing meta-data about the module. */
+  DictAttrs attrs;
+  /*! \brief Globally static object that are referred by the IR itself */
+  ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos;
+  /*!
+   * \brief A map from string names to global variables that
+   * ensures global uniqueness.
+   */
+  ffi::Map<ffi::String, GlobalVar> global_var_map_;
+
+  /*!
+   * \brief Get a module attribute.
+   *
+   * \param attr_key The attribute key.
+   * \param default_value The default value if the key does not exist, defaults to nullptr.
+   *
+   * \return The result
+   *
+   * \tparam TObjectRef the expected object type.
+   * \throw Error if the key exists but the value does not match TObjectRef
+   *
+   * \code
+   *
+   *  void GetAttrExample(const IRModule& mod) {
+   *    auto value = f->GetAttr<int64_t>("AttrKey", 0);
+   *  }
+   *
+   * \endcode
+   */
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetAttr(
+      const std::string& attr_key,
+      ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
+    return attrs.GetAttr(attr_key, default_value);
+  }
+  // variant that uses TObjectRef to enable implicit conversion to default value.
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
+    return GetAttr<TObjectRef>(attr_key, ffi::Optional<TObjectRef>(default_value));
+  }
+
+  /*!
+   * \brief Get the metadata attributes.
+   * \returns The additional meta-data attributes
+   */
+  DictAttrs GetAttrs() const { return attrs; }
+
+  /*!
+   * \brief Check whether the module has an non-zero integer attr.
+   *
+   * This function can be used to check whether an optional
+   * attribute mark(e.g. inline) exists.
+   *
+   * \param attr_key The key to the attribute.
+   * \return The check result.
+   *
+   * \code
+   *
+   *  void HasNonzeroAttrExample(const IRModule& mod) {
+   *    if (mod->HasNonzeroAttr(attr::kInline)) {
+   *      // inline the function.
+   *    }
+   *  }
+   *
+   * \endcode
+   */
+  bool HasNonzeroAttr(const std::string& attr_key) const { return attrs.HasNonzeroAttr(attr_key); }
+
+  IRModuleNode() : source_map() {}
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<IRModuleNode>()
+        .def_ro("functions", &IRModuleNode::functions)
+        .def_ro("global_var_map_", &IRModuleNode::global_var_map_)
+        .def_ro("source_map", &IRModuleNode::source_map)
+        .def_ro("attrs", &IRModuleNode::attrs)
+        .def_ro("global_infos", &IRModuleNode::global_infos);
+    // register custom structural equal and hash.
+    refl::TypeAttrDef<IRModuleNode>()
+        .def("__s_equal__", &IRModuleNode::SEqual)
+        .def("__s_hash__", &IRModuleNode::SHash);
+  }
+
+  TVM_DLL bool SEqual(const IRModuleNode* other,
+                      ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const;
+  TVM_DLL int64_t SHash(int64_t init_hash,
+                        ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const;
+
+  /*!
+   * \brief Add a function to the global environment.
+   * \param var The var of the global function.
+   * \param func The function.
+   * \param update Controls whether you can replace a definition in the
+   * environment.
+   */
+  TVM_DLL void Add(const GlobalVar& var, const BaseFunc& func, bool update = false);
+
+  /*!
+   * \brief Add a function to the global environment.
+   * \param var The name of the global function.
+   * \param func The function.
+   *
+   * It does not do type inference as Add does.
+   */
+  TVM_DLL void AddUnchecked(const GlobalVar& var, const BaseFunc& func);
+
+  /*!
+   * \brief Update a function in the global environment.
+   * \param var The name of the global function to update.
+   * \param func The new function.
+   */
+  TVM_DLL void Update(const GlobalVar& var, const BaseFunc& func);
+
+  /*!
+   * \brief Update an array of global infos in the global environment.
+   * \param name The name of the global info.
+   * \param info The new array of global infos.
+   */
+  TVM_DLL void UpdateGlobalInfo(const ffi::String& name, const ffi::Array<GlobalInfo>& info);
+
+  /*!
+   * \brief Remove a function from the global environment.
+   * \param var The name of the global function to update.
+   */
+  TVM_DLL void Remove(const GlobalVar& var);
+
+  /*!
+   * \brief Check if the global_var_map_ contains a global variable.
+   * \param name The variable name.
+   * \returns true if contains, otherise false.
+   */
+  TVM_DLL bool ContainGlobalVar(const ffi::String& name) const;
+
+  /*!
+   * \brief Lookup a global function by its variable.
+   * \param str The unique string specifying the global variable.
+   * \returns The global variable.
+   */
+  TVM_DLL GlobalVar GetGlobalVar(const ffi::String& str) const;
+
+  /*!
+   * \brief Collect all global vars defined in this module, ordered by
+   *        the global variable name.
+   * \returns An array of global vars
+   */
+  TVM_DLL ffi::Array<GlobalVar> GetGlobalVars() const;
+
+  /*!
+   * \brief Look up a global function by its variable.
+   * \param var The global var to lookup.
+   * \returns The function named by the variable argument.
+   */
+  TVM_DLL BaseFunc Lookup(const GlobalVar& var) const;
+
+  /*!
+   * \brief Look up a global function by its string name
+   * \param name The name of the function.
+   * \returns The function named by the argument.
+   */
+  TVM_DLL BaseFunc Lookup(const ffi::String& name) const;
+
+  /*!
+   * \brief Update the functions inside this environment by
+   *        functions in another environment.
+   * \param other The other environment.
+   */
+  TVM_DLL void Update(const IRModule& other);
+
+  /*!
+   * \brief Create a shallow copy of this IRModule.
+   * \returns The shallow copy of the IRModule.
+   */
+  TVM_DLL IRModule ShallowCopy();
+  /*!
+   * \brief The set of imported files.
+   */
+  TVM_DLL std::unordered_set<ffi::String> Imports() const;
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IRModule", IRModuleNode, ffi::Object);
+
+ private:
+  friend class IRModule;
+};
+
+/*!
+ * \brief Managed reference class to IRModuleNode.
+ * \sa IRModuleNode
+ */
+class IRModule : public ffi::ObjectRef {
+ public:
+  /*!
+   * \brief constructor
+   * \param functions Functions in the module.
+   * \param map The module source map.
+   * \param attrs The module meta-data attributes.
+   * \param global_infos Global infos in the module.
+   */
+  TVM_DLL explicit IRModule(ffi::Map<GlobalVar, BaseFunc> functions, SourceMap map = {},
+                            DictAttrs attrs = DictAttrs(),
+                            ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos = {});
+
+  /*! \brief default constructor */
+  IRModule() : IRModule(ffi::Map<GlobalVar, BaseFunc>({})) {}
+  /*!
+   * \brief constructor
+   * \param n The object pointer.
+   */
+  explicit IRModule(ffi::ObjectPtr<IRModuleNode> n) : ffi::ObjectRef(n) {}
+  /*!
+   * \brief constructor with UnsafeInit
+   */
+  explicit IRModule(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
+  /*! \return mutable pointers to the node. */
+  IRModuleNode* operator->() const {
+    auto* ptr = get_mutable();
+    TVM_FFI_ICHECK(ptr != nullptr);
+    return static_cast<IRModuleNode*>(ptr);
+  }
+
+  /*!
+   * \brief As for \p FromExprInContext, but assuming \p expr is bound to 'main' and no
+   * imports.
+   */
+  TVM_DLL static IRModule FromExpr(const Expr& expr,
+                                   const ffi::Map<GlobalVar, BaseFunc>& global_funcs = {});
+
+  /*!
+   * \brief Create a shallow copy of an IRModule.
+   * \param mod The module to copy.
+   * \return The copied module.
+   */
+  IRModule ShallowCopyIRModule(IRModule mod);
+
+  /*! \brief Declare the container type. */
+  using ContainerType = IRModuleNode;
+
+  // allow copy on write.
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(IRModuleNode);
+};
+
+namespace attr {
+
+// Following are attributes for IRModule only.
+
+/*!
+ * \brief Name of the module
+ *
+ * Type: String
+ */
+constexpr const char* kModuleName = "mod_name";
+
+/*!
+ * \brief All the runtime::Modules accumulated during compilation by external codegen. These
+ * modules must be either directly linked or captured in the final compilation artifact.
+ *
+ * Type: ffi::Array<runtime::Module>
+ */
+constexpr const char* kExternalMods = "external_mods";
+
+/*!
+ * \brief A prefix for generating C symbols  system lib creation.
+ *
+ * This prefix guides passes that creates global_symbol for internal functions
+ * that may have c linkage (e.g. TIR functions and some BYOC functions). It also affects
+ * the symbol of the fat bin blob during module export.
+ *
+ * This attribute is used to avoid symbol conflict when we
+ * generate and combine multiple system libs that get linked into one.
+ *
+ * Rationale: mechanisms like BYOC rely on the common global symbol
+ * and each external compiler also has its own mechanism of mangling.
+ * As a result, we cannot rely on other mechanisms on setting a global_symbol and then renaming,
+ * because the external compiler already agreed on the name.
+ *
+ * system_lib_prefix provides a way to hint at the passes to allow names to
+ * avoid name conflict at the beginning.
+ *
+ * Note that users can still directly specify global symbols that may conflict.
+ * It is up to the downstream toolchain to manage those external-facing functions.
+ *
+ * This does not affect non-C linkage functions it is less of an issue because
+ * they will be embedded into fatbin that in different symbols,
+ * The system lib loader can pick the right prefix for a given prefix.
+ *
+ * Having this attribute implies system lib generation linkage.
+ */
+constexpr const char* kSystemLibPrefix = "system_lib_prefix";
+
+/*!
+ * \brief All the named runtime::Tensors accumulated during compilation by external codegen.
+ * Generally the associated runtime::Module will indicate it requires bindings for these names,
+ * and during module initialization these bindings will be recovered from a ConstLoaderModule.
+ *
+ * Type: ffi::Map<ffi::String, runtime::Tensor>
+ */
+constexpr const char* kConstNameToConstant = "const_name_to_constant";
+
+}  // namespace attr
+}  // namespace tvm
+#endif  // TVM_IR_MODULE_H_

include/tvm/ir/node_functor.h

@@ -0,0 +1,194 @@
+/*
+ * 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.
+ */
+/*!
+ * \file tvm/ir/node_functor.h
+ * \brief Defines the Functor data structures.
+ */
+#ifndef TVM_IR_NODE_FUNCTOR_H_
+#define TVM_IR_NODE_FUNCTOR_H_
+
+#include <tvm/ffi/error.h>
+
+#include <cstring>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+
+/*!
+ * \brief A dynamically dispatched functor on the type of the first argument.
+ *
+ * This is a class that is useful to construct polymorphic dispatching
+ * base on the AST/IR node's type.
+ *
+ * \code
+ *   NodeFunctor<std::string (const ffi::ObjectRef& n, std::string prefix)> tostr;
+ *   tostr.set_dispatch<Add>([](const ffi::ObjectRef& op, std::string prefix) {
+ *     return prefix + "Add";
+ *   });
+ *   tostr.set_dispatch<IntImm>([](const ffi::ObjectRef& op, std::string prefix) {
+ *     return prefix + "IntImm"
+ *   });
+ *
+ *   Expr x = MakeConst(1);
+ *   Expr y = x + x;
+ *   // dispatch to IntImm, outputs "MyIntImm"
+ *   LOG(INFO) << tostr(x, "My");
+ *   // dispatch to IntImm, outputs "MyAdd"
+ *   LOG(INFO) << tostr(y, "My");
+ * \endcode
+ *
+ * \tparam FType function signiture
+ *  This type if only defined for FType with function signature
+ */
+template <typename FType>
+class NodeFunctor;
+
+template <typename R, typename... Args>
+class NodeFunctor<R(const ffi::ObjectRef& n, Args...)> {
+ private:
+  /*! \brief internal function pointer type */
+  typedef R (*FPointer)(const ffi::ObjectRef& n, Args...);
+  /*! \brief refer to itself. */
+  using TSelf = NodeFunctor<R(const ffi::ObjectRef& n, Args...)>;
+  /*! \brief internal function table */
+  std::vector<FPointer> func_;
+  /*! \brief start range of func index */
+  uint32_t begin_type_index_{0};
+
+ public:
+  /*! \brief the result type of this functor */
+  using result_type = R;
+  /*!
+   * \brief Whether the functor can dispatch the corresponding Node
+   * \param n The node to be dispatched
+   * \return Whether dispatching function is registered for n's type.
+   */
+  bool can_dispatch(const ffi::ObjectRef& n) const {
+    uint32_t type_index = n->type_index();
+    if (type_index < begin_type_index_) return false;
+    type_index -= begin_type_index_;
+    return type_index < func_.size() && func_[type_index] != nullptr;
+  }
+  /*!
+   * \brief invoke the functor, dispatch on type of n
+   * \param n The Node argument
+   * \param args The additional arguments
+   * \return The result.
+   */
+  R operator()(const ffi::ObjectRef& n, Args... args) const {
+    TVM_FFI_ICHECK(can_dispatch(n))
+        << "NodeFunctor calls un-registered function on type " << n->GetTypeKey();
+    return (*func_[n->type_index() - begin_type_index_])(n, std::forward<Args>(args)...);
+  }
+  /*!
+   * \brief set the dispatcher for type TNode
+   * \param f The function to be set.
+   * \tparam TNode the type of Node to be dispatched.
+   * \return reference to self.
+   */
+  template <typename TNode>
+  TSelf& set_dispatch(FPointer f) {  // NOLINT(*)
+    uint32_t tindex = TNode::RuntimeTypeIndex();
+    if (func_.size() <= tindex) {
+      func_.resize(tindex + 1, nullptr);
+    }
+    TVM_FFI_ICHECK(func_[tindex] == nullptr)
+        << "Dispatch for " << TNode::_type_key << " is already set";
+    TVM_FFI_ICHECK_EQ(begin_type_index_, 0) << " Cannot call set_dispatch after calling Finalize";
+    func_[tindex] = f;
+    return *this;
+  }
+  /*!
+   * \brief unset the dispatcher for type TNode
+   *
+   * \tparam TNode the type of Node to be dispatched.
+   * \return reference to self.
+   */
+  template <typename TNode>
+  TSelf& clear_dispatch() {  // NOLINT(*)
+    uint32_t tindex = TNode::RuntimeTypeIndex();
+    TVM_FFI_ICHECK_LT(tindex, func_.size()) << "clear_dispatch: index out of range";
+    TVM_FFI_ICHECK_EQ(begin_type_index_, 0) << " Cannot call clear_dispatch after calling Finalize";
+    func_[tindex] = nullptr;
+    return *this;
+  }
+  /*!
+   * \brief Finalize the functor after calling sequence of set_dispatch
+   * This function will attempt to find the min type index that is not null
+   * and optimize the space of the func table so it is more compact
+   */
+  void Finalize() {
+    TVM_FFI_ICHECK_EQ(begin_type_index_, 0) << "Can only call Finalize once";
+    while (begin_type_index_ < func_.size() && func_[begin_type_index_] == nullptr) {
+      ++begin_type_index_;
+    }
+    // shift up the function value
+    size_t new_ftable_size = func_.size() - begin_type_index_;
+    if (begin_type_index_ != 0) {
+      std::memmove(func_.data(), func_.data() + begin_type_index_,
+                   new_ftable_size * sizeof(FPointer));
+    }
+    func_.resize(new_ftable_size);
+    func_.shrink_to_fit();
+  }
+};
+
+#define TVM_REG_FUNC_VAR_DEF(ClsName) [[maybe_unused]] static auto& __make_functor##_##ClsName
+
+/*!
+ * \brief Useful macro to set NodeFunctor dispatch in a global static field.
+ *
+ * \code
+ *  // Use NodeFunctor to implement TVMScriptPrinter similar to Visitor Pattern.
+ *  // vtable allows easy patch of new Node types, without changing
+ *  // the interface of TVMScriptPrinter.
+ *
+ *  class TVMScriptPrinter {
+ *   public:
+ *    // the dispatch function.
+ *    static std::string Script(const ffi::ObjectRef& node, const PrinterConfig& cfg) {
+ *      return vtable()(node, cfg);
+ *    }
+ *    using FType = NodeFunctor<std::string(const ffi::ObjectRef&, const PrinterConfig&)>;
+ *    // function to return global function table
+ *    static FType& vtable();
+ *  };
+ *
+ *  // in cpp/cc file
+ *  TVMScriptPrinter::FType& TVMScriptPrinter::vtable() {
+ *    static FType inst; return inst;
+ *  }
+ *
+ *  TVM_STATIC_IR_FUNCTOR(TVMScriptPrinter, vtable)
+ *  .set_dispatch<AddNode>([](const ffi::ObjectRef& ref, const PrinterConfig& cfg) {
+ *    auto* n = static_cast<const AddNode*>(ref.get());
+ *    return Script(n->a, cfg) + " + " + Script(n->b, cfg);
+ *  });
+ *
+ * \endcode
+ *
+ * \param ClsName The name of the class
+ * \param FField The static function that returns a singleton of NodeFunctor.
+ */
+#define TVM_STATIC_IR_FUNCTOR(ClsName, FField) \
+  TVM_FFI_STR_CONCAT(TVM_REG_FUNC_VAR_DEF(ClsName), __COUNTER__) = ClsName::FField()
+}  // namespace tvm
+#endif  // TVM_IR_NODE_FUNCTOR_H_

include/tvm/ir/op.h

@@ -0,0 +1,423 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/op.h
+ * \brief Primitive operators(builtin intrinsics)
+ *        and registry for them.
+ */
+#ifndef TVM_IR_OP_H_
+#define TVM_IR_OP_H_
+
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/attr_registry_map.h>
+#include <tvm/ir/attrs.h>
+#include <tvm/ir/env_func.h>
+#include <tvm/ir/expr.h>
+#include <tvm/ir/type.h>
+
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+
+// forward declare name.
+template <typename>
+class OpAttrMap;
+
+/*!
+ * \brief Information about an input field of an Op (name, type, description).
+ *
+ *  Populated via OpRegEntry::add_argument and consumed both by
+ *  internal sanity checks / error messages and by external tooling
+ *  that wants to introspect an Op's argument schema.
+ */
+class ArgumentInfoNode : public ffi::Object {
+ public:
+  /*! \brief name of the field */
+  ffi::String name;
+  /*! \brief type docstring information in str. */
+  ffi::String type_info;
+  /*! \brief detailed description of the type */
+  ffi::String description;
+
+  static void RegisterReflection() {
+    namespace rfl = ffi::reflection;
+    rfl::ObjectDef<ArgumentInfoNode>()
+        .def_ro("name", &ArgumentInfoNode::name)
+        .def_ro("type_info", &ArgumentInfoNode::type_info)
+        .def_ro("description", &ArgumentInfoNode::description);
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.ArgumentInfo", ArgumentInfoNode, ffi::Object);
+};
+
+/*! \brief Managed reference to ArgumentInfoNode. */
+class ArgumentInfo : public ffi::ObjectRef {
+ public:
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ArgumentInfo, ffi::ObjectRef, ArgumentInfoNode);
+};
+
+// TODO(tvm-team): migrate low-level intrinsics to use Op
+/*!
+ * \brief Primitive Op(builtin intrinsics)
+ *
+ * This data structure stores the meta-data
+ * about primitive operators that can be invoked via Call.
+ *
+ * Low-level IR intrinsics(such as libc.expf) are also
+ * implemented via Op.
+ *
+ * \sa Op
+ */
+class OpNode : public ExprNode {
+ public:
+  /*! \brief name of the operator */
+  ffi::String name;
+  /*!
+   * \brief detailed description of the operator
+   *  This can be used to generate docstring automatically for the operator.
+   */
+  ffi::String description;
+  /* \brief Information of input arguments to the operator */
+  ffi::Array<ArgumentInfo> arguments;
+  /*!
+   * \brief The type key of the attribute field
+   *  This can be empty, in which case it defaults to anything.
+   */
+  ffi::String attrs_type_key;
+  /*!
+   * \brief attribute type index,
+   * this field varies in each run and is not exposed to frontend.
+   */
+  uint32_t attrs_type_index{0};
+  /*!
+   * \brief number of input arguments to the operator,
+   * -1 means it is variable length
+   */
+  int32_t num_inputs = -1;
+  /*!
+   * \brief support level of the operator,
+   *  The lower the more priority it contains.
+   *  This is in analogies to BLAS levels.
+   */
+  int32_t support_level = 10;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<OpNode>()
+        .def_ro("name", &OpNode::name)
+        .def_ro("description", &OpNode::description, refl::AttachFieldFlag::SEqHashIgnore())
+        .def_ro("arguments", &OpNode::arguments, refl::AttachFieldFlag::SEqHashIgnore())
+        .def_ro("attrs_type_key", &OpNode::attrs_type_key, refl::AttachFieldFlag::SEqHashIgnore())
+        .def_ro("num_inputs", &OpNode::num_inputs, refl::AttachFieldFlag::SEqHashIgnore())
+        .def_ro("support_level", &OpNode::support_level, refl::AttachFieldFlag::SEqHashIgnore());
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Op", OpNode, ExprNode);
+
+ private:
+  /*! \return the internal attr registry index. */
+  uint32_t AttrRegistryIndex() const { return index_; }
+  /*! \brief repr to be printed in registry*/
+  std::string AttrRegistryName() const { return name; }
+
+  // friend class
+  template <typename>
+  friend class AttrRegistryMapContainerMap;
+  template <typename, typename>
+  friend class AttrRegistry;
+  friend class OpRegEntry;
+
+  // Program internal unique index of operator.
+  // Used to help index the program.
+  uint32_t index_{0};
+};
+
+/*!
+ * \brief Managed reference class to OpNode.
+ * \sa OpNode
+ */
+class Op : public Expr {
+ public:
+  explicit Op(ffi::ObjectPtr<OpNode> n) : Expr(std::move(n)) {
+    TVM_FFI_CHECK(defined(), ValueError) << "Op expects a defined OpNode";
+  }
+
+  /*!
+   * \brief Get additional registered attribute about operators.
+   *  If nothing has been registered, an empty OpAttrMap will be returned.
+   * \param attr_name The name of the attribute.
+   * \return An OpAttrMap of specified attr_name.
+   * \tparam ValueType The type of the attribute.
+   */
+  template <typename ValueType>
+  inline static OpAttrMap<ValueType> GetAttrMap(const ffi::String& attr_name);
+  /*!
+   * \brief Checks if an attr map is present in the registry.
+   * \param attr_name The name of the attribute.
+   * \return bool True if the attr is present.
+   */
+  TVM_DLL static bool HasAttrMap(const ffi::String& attr_name);
+  /*!
+   * \brief Get an Op for a given operator name.
+   *  Will raise an error if the op has not been registered.
+   * \param op_name Name of the operator.
+   * \return Pointer to a Op, valid throughout program lifetime.
+   */
+  TVM_DLL static const Op& Get(const ffi::String& op_name);
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Op, Expr, OpNode);
+
+ private:
+  /*!
+   * \brief Get generic attrmap given attr name
+   * \param key The attribute key
+   * \return The attr map.
+   */
+  TVM_DLL static const AttrRegistryMapContainerMap<Op>& GetAttrMapContainer(const ffi::String& key);
+};
+
+/*!
+ * \brief Helper structure to register operators
+ * \sa TVM_REGISTER_OP
+ */
+class OpRegEntry {
+ public:
+  /*! \return the operator */
+  const Op& op() const { return op_; }
+  /*!
+   * \brief setter function during registration
+   *  Set the description of operator
+   * \param descr the description string.
+   * \return reference to self.
+   */
+  inline OpRegEntry& describe(const std::string& descr);  // NOLINT(*)
+  /*!
+   * \brief Add argument information to the function.
+   * \param name Name of the argument.
+   * \param type Type of the argument.
+   * \param description Description of the argument.
+   * \return reference to self.
+   */
+  inline OpRegEntry& add_argument(const std::string& name, const std::string& type,
+                                  const std::string& description);
+  /*!
+   * \brief Set the attrs type key and index to be AttrsType.
+   * \tparam AttrsType the attribute type to b set.
+   * \return reference to self.
+   */
+  template <typename AttrsType>
+  inline OpRegEntry& set_attrs_type();
+  /*!
+   * \brief Set the attrs type key and index to be AttrsType.
+   * \param key The attribute type key to be set.
+   * \return reference to self.
+   */
+  inline OpRegEntry& set_attrs_type_key(const ffi::String& key);
+  /*!
+   * \brief Set the num_inputs
+   * \param n The number of inputs to be set.
+   * \return reference to self.
+   */
+  inline OpRegEntry& set_num_inputs(int32_t n);  // NOLINT(*)
+  /*!
+   * \brief Set the support level of op.
+   * \param level The support level.
+   * \return reference to self.
+   */
+  inline OpRegEntry& set_support_level(int32_t level);  // NOLINT(*)
+  /*!
+   * \brief Register additional attributes to operator.
+   * \param attr_name The name of the attribute.
+   * \param value The value to be set.
+   * \param plevel The priority level of this set,
+   *  an higher priority level attribute
+   *  will replace lower priority level attribute.
+   *  Must be bigger than 0.
+   *
+   *  Cannot set with same plevel twice in the code.
+   *
+   * \tparam ValueType The type of the value to be set.
+   */
+  template <typename ValueType>
+  inline OpRegEntry& set_attr(const std::string& attr_name,  // NOLINT(*)
+                              const ValueType& value, int plevel = 10);
+
+  /*!
+   * \brief Resets an attr of the registry.
+   * \param attr_name The name of the attribute.
+   */
+  inline void reset_attr(const std::string& attr_name);
+
+  // set the name of the op to be the same as registry
+  inline OpRegEntry& set_name() {  // NOLINT(*)
+    if (get()->name.length() == 0) {
+      get()->name = name;
+    }
+    return *this;
+  }
+  /*!
+   * \brief Register or get a new entry.
+   * \param name The name of the operator.
+   * \return the corresponding entry.
+   */
+  TVM_DLL static OpRegEntry& RegisterOrGet(const ffi::String& name);
+
+ private:
+  template <typename, typename>
+  friend class AttrRegistry;
+  // the name
+  std::string name;
+  /*! \brief The operator */
+  Op op_;
+  /*! \brief Construct the non-null Op for this registry entry. */
+  static Op MakeOp(uint32_t reg_index);
+  // private constructor
+  TVM_DLL OpRegEntry(uint32_t reg_index);
+  // return internal pointer to op.
+  inline OpNode* get();
+  // update the attribute OpAttrMap
+  TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
+};
+
+/*!
+ * \brief ffi::Map<Op,ValueType> used to store meta-information about Op.
+ * \tparam ValueType The type of the value stored in map.
+ */
+template <typename ValueType>
+class OpAttrMap : public AttrRegistryMap<Op, ValueType> {
+ public:
+  /*!
+   * \brief get the corresponding value element at op with default value.
+   * \param expr The key to the map
+   * \param def_value The default value when the key does not exist
+   *         or if expr is not an Op.
+   * \return the const reference to the content value.
+   */
+  inline ValueType get(const Expr& expr, ValueType def_value) const;
+
+  using TParent = AttrRegistryMap<Op, ValueType>;
+  using TParent::count;
+  using TParent::get;
+  using TParent::operator[];
+
+ private:
+  friend class Op;
+  // constructor
+  explicit OpAttrMap(const AttrRegistryMapContainerMap<Op>& map) : TParent(map) {}
+};
+
+// internal macros to make
+#define TVM_OP_REGISTER_VAR_DEF [[maybe_unused]] static ::tvm::OpRegEntry& __make_##Op
+
+/*!
+ * \def TVM_REGISTER_OP
+ * \brief Register a new operator, or set attribute of the corresponding op.
+ *
+ * \param OpName The name of registry
+ *
+ * \code
+ *
+ *  TVM_REGISTER_OP("add")
+ *  .describe("add two inputs together")
+ *  .set_num_inputs(2)
+ *  .set_attr<OpKernel>("gpu_kernel", AddKernel);
+ *
+ * \endcode
+ */
+#define TVM_REGISTER_OP(OpName)                              \
+  TVM_FFI_STR_CONCAT(TVM_OP_REGISTER_VAR_DEF, __COUNTER__) = \
+      ::tvm::OpRegEntry::RegisterOrGet(OpName).set_name()
+
+// implementations
+
+template <typename ValueType>
+inline OpAttrMap<ValueType> Op::GetAttrMap(const ffi::String& key) {
+  return OpAttrMap<ValueType>(Op::GetAttrMapContainer(key));
+}
+
+inline OpNode* OpRegEntry::get() { return const_cast<OpNode*>(op_.operator->()); }
+
+inline OpRegEntry& OpRegEntry::describe(const std::string& descr) {  // NOLINT(*)
+  get()->description = descr;
+  return *this;
+}
+
+inline OpRegEntry& OpRegEntry::add_argument(const std::string& name, const std::string& type,
+                                            const std::string& description) {
+  auto n = ffi::make_object<ArgumentInfoNode>();
+  n->name = name;
+  n->type_info = type;
+  n->description = description;
+  get()->arguments.push_back(ArgumentInfo(n));
+  return *this;
+}
+
+inline OpRegEntry& OpRegEntry::set_num_inputs(int32_t n) {  // NOLINT(*)
+  get()->num_inputs = n;
+  return *this;
+}
+
+template <typename AttrsType>
+inline OpRegEntry& OpRegEntry::set_attrs_type() {  // NOLINT(*)
+  get()->attrs_type_key = AttrsType::_type_key;
+  get()->attrs_type_index = AttrsType::RuntimeTypeIndex();
+  return *this;
+}
+
+inline OpRegEntry& OpRegEntry::set_attrs_type_key(const ffi::String& key) {  // NOLINT(*)
+  get()->attrs_type_key = key;
+  get()->attrs_type_index = tvm::ffi::TypeKeyToIndex(key.c_str());
+  return *this;
+}
+
+inline OpRegEntry& OpRegEntry::set_support_level(int32_t n) {  // NOLINT(*)
+  get()->support_level = n;
+  return *this;
+}
+
+template <typename ValueType>
+inline OpRegEntry& OpRegEntry::set_attr(  // NOLINT(*)
+    const std::string& attr_name, const ValueType& value, int plevel) {
+  TVM_FFI_ICHECK_GT(plevel, 0) << "plevel in set_attr must be greater than 0";
+  UpdateAttr(attr_name, Any(value), plevel);
+  return *this;
+}
+
+// member functions of OpAttrMap
+
+template <typename ValueType>
+inline ValueType OpAttrMap<ValueType>::get(const Expr& expr, ValueType def_value) const {
+  TVM_FFI_ICHECK(expr.defined());
+  if (const OpNode* op = expr.as<OpNode>()) {
+    return this->map_.get(ffi::GetRef<Op>(op), def_value);
+  } else {
+    return def_value;
+  }
+}
+
+}  // namespace tvm
+#endif  // TVM_IR_OP_H_

include/tvm/ir/scope_stack.h

@@ -0,0 +1,123 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/scope_stack.h
+ * \brief A generic scope stack for managing hierarchical state during IR visiting.
+ */
+#ifndef TVM_IR_SCOPE_STACK_H_
+#define TVM_IR_SCOPE_STACK_H_
+
+#include <tvm/ffi/error.h>
+
+#include <deque>
+#include <type_traits>
+
+namespace tvm {
+
+/*!
+ * \brief A scope stack for maintaining hierarchical state during IR visiting.
+ *
+ * During IR tree traversal, visitors often need to track scope-local state
+ * (e.g., active constraints, variable bindings) that should be automatically
+ * cleaned up when leaving a scope. ScopeStack provides this via WithNewScope,
+ * which pushes a new element on entry and pops it on exit.
+ *
+ * \code
+ *   ScopeStack<WithGroup<ConstraintContext>> constraints;
+ *
+ *   // In VisitStmt_(ForNode):
+ *   return constraints.WithNewScope([&]() -> Stmt {
+ *     constraints.Current().Emplace(analyzer, condition);
+ *     return StmtExprMutator::VisitStmt_(op);
+ *   });
+ * \endcode
+ *
+ * \tparam T The element type stored on the stack. Must be default-constructible.
+ */
+template <typename T>
+class ScopeStack {
+ public:
+  /*! \brief Construct with one initial scope level. */
+  ScopeStack() { stack_.emplace_back(); }
+
+  /*! \brief Return the number of active scopes. */
+  size_t size() const { return stack_.size(); }
+
+  /*! \brief Return true if no scopes are active. */
+  bool empty() const { return stack_.empty(); }
+
+  /*!
+   * \brief Access the current (innermost) scope element.
+   *
+   * The returned reference is stable across push_back/pop_back because
+   * std::deque guarantees pointer stability for these operations.
+   *
+   * \return Mutable reference to the top element.
+   */
+  T& Current() {
+    TVM_FFI_ICHECK(!stack_.empty());
+    return stack_.back();
+  }
+
+  /*! \brief Const access to the current (innermost) scope element. */
+  const T& Current() const {
+    TVM_FFI_ICHECK(!stack_.empty());
+    return stack_.back();
+  }
+
+  /*!
+   * \brief Execute body within a new scope.
+   *
+   * Pushes a new T onto the stack, executes the body, then pops it.
+   *
+   * \param body A callable to execute within the scope.
+   * \return The return value of body(), if non-void.
+   */
+  template <typename F>
+  auto WithNewScope(F&& body) -> decltype(body()) {
+    stack_.emplace_back();
+    struct Guard {
+      std::deque<T>* stack;
+      ~Guard() noexcept(false) { stack->pop_back(); }
+    } guard{&stack_};
+    if constexpr (std::is_void_v<decltype(body())>) {
+      body();
+    } else {
+      return body();
+    }
+  }
+
+ private:
+  /*!
+   * \brief The scope stack.
+   *
+   * We use std::deque rather than std::vector for pointer stability:
+   * references returned by Current() remain valid across push/pop operations.
+   * This is critical because methods called on Current() (e.g., Emplace on
+   * a WithGroup) may trigger re-entrant code that pushes new scopes onto
+   * the same stack. With std::vector the internal buffer reallocation would
+   * invalidate the reference, causing use-after-free.
+   */
+  std::deque<T> stack_;
+};
+
+}  // namespace tvm
+
+#endif  // TVM_IR_SCOPE_STACK_H_

include/tvm/ir/source_map.h

@@ -0,0 +1,226 @@
+/*
+ * 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.
+ */
+/*!
+ * \file source_map.h
+ * \brief A map from source names to source code.
+ */
+#ifndef TVM_IR_SOURCE_MAP_H_
+#define TVM_IR_SOURCE_MAP_H_
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/runtime/base.h>
+
+#include <fstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+
+/*!
+ * \brief The source name in the Span
+ * \sa SourceNameNode, Span
+ */
+class SourceName;
+/*!
+ * \brief The name of a source fragment.
+ */
+class SourceNameNode : public ffi::Object {
+ public:
+  /*! \brief The source name. */
+  ffi::String name;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<SourceNameNode>().def_ro("name", &SourceNameNode::name);
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.SourceName", SourceNameNode, ffi::Object);
+};
+
+/*!
+ * \brief The source name of a file span.
+ * \sa SourceNameNode, Span
+ */
+class SourceName : public ffi::ObjectRef {
+ public:
+  /*!
+   * \brief Get an SourceName for a given operator name.
+   *  Will raise an error if the source name has not been registered.
+   * \param name Name of the operator.
+   * \return SourceName valid throughout program lifetime.
+   */
+  TVM_DLL static SourceName Get(const ffi::String& name);
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SourceName, ffi::ObjectRef, SourceNameNode);
+};
+
+/*!
+ * \brief Span information for debugging purposes
+ */
+class Span;
+/*!
+ * \brief Stores locations in frontend source that generated a node.
+ */
+class SpanNode : public ffi::Object {
+ public:
+  /*! \brief The source name. */
+  SourceName source_name;
+  /*! \brief The line number. */
+  int line;
+  /*! \brief The column offset. */
+  int column;
+  /*! \brief The end line number. */
+  int end_line;
+  /*! \brief The end column number. */
+  int end_column;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<SpanNode>()
+        .def_ro("source_name", &SpanNode::source_name)
+        .def_ro("line", &SpanNode::line)
+        .def_ro("column", &SpanNode::column)
+        .def_ro("end_line", &SpanNode::end_line)
+        .def_ro("end_column", &SpanNode::end_column);
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+  TVM_FFI_DECLARE_OBJECT_INFO("ir.Span", SpanNode, ffi::Object);
+};
+
+class Span : public ffi::ObjectRef {
+ public:
+  TVM_DLL Span(SourceName source_name, int line, int end_line, int column, int end_column);
+
+  /*! \brief Merge two spans into one which captures the combined regions. */
+  TVM_DLL Span Merge(const Span& other) const;
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Span, ffi::ObjectRef, SpanNode);
+};
+
+/*!
+ * \brief Store a list of spans for an expr generated from mulitple source exprs
+ */
+class SequentialSpanNode : public SpanNode {
+ public:
+  /*! \brief The original source list of spans to construct a sequential span. */
+  ffi::Array<Span> spans;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<SequentialSpanNode>().def_ro("spans", &SequentialSpanNode::spans);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.SequentialSpan", SequentialSpanNode, SpanNode);
+};
+
+/*!
+ * \brief Reference class of SequentialSpanNode.
+ * \sa SequentialSpanNode
+ */
+class SequentialSpan : public Span {
+ public:
+  TVM_DLL SequentialSpan(ffi::Array<Span> spans);
+
+  TVM_DLL SequentialSpan(std::initializer_list<Span> init);
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SequentialSpan, Span, SequentialSpanNode);
+};
+
+/*! \brief A program source in any language.
+ *
+ * Could represent the source from an ML framework or a source
+ * representing a tvm::IRModule.
+ */
+class Source;
+
+class SourceNode : public ffi::Object {
+ public:
+  /*! \brief The source name. */
+  SourceName source_name;
+
+  /*! \brief The raw source. */
+  ffi::String source;
+
+  /*! \brief A mapping of line breaks into the raw source. */
+  std::vector<std::pair<int, int>> line_map;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<SourceNode>()
+        .def_ro("source_name", &SourceNode::source_name)
+        .def_ro("source", &SourceNode::source);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Source", SourceNode, ffi::Object);
+};
+
+class Source : public ffi::ObjectRef {
+ public:
+  TVM_DLL Source(SourceName src_name, std::string source);
+  TVM_DLL tvm::ffi::String GetLine(int line);
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Source, ffi::ObjectRef, SourceNode);
+};
+
+/*!
+ * \brief A mapping from a unique source name to source fragment.
+ */
+class SourceMap;
+/*!
+ * \brief Stores locations in frontend source that generated a node.
+ */
+class SourceMapObj : public ffi::Object {
+ public:
+  /*! \brief The source mapping. */
+  ffi::Map<SourceName, Source> source_map;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<SourceMapObj>().def_ro("source_map", &SourceMapObj::source_map);
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.SourceMap", SourceMapObj, ffi::Object);
+};
+
+class SourceMap : public ffi::ObjectRef {
+ public:
+  explicit SourceMap(ffi::Map<SourceName, Source> source_map);
+
+  explicit SourceMap(std::initializer_list<std::pair<SourceName, Source>> source_map)
+      : SourceMap(ffi::Map<SourceName, Source>(source_map)) {}
+
+  SourceMap() : SourceMap(ffi::Map<SourceName, Source>()) {}
+
+  void Add(const Source& source);
+
+  SourceMapObj* operator->() {
+    TVM_FFI_ICHECK(get() != nullptr);
+    return static_cast<SourceMapObj*>(get_mutable());
+  }
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SourceMap, ffi::ObjectRef, SourceMapObj);
+};
+
+}  // namespace tvm
+
+#endif  // TVM_IR_SOURCE_MAP_H_

include/tvm/ir/transform.h

@@ -0,0 +1,584 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/transform.h
+ *
+ * This file implements a pass manager. The pass manager manages a sequence
+ * of IRModule -> IRModule transformation passes over a particlar unit of AST. The
+ * design is largely inspired from LLVM's pass manager and modern deep learning
+ * frameworks that perform tensor->tensor transformations.
+ *
+ * The responsibilities of a traditional compiler pass manager usually involves:
+ *  - Organizing the execution order of optimization passes though not
+ * necessarily in the optimal sequence.
+ *  - Collecting required analysis information and keep them up-to-date.
+ *  - Reducing the effort required to implement new passes for compiler
+ * developers, etc.
+ *
+ * Similar to LLVM's pass manager, we designed the Relax pass manager to work
+ * different granularity, i.e. module level, function level, and even sequential
+ * passe that contains a host of passes.
+ *
+ * However, we also extend the functionality of the traditional pass manager
+ * with the consideration of requirements/convention from deep learning
+ * frameworks, such as Pytorch and Gluon, etc. Each pass in the Relax pass
+ * manager performs the IRModule -> IRModule transformation. All
+ * different types of passes, including the sequential-level pass object, are
+ * essentially pass objects. This design, therefore, effectively provides users
+ * a consistent and convenient interface, i.e. Pass, to play with. It offers a
+ * means to ease the development and testing of Relax passes. For example, with
+ * the pass manager, external users will be able to have custom passes correctly
+ * scheduled without having to modify a single handcrafted pass order.
+ *
+ * In the future we need to describe constraints between passes. For example,
+ * we may want to preserve dependencies between different passes and validate
+ * them on the completion of a certain pass.
+ *
+ * We also need to store side information and import the error reporting system.
+ */
+#ifndef TVM_IR_TRANSFORM_H_
+#define TVM_IR_TRANSFORM_H_
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/creator.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/string.h>
+#include <tvm/ir/instrument.h>
+#include <tvm/ir/module.h>
+#include <tvm/ir/with_context.h>
+
+#include <string>
+#include <type_traits>
+#include <utility>
+
+namespace tvm {
+namespace transform {
+
+/*!
+ * \brief PassContextNode contains the information that a pass can rely on,
+ * such as analysis results.
+ * \sa PassContext
+ */
+class PassContextNode : public ffi::Object {
+ public:
+  /*! \brief The default optimization level. */
+  int opt_level{2};
+
+  /*! \brief The list of required passes. */
+  ffi::Array<ffi::String> required_pass;
+  /*! \brief The list of disabled passes. */
+  ffi::Array<ffi::String> disabled_pass;
+  /*! \brief Pass specific configurations. */
+  ffi::Map<ffi::String, Any> config;
+
+  /*! \brief A list of pass instrument implementations. */
+  ffi::Array<instrument::PassInstrument> instruments;
+
+  PassContextNode() = default;
+
+  /*!
+   * \brief Get a config value from the pass context.
+   *
+   * \param key The config key.
+   * \param default_value The default value if the key does not exist, defaults to nullptr.
+   *
+   * \return The result
+   *
+   * \tparam TOBjectRef the expected object type.
+   * \throw Error if the key exists but the value does not match TObjectRef.
+   */
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetConfig(
+      const std::string& key,
+      ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
+    if (!config.defined()) return default_value;
+    auto it = config.find(key);
+    if (it != config.end()) {
+      return (*it).second.as_or_throw<ffi::Optional<TObjectRef>>();
+    } else {
+      return default_value;
+    }
+  }
+  // variant that uses TObjectRef to enable implicit conversion to default value.
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetConfig(const std::string& key, TObjectRef default_value) const {
+    return GetConfig<TObjectRef>(key, ffi::Optional<TObjectRef>(default_value));
+  }
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<PassContextNode>()
+        .def_ro("opt_level", &PassContextNode::opt_level)
+        .def_ro("required_pass", &PassContextNode::required_pass)
+        .def_ro("disabled_pass", &PassContextNode::disabled_pass)
+        .def_ro("instruments", &PassContextNode::instruments)
+        .def_ro("config", &PassContextNode::config);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.PassContext", PassContextNode, ffi::Object);
+};
+
+/*!
+ * \brief PassContext that is used to configure the pass behavior.
+ *
+ * \code
+ *
+ *  auto new_ctx = PassContext::Create();
+ *  ctx->opt_level = 2;
+ *  With<PassContext> scope(ctx);
+ *  // pass context in effect.
+ *
+ * \endcode
+ * \sa PassContextNode
+ */
+class PassContext : public ffi::ObjectRef {
+ public:
+  PassContext() {}
+  /*!
+   * \brief constructor with UnsafeInit
+   */
+  explicit PassContext(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
+  /*!
+   * \brief constructor with ffi::ObjectPtr
+   */
+  explicit PassContext(ffi::ObjectPtr<PassContextNode> n) : ffi::ObjectRef(n) {}
+  /*!
+   * \brief const accessor.
+   * \return const access pointer.
+   */
+  const PassContextNode* operator->() const {
+    TVM_FFI_ICHECK(get() != nullptr);
+    return static_cast<const PassContextNode*>(get());
+  }
+  /*!
+   * \brief mutable accessor.
+   * \return mutable access pointer.
+   */
+  PassContextNode* operator->() {
+    TVM_FFI_ICHECK(get() != nullptr);
+    return static_cast<PassContextNode*>(get_mutable());
+  }
+
+  /*!
+   * \brief Construct a PassContext containing the default configurations.
+   * \return The new PassContext.
+   */
+  TVM_DLL static PassContext Create();
+  /*!
+   * \brief Get the default pass context in the current scope.
+   * \return The pass context.
+   */
+  TVM_DLL static PassContext Current();
+
+  /*!
+   * \brief Get all supported configuration names and metadata, registered within the PassContext.
+   * \return Map indexed by the config name, pointing to the metadata map as key-value
+   */
+  TVM_DLL static ffi::Map<ffi::String, ffi::Map<ffi::String, ffi::String>> ListConfigs();
+
+  /*!
+   * \brief Call instrument implementations' callbacks when entering PassContext.
+   *        The callbacks are called in order, and if one raises an exception, the rest will not be
+   *        called.
+   */
+  TVM_DLL void InstrumentEnterPassContext();
+
+  /*!
+   * \brief Call instrument implementations' callbacks when exiting PassContext.
+   *        The callbacks are called in order, and if one raises an exception, the rest will not be
+   *        called.
+   */
+  TVM_DLL void InstrumentExitPassContext();
+
+  /*!
+   * \brief Call instrument implementations' callbacks before a pass run.
+   *        The callbacks are called in order, and if one raises an exception, the rest will not be
+   *        called.
+   *
+   * \param mod The module that an optimization pass runs on.
+   * \param info The pass information.
+   *
+   * \return false: the pass is skipped; true: the pass runs.
+   */
+  TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
+
+  /*!
+   * \brief Call instrument implementations callbacks after a pass run.
+   *        The callbacks are called in order, and if one raises an exception, the rest will not be
+   *        called.
+   *
+   * \param mod The module that an optimization pass runs on.
+   * \param info The pass information.
+   */
+  TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
+
+  /*!
+   * \brief Check whether a pass is enabled.
+   * \param info The pass information.
+   * \return true if the pass is enabled. Otherwise, false.
+   */
+  TVM_DLL bool PassEnabled(const PassInfo& info) const;
+
+  /*!
+   * \brief Register a valid configuration option and its ValueType for validation.
+   *
+   * \param key The configuration key.
+   * \tparam ValueType The value type to be registered
+   */
+  template <typename ValueType>
+  static int32_t RegisterConfigOption(const char* key) {
+    // NOTE: we could further update the function later.
+    if constexpr (std::is_base_of_v<ffi::ObjectRef, ValueType>) {
+      int32_t tindex = ffi::TypeToRuntimeTypeIndex<ValueType>::v();
+      auto type_key = ffi::TypeIndexToTypeKey(tindex);
+      auto legalization = [=](ffi::Any value) -> ffi::Any {
+        if (auto opt_map = value.try_cast<ffi::Map<ffi::String, ffi::Any>>()) {
+          return ffi::reflection::ObjectCreator(type_key)(opt_map.value());
+        } else {
+          auto opt_val = value.try_cast<ValueType>();
+          if (!opt_val.has_value()) {
+            TVM_FFI_THROW(AttributeError)
+                << "Expect config " << key << " to have type " << type_key << ", but instead get "
+                << ffi::details::AnyUnsafe::GetMismatchTypeInfo<ValueType>(value);
+          }
+          return *opt_val;
+        }
+      };
+      RegisterConfigOption(key, type_key, legalization);
+    } else {
+      // non-object type, do not support implicit conversion from map
+      std::string type_str = ffi::TypeTraits<ValueType>::TypeStr();
+      auto legalization = [=](ffi::Any value) -> ffi::Any {
+        auto opt_val = value.try_cast<ValueType>();
+        if (!opt_val.has_value()) {
+          TVM_FFI_THROW(AttributeError)
+              << "Expect config " << key << " to have type " << type_str << ", but instead get "
+              << ffi::details::AnyUnsafe::GetMismatchTypeInfo<ValueType>(value);
+        } else {
+          return *opt_val;
+        }
+      };
+      RegisterConfigOption(key, type_str, legalization);
+    }
+    return 0;
+  }
+
+  // accessor.
+  using ContainerType = PassContextNode;
+  class Internal;
+
+ private:
+  // The entry of a pass context scope.
+  TVM_DLL void EnterWithScope();
+  // The exit of a pass context scope.
+  TVM_DLL void ExitWithScope();
+  // Register configuration key value type.
+  TVM_DLL static void RegisterConfigOption(const char* key, ffi::String value_type_str,
+                                           std::function<ffi::Any(ffi::Any)> legalization);
+
+  // Classes to get the Python `with` like syntax.
+  friend class Internal;
+  friend class With<PassContext>;
+};
+
+/*!
+ * \brief Create a pass-config object with all default values, using the
+ *        reflection defaults.
+ * \tparam TConfig the ObjectRef type to be created.
+ * \return An instance with all reflection-defined default values applied.
+ */
+template <typename TConfig>
+inline TConfig PassConfigWithDefaults() {
+  static_assert(std::is_base_of_v<ffi::ObjectRef, TConfig>,
+                "Can only create ObjectRef-derived types");
+  using ContainerType = typename TConfig::ContainerType;
+  static auto finit_object = ffi::Function::GetGlobalRequired("ffi.MakeObjectFromPackedArgs");
+  ffi::AnyView packed_args[1];
+  packed_args[0] = ContainerType::RuntimeTypeIndex();
+  ffi::Any rv;
+  finit_object.CallPacked(ffi::PackedArgs(packed_args, 1), &rv);
+  return rv.cast<TConfig>();
+}
+
+#define TVM_PASS_CTX_CONFIG_VAR_DEF [[maybe_unused]] static uint32_t __make_PassContext_tid
+
+/*!
+ * \brief Helper macro to register the object type to runtime.
+ *  Makes sure that the runtime type table is correctly populated.
+ *
+ *  Use this macro in the cc file for each terminal class.
+ */
+#define TVM_REGISTER_PASS_CONFIG_OPTION(Key, ValueType)          \
+  TVM_FFI_STR_CONCAT(TVM_PASS_CTX_CONFIG_VAR_DEF, __COUNTER__) = \
+      ::tvm::transform::PassContext::RegisterConfigOption<ValueType>(Key)
+
+/*!
+ * \brief Meta data that will be used to help optimization and analysis.
+ * \sa PassInfo
+ */
+class PassInfoNode : public ffi::Object {
+ public:
+  /*! \brief The minimal optimization level that this pass will be enabled. */
+  int opt_level;
+
+  /*! \brief The name of an optimization/analysis pass. */
+  ffi::String name;
+
+  /*! \brief Boolean that tells whether this pass will be traced or not. */
+  bool traceable;
+
+  /*! \brief The passes that are required to perform the current pass. */
+  ffi::Array<ffi::String> required;
+
+  PassInfoNode() = default;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<PassInfoNode>()
+        .def_ro("opt_level", &PassInfoNode::opt_level)
+        .def_ro("name", &PassInfoNode::name)
+        .def_ro("required", &PassInfoNode::required)
+        .def_ro("traceable", &PassInfoNode::traceable);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.PassInfo", PassInfoNode, ffi::Object);
+};
+
+/*!
+ * \brief Managed reference class for PassInfoNode
+ * \sa PassInfoNode
+ */
+class PassInfo : public ffi::ObjectRef {
+ public:
+  /*!
+   * \brief Constructor
+   * \param opt_level The optimization level
+   * \param name Name of the pass.
+   * \param required  The passes that are required to perform the current pass.
+   * \param traceable Boolean that tells whether the pass is traceable.
+   */
+  TVM_DLL PassInfo(int opt_level, ffi::String name, ffi::Array<ffi::String> required,
+                   bool traceable);
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PassInfo, ffi::ObjectRef, PassInfoNode);
+};
+
+/*!
+ * \brief PassNode is the base type of differnt types of optimization passes.
+ * It is designed as a pure class and implemented by different pass subclasses
+ * at different granularity of Relax nodes.
+ */
+class PassNode : public ffi::Object {
+ public:
+  virtual ~PassNode() {}
+  /*!
+   * \brief Get the pass information/meta data. */
+  virtual PassInfo Info() const = 0;
+
+  /*!
+   * \brief Transform mod using the default PassContext in the current scope.
+   *
+   * \param mod The module that an optimization pass runs on.
+   *
+   * \return The transformed module.
+   */
+  IRModule operator()(IRModule mod) const {
+    return this->operator()(std::move(mod), PassContext::Current());
+  }
+
+  /*!
+   * \brief Transform mod using a functor under a given pass context.
+   *
+   * \param mod The module that an optimization pass runs on.
+   * \param pass_ctx The pass context that can provide information for the optimization.
+   *
+   * \return The transformed module.
+   */
+  virtual IRModule operator()(IRModule mod, const PassContext& pass_ctx) const = 0;
+  TVM_FFI_DECLARE_OBJECT_INFO("transform.Pass", PassNode, ffi::Object);
+};
+
+class Pass : public ffi::ObjectRef {
+ public:
+  /*!
+   * \brief Transform mod using the default PassContext in the current scope.
+   *
+   * \code
+   *
+   * // If you do no longer need the input module
+   * // it is recommended to use std::move to move your input module.
+   * mod = pass(std::move(mod));
+   *
+   * \endcode
+   *
+   * \param mod The module that an optimization pass runs on.
+   *
+   * \return The transformed module.
+   */
+  IRModule operator()(IRModule mod) const;
+
+  /*!
+   * \brief Transform mod using a functor under a given pass context.
+   *
+   * \param mod The module that an optimization pass runs on.
+   * \param pass_ctx The pass context that can provide information for the optimization.
+   *
+   * \return The transformed module.
+   */
+  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const;
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Pass, ffi::ObjectRef, PassNode);
+
+ private:
+  IRModule static AssertImmutableModule(const IRModule& mod, const PassNode* node,
+                                        const PassContext& pass_ctx);
+};
+
+/*!
+ * \brief The SequentialNode contains a set of passes that transform Relax
+ * programs from one AST to another semantically equivalent one.
+ *
+ * One example of this level of pass is that the pass manager needs to correctly
+ * perform a host of optimizations with a given optimization level and disabled
+ * passes.
+ */
+class SequentialNode : public PassNode {
+ public:
+  /* \brief The pass meta data.*/
+  PassInfo pass_info;
+
+  /*! \brief A list of passes that used to compose a sequential pass. */
+  tvm::ffi::Array<Pass> passes;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<SequentialNode>()
+        .def_ro("pass_info", &SequentialNode::pass_info)
+        .def_ro("passes", &SequentialNode::passes);
+  }
+
+  /*!
+   * \brief Get the pass information/meta data.
+   */
+  PassInfo Info() const override { return pass_info; }
+
+  /*!
+   * \brief Resolve the pass dependency. It globs all required passes by
+   *        a given pass and executes them.
+   *
+   * \param mod The module that an optimization pass runs on.
+   *
+   * TODO(zhiics) Build a dependency graph among the passes using provided
+   * metadata, i.e. required_passes. Likely, we can have a data structure, i.e.
+   * PassInfo, to store the relevant information including the parent passes.
+   */
+  void ResolveDependency(const IRModule& mod);
+
+  /*!
+   * \brief Perform optimizations on a series of passes. The aforementioned
+   *        typical pass manager jobs could be done by it. This function could
+   *        be overloaded to focus on different metrics, i.e. performance,
+   *        memory footprint, etc.
+   *
+   * \param mod The module that these passes are applied on.
+   * \param pass_ctx The context that these passes execute on.
+   *
+   * \return Return the updated module.
+   */
+  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.Sequential", SequentialNode, PassNode);
+};
+
+class Sequential : public Pass {
+ public:
+  /*!
+   * \brief The constructor of `Sequential`.
+   *
+   * \param passes The passes to apply.
+   * \param pass_info The pass metadata.
+   */
+  TVM_DLL Sequential(ffi::Array<Pass> passes, PassInfo pass_info);
+
+  /*!
+   * \brief The constructor of `Sequential`.
+   *
+   * \param passes The passes to apply.
+   * \param name The name of a sequential pass. It's defaulted to "sequential".
+   *        This allows users to only provide a list of passes and execute them
+   *        under a given context.
+   */
+  TVM_DLL Sequential(ffi::Array<Pass> passes, ffi::String name = "sequential");
+
+  Sequential() = default;
+  explicit Sequential(ffi::ObjectPtr<SequentialNode> n) : Pass(n) {}
+
+  const SequentialNode* operator->() const;
+  using ContainerType = SequentialNode;
+};
+
+/*
+ * \brief Create a module pass.
+ *
+ * \param pass_func The packed function that contains the optimization.
+ * \param opt_level The optimization level of the module pass.
+ * \param name The name of the module pass.
+ * \param required The list of the passes that the module pass is dependent on.
+ *
+ * \return The created module pass.
+ */
+TVM_DLL Pass CreateModulePass(std::function<IRModule(IRModule, PassContext)> pass_func,
+                              int opt_level, ffi::String name, ffi::Array<ffi::String> required,
+                              bool traceable = false);
+
+/*!
+ * \brief A special trace pass that prints the header and IR to LOG(INFO).
+ * \param header The header to be attached to the output.
+ * \return The pass.
+ */
+TVM_DLL Pass PrintIR(ffi::String header = "");
+
+/*!
+ * \brief Enrich a pass-time error with a TVMScript-rendered, underlined source
+ *        location derived from the error's embedded VisitErrorContext.
+ *
+ * Returns an ffi::Error that preserves err's kind, original message, and
+ * backtrace, and appends the failing pass name plus the offending location
+ * rendered as TVMScript (the whole \p mod, or local to \p func when provided).
+ * The returned error drops the VisitErrorContext payload, so an outer catch
+ * that re-enriches finds no context and returns the error unchanged.
+ *
+ * Pure and total: never throws; returns \p err unchanged when there is no
+ * context, the path is unresolvable, or rendering fails.
+ *
+ * \param err The error thrown by the pass body.
+ * \param mod The IRModule the pass ran on (the access-path root, or the
+ *            container of \p func when \p func is provided).
+ * \param pass_name The name of the failing pass, shown in the message.
+ * \param func When set, resolve and render the location local to
+ *             \p mod->functions[func]; otherwise use the whole module.
+ * \return The enriched (or, on any fallback, the original) error.
+ */
+TVM_DLL ffi::Error EnrichPassErrorWithContext(
+    const ffi::Error& err, const IRModule& mod, ffi::String pass_name,
+    ffi::Optional<GlobalVar> func = ffi::Optional<GlobalVar>(std::nullopt));
+
+}  // namespace transform
+}  // namespace tvm
+
+#endif  // TVM_IR_TRANSFORM_H_

include/tvm/ir/type.h

@@ -0,0 +1,207 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/type.h
+ * \brief IR/AST nodes for TVM types shared across IR variants.
+ */
+#ifndef TVM_IR_TYPE_H_
+#define TVM_IR_TYPE_H_
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/base_expr.h>
+#include <tvm/ir/source_map.h>
+
+#include <string>
+
+namespace tvm {
+
+/*!
+ * \brief Low-level raw pointer type.
+ *
+ *  PointerType represents type hints in the TIR to be
+ *  passed to the final code generator.
+ *
+ *  PointerType should not occur in the high-level analysis.
+ *
+ * \sa PointerType
+ */
+class PointerTypeNode : public TypeNode {
+ public:
+  /*!
+   * \brief The type of the element which the pointer points to.
+   */
+  Type element_type = PrimType::Void();
+  /*!
+   * \brief The storage scope of the pointer
+   */
+  ffi::String storage_scope;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<PointerTypeNode>()
+        .def_ro("element_type", &PointerTypeNode::element_type)
+        .def_ro("storage_scope", &PointerTypeNode::storage_scope);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.PointerType", PointerTypeNode, TypeNode);
+};
+
+/*
+ * \brief Managed reference to PointerTypeNode.
+ * \sa PointerTypeNode
+ */
+class PointerType : public Type {
+ public:
+  /*!
+   * \brief Constructor
+   * \param element_type The type of the element which the pointer points to.
+   * \param storage_scope The storage scope into which the pointer addresses
+   */
+  TVM_DLL explicit PointerType(Type element_type, ffi::String storage_scope = "");
+
+  /*! \brief Construct an opaque pointer with void element type. */
+  TVM_DLL static PointerType VoidPointerTy(ffi::String storage_scope = "");
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PointerType, Type, PointerTypeNode);
+};
+
+/*!
+ * \brief The type of tuple values.
+ * \sa TupleType
+ */
+class TupleTypeNode : public TypeNode {
+ public:
+  /*! \brief The type of each field in the tuple. */
+  ffi::Array<Type> fields;
+
+  TupleTypeNode() {}
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<TupleTypeNode>().def_ro("fields", &TupleTypeNode::fields);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TupleType", TupleTypeNode, TypeNode);
+};
+
+/*!
+ * \brief Managed reference to TupleTypeNode.
+ * \sa TupleTypeNode.
+ */
+class TupleType : public Type {
+ public:
+  /*!
+   * \brief Constructor
+   * \param fields Fields in the tuple.
+   * \param span The span of the type.
+   */
+  TVM_DLL explicit TupleType(ffi::Array<Type> fields, Span span = Span());
+
+  /*!
+   * \brief Create an empty tuple type that constains nothing.
+   * \return A empty tuple type.
+   */
+  TVM_DLL TupleType static Empty();
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleType, Type, TupleTypeNode);
+};
+
+/*!
+ * \return a type that represents void.
+ */
+inline Type VoidType() { return TupleType::Empty(); }
+
+/*!
+ * \brief Check whether the tyep represents void.
+ * \return The check result.
+ */
+inline bool IsVoidType(const Type& type) {
+  auto* n = type.as<TupleTypeNode>();
+  return n && n->fields.size() == 0;
+}
+
+/*!
+ * \brief Function type.
+ *
+ * We support polymorphic function type.
+ * This can be roughly viewed as template function in C++.
+ *
+ * \sa FuncType, TypeConstraint
+ */
+class FuncTypeNode : public TypeNode {
+ public:
+  /*! \brief type type of arguments */
+  ffi::Array<Type> arg_types;
+  /*! \brief The type of return value. */
+  Type ret_type = VoidType();
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<FuncTypeNode>()
+        .def_ro("arg_types", &FuncTypeNode::arg_types)
+        .def_ro("ret_type", &FuncTypeNode::ret_type);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.FuncType", FuncTypeNode, TypeNode);
+};
+
+/*!
+ * \brief Managed reference to FuncTypeNode.
+ * \sa FuncTypeNode
+ */
+class FuncType : public Type {
+ public:
+  /*!
+   * \brief Constructor
+   * \param arg_types The types of the arguments.
+   * \param ret_type The type of the return value.
+   * \param span The span information.
+   * \sa FuncTypeNode for more docs about these fields.
+   */
+  TVM_DLL FuncType(ffi::Array<Type> arg_types, Type ret_type, Span span = Span());
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeNode);
+};
+
+/*!
+ * \brief The type of tensor map.
+ * \sa TensorMapType
+ */
+class TensorMapTypeNode : public TypeNode {
+ public:
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<TensorMapTypeNode>();
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TensorMapType", TensorMapTypeNode, TypeNode);
+};
+
+/*!
+ * \brief Managed reference to TensorMapTypeNode.
+ * \sa TensorMapTypeNode
+ */
+class TensorMapType : public Type {
+ public:
+  TVM_DLL TensorMapType(Span span = Span());
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorMapType, Type, TensorMapTypeNode);
+};
+
+}  // namespace tvm
+#endif  // TVM_IR_TYPE_H_

include/tvm/ir/unique_name_supply.h

@@ -0,0 +1,143 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/unique_name_supply.h
+ * \brief UniqueNameSupply that can be used to generate unique variable names.
+ */
+#ifndef TVM_IR_UNIQUE_NAME_SUPPLY_H_
+#define TVM_IR_UNIQUE_NAME_SUPPLY_H_
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/expr.h>
+
+#include <cstdint>
+#include <string>
+#include <utility>
+
+namespace tvm {
+
+/*!
+ * \brief UniqueNameSupply can be used to generate unique names.
+ */
+class UniqueNameSupplyNode : public ffi::Object {
+ public:
+  /*!
+   * \brief Empty constructor. Needed by the TVM_REGISTER_NODE_TYPE macro.
+   */
+  UniqueNameSupplyNode() = default;
+
+  /*!
+   * \brief Constructor.
+   * \param prefix The prefix to be used with this UniqueNameSupply.
+   * \param name_map The map used to guarantee uniqueness.
+   */
+  UniqueNameSupplyNode(const ffi::String& prefix, ffi::Map<ffi::String, int64_t> name_map)
+      : prefix_(prefix), name_map(std::move(name_map)) {}
+
+  /*!
+   * \brief Generates a unique name from this UniqueNameSupply.
+   * \param name The name from which the generated name is derived.
+   * \param add_prefix If set to true, then the prefix of this UniqueNameSupply will be prepended to
+   * the name.
+   * \param add_underscore If set to true, add '_' between prefix and a digit.
+   * \return A unique name.
+   */
+  ffi::String FreshName(const ffi::String& name, bool add_prefix = true,
+                        bool add_underscore = true);
+
+  /*!
+   * \brief Reserves an existing name with this UniqueNameSupply.
+   * \param name The name to be reserved.
+   * \param add_prefix If set to true, then the prefix of this UniqueNameSupply will be prepended to
+   * the name before reserving it. \return The name that was reserved with the UniqueNameSupply. It
+   * can be different if a prefix is added.
+   */
+  ffi::String ReserveName(const ffi::String& name, bool add_prefix = true);
+
+  /*!
+   * \brief Checks if this UniqueNameSupply already generated a name.
+   * \param name The name to check.
+   * \param add_prefix If set to true, then the prefix of this UniqueNameSupply will be prepended to
+   * the name before checking for it. \return True if the name has already been generated. False
+   * otherwise.
+   */
+  bool ContainsName(const ffi::String& name, bool add_prefix = true);
+
+  // Prefix for all GlobalVar names. It can be empty.
+  std::string prefix_;
+
+  static constexpr const bool _type_mutable = true;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<UniqueNameSupplyNode>();
+  }
+
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.UniqueNameSupply", UniqueNameSupplyNode, ffi::Object);
+
+ private:
+  /*! \brief Helper function to add the UniqueNameSupply prefix to the name. */
+  ffi::String AddPrefixToName(const ffi::String& name);
+
+  /*!
+   * \brief Function that will generate a unique name.
+   * \param name The name to be used as a base.
+   * \param add_underscore If set to true, add '_' between prefix and a digit.
+   * \return A unique name.
+   */
+  std::string GetUniqueName(std::string name, bool add_underscore = true);
+
+  /*! \brief A map that is used to generate unique names. */
+  ffi::Map<ffi::String, int64_t> name_map;
+};
+
+/*!
+ * \brief Managed reference class to UniqueNameSupplyNode.
+ * \sa UniqueNameSupplyNode
+ */
+class UniqueNameSupply : public ffi::ObjectRef {
+ public:
+  /*!
+   * \brief Constructor.
+   * \param prefix The prefix to be used with this UniqueNameSupply.
+   * \param name_map An optional map.
+   */
+  TVM_DLL explicit UniqueNameSupply(const ffi::String& prefix = "",
+                                    ffi::Map<ffi::String, int64_t> name_map = {});
+
+  /*!
+   * \brief Construct UniqueNameSupply by reserving names from the given iterator range.
+   *
+   * The functor should return the name of the dereferenced object.
+   */
+  template <typename Iter, typename Lambda>
+  TVM_DLL UniqueNameSupply(Iter begin, Iter end, Lambda f) : UniqueNameSupply("") {
+    for (auto it = begin; it != end; ++it) {
+      this->operator->()->ReserveName(f(*it), false);
+    }
+  }
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(UniqueNameSupply, ffi::ObjectRef,
+                                                UniqueNameSupplyNode);
+};
+
+}  // namespace tvm
+
+#endif  // TVM_IR_UNIQUE_NAME_SUPPLY_H_

include/tvm/ir/with_context.h

@@ -0,0 +1,168 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/ir/with_context.h
+ * \brief RAII wrapper function to enter and exit a context object
+ *        similar to python's with syntax.
+ */
+#ifndef TVM_IR_WITH_CONTEXT_H_
+#define TVM_IR_WITH_CONTEXT_H_
+
+#include <exception>
+#include <functional>
+#include <memory>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+
+/*!
+ * \brief RAII wrapper function to enter and exit a context object
+ *        similar to python's with syntax.
+ *
+ * \code
+ * // context class
+ * class MyContext {
+ *  private:
+ *    friend class With<MyContext>;
+      MyContext(arguments);
+ *    void EnterWithScope();
+ *    void ExitWithScope();
+ * };
+ *
+ * {
+ *   With<MyContext> scope(arguments);
+ *   // effect take place.
+ * }
+ * \endcode
+ *
+ * \tparam ContextType Type of the context object.
+ */
+template <typename ContextType>
+class With {
+ public:
+  /*!
+   * \brief constructor.
+   *  Enter the scope of the context.
+   */
+  template <typename... Args>
+  explicit With(Args&&... args) : ctx_(std::forward<Args>(args)...) {
+    ctx_.EnterWithScope();
+  }
+  /*! \brief destructor, leaves the scope of the context. */
+  ~With() noexcept(false) { ctx_.ExitWithScope(); }
+
+  // Disable copy and move construction.  `With` is intended only for
+  // use in nested contexts that are exited in the reverse order of
+  // entry.  Allowing context to be copied or moved would break this
+  // expectation.
+  With(const With& other) = delete;
+  With& operator=(const With& other) = delete;
+  With(With&& other) = delete;
+  With& operator=(With&& other) = delete;
+
+  ContextType* get() { return &ctx_; }
+  const ContextType* get() const { return &ctx_; }
+
+  ContextType* operator->() { return get(); }
+  const ContextType* operator->() const { return get(); }
+  ContextType& operator*() { return *get(); }
+  const ContextType* operator*() const { return *get(); }
+
+  ContextType operator()() { return ctx_; }
+
+ private:
+  /*! \brief internal context type. */
+  ContextType ctx_;
+};
+
+/*!
+ * \brief A group of RAII contexts managed together.
+ *
+ * Allows dynamically emplacing multiple context objects that are
+ * all exited (in reverse order) when the group is destroyed.
+ * ContextType must declare `friend class With<ContextType>`
+ * and provide EnterWithScope() / ExitWithScope() methods.
+ *
+ * \code
+ *   WithGroup<ConstraintContext> group;
+ *   group.Emplace(analyzer, cond1);  // constructs and enters
+ *   group.Emplace(analyzer, cond2);  // constructs and enters
+ *   // destructor: exits cond2, then cond1
+ * \endcode
+ *
+ * \tparam ContextType The context type with EnterWithScope/ExitWithScope.
+ */
+template <typename ContextType>
+class WithGroup {
+ public:
+  WithGroup() = default;
+  WithGroup(WithGroup&&) = default;
+  WithGroup& operator=(WithGroup&&) = default;
+  WithGroup(const WithGroup&) = delete;
+  WithGroup& operator=(const WithGroup&) = delete;
+
+  /*!
+   * \brief Construct a context and enter its scope.
+   * \param args Arguments forwarded to ContextType constructor.
+   */
+  template <typename... Args>
+  void Emplace(Args&&... args) {
+    entries_.push_back(std::make_unique<With<ContextType>>(std::forward<Args>(args)...));
+  }
+
+  /*! \brief Number of active contexts in this group. */
+  size_t size() const { return entries_.size(); }
+
+  /*!
+   * \brief Destructor — exits all contexts in reverse order.
+   *
+   * On normal exit: if any ExitWithScope throws, the remaining
+   * contexts are still cleaned up, then the first exception
+   * is re-thrown.
+   *
+   * During stack unwinding: all exceptions are swallowed
+   * to avoid std::terminate.
+   */
+  ~WithGroup() noexcept(false) {
+    bool unwinding = std::uncaught_exceptions() > 0;
+    std::exception_ptr first_exc;
+    while (!entries_.empty()) {
+      // Move the last entry out of the vector first, then destroy it.
+      // This ensures entries_ shrinks even if ~With() throws.
+      auto entry = std::move(entries_.back());
+      entries_.pop_back();
+      try {
+        entry.reset();  // calls ~With<ContextType>() -> ExitWithScope()
+      } catch (...) {
+        if (!unwinding && !first_exc) {
+          first_exc = std::current_exception();
+        }
+      }
+    }
+    if (first_exc) std::rethrow_exception(first_exc);
+  }
+
+ private:
+  std::vector<std::unique_ptr<With<ContextType>>> entries_;
+};
+
+}  // namespace tvm
+#endif  // TVM_IR_WITH_CONTEXT_H_