[TARGET] Fix round-trip reconstruction of targets with canonicalizer-generated feature.* attrs

src/target/target.cc

@@ -330,11 +330,38 @@ ObjectPtr<TargetNode> TargetInternal::FromConfig(ffi::Map<ffi::String, ffi::Any>
     target->host = std::nullopt;
   }
 
-  // Step 3: Use ConfigSchema to validate types, apply defaults, and run canonicalizer
+  // Step 3: Extract any "feature.*" keys from the input config before schema validation.
+  // These are canonicalizer-owned metadata that may appear in exported configs (via ToConfig())
+  // but are not part of the target kind's declared schema. We preserve them across round-trip
+  // and let the canonicalizer's output take priority if it re-emits the same key.
+  std::unordered_map<std::string, ffi::Any> saved_features;
+  {
+    std::vector<ffi::String> feature_keys;
+    for (const auto& kv : config) {
+      std::string key_str(kv.first);
+      if (key_str.size() > 8 && key_str.compare(0, 8, "feature.") == 0) {
+        saved_features[key_str] = kv.second;
+        feature_keys.push_back(kv.first);
+      }
+    }
+    for (const auto& k : feature_keys) {
+      config.erase(k);
+    }
+  }
+
+  // Step 4: Use ConfigSchema to validate types, apply defaults, and run canonicalizer
   // Note: structural keys (kind, tag, keys, device) pass through to canonicalizer
   ffi::Map<ffi::String, ffi::Any> resolved = target->kind->schema_.Resolve(config);
 
-  // Step 4: Extract structural fields from resolved config
+  // Step 5: Merge back preserved feature.* keys. Canonicalizer output is authoritative:
+  // only restore a saved feature if the canonicalizer did not re-emit it.
+  for (const auto& kv : saved_features) {
+    if (resolved.find(ffi::String(kv.first)) == resolved.end()) {
+      resolved.Set(ffi::String(kv.first), kv.second);
+    }
+  }
+
+  // Step 6: Extract structural fields from resolved config
   if (resolved.count(kTag)) {
     if (auto tag = resolved[kTag].try_cast<ffi::String>()) {
       target->tag = tag.value();
@@ -367,14 +394,14 @@ ObjectPtr<TargetNode> TargetInternal::FromConfig(ffi::Map<ffi::String, ffi::Any>
     resolved.erase(kKeys);
   }
 
-  // Step 5: Build attrs from resolved entries (excluding structural keys)
+  // Step 7: Build attrs from resolved entries (excluding structural keys)
   resolved.erase(kKind);
   std::unordered_map<ffi::String, ffi::Any> attrs;
   for (const auto& kv : resolved) {
     attrs[kv.first] = kv.second;
   }
 
-  // Step 6: If requested, query attributes from the device. User-specified
+  // Step 8: If requested, query attributes from the device. User-specified
   // parameters take precedence over queried parameters.
   int64_t from_device_id = -1;
   if (auto it = attrs.find(kFromDevice); it != attrs.end()) {

tests/cpp/target_test.cc

@@ -222,6 +222,63 @@ TEST(TargetCreation, TargetParserProcessing) {
   ASSERT_EQ(test_target->GetAttr<ffi::String>("mattr").value(), "cake");
 }
 
+TEST(TargetCreation, RoundTripCanonicalizerFeatures) {
+  // Construct a target whose canonicalizer sets feature.test and transforms mcpu
+  Target original(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestTargetParser")},
+      {"mcpu", ffi::String("woof")},
+  });
+  ASSERT_EQ(original->GetAttr<ffi::String>("mcpu").value(), "super_woof");
+  ASSERT_EQ(original->GetAttr<bool>("feature.test").value(), true);
+
+  // Export to config and reconstruct
+  ffi::Map<ffi::String, ffi::Any> exported = original->ToConfig();
+  Target reconstructed(exported);
+
+  // Canonicalized attrs must survive the round-trip
+  // Note: mcpu gets canonicalized again (super_super_woof) because the canonicalizer runs
+  ASSERT_TRUE(reconstructed->GetAttr<ffi::String>("mcpu").has_value());
+  ASSERT_EQ(reconstructed->GetAttr<bool>("feature.test").value(), true);
+  ASSERT_EQ(reconstructed->keys.size(), 1);
+  ASSERT_EQ(reconstructed->keys[0], "super");
+}
+
+TEST(TargetCreation, RoundTripCanonicalizerFeaturesNestedHost) {
+  // Construct a host target whose canonicalizer sets feature.test
+  Target host(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestTargetParser")},
+      {"mcpu", ffi::String("woof")},
+  });
+  ASSERT_EQ(host->GetAttr<bool>("feature.test").value(), true);
+
+  // Attach it as host to another target
+  Target outer(ffi::Map<ffi::String, ffi::Any>{
+      {"kind", ffi::String("TestTargetKind")},
+      {"my_bool", true},
+  });
+  Target combined(outer, host);
+
+  // Export the outer target (includes nested host) and reconstruct
+  ffi::Map<ffi::String, ffi::Any> exported = combined->ToConfig();
+  Target reconstructed(exported);
+
+  // The nested host must reconstruct successfully with feature.* preserved
+  ffi::Optional<Target> reconstructed_host = reconstructed->GetHost();
+  ASSERT_TRUE(reconstructed_host.defined());
+  ASSERT_EQ(reconstructed_host.value()->GetAttr<bool>("feature.test").value(), true);
+  ASSERT_TRUE(reconstructed_host.value()->GetAttr<ffi::String>("mcpu").has_value());
+}
+
+TEST(TargetCreationFail, UnknownNonFeatureKeyStillFails) {
+  // Verify that unknown non-feature.* keys still fail schema validation
+  ffi::Map<ffi::String, ffi::Any> config = {
+      {"kind", ffi::String("TestTargetParser")},
+      {"mcpu", ffi::String("woof")},
+      {"unknown_key", ffi::String("bad")},
+  };
+  ASSERT_THROW({ Target{config}; }, tvm::Error);
+}
+
 TVM_REGISTER_TARGET_KIND("TestStringKind", kDLCPU)
     .add_attr_option<ffi::String>("single")
     .add_attr_option<ffi::Array<ffi::String>>("array")