diff --git a/src/target/target.cc b/src/target/target.cc index 91a5854b3934..7d03d5e14c69 100644 --- a/src/target/target.cc +++ b/src/target/target.cc @@ -330,11 +330,38 @@ ObjectPtr TargetInternal::FromConfig(ffi::Map 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 saved_features; + { + std::vector 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 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()) { target->tag = tag.value(); @@ -367,14 +394,14 @@ ObjectPtr TargetInternal::FromConfig(ffi::Map 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 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()) { diff --git a/tests/cpp/target_test.cc b/tests/cpp/target_test.cc index df592f760866..342870e88f55 100644 --- a/tests/cpp/target_test.cc +++ b/tests/cpp/target_test.cc @@ -222,6 +222,63 @@ TEST(TargetCreation, TargetParserProcessing) { ASSERT_EQ(test_target->GetAttr("mattr").value(), "cake"); } +TEST(TargetCreation, RoundTripCanonicalizerFeatures) { + // Construct a target whose canonicalizer sets feature.test and transforms mcpu + Target original(ffi::Map{ + {"kind", ffi::String("TestTargetParser")}, + {"mcpu", ffi::String("woof")}, + }); + ASSERT_EQ(original->GetAttr("mcpu").value(), "super_woof"); + ASSERT_EQ(original->GetAttr("feature.test").value(), true); + + // Export to config and reconstruct + ffi::Map 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("mcpu").has_value()); + ASSERT_EQ(reconstructed->GetAttr("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{ + {"kind", ffi::String("TestTargetParser")}, + {"mcpu", ffi::String("woof")}, + }); + ASSERT_EQ(host->GetAttr("feature.test").value(), true); + + // Attach it as host to another target + Target outer(ffi::Map{ + {"kind", ffi::String("TestTargetKind")}, + {"my_bool", true}, + }); + Target combined(outer, host); + + // Export the outer target (includes nested host) and reconstruct + ffi::Map exported = combined->ToConfig(); + Target reconstructed(exported); + + // The nested host must reconstruct successfully with feature.* preserved + ffi::Optional reconstructed_host = reconstructed->GetHost(); + ASSERT_TRUE(reconstructed_host.defined()); + ASSERT_EQ(reconstructed_host.value()->GetAttr("feature.test").value(), true); + ASSERT_TRUE(reconstructed_host.value()->GetAttr("mcpu").has_value()); +} + +TEST(TargetCreationFail, UnknownNonFeatureKeyStillFails) { + // Verify that unknown non-feature.* keys still fail schema validation + ffi::Map 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("single") .add_attr_option>("array")