(tvm) branch main updated: Fix GPU detection in PerStoreFeatureNode (#17593)

[yongwww]

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

yongwww pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new da2e89ac19 Fix GPU detection in PerStoreFeatureNode (#17593)
da2e89ac19 is described below

commit da2e89ac19d2986b6e7ee8874952e80dd8c72dde
Author: Mingyu Xiong <102020096+moonque...@users.noreply.github.com>
AuthorDate: Sun Jan 19 11:12:37 2025 +0800

    Fix GPU detection in PerStoreFeatureNode (#17593)
---
 src/meta_schedule/feature_extractor/per_store_feature.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/meta_schedule/feature_extractor/per_store_feature.cc 
b/src/meta_schedule/feature_extractor/per_store_feature.cc
index 82bc7c2de0..955e39df4f 100644
--- a/src/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/meta_schedule/feature_extractor/per_store_feature.cc
@@ -1392,7 +1392,8 @@ class PerStoreFeatureNode : public FeatureExtractorNode {
 
   Array<runtime::NDArray> ExtractFrom(const TuneContext& tune_context,
                                       const Array<MeasureCandidate>& 
candidates) {
-    bool is_gpu = tune_context->target.value()->kind->name == "cuda";
+    auto& target_keys = tune_context->target.value()->keys;
+    bool is_gpu = std::find(target_keys.begin(), target_keys.end(), "gpu") != 
target_keys.end();
     std::vector<runtime::NDArray> results;
     results.resize(candidates.size());
     std::unique_ptr<tir::group6::Feature> feature_group6 = nullptr;