chore: import upstream snapshot with attribution

include/tvm/script/ir_builder/base.h

@@ -0,0 +1,309 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_SCRIPT_IR_BUILDER_BASE_H_
+#define TVM_SCRIPT_IR_BUILDER_BASE_H_
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/expr.h>
+#include <tvm/ir/function.h>
+#include <tvm/ir/node_functor.h>
+
+#include <vector>
+
+namespace tvm {
+namespace script {
+namespace ir_builder {
+
+////////////////////////////// IRBuilderFrame //////////////////////////////
+
+/*!
+ * \brief A stack frame of the IRBuilder used to keep track of the current scope.
+ * Furthermore, the information stored in each stack frame can be useful for context-dependent
+ * IR construction.
+ *
+ * \example
+ *
+ * The `T::MatchBuffer` below adds an element in `PrimFuncNode::buffer_map`:
+ *
+ * \code {.cpp}
+ *
+ * using T = tvm::script::ir_builder::tirx;
+ * With <PrimFuncFrame> _(...);
+ * Buffer buffer = T::MatchBuffer(...);
+ *
+ * \endcode
+ *
+ * The `T::MatchBuffer` below instead generates `MatchBufferRegion` in a TIR block:
+ *
+ * \code {.cpp}
+ *
+ * using T = tvm::script::ir_builder::tirx;
+ * With <PrimFuncFrame> _(...);
+ * {
+ *   With<SBlockFrame> _2(...);
+ *   Buffer buffer = T::MatchBuffer(...);
+ * }
+ *
+ * \endcode
+ */
+class IRBuilderFrameNode : public ffi::Object {
+ public:
+  /*! \brief A list of callbacks used when exiting the frame. */
+  std::vector<ffi::TypedFunction<void()>> callbacks;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<IRBuilderFrameNode>();
+    // `callbacks` is not registered as it's not visited.
+  }
+
+  static constexpr const bool _type_mutable = true;
+  TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.IRBuilderFrame", IRBuilderFrameNode, ffi::Object);
+
+ public:
+  /*! \brief Default destructor. */
+  virtual ~IRBuilderFrameNode() = default;
+  /*!
+   * \brief The method called when entering RAII scope.
+   * \sa tvm::support::With
+   */
+  virtual void EnterWithScope();
+  /*!
+   * \brief The method called when exiting RAII scope.
+   * \sa tvm::support::With
+   */
+  virtual void ExitWithScope();
+  /*!
+   * \brief Add a callback method invoked when exiting the RAII scope.
+   * \param callback The callback to be added.
+   */
+  void AddCallback(ffi::TypedFunction<void()> callback);
+};
+
+/*!
+ * \brief Managed reference to an IRBuilderFrameNode.
+ * \sa IRBuilderFrameNode
+ */
+class IRBuilderFrame : public ffi::ObjectRef {
+ public:
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRBuilderFrame, ffi::ObjectRef, IRBuilderFrameNode);
+
+ protected:
+  /*! \brief Disallow direct construction of this object. */
+  IRBuilderFrame() = default;
+  explicit IRBuilderFrame(ffi::ObjectPtr<IRBuilderFrameNode> data) : ffi::ObjectRef(data) {}
+
+ public:
+  /*!
+   * \brief Redirected to `IRBuilderFrameNode::EnterWithScope`.
+   * \sa IRBuilderFrameNode::EnterWithScope
+   */
+  inline void EnterWithScope() {
+    TVM_FFI_ICHECK(data_ != nullptr);
+    static_cast<IRBuilderFrameNode*>(data_.get())->EnterWithScope();
+  }
+  /*!
+   * \brief Redirected to `IRBuilderFrameNode::ExitWithScope`.
+   * \sa IRBuilderFrameNode::ExitWithScope
+   */
+  inline void ExitWithScope() {
+    TVM_FFI_ICHECK(data_ != nullptr);
+    static_cast<IRBuilderFrameNode*>(data_.get())->ExitWithScope();
+    data_.reset();
+  }
+};
+
+////////////////////////////// IRBuilder //////////////////////////////
+
+/*!
+ * \brief A dialect-agnostic IRBuilder that constructs any IR of TVM.
+ * An idiomatic use of this class is to put this inside the RAII with-scope,
+ * call dialect-specific methods accordingly. Upon exiting the scope.
+ *
+ * \code
+ *
+ * PrimFunc ConstructPrimFunc() {
+ *   using tvm::script::ir_builder::IRBuilder;
+ *   using T = tvm::script::ir_builder::tirx;
+ *   IRBuilder builder;
+ *   // Step 1. Place IRBuilder inside the with-scope.
+ *   {
+ *     With<IRBuilder> _(builder);
+ *     // Step 2. Call dialect-specific methods.
+ *     With<T::PrimFuncFrame> _2(...);
+ *     T::MatchBuffer(...);
+ *   }
+ *   // Step 3. Return the constructed PrimFunc.
+ *   return builder->Get<PrimFunc>();
+ * }
+ *
+ * \endcode
+ */
+class IRBuilderNode : public ffi::Object {
+ public:
+  /*! \brief A stack of context frames in the IRBuilder */
+  ffi::Array<IRBuilderFrame> frames;
+  /*! \brief The outcome of IR construction */
+  ffi::Optional<ffi::ObjectRef> result;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<IRBuilderNode>()
+        .def_ro("frames", &IRBuilderNode::frames)
+        .def_ro("result", &IRBuilderNode::result);
+  }
+
+  static constexpr const bool _type_mutable = true;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.IRBuilder", IRBuilderNode, ffi::Object);
+
+ public:
+  /*!
+   * \brief Find a frame of the given type in the stack `this->frames` from top to bottom.
+   * \tparam T The type of the frame to find.
+   * \return The frame if found, otherwise std::nullopt.
+   */
+  template <typename TFrame>
+  inline ffi::Optional<TFrame> FindFrame() const;
+  /*!
+   * \brief Get the frame on top of the stack `this->frames` if its type is `TFrame`.
+   * \tparam TFrame The assumed type of the last frame on stack.
+   * \return The frame if the stack is non-empty and the top of the stack is of type `TFrame`.
+   * Otherwise std::nullopt.
+   */
+  template <typename TFrame>
+  inline ffi::Optional<TFrame> GetLastFrame() const;
+  /*!
+   * \brief Get the IR being constructed.
+   * \tparam TObjectRef The type of the IR being constructed.
+   * \return The resulting IR. Throw an exception if the IR is not constructed yet.
+   */
+  template <typename TObjectRef>
+  inline TObjectRef Get() const;
+};
+
+/*!
+ * \brief Managed reference to an IRBuilderNode.
+ * \sa IRBuilderNode
+ */
+class IRBuilder : public ffi::ObjectRef {
+ public:
+  /*! \brief Creates an IRBuilder. */
+  IRBuilder();
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRBuilder, ffi::ObjectRef, IRBuilderNode);
+
+ public:
+  /*!
+   * \brief Puts the current IRBuilder into a thread-local scope, which can be retrieved using
+   * `IRBuilder::Current()`.
+   *
+   * \code {.cpp}
+   * IRBuilder builder;
+   * {
+   *   With<IRBuilder> _(builder);
+   *   // IRBuilder::Current() == builder
+   * }
+   * // IRBuilder::Current() == nullptr
+   * \endcode
+   *
+   * \sa IRBuilder::Current
+   * \sa IRBuilder::ExitWithScope
+   * \sa tvm::support::With
+   */
+  void EnterWithScope();
+  /*!
+   * \brief Exit the RAII scope.
+   * \sa IRBuilder::EnterWithScope
+   * \sa IRBuilder::Current
+   * \sa tvm::support::With
+   */
+  void ExitWithScope();
+  /*!
+   * \brief Get the current IRBuilder in the current thread-local scope.
+   * \return The current IRBuilder.
+   * \sa IRBuilder::EnterWithScope
+   * \sa IRBuilder::ExitWithScope
+   * \sa tvm::support::With
+   */
+  static IRBuilder Current();
+  /*! \brief See if the current thread-local scope has an IRBuilder. */
+  static bool IsInScope();
+  /*!
+   * \brief Give a string name to the `obj`
+   * \tparam TObjectRef The type of the object to name.
+   * \param name The name to give to the object.
+   * \param obj The object to name.
+   */
+  template <class TObjectRef>
+  inline static TObjectRef Name(ffi::String name, TObjectRef obj);
+};
+
+////////////////////////////// Details //////////////////////////////
+
+namespace details {
+
+class Namer {
+ public:
+  using FType = NodeFunctor<void(const ffi::ObjectRef&, ffi::String)>;
+  static FType& vtable();
+  static void Name(ffi::ObjectRef node, ffi::String name);
+};
+
+}  // namespace details
+
+template <class TObjectRef>
+inline TObjectRef IRBuilder::Name(ffi::String name, TObjectRef obj) {
+  details::Namer::Name(obj, name);
+  return obj.template as_or_throw<TObjectRef>();
+}
+
+template <typename TFrame>
+inline ffi::Optional<TFrame> IRBuilderNode::FindFrame() const {
+  using TFrameNode = typename TFrame::ContainerType;
+  for (auto it = frames.rbegin(); it != frames.rend(); ++it) {
+    if (const TFrameNode* p = (*it).template as<TFrameNode>()) {
+      return ffi::GetRef<TFrame>(p);
+    }
+  }
+  return std::nullopt;
+}
+
+template <typename TFrame>
+inline ffi::Optional<TFrame> IRBuilderNode::GetLastFrame() const {
+  using TFrameNode = typename TFrame::ContainerType;
+  if (!frames.empty() && frames.back()->IsInstance<TFrameNode>()) {
+    return frames.back().as_or_throw<TFrame>();
+  }
+  return std::nullopt;
+}
+
+template <typename TObjectRef>
+inline TObjectRef IRBuilderNode::Get() const {
+  using TObject = typename TObjectRef::ContainerType;
+  TVM_FFI_CHECK(result.has_value(), IndexError) << "No result exists in IRBuilder yet";
+  const auto* n = result.as<TObject>();
+  TVM_FFI_CHECK(n != nullptr, TypeError)
+      << "IRBuilder result is not of type: " << TObject::_type_key;
+  return ffi::GetRef<TObjectRef>(n);
+}
+
+}  // namespace ir_builder
+}  // namespace script
+}  // namespace tvm
+
+#endif  // TVM_SCRIPT_IR_BUILDER_BASE_H_

