This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch tqchen/printer-refinement
in repository https://gitbox.apache.org/repos/asf/tvm.git

commit da005400de5ab4cfd37217031f510c2c738334d4
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Oct 1 20:34:28 2026 +0000

    Preserve generic attribute keys and typed constant values
---
 src/script/printer/ir/misc.cc | 38 +++++++++++++++++++++++++++++++-------
 1 file changed, 31 insertions(+), 7 deletions(-)

diff --git a/src/script/printer/ir/misc.cc b/src/script/printer/ir/misc.cc
index bc98604d75..58427196ea 100644
--- a/src/script/printer/ir/misc.cc
+++ b/src/script/printer/ir/misc.cc
@@ -134,16 +134,19 @@ ExprDoc AnyValue(DocTranslatorObj* d, ffi::AnyView value) 
{
     return ListDoc(elements);
   }
   if (auto map = value.as<ffi::Map<ffi::Any, ffi::Any>>()) {
-    std::vector<std::pair<ffi::String, ffi::Any>> entries;
-    for (const auto& [key, item] : map.value()) {
-      entries.emplace_back(key.as_or_throw<ffi::String>(), item);
+    std::vector<std::pair<ffi::Any, ffi::Any>> entries(map->begin(), 
map->end());
+    if (std::all_of(entries.begin(), entries.end(), [](const auto& entry) {
+          return entry.first.template as<ffi::String>().has_value();
+        })) {
+      std::sort(entries.begin(), entries.end(), [](const auto& a, const auto& 
b) {
+        return a.first.template as_or_throw<ffi::String>() <
+               b.first.template as_or_throw<ffi::String>();
+      });
     }
-    std::sort(entries.begin(), entries.end(),
-              [](const auto& a, const auto& b) { return a.first < b.first; });
     ffi::Array<ExprDoc> keys;
     ffi::Array<ExprDoc> values;
     for (const auto& [key, item] : entries) {
-      keys.push_back(LiteralDoc::Str(key, std::nullopt));
+      keys.push_back(AnyValue(d, key));
       values.push_back(AnyValue(d, item));
     }
     return DictDoc(keys, values);
@@ -175,8 +178,29 @@ ExprDoc AnyValue(DocTranslatorObj* d, ffi::AnyView value) {
         {LiteralDoc::Int(vdevice.value()->vdevice_id, std::nullopt),
          LiteralDoc::Str(vdevice.value()->memory_scope, std::nullopt)});
   }
+  // Any-valued fields must retain IR objects; Expr-valued fields convert 
literals.
+  if (const auto* integer = value.as<IntImmNode>()) {
+    ExprDoc doc = NamespaceDoc("tirx")->Attr("IntImm")->Call(
+        {LiteralDoc::DataType(integer->ty.as_or_throw<PrimType>()->dtype, 
std::nullopt),
+         LiteralDoc::Int(ffi::GetRef<IntImm>(integer), std::nullopt)});
+    d->RecordOrigin(doc, ffi::GetRef<IntImm>(integer));
+    return doc;
+  }
+  if (const auto* floating = value.as<FloatImmNode>()) {
+    ExprDoc literal = LiteralDoc::Float(floating->value, std::nullopt);
+    if (!std::isfinite(floating->value)) {
+      // Nonfinite literals render as strings; the scalar helper converts them 
to a number.
+      literal = 
NamespaceDoc("tirx")->Attr("float64")->Call({literal})->Attr("value");
+    }
+    ExprDoc doc =
+        NamespaceDoc("tirx")
+            ->Attr("FloatImm")
+            
->Call({LiteralDoc::DataType(floating->ty.as_or_throw<PrimType>()->dtype, 
std::nullopt),
+                    literal});
+    d->RecordOrigin(doc, ffi::GetRef<FloatImm>(floating));
+    return doc;
+  }
   if (const auto* string = value.as<StringImmNode>()) {
-    // Any-valued fields must retain the IR object; Expr-valued fields convert 
literals.
     ExprDoc doc =
         
NamespaceDoc("ir")->Attr("StringImm")->Call({LiteralDoc::Str(string->value, 
std::nullopt)});
     d->RecordOrigin(doc, ffi::GetRef<StringImm>(string));

Reply via email to