wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tvm/script/printer/config.h
+164
View File
@@ -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/doc.h
+1394
View File
File diff suppressed because it is too large Load Diff
include/tvm/script/printer/ir_docsifier.h
+364
View File
@@ -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
+199
View File
@@ -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
+71
View File
@@ -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_