include/tvm/script/ir_builder/ir/frame.h

@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_
+#define TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/expr.h>
+#include <tvm/ir/function.h>
+#include <tvm/ir/module.h>
+#include <tvm/script/ir_builder/base.h>
+
+#include <vector>
+
+namespace tvm {
+namespace script {
+namespace ir_builder {
+namespace ir {
+
+/*!
+ * \brief A frame that represents the IRModule frame with functions and global variables.
+ *
+ * \sa IRModuleFrame
+ */
+class IRModuleFrameNode : public IRBuilderFrameNode {
+ public:
+  /*! \brief A map from string names to global variables that ensures global uniqueness. */
+  ffi::Map<ffi::String, GlobalVar> global_var_map;
+  /*!
+   * \brief A map from GlobalVar to all global functions.
+   * \note Only defined functions are in the map, while declared functions are not included.
+   */
+  ffi::Map<GlobalVar, BaseFunc> functions;
+  /*! \brief IRModule's attributes. */
+  ffi::Map<ffi::String, Any> attrs;
+  /*! \brief IRModule's global_infos */
+  ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<IRModuleFrameNode>()
+        .def_ro("global_vars", &IRModuleFrameNode::global_var_map)
+        .def_ro("functions", &IRModuleFrameNode::functions)
+        .def_ro("attrs", &IRModuleFrameNode::attrs)
+        .def_ro("global_infos", &IRModuleFrameNode::global_infos);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.IRModuleFrame", IRModuleFrameNode,
+                                    IRBuilderFrameNode);
+
+ public:
+  void ExitWithScope() final;
+};
+
+/*!
+ * \brief Managed reference to IRModuleFrameNode.
+ *
+ * \sa IRModuleFrameNode
+ */
+class IRModuleFrame : public IRBuilderFrame {
+ public:
+  explicit IRModuleFrame(ffi::ObjectPtr<IRModuleFrameNode> data) : IRBuilderFrame(data) {
+    TVM_FFI_ICHECK(data != nullptr);
+  }
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRModuleFrame, IRBuilderFrame, IRModuleFrameNode);
+};
+
+}  // namespace ir
+}  // namespace ir_builder
+}  // namespace script
+}  // namespace tvm
+
+#endif  // TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_

