This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tqchen/printer-doc-translation in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 5fdfb5750bd74c07468392cf794c6f505144e6d7 Author: Tianqi Chen <[email protected]> AuthorDate: Wed Sep 30 20:40:35 2026 +0000 [REFACTOR][Printer] Inline translator operations and localize metadata --- include/tvm/printer/doc_translator.h | 49 +++++++++++++++++++------- src/printer/doc_translator.cc | 67 ++++++++---------------------------- 2 files changed, 52 insertions(+), 64 deletions(-) diff --git a/include/tvm/printer/doc_translator.h b/include/tvm/printer/doc_translator.h index 2bef66bf12..e1cb292e07 100644 --- a/include/tvm/printer/doc_translator.h +++ b/include/tvm/printer/doc_translator.h @@ -173,7 +173,10 @@ class DocTranslatorObj : public ffi::Object { * \param doc The emitted Doc. * \param origin The original IR object represented by this emission. */ - TVM_DLL void Emit(const Doc& doc, ffi::Optional<ffi::ObjectRef> origin = std::nullopt); + void Emit(const Doc& doc, ffi::Optional<ffi::ObjectRef> origin = std::nullopt) { + CurrentScopeDocs().push_back(doc); + if (origin.has_value()) RecordOrigin(doc, origin.value()); + } /*! * \brief Reuse or allocate a variable's stable identifier. @@ -263,17 +266,25 @@ class DocTranslatorObj : public ffi::Object { /*! * \brief Push an empty Doc collection. */ - TVM_DLL void BeginDocScope(); + void BeginDocScope() { doc_scopes_.push_back(ffi::List<Doc>()); } /*! * \brief Pop the latest Doc collection. * \return The collected Docs. */ - TVM_DLL ffi::List<Doc> EndDocScope(); + ffi::List<Doc> EndDocScope() { + auto docs = CurrentScopeDocs(); + doc_scopes_.pop_back(); + return docs; + } /*! * \brief Get the current Doc collection without popping it. * \return A mutable handle sharing the active collection's storage. */ - TVM_DLL ffi::List<Doc> CurrentScopeDocs() const; + ffi::List<Doc> CurrentScopeDocs() const { + TVM_FFI_CHECK(!doc_scopes_.empty(), ValueError) + << "printer emission requires a collection scope"; + return doc_scopes_.back(); + } /*! * \brief Run a callback in a balanced Doc collection scope. * \tparam Callback The callable type. @@ -306,20 +317,34 @@ class DocTranslatorObj : public ffi::Object { * \param origin The original IR object. */ void RecordOrigin(const Doc& doc, const ffi::ObjectRef& origin) { doc_origins_.Set(doc, origin); } - /*! - * \brief Intern a metadata value. - * \param value The retained metadata value. - * \return An ordinary Doc indexing the metadata registry. - */ - TVM_DLL ExprDoc AddMetadata(ffi::Any value); /*! * \brief Dispatch directly to the registered native or packed hook. * \param value The input value. * \param destination A borrowed Var, or nullptr when no binding is requested. * \return An expression for the caller, or None after completed emission. */ - TVM_DLL ffi::Optional<ExprDoc> DefaultTranslate(ffi::AnyView value, - const ffi::Object* destination = nullptr); + ffi::Optional<ExprDoc> DefaultTranslate(ffi::AnyView value, + const ffi::Object* destination = nullptr) { + static ffi::reflection::TypeAttrColumn column(kDocTranslate); + ffi::AnyView attr = column[value.type_index()]; + if (attr.type_index() == ffi::TypeIndex::kTVMFFIOpaquePtr) { + return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Optional<ExprDoc>>( + reinterpret_cast<decltype(DocTranslatorVTable::translate)>(attr.cast<void*>())( + this, value, destination)) + .value(); + } + if (attr.type_index() == ffi::TypeIndex::kTVMFFIFunction) { + ffi::Any destination_arg = nullptr; + if (destination) destination_arg = ffi::GetRef<ffi::ObjectRef>(destination); + return attr.cast<ffi::Function>() + .CallExpected<ffi::Optional<ExprDoc>>(this, value, destination_arg) + .value(); + } + TVM_FFI_THROW(TypeError) << (attr.type_index() == ffi::TypeIndex::kTVMFFINone + ? std::string("printer has no Doc hook for ") + + value.GetTypeKey() + : "printer Doc hook must be a native pointer or ffi.Function"); + } static constexpr bool _type_mutable = true; TVM_FFI_DECLARE_OBJECT_INFO("printer.DocTranslator", DocTranslatorObj, ffi::Object); diff --git a/src/printer/doc_translator.cc b/src/printer/doc_translator.cc index b4e7e52a1d..5818b4cdc1 100644 --- a/src/printer/doc_translator.cc +++ b/src/printer/doc_translator.cc @@ -71,43 +71,6 @@ namespace tvm { namespace printer { -ffi::Optional<ExprDoc> DocTranslatorObj::DefaultTranslate(ffi::AnyView value, - const ffi::Object* destination) { - static ffi::reflection::TypeAttrColumn column(kDocTranslate); - ffi::AnyView attr = column[value.type_index()]; - if (attr.type_index() == ffi::TypeIndex::kTVMFFIOpaquePtr) { - return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Optional<ExprDoc>>( - reinterpret_cast<decltype(DocTranslatorVTable::translate)>(attr.cast<void*>())( - this, value, destination)) - .value(); - } - if (attr.type_index() == ffi::TypeIndex::kTVMFFIFunction) { - ffi::Any destination_arg = nullptr; - if (destination) destination_arg = ffi::GetRef<ffi::ObjectRef>(destination); - return attr.cast<ffi::Function>() - .CallExpected<ffi::Optional<ExprDoc>>(this, value, destination_arg) - .value(); - } - TVM_FFI_THROW(TypeError) << (attr.type_index() == ffi::TypeIndex::kTVMFFINone - ? std::string("printer has no Doc hook for ") + - value.GetTypeKey() - : "printer Doc hook must be a native pointer or ffi.Function"); -} - -void DocTranslatorObj::BeginDocScope() { doc_scopes_.push_back(ffi::List<Doc>()); } -ffi::List<Doc> DocTranslatorObj::CurrentScopeDocs() const { - TVM_FFI_CHECK(!doc_scopes_.empty(), ValueError) << "printer emission requires a collection scope"; - return doc_scopes_.back(); -} -void DocTranslatorObj::Emit(const Doc& doc, ffi::Optional<ffi::ObjectRef> origin) { - CurrentScopeDocs().push_back(doc); - if (origin.has_value()) doc_origins_.Set(doc, origin.value()); -} -ffi::List<Doc> DocTranslatorObj::EndDocScope() { - auto docs = CurrentScopeDocs(); - doc_scopes_.pop_back(); - return docs; -} ffi::Array<StmtDoc> AsStatements(const ffi::List<Doc>& docs) { ffi::Array<StmtDoc> statements; for (const auto& doc : docs) { @@ -120,10 +83,12 @@ ffi::Array<StmtDoc> AsStatements(const ffi::List<Doc>& docs) { return statements; } -ExprDoc DocTranslatorObj::AddMetadata(ffi::Any value) { +namespace { + +ExprDoc AddMetadata(DocTranslatorObj* d, ffi::Any value) { TVM_FFI_CHECK(value != nullptr, TypeError) << "Metadata cannot contain None"; using MetadataMap = ffi::Dict<ffi::String, ffi::List<ffi::Any>>; - auto metadata = GetOrCreateExtraState<MetadataMap>("ir.metadata_map"); + auto metadata = d->GetOrCreateExtraState<MetadataMap>("ir.metadata_map"); ffi::String key = value.GetTypeKey(); auto entries = metadata.Get(key).value_or(ffi::List<ffi::Any>{}); size_t index = 0; @@ -134,15 +99,13 @@ ExprDoc DocTranslatorObj::AddMetadata(ffi::Any value) { entries.push_back(value); metadata.Set(key, entries); } - ExprDoc entries_doc = IdDoc(GetExtraConfig<ffi::String>( + ExprDoc entries_doc = IdDoc(d->GetExtraConfig<ffi::String>( "script.metadata_name", "metadata"))[{LiteralDoc::Str(key, std::nullopt)}]; ExprDoc doc = entries_doc[{LiteralDoc::Int(index, std::nullopt)}]; - if (auto origin = value.as<ffi::ObjectRef>()) RecordOrigin(doc, origin.value()); + if (auto origin = value.as<ffi::ObjectRef>()) d->RecordOrigin(doc, origin.value()); return doc; } -namespace { - using GlobalInfoMap = ffi::Dict<ffi::String, ffi::List<GlobalInfo>>; // Selectors are derived from the module's forward registry, never cached as @@ -298,7 +261,7 @@ ExprDoc TranslateFloat(DocTranslatorObj* d, const FloatImmNode* imm) { } ExprDoc TranslateString(DocTranslatorObj* d, const StringImmNode* imm) { if (imm->value.find("\n") != ffi::String::npos) { - return d->AddMetadata(ffi::GetRef<StringImm>(imm)); + return AddMetadata(d, ffi::GetRef<StringImm>(imm)); } return NamespaceDoc("ir")->Attr("StringImm")->Call({LiteralDoc::Str(imm->value, std::nullopt)}); } @@ -346,7 +309,7 @@ ExprDoc TranslateGenericConst(DocTranslatorObj* d, const GenericConstNode* const return LiteralDoc::DataType(dtype.value(), std::nullopt); } auto tensor = constant->value.as<runtime::Tensor>(); - if (!tensor.has_value()) return d->AddMetadata(ffi::GetRef<GenericConst>(constant)); + if (!tensor.has_value()) return AddMetadata(d, ffi::GetRef<GenericConst>(constant)); const runtime::Tensor& data = tensor.value(); DLDataType dtype = data.DataType(); auto print_const = [&](ExprDoc value) -> ExprDoc { @@ -358,17 +321,17 @@ ExprDoc TranslateGenericConst(DocTranslatorObj* d, const GenericConstNode* const if (data->ndim != 0) { if (d->GetExtraConfig<bool>("script.show_meta", false) || d->GetExtraConfig<bool>("script.legacy_display", false)) { - return d->AddMetadata(ffi::GetRef<GenericConst>(constant)); + return AddMetadata(d, ffi::GetRef<GenericConst>(constant)); } // Nested Python lists cannot encode dimensions following an empty axis. for (int i = 0; i + 1 < data->ndim; ++i) { - if (data->shape[i] == 0) return d->AddMetadata(ffi::GetRef<GenericConst>(constant)); + if (data->shape[i] == 0) return AddMetadata(d, ffi::GetRef<GenericConst>(constant)); } bool supported = dtype == (DLDataType{kDLFloat, 32, 1}) || dtype == (DLDataType{kDLFloat, 64, 1}) || dtype == (DLDataType{kDLInt, 32, 1}) || dtype == (DLDataType{kDLInt, 64, 1}) || dtype == (DLDataType{kDLBool, 8, 1}); - if (!supported) return d->AddMetadata(ffi::GetRef<GenericConst>(constant)); + if (!supported) return AddMetadata(d, ffi::GetRef<GenericConst>(constant)); size_t count = 1; for (int i = 0; i < data->ndim; ++i) { TVM_FFI_CHECK( @@ -400,9 +363,9 @@ ExprDoc TranslateGenericConst(DocTranslatorObj* d, const GenericConstNode* const return print_const(make_list(0)); } if (data->device.device_type != kDLCPU) - return d->AddMetadata(ffi::GetRef<GenericConst>(constant)); + return AddMetadata(d, ffi::GetRef<GenericConst>(constant)); if (dtype.lanes != 1 || dtype.bits == 0 || dtype.bits > 64 || dtype.bits % 8 != 0) { - return d->AddMetadata(ffi::GetRef<GenericConst>(constant)); + return AddMetadata(d, ffi::GetRef<GenericConst>(constant)); } alignas(double) uint8_t scalar_bytes[sizeof(double)]{}; data.CopyToBytes(scalar_bytes, dtype.bits / 8); @@ -438,7 +401,7 @@ ExprDoc TranslateGenericConst(DocTranslatorObj* d, const GenericConstNode* const } else if (dtype == (DLDataType{kDLBool, 8, 1})) { scalar = LiteralDoc::Boolean(read_scalar(uint8_t{}), std::nullopt); } else { - return d->AddMetadata(ffi::GetRef<GenericConst>(constant)); + return AddMetadata(d, ffi::GetRef<GenericConst>(constant)); } return print_const(scalar); } @@ -884,7 +847,7 @@ ExprDoc CallAttrsValue(DocTranslatorObj* d, const Attrs& attrs) { ffi::Array<ffi::String> keys; ffi::Array<ExprDoc> values; for (const auto& [key, value] : fields) { - if (key == "type_key") return d->AddMetadata(attrs); + if (key == "type_key") return AddMetadata(d, attrs); keys.push_back(key); values.push_back(AnyValue(d, value)); }
