Skip to content

Commit e1255c9

Browse files
junrushaospectrometerHBHMasterJH5574jinhongyiivinx13
authored
[MetaSchedule] Random Feature Extractor (#9760)
Co-authored-by: Junru Shao <[email protected]> Co-authored-by: Bohan Hou <[email protected]> Co-authored-by: Ruihang Lai <[email protected]> Co-authored-by: Hongyi Jin <[email protected]> Co-authored-by: Wuwei Lin <[email protected]> Co-authored-by: Siyuan Feng <[email protected]> Co-authored-by: Bohan Hou <[email protected]> Co-authored-by: Ruihang Lai <[email protected]> Co-authored-by: Hongyi Jin <[email protected]> Co-authored-by: Wuwei Lin <[email protected]> Co-authored-by: Siyuan Feng <[email protected]>
1 parent a374cdd commit e1255c9

File tree

16 files changed

+1137
-3
lines changed

16 files changed

+1137
-3
lines changed

‎include/tvm/meta_schedule/cost_model.h‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#ifndef TVM_META_SCHEDULE_COST_MODEL_H_
21+
#define TVM_META_SCHEDULE_COST_MODEL_H_
22+
23+
#include <tvm/meta_schedule/search_strategy.h>
24+
25+
#include <vector>
26+
27+
namespace tvm {
28+
namespace meta_schedule {
29+
30+
class TuneContext;
31+
32+
/*! \brief Cost model. */
33+
class CostModelNode : public runtime::Object {
34+
public:
35+
/*! \brief Virtual destructor. */
36+
virtual ~CostModelNode() = default;
37+
38+
void VisitAttrs(tvm::AttrVisitor* v) {}
39+
40+
/*!
41+
* \brief Load the cost model from given file location.
42+
* \param path The file path.
43+
*/
44+
virtual void Load(const String& path) = 0;
45+
46+
/*!
47+
* \brief Save the cost model to given file location.
48+
* \param path The file path.
49+
*/
50+
virtual void Save(const String& path) = 0;
51+
52+
/*!
53+
* \brief Update the cost model given running results.
54+
* \param tune_context The tuning context.
55+
* \param candidates The measure candidates.
56+
* \param results The running results of the measure candidates.
57+
*/
58+
virtual void Update(const TuneContext& tune_context, const Array<MeasureCandidate>& candidates,
59+
const Array<RunnerResult>& results) = 0;
60+
61+
/*!
62+
* \brief Predict the normalized score (the larger the better) of given measure candidates.
63+
* \param tune_context The tuning context.
64+
* \param candidates The measure candidates.
65+
* \return The predicted normalized score.
66+
*/
67+
virtual std::vector<double> Predict(const TuneContext& tune_context,
68+
const Array<MeasureCandidate>& candidates) = 0;
69+
70+
static constexpr const char* _type_key = "meta_schedule.CostModel";
71+
TVM_DECLARE_BASE_OBJECT_INFO(CostModelNode, Object);
72+
};
73+
74+
/*! \brief The cost model with customized methods on the python-side. */
75+
class PyCostModelNode : public CostModelNode {
76+
public:
77+
/*!
78+
* \brief Load the cost model from given file location.
79+
* \param path The file path.
80+
*/
81+
using FLoad = runtime::TypedPackedFunc<void(String)>;
82+
/*!
83+
* \brief Save the cost model to given file location.
84+
* \param path The file path.
85+
*/
86+
using FSave = runtime::TypedPackedFunc<void(String)>;
87+
/*!
88+
* \brief Update the cost model given running results.
89+
* \param tune_context The tuning context.
90+
* \param candidates The measure candidates.
91+
* \param results The running results of the measure candidates.
92+
* \return Whether cost model was updated successfully.
93+
*/
94+
using FUpdate = runtime::TypedPackedFunc<void(const TuneContext&, const Array<MeasureCandidate>&,
95+
const Array<RunnerResult>&)>;
96+
/*!
97+
* \brief Predict the running results of given measure candidates.
98+
* \param tune_context The tuning context.
99+
* \param candidates The measure candidates.
100+
* \param p_addr The address to save the the estimated running results.
101+
*/
102+
using FPredict = runtime::TypedPackedFunc<void(const TuneContext&, const Array<MeasureCandidate>&,
103+
void* p_addr)>;
104+
/*!
105+
* \brief Get the cost model as string with name.
106+
* \return The string representation of the cost model.
107+
*/
108+
using FAsString = runtime::TypedPackedFunc<String()>;
109+
110+
/*! \brief The packed function to the `Load` function. */
111+
FLoad f_load;
112+
/*! \brief The packed function to the `Save` function. */
113+
FSave f_save;
114+
/*! \brief The packed function to the `Update` function. */
115+
FUpdate f_update;
116+
/*! \brief The packed function to the `Predict` function. */
117+
FPredict f_predict;
118+
/*! \brief The packed function to the `AsString` function. */
119+
FAsString f_as_string;
120+
121+
void VisitAttrs(tvm::AttrVisitor* v) {
122+
// `f_load` is not visited
123+
// `f_save` is not visited
124+
// `f_update` is not visited
125+
// `f_predict` is not visited
126+
// `f_as_string` is not visited
127+
}
128+
129+
void Load(const String& path) {
130+
ICHECK(f_load != nullptr) << "PyCostModel's Load method not implemented!";
131+
f_load(path);
132+
}
133+
134+
void Save(const String& path) {
135+
ICHECK(f_save != nullptr) << "PyCostModel's Save method not implemented!";
136+
f_save(path);
137+
}
138+
void Update(const TuneContext& tune_context, const Array<MeasureCandidate>& candidates,
139+
const Array<RunnerResult>& results) {
140+
ICHECK(f_update != nullptr) << "PyCostModel's Update method not implemented!";
141+
f_update(tune_context, candidates, results);
142+
}
143+
144+
std::vector<double> Predict(const TuneContext& tune_context,
145+
const Array<MeasureCandidate>& candidates) {
146+
ICHECK(f_predict != nullptr) << "PyCostModel's Predict method not implemented!";
147+
std::vector<double> result(candidates.size(), 0.0);
148+
f_predict(tune_context, candidates, result.data());
149+
return result;
150+
}
151+
152+
static constexpr const char* _type_key = "meta_schedule.PyCostModel";
153+
TVM_DECLARE_FINAL_OBJECT_INFO(PyCostModelNode, CostModelNode);
154+
};
155+
156+
/*!
157+
* \brief Managed reference to CostModelNode
158+
* \sa CostModelNode
159+
*/
160+
class CostModel : public runtime::ObjectRef {
161+
public:
162+
/*!
163+
* \brief Create a feature extractor with customized methods on the python-side.
164+
* \param f_load The packed function of `Load`.
165+
* \param f_save The packed function of `Save`.
166+
* \param f_update The packed function of `Update`.
167+
* \param f_predict The packed function of `Predict`.
168+
* \param f_as_string The packed function of `AsString`.
169+
* \return The feature extractor created.
170+
*/
171+
TVM_DLL static CostModel PyCostModel(PyCostModelNode::FLoad f_load, //
172+
PyCostModelNode::FSave f_save, //
173+
PyCostModelNode::FUpdate f_update, //
174+
PyCostModelNode::FPredict f_predict, //
175+
PyCostModelNode::FAsString f_as_string);
176+
TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(CostModel, ObjectRef, CostModelNode);
177+
};
178+
179+
} // namespace meta_schedule
180+
} // namespace tvm
181+
182+
#endif // TVM_META_SCHEDULE_COST_MODEL_H_

‎include/tvm/meta_schedule/feature_extractor.h‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#ifndef TVM_META_SCHEDULE_FEATURE_EXTRACTOR_H_
21+
#define TVM_META_SCHEDULE_FEATURE_EXTRACTOR_H_
22+
23+
#include <tvm/meta_schedule/search_strategy.h>
24+
25+
namespace tvm {
26+
namespace meta_schedule {
27+
28+
class TuneContext;
29+
30+
/*! \brief Extractor for features from measure candidates for use in cost model. */
31+
class FeatureExtractorNode : public runtime::Object {
32+
public:
33+
/*! \brief Virtual destructor. */
34+
virtual ~FeatureExtractorNode() = default;
35+
36+
void VisitAttrs(tvm::AttrVisitor* v) {}
37+
38+
/*!
39+
* \brief Extract features from the given measure candidate.
40+
* \param tune_context The tuning context for feature extraction.
41+
* \param candidates The measure candidates to extract features from.
42+
* \return The feature ndarray extracted.
43+
*/
44+
virtual Array<tvm::runtime::NDArray> ExtractFrom(const TuneContext& tune_context,
45+
const Array<MeasureCandidate>& candidates) = 0;
46+
47+
static constexpr const char* _type_key = "meta_schedule.FeatureExtractor";
48+
TVM_DECLARE_BASE_OBJECT_INFO(FeatureExtractorNode, Object);
49+
};
50+
51+
/*! \brief The feature extractor with customized methods on the python-side. */
52+
class PyFeatureExtractorNode : public FeatureExtractorNode {
53+
public:
54+
/*!
55+
* \brief Extract features from the given measure candidate.
56+
* \param tune_context The tuning context for feature extraction.
57+
* \param candidates The measure candidates to extract features from.
58+
* \return The feature ndarray extracted.
59+
*/
60+
using FExtractFrom = runtime::TypedPackedFunc<Array<tvm::runtime::NDArray>(
61+
const TuneContext& tune_context, const Array<MeasureCandidate>& candidates)>;
62+
/*!
63+
* \brief Get the feature extractor as string with name.
64+
* \return The string of the feature extractor.
65+
*/
66+
using FAsString = runtime::TypedPackedFunc<String()>;
67+
68+
/*! \brief The packed function to the `ExtractFrom` function. */
69+
FExtractFrom f_extract_from;
70+
/*! \brief The packed function to the `AsString` function. */
71+
FAsString f_as_string;
72+
73+
void VisitAttrs(tvm::AttrVisitor* v) {
74+
// `f_extract_from` is not visited
75+
// `f_as_string` is not visited
76+
}
77+
78+
Array<tvm::runtime::NDArray> ExtractFrom(const TuneContext& tune_context,
79+
const Array<MeasureCandidate>& candidates) {
80+
ICHECK(f_extract_from != nullptr) << "PyFeatureExtractor's ExtractFrom method not implemented!";
81+
return f_extract_from(tune_context, candidates);
82+
}
83+
84+
static constexpr const char* _type_key = "meta_schedule.PyFeatureExtractor";
85+
TVM_DECLARE_FINAL_OBJECT_INFO(PyFeatureExtractorNode, FeatureExtractorNode);
86+
};
87+
88+
/*!
89+
* \brief Managed reference to FeatureExtractorNode
90+
* \sa FeatureExtractorNode
91+
*/
92+
class FeatureExtractor : public runtime::ObjectRef {
93+
public:
94+
/*!
95+
* \brief Create a feature extractor that extracts features from each BufferStore
96+
* \param buffers_per_store The number of buffers in each BufferStore; Pad or truncate if
97+
* necessary.
98+
* \param arith_intensity_curve_num_samples The number of samples used in the arithmetic intensity
99+
* curve.
100+
* \param cache_line_bytes The number of bytes in a cache line.
101+
* \return The feature extractor created.
102+
*/
103+
TVM_DLL static FeatureExtractor PerStoreFeature(int buffers_per_store = 5,
104+
int arith_intensity_curve_num_samples = 10,
105+
int cache_line_bytes = 64);
106+
/*!
107+
* \brief Create a feature extractor with customized methods on the python-side.
108+
* \param f_extract_from The packed function of `ExtractFrom`.
109+
* \param f_as_string The packed function of `AsString`.
110+
* \return The feature extractor created.
111+
*/
112+
TVM_DLL static FeatureExtractor PyFeatureExtractor(
113+
PyFeatureExtractorNode::FExtractFrom f_extract_from,
114+
PyFeatureExtractorNode::FAsString f_as_string);
115+
TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(FeatureExtractor, ObjectRef, FeatureExtractorNode);
116+
};
117+
118+
} // namespace meta_schedule
119+
} // namespace tvm
120+
121+
#endif // TVM_META_SCHEDULE_FEATURE_EXTRACTOR_H_

‎python/tvm/meta_schedule/cost_model/__init__.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""
18+
The tvm.meta_schedule.cost_model package.
19+
"""
20+
from .cost_model import CostModel, PyCostModel
21+
from .random_model import RandomModel

0 commit comments

Comments
 (0)