include/tvm/script/ir_builder/ir/ir.h

@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_SCRIPT_IR_BUILDER_IR_IR_H_
+#define TVM_SCRIPT_IR_BUILDER_IR_IR_H_
+
+#include <tvm/ir/expr.h>
+#include <tvm/ir/function.h>
+#include <tvm/script/ir_builder/ir/frame.h>
+
+#include <vector>
+
+namespace tvm {
+namespace script {
+namespace ir_builder {
+namespace ir {
+
+/*!
+ * \brief The IRModule declaration statement.
+ * \return The IRModuleFrame.
+ */
+TVM_DLL IRModuleFrame IRModule();
+
+/*!
+ * \brief Declare a Function without given the specific function implementation.
+ * \note It is usually used in cross-function call. And we can specify the function by `DefFunction`
+ * \param func_name The function unique name.
+ * \param func_signature A Function w/o body, which used to specify the function signature
+ *                       (i.e. func params and func return type/shape).
+ * \return The corresponding GlobalVar.
+ */
+TVM_DLL GlobalVar DeclFunction(const ffi::String& func_name, const BaseFunc& func_signature);
+
+/*!
+ * \brief Define the function which is declared before.
+ * \param func_name The function unique name.
+ * \param func The given function implementation
+ */
+TVM_DLL void DefFunction(const ffi::String& func_name, const BaseFunc& func);
+
+}  // namespace ir
+}  // namespace ir_builder
+}  // namespace script
+}  // namespace tvm
+
+#endif  // TVM_SCRIPT_IR_BUILDER_IR_IR_H_

include/tvm/script/printer/config.h

@@ -0,0 +1,164 @@
+/*
+ * 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/script/printer/config.h
+ * \brief Configuration object for the TVMScript printer.
+ *
+ * Contains PrinterConfig / PrinterConfigNode, GetBuiltinKeywords, GetExtraConfig,
+ * and RedirectedReprPrinterMethod.  The entry-point free function tvm::Script()
+ * and the dispatch vtable TVMScriptPrinter live in printer.h.
+ */
+#ifndef TVM_SCRIPT_PRINTER_CONFIG_H_
+#define TVM_SCRIPT_PRINTER_CONFIG_H_
+
+#include <tvm/ffi/any.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/reflection/access_path.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/string.h>
+#include <tvm/runtime/base.h>
+
+#include <string>
+
+namespace tvm {
+
+class PrinterConfigNode : public ffi::Object {
+ public:
+  /*! \brief A stack that tracks the names of the binding hierarchy */
+  ffi::Array<ffi::String> binding_names = {};
+  /*! \brief Whether or not to show metadata. */
+  bool show_meta = false;
+  /*! \brief The prefix of IR nodes */
+  ffi::String ir_prefix = "I";
+  /*!
+   * \brief The alias of the current module at cross-function call
+   * \note Directly use module name if it's empty.
+   */
+  ffi::String module_alias = "cls";
+  /*! \brief Default buffer dtype */
+  DLDataType buffer_dtype = DLDataType{kDLFloat, 32, 1};
+  /*! \brief Default data type of integer literals */
+  DLDataType int_dtype = DLDataType{kDLInt, 32, 1};
+  /*!
+   * \brief Default data type of float literals. Right now we always print out the explicit type
+   * of floating point values, so setting it to Void means we do not print without the
+   * T.float32/T.float64 wrapper.
+   */
+  DLDataType float_dtype = DLDataType{kDLOpaqueHandle, 0, 0};
+  /*! \brief Whether or not to verbose print expressions. */
+  bool verbose_expr = false;
+  /*! \brief Number of spaces used for indentation*/
+  int indent_spaces = 4;
+  /*! \brief Whether to print line numbers */
+  bool print_line_numbers = false;
+  /*! \brief Number of context lines to print around the underlined text */
+  int num_context_lines = -1;
+  /*! \brief Whether to output with syntax sugar, set false for complete printing. */
+  bool syntax_sugar = true;
+  /*! \brief Whether variable names should include the object's address */
+  bool show_object_address = false;
+
+  /*! \brief Whether to render access-path context for invisible underlined paths. Defaults true. */
+  bool render_invisible_path_info = true;
+
+  /* \brief ffi::Object path to be underlined */
+  ffi::Array<ffi::reflection::AccessPath> path_to_underline;
+  /*! \brief ffi::Object path to be annotated. */
+  ffi::Map<ffi::reflection::AccessPath, ffi::String> path_to_annotate;
+  /*! \brief ffi::Object to be underlined. */
+  ffi::Array<ffi::ObjectRef> obj_to_underline = ffi::Array<ffi::ObjectRef>();
+  /*! \brief ffi::Object to be annotated. */
+  ffi::Map<ffi::ObjectRef, ffi::String> obj_to_annotate = ffi::Map<ffi::ObjectRef, ffi::String>();
+
+  /*!
+   * \brief Generic extension map for dialect-specific config knobs.
+   *
+   * Keys are conventionally namespaced as "<dialect>.<knob>", e.g.:
+   *   "tirx.prefix"              — the TIR prefix (default "T")
+   *   "relax.prefix"             — the Relax prefix (default "R")
+   *   "relax.show_all_ty" — whether to show all Relax type annotations (default true)
+   *
+   * Use GetExtraConfig<T>(key, fallback) to read values with a typed fallback.
+   */
+  ffi::Map<ffi::String, ffi::Any> extra_config;
+
+  /*!
+   * \brief Look up a value in extra_config with type cast and fallback.
+   *
+   * Keys are conventionally namespaced as "<dialect>.<knob>"
+   * (e.g. "tirx.prefix", "relax.show_all_ty").
+   */
+  template <typename T>
+  T GetExtraConfig(const ffi::String& key, T fallback) const {
+    auto it = extra_config.find(key);
+    if (it == extra_config.end()) return fallback;
+    return (*it).second.as_or_throw<T>();
+  }
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<PrinterConfigNode>()
+        .def_ro("binding_names", &PrinterConfigNode::binding_names)
+        .def_ro("show_meta", &PrinterConfigNode::show_meta)
+        .def_ro("ir_prefix", &PrinterConfigNode::ir_prefix)
+        .def_ro("module_alias", &PrinterConfigNode::module_alias)
+        .def_ro("buffer_dtype", &PrinterConfigNode::buffer_dtype)
+        .def_ro("int_dtype", &PrinterConfigNode::int_dtype)
+        .def_ro("float_dtype", &PrinterConfigNode::float_dtype)
+        .def_ro("verbose_expr", &PrinterConfigNode::verbose_expr)
+        .def_ro("indent_spaces", &PrinterConfigNode::indent_spaces)
+        .def_ro("print_line_numbers", &PrinterConfigNode::print_line_numbers)
+        .def_ro("num_context_lines", &PrinterConfigNode::num_context_lines)
+        .def_ro("syntax_sugar", &PrinterConfigNode::syntax_sugar)
+        .def_ro("show_object_address", &PrinterConfigNode::show_object_address)
+        .def_ro("render_invisible_path_info", &PrinterConfigNode::render_invisible_path_info)
+        .def_ro("path_to_underline", &PrinterConfigNode::path_to_underline)
+        .def_ro("path_to_annotate", &PrinterConfigNode::path_to_annotate)
+        .def_ro("obj_to_underline", &PrinterConfigNode::obj_to_underline)
+        .def_ro("obj_to_annotate", &PrinterConfigNode::obj_to_annotate)
+        .def_ro("extra_config", &PrinterConfigNode::extra_config);
+  }
+
+  ffi::Array<ffi::String> GetBuiltinKeywords();
+
+  static constexpr const bool _type_mutable = true;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.PrinterConfig", PrinterConfigNode, ffi::Object);
+};
+
+class TVM_DLL PrinterConfig : public ffi::ObjectRef {
+ public:
+  explicit PrinterConfig(
+      ffi::Map<ffi::String, ffi::Any> config_dict = ffi::Map<ffi::String, ffi::Any>());
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrinterConfig, ffi::ObjectRef, PrinterConfigNode);
+};
+
+/*!
+ * \brief The fallback body used by TVM_REGISTER_SCRIPT_AS_REPR (defined in printer.h).
+ *
+ * Tries to format \p obj via tvm::Script; on error falls back to a plain
+ * address string.  Defined in src/script/printer/config.cc so that
+ * <tvm/runtime/logging.h> is not pulled into this public header.
+ */
+TVM_DLL std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj);
+
+}  // namespace tvm
+#endif  // TVM_SCRIPT_PRINTER_CONFIG_H_

include/tvm/script/printer/ir_docsifier.h

@@ -0,0 +1,364 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_
+#define TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_
+
+#include <tvm/ffi/reflection/access_path.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/module.h>
+#include <tvm/script/printer/config.h>
+#include <tvm/script/printer/doc.h>
+#include <tvm/script/printer/ir_docsifier_functor.h>
+
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace script {
+namespace printer {
+
+using AccessPath = ffi::reflection::AccessPath;
+
+//////////////////////// Frame ////////////////////////
+
+class IRDocsifier;
+class IRDocsifierNode;
+
+/*!
+ * Frame is the core data structure for semantic information
+ * when printing IR graph into TVMScript code.
+ */
+class FrameNode : public ffi::Object {
+ public:
+  /*! The docs generated in the frame */
+  ffi::Array<StmtDoc> stmts;
+  /*! The corresponding IRDocsifier */
+  IRDocsifierNode* d;
+  /*! The callbacks that are going to be invoked when the frame exits */
+  std::vector<std::function<void()>> callbacks;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<FrameNode>().def_ro("stmts", &FrameNode::stmts);
+  }
+
+  static constexpr const bool _type_mutable = true;
+  TVM_FFI_DECLARE_OBJECT_INFO("script.printer.Frame", FrameNode, ffi::Object);
+
+ public:
+  virtual ~FrameNode() = default;
+
+  /*!
+   * \brief Add a callback function to be called when this frame exits.
+   * \param cb The callback function. It should have signature void().
+   */
+  template <typename TCallback>
+  void AddExitCallback(TCallback&& cb) {
+    callbacks.emplace_back(std::forward<TCallback>(cb));
+  }
+  /*!
+   * \brief Add a dispatch token to the docsifier, and a callback that pops the token when this
+   * frame exits.
+   * \param d The docsifier.
+   * \param token The token to be added.
+   */
+  void AddDispatchToken(const IRDocsifier& d, const ffi::String& token);
+  /*!
+   * \brief Method that's called when Frame enters the scope.
+   */
+  virtual void EnterWithScope();
+  /*!
+   * \brief Method that's called when Frame exits the scope.
+   */
+  virtual void ExitWithScope();
+};
+
+/*!
+ * \brief Reference type of FrameNode
+ */
+class Frame : public ffi::ObjectRef {
+ protected:
+  Frame() = default;
+
+ public:
+  virtual ~Frame() = default;
+
+  /*! \brief Method that's called when Frame enters the scope. */
+  void EnterWithScope() { get()->EnterWithScope(); }
+
+  /*! \brief Method that's called when Frame exits the scope. */
+  void ExitWithScope() { get()->ExitWithScope(); }
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Frame, ffi::ObjectRef, FrameNode);
+};
+
+//////////////////////// IRDocsifier ////////////////////////
+
+/*!
+ * \brief IRDocsifier is the top-level interface in the IR->Doc process.
+ *
+ * It provides methods to convert IR node object to Doc, operate on Frame
+ * objects and change dispatch tokens.
+ */
+class IRDocsifierNode : public ffi::Object {
+ public:
+  /*! \brief A function that creates the doc for a variable */
+  using DocCreator = std::function<ExprDoc()>;
+  /*! \brief Information about a variable, including its optional name and its doc creator */
+  struct VariableInfo {
+    /*! \brief The creator */
+    DocCreator creator;
+    /*! \brief The name of the variable */
+    ffi::Optional<ffi::String> name;
+  };
+  /*! \brief The configuration of the printer */
+  PrinterConfig cfg{ffi::UnsafeInit()};
+  /*!
+   * \brief The stack of frames.
+   * \sa FrameNode
+   */
+  ffi::Array<Frame> frames;
+  /*!
+   * \brief The stack of dispatch tokens.
+   *
+   * The dispatch token on the top decides which dispatch function to use
+   * when converting IR node object to Doc.
+   */
+  ffi::Array<ffi::String> dispatch_tokens;
+  /*! \brief Mapping from a var to its info */
+  std::unordered_map<ffi::ObjectRef, VariableInfo, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
+      obj2info;
+  /*! \brief Metadata printing */
+  std::unordered_map<ffi::String, ffi::Array<ffi::Any>> metadata;
+  /*! \brief GlobalInfo printing */
+  std::unordered_map<ffi::String, ffi::Array<GlobalInfo>> global_infos;
+  /*! \brief The variable names used already */
+  std::unordered_set<ffi::String> defined_names;
+  /*! \brief Common prefixes of variable usages */
+  std::unordered_map<const ffi::Object*, std::vector<const ffi::Object*>> common_prefix;
+  /*! \brief The IR usages for headers printing */
+  std::unordered_set<std::string> ir_usage;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<IRDocsifierNode>()
+        .def_ro("frames", &IRDocsifierNode::frames)
+        .def_ro("dispatch_tokens", &IRDocsifierNode::dispatch_tokens);
+  }
+
+  static constexpr const bool _type_mutable = true;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.IRDocsifier", IRDocsifierNode, ffi::Object);
+
+ public:
+  /*!
+   * \brief Define variable by name.
+   * \param obj The variable object.
+   * \param frame The frame that this variable is defined in.
+   * \param name_hint The hint for variable name.
+   *
+   * \return The id doc for this variable.
+   *
+   * This function will rename the variable to avoid name conflict with other variables
+   * in the table.
+   */
+  IdDoc Define(const ffi::ObjectRef& obj, const Frame& frame, const ffi::String& name_hint);
+
+  /*!
+   * \brief Define variable by doc factory.
+   * \param obj The variable object.
+   * \param frame The frame that this variable is defined in.
+   * \param doc_factory The function to return an ExprDoc object for this variable.
+   *
+   * This function is a special form of `Define`. Variable is mapped to ExprDoc rather
+   * than IdDoc. It's useful when a variable is implicitly defined without a name, like
+   * the buf->data in TIR, which should be mapped to `AttrDoc(IdDoc("<buffer_name>"), "data")`.
+   *
+   * This function takes a DocFactory instead of Doc. It's because GetVarDoc needs to
+   * return a new Doc object every time it's called, as the returned doc will have
+   * different `source_path`. Currently there isn't a good way to deep copy a TVMObject
+   * so VarTable needs to call a factory function to get a freshly-constructed Doc object
+   * every time GetVarDoc is called.
+   */
+  void Define(const ffi::ObjectRef& obj, const Frame& frame, DocCreator doc_factory);
+
+  /*!
+   * \brief Get the doc for variable.
+   * \param obj The variable object.
+   *
+   * \return The doc for variable, if it exists in the table. Otherwise it returns std::nullopt.
+   */
+  ffi::Optional<ExprDoc> GetVarDoc(const ffi::ObjectRef& obj) const;
+  /*! \brief Add a TVM object to the metadata section*/
+  ExprDoc AddMetadata(const ffi::Any& obj);
+  /*! \brief Add a GlobalInfo to the global_infos map.
+   * \param name The name of key of global_infos.
+   * \param ginfo The GlobalInfo to be added.
+   */
+  void AddGlobalInfo(const ffi::String& name, const GlobalInfo& ginfo);
+  /*!
+   * \brief Check if a variable exists in the table.
+   * \param obj The variable object.
+   *
+   * \return a boolean for whether variable exists.
+   */
+  bool IsVarDefined(const ffi::ObjectRef& obj) const;
+  /*! \brief Remove the variable defined */
+  void RemoveVar(const ffi::ObjectRef& obj);
+  /*!
+   * \brief Set the common prefix information of variable usage.
+   * \param root The root of the AST.
+   * \param is_var A function that returns true if the given object is considered a variable.
+   */
+  void SetCommonPrefix(const ffi::ObjectRef& root, ffi::TypedFunction<bool(ffi::ObjectRef)> is_var);
+  /*!
+   * \brief Transform the input object into TDoc.
+   * \param obj The object to be transformed.
+   * \param path The path to this object.
+   *
+   * \return The Doc object.
+   */
+  template <class TDoc = Doc>
+  inline TDoc AsDoc(const Any& obj, const AccessPath& path) const;
+};
+
+/*!
+ * \brief Reference type of IRDocsifierNode.
+ */
+class IRDocsifier : public ffi::ObjectRef {
+ public:
+  using FType = IRDocsifierFunctor<printer::Doc, AccessPath, IRDocsifier>;
+  /*! \brief Create a IRDocsifier. */
+  explicit IRDocsifier(const PrinterConfig& cfg);
+  /*! \brief The registration table for IRDocsifier. */
+  TVM_DLL static FType& vtable();
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRDocsifier, ffi::ObjectRef, IRDocsifierNode);
+};
+
+//////////////////////// Implementation ////////////////////////
+
+inline void FrameNode::EnterWithScope() {
+  if (d != nullptr) {
+    d->frames.push_back(ffi::GetRef<Frame>(this));
+  }
+}
+
+inline void FrameNode::ExitWithScope() {
+  for (const std::function<void()>& callback : callbacks) {
+    callback();
+  }
+  callbacks.clear();
+  if (d != nullptr) {
+    d->frames.pop_back();
+  }
+}
+
+template <class TDoc>
+inline static void AddDocDecoration(const Doc& d, const ffi::ObjectRef& obj, const AccessPath& path,
+                                    const PrinterConfig& cfg) {
+  if (cfg->obj_to_annotate.count(obj)) {
+    if (const auto* stmt = d.as<StmtDocNode>()) {
+      if (stmt->comment.has_value()) {
+        stmt->comment = stmt->comment.value() + "\n" + cfg->obj_to_annotate.at(obj);
+      } else {
+        stmt->comment = cfg->obj_to_annotate.at(obj);
+      }
+    } else {
+      LOG(WARNING) << "Expect StmtDoc to be annotated for object " << obj << ", but got "
+                   << d.as_or_throw<TDoc>()->_type_key;
+    }
+  }
+  for (const ffi::ObjectRef& o : cfg->obj_to_underline) {
+    if (o.same_as(obj)) {
+      cfg->path_to_underline.push_back(path);
+    }
+  }
+  for (const auto& pair : cfg->path_to_annotate) {
+    AccessPath p = pair.first;
+    ffi::String attn = pair.second;
+    if (p->IsPrefixOf(path) && path->IsPrefixOf(p)) {
+      if (const auto* stmt = d.as<StmtDocNode>()) {
+        if (stmt->comment.has_value()) {
+          stmt->comment = stmt->comment.value() + "\n" + attn;
+        } else {
+          stmt->comment = attn;
+        }
+      } else {
+        LOG(WARNING) << "Expect StmtDoc to be annotated at object path " << p << ", but got "
+                     << d.as_or_throw<TDoc>()->_type_key;
+      }
+    }
+  }
+}
+
+template <class TDoc>
+inline TDoc IRDocsifierNode::AsDoc(const Any& value, const AccessPath& path) const {
+  switch (value.type_index()) {
+    case ffi::TypeIndex::kTVMFFINone:
+      return LiteralDoc::None(path).as_or_throw<TDoc>();
+    case ffi::TypeIndex::kTVMFFIBool:
+      return LiteralDoc::Boolean(value.as<bool>().value(), path).as_or_throw<TDoc>();
+    case ffi::TypeIndex::kTVMFFIInt:
+      return LiteralDoc::Int(value.as<int64_t>().value(), path).as_or_throw<TDoc>();
+    case ffi::TypeIndex::kTVMFFIFloat:
+      return LiteralDoc::Float(value.as<double>().value(), path).as_or_throw<TDoc>();
+    case ffi::TypeIndex::kTVMFFISmallStr:
+    case ffi::TypeIndex::kTVMFFIStr: {
+      std::string string_value = value.cast<std::string>();
+      bool has_multiple_lines = string_value.find_first_of('\n') != std::string::npos;
+      if (has_multiple_lines) {
+        Doc d = const_cast<IRDocsifierNode*>(this)->AddMetadata(string_value);
+        // TODO(tqchen): cross check AddDocDecoration
+        return d.as_or_throw<TDoc>();
+      }
+      return LiteralDoc::Str(string_value, path).as_or_throw<TDoc>();
+    }
+    case ffi::TypeIndex::kTVMFFIDataType:
+      return LiteralDoc::DataType(value.as<DLDataType>().value(), path).as_or_throw<TDoc>();
+    case ffi::TypeIndex::kTVMFFIDevice:
+      return LiteralDoc::Device(value.as<DLDevice>().value(), path).as_or_throw<TDoc>();
+    default: {
+      if (auto opt_obj = value.as<ffi::ObjectRef>()) {
+        ffi::ObjectRef obj = opt_obj.value();
+        Doc d = IRDocsifier::vtable()(dispatch_tokens.back(), obj, path,
+                                      ffi::GetRef<IRDocsifier>(this));
+        d->source_paths.push_back(path);
+        AddDocDecoration<TDoc>(d, obj, path, cfg);
+        return d.as_or_throw<TDoc>();
+      } else {
+        TVM_FFI_THROW(TypeError) << "Cannot handle Any type: `" << value.GetTypeKey() << "`";
+        TVM_FFI_UNREACHABLE();
+      }
+    }
+  }
+}
+
+inline void FrameNode::AddDispatchToken(const IRDocsifier& d, const ffi::String& token) {
+  d->dispatch_tokens.push_back(token);
+  this->AddExitCallback([doc = d.get()]() { doc->dispatch_tokens.pop_back(); });
+}
+
+}  // namespace printer
+}  // namespace script
+}  // namespace tvm
+
+#endif  // TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_

include/tvm/script/printer/ir_docsifier_functor.h

@@ -0,0 +1,199 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_SCRIPT_PRINTER_IR_DOCSIFIER_FUNCTOR_H_
+#define TVM_SCRIPT_PRINTER_IR_DOCSIFIER_FUNCTOR_H_
+
+#include <tvm/ffi/function.h>
+#include <tvm/runtime/logging.h>
+
+#include <optional>
+#include <string>
+#include <type_traits>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace script {
+namespace printer {
+
+/*!
+ * \brief Dynamic dispatch functor based on AccessPath.
+ *
+ * This functor dispatches based on the type of object and the input dispatch token.
+ */
+template <typename R, typename... Args>
+class IRDocsifierFunctor {
+ private:
+  using TSelf = IRDocsifierFunctor<R, Args...>;
+
+  template <class TObjectRef, class TCallable>
+  using IsDispatchFunction =
+      typename std::is_convertible<TCallable, std::function<R(TObjectRef, Args...)>>;
+
+ public:
+  /*!
+   * \brief Call the dispatch function.
+   * \param token The dispatch token.
+   * \param obj The object.
+   * \param args Other args.
+   *
+   * \return The return value of the dispatch function
+   *
+   * If the TObjectRef isn't registered with the token, it will try to find
+   * dispatch function for TObjectRef with the default dispatch token (empty string).
+   */
+  template <class TObjectRef>
+  R operator()(const ffi::String& token, TObjectRef obj, Args... args) const {
+    uint32_t type_index = obj.defined() ? obj->type_index() : 0;
+    const ffi::Function* pf = nullptr;
+    if ((pf = LookupDispatchTable(token, type_index)) != nullptr) {
+      return (*pf)(obj, args...).template cast<R>();
+    }
+    if ((pf = LookupDispatchTable("", type_index)) != nullptr) {
+      return (*pf)(obj, args...).template cast<R>();
+    }
+    if ((pf = LookupFallback()) != nullptr) {
+      return (*pf)(obj, args...).template cast<R>();
+    }
+
+    LOG(WARNING) << "ObjectFunctor calls un-registered function on type: "
+                 << ffi::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
+                 << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
+    TVM_FFI_ICHECK(false) << "ObjectFunctor calls un-registered function on type: "
+                          << ffi::Object::TypeIndex2Key(type_index) << " (token: " << token << ")"
+                          << ". ObjectType: " << obj->GetTypeKey() << ". Object: " << obj;
+    TVM_FFI_UNREACHABLE();
+  }
+
+  /*!
+   * \brief Set the dispatch function
+   * \param token The dispatch token.
+   * \param type_index The TVM object type index for this dispatch function.
+   * \param f The dispatch function.
+   *
+   * This takes a type-erased packed function as input. It should be used
+   * through FFI boundary, for example, registering dispatch function from Python.
+   */
+  TSelf& set_dispatch(ffi::String token, uint32_t type_index, ffi::Function f) {
+    std::vector<ffi::Function>* table = &dispatch_table_[token];
+    if (table->size() <= type_index) {
+      table->resize(type_index + 1, nullptr);
+    }
+    ffi::Function& slot = (*table)[type_index];
+    if (slot != nullptr) {
+      TVM_FFI_ICHECK(false) << "Dispatch for type is already registered: "
+                            << ffi::Object::TypeIndex2Key(type_index);
+    }
+    slot = f;
+    return *this;
+  }
+
+  TSelf& set_fallback(ffi::Function f) {
+    TVM_FFI_ICHECK(!dispatch_fallback_.has_value()) << "Fallback is already defined";
+    dispatch_fallback_ = f;
+    return *this;
+  }
+
+  void remove_fallback() { dispatch_fallback_ = std::nullopt; }
+
+  /*!
+   * \brief Set the dispatch function
+   * \param token The dispatch token.
+   * \param f The dispatch function.
+   */
+  template <typename TObjectRef, typename TCallable,
+            typename = std::enable_if_t<IsDispatchFunction<TObjectRef, TCallable>::value>>
+  TSelf& set_dispatch(ffi::String token, TCallable f) {
+    return set_dispatch(token, TObjectRef::ContainerType::RuntimeTypeIndex(),
+                        ffi::TypedFunction<R(TObjectRef, Args...)>(f));
+  }
+
+  template <typename TCallable,
+            typename = std::enable_if_t<IsDispatchFunction<ffi::ObjectRef, TCallable>::value>>
+  TSelf& set_fallback(TCallable f) {
+    ffi::Function func = ffi::TypedFunction<R(ffi::ObjectRef, Args...)>(f);
+    return set_fallback(func);
+  }
+
+  /*!
+   * \brief Remove dispatch function
+   * \param token The dispatch token.
+   * \param type_index The TVM object type index for the dispatch function to be removed.
+   *
+   * This is useful when dispatch function comes from other language's runtime, and
+   * those function should be removed before that language runtime shuts down.
+   */
+  void remove_dispatch(ffi::String token, uint32_t type_index) {
+    std::vector<ffi::Function>* table = &dispatch_table_[token];
+    if (table->size() <= type_index) {
+      return;
+    }
+    (*table)[type_index] = nullptr;
+  }
+
+ private:
+  /*!
+   * \brief Look up the dispatch table for the given token and type_index.
+   * \param token The dispatch token.
+   * \param type_index The TVM object type index.
+   * \return Returns the functor if the lookup succeeds, nullptr otherwise.
+   */
+  const ffi::Function* LookupDispatchTable(const ffi::String& token, uint32_t type_index) const {
+    auto it = dispatch_table_.find(token);
+    if (it == dispatch_table_.end()) {
+      return nullptr;
+    }
+    const std::vector<ffi::Function>& tab = it->second;
+    if (type_index >= tab.size()) {
+      return nullptr;
+    }
+    const ffi::Function* f = &tab[type_index];
+    if (f->defined()) {
+      return f;
+    } else {
+      return nullptr;
+    }
+  }
+
+  /*!
+   * \brief Look up the fallback to be used if no handler is registered
+   */
+  const ffi::Function* LookupFallback() const {
+    if (dispatch_fallback_.has_value()) {
+      return &*dispatch_fallback_;
+    } else {
+      return nullptr;
+    }
+  }
+
+  /*
+   * This type alias and the following free functions are created to reduce the binary bloat
+   * from template and also hide implementation details from this header
+   */
+  using DispatchTable = std::unordered_map<std::string, std::vector<ffi::Function>>;
+  /*! \brief The dispatch table. */
+  DispatchTable dispatch_table_;
+  std::optional<ffi::Function> dispatch_fallback_;
+};
+
+}  // namespace printer
+}  // namespace script
+}  // namespace tvm
+#endif  // TVM_SCRIPT_PRINTER_IR_DOCSIFIER_FUNCTOR_H_

include/tvm/script/printer/printer.h

@@ -0,0 +1,71 @@
+/*
+ * 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/script/printer/printer.h
+ * \brief Entry-point header for TVMScript printing.
+ *
+ * Declares the free function `tvm::Script(node, optional_config)` and the
+ * dispatch vtable `TVMScriptPrinter::vtable()` used by per-dialect printers.
+ * `PrinterConfig` and its dataclass helpers live in config.h; this header is
+ * what callers include to invoke printing.
+ */
+#ifndef TVM_SCRIPT_PRINTER_PRINTER_H_
+#define TVM_SCRIPT_PRINTER_PRINTER_H_
+
+#include <tvm/ir/node_functor.h>
+#include <tvm/script/printer/config.h>
+
+namespace tvm {
+
+/*! \brief Print \p node as TVMScript with the given \p config.
+ *
+ *  Falls back to ffi::ReprPrint for types not registered with TVMScriptPrinter.
+ */
+TVM_DLL std::string Script(const ffi::ObjectRef& node,
+                           const ffi::Optional<PrinterConfig>& config = std::nullopt);
+
+/*! \brief Dispatch vtable used by per-dialect printers to register their
+ *         object-type printing functions.  Internal, but exposed here because
+ *         TVM_REGISTER_SCRIPT_AS_REPR refers to it.
+ */
+class TVMScriptPrinter {
+ public:
+  using FType = NodeFunctor<std::string(const ffi::ObjectRef&, const PrinterConfig&)>;
+  TVM_DLL static FType& vtable();
+};
+
+/*!
+ * \brief Register Script as the kRepr callback for ObjectType and install
+ *        the per-type dispatch entry in TVMScriptPrinter::vtable().
+ *
+ * \param ObjectType  The concrete object node type (e.g. tirx::VarNode).
+ * \param Method      The TVMScriptPrinter vtable dispatch function.
+ */
+#define TVM_REGISTER_SCRIPT_AS_REPR(ObjectType, Method)                                        \
+  TVM_FFI_STATIC_INIT_BLOCK() {                                                                \
+    namespace refl = tvm::ffi::reflection;                                                     \
+    refl::TypeAttrDef<ObjectType>().def(refl::type_attr::kRepr,                                \
+                                        [](ffi::ObjectRef obj, ffi::Function) -> ffi::String { \
+                                          return RedirectedReprPrinterMethod(obj);             \
+                                        });                                                    \
+  }                                                                                            \
+  TVM_STATIC_IR_FUNCTOR(TVMScriptPrinter, vtable).set_dispatch<ObjectType>(Method)
+
+}  // namespace tvm
+#endif  // TVM_SCRIPT_PRINTER_PRINTER_H_