forked from apache/tvm
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Meta Schedule] Feature Extractor & Cost Model (apache#510)
* Fix sttr func & schedule naming. * Fix schedule -> sch. * Add feature extractor. * Fix init. * Revert wrong description. * Add cost model. * Remove unused include. * Fix issues. * Fix init. Co-authored-by: Junru Shao <[email protected]>
- Loading branch information
Showing
22 changed files
with
884 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,185 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, | ||
| * software distributed under the License is distributed on an | ||
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| * KIND, either express or implied. See the License for the | ||
| * specific language governing permissions and limitations | ||
| * under the License. | ||
| */ | ||
|
|
||
| #ifndef TVM_META_SCHEDULE_COST_MODEL_H_ | ||
| #define TVM_META_SCHEDULE_COST_MODEL_H_ | ||
|
|
||
| #include <tvm/meta_schedule/search_strategy.h> | ||
|
|
||
| namespace tvm { | ||
| namespace meta_schedule { | ||
|
|
||
| class TuneContext; | ||
|
|
||
| /*! \brief Cost model. */ | ||
| class CostModelNode : public runtime::Object { | ||
| public: | ||
| /*! \brief Virtual destructor. */ | ||
| virtual ~CostModelNode() = default; | ||
|
|
||
| void VisitAttrs(tvm::AttrVisitor* v) {} | ||
|
|
||
| /*! | ||
| * \brief Load the cost model from given file location. | ||
| * \param file_location The file location. | ||
| * \return Whether cost model was loaded successfully. | ||
| */ | ||
| virtual bool Load(const String& file_location) = 0; | ||
|
|
||
| /*! | ||
| * \brief Save the cost model to given file location. | ||
| * \param file_location The file location. | ||
| * \return Whether cost model was saved successfully. | ||
| */ | ||
| virtual bool Save(const String& file_location) = 0; | ||
|
|
||
| /*! | ||
| * \brief Update the cost model given running results. | ||
| * \param tune_context The tuning context. | ||
| * \param candidates The measure candidates. | ||
| * \param results The running results of the measure candidates. | ||
| */ | ||
| virtual void Update(const TuneContext& tune_context, const Array<MeasureCandidate>& candidates, | ||
| const Array<RunnerResult>& results) = 0; | ||
|
|
||
| /*! | ||
| * \brief Predict the running results of given measure candidates. | ||
| * \param tune_context The tuning context. | ||
| * \param candidates The measure candidates. | ||
| * \return The predicted running results. | ||
| */ | ||
| virtual std::vector<double> Predict(const TuneContext& tune_context, | ||
| const Array<MeasureCandidate>& candidates) = 0; | ||
|
|
||
| static constexpr const char* _type_key = "meta_schedule.CostModel"; | ||
| TVM_DECLARE_BASE_OBJECT_INFO(CostModelNode, Object); | ||
| }; | ||
|
|
||
| /*! \brief The cost model with customized methods on the python-side. */ | ||
| class PyCostModelNode : public CostModelNode { | ||
| public: | ||
| /*! | ||
| * \brief Load the cost model from given file location. | ||
| * \param file_location The file location. | ||
| * \return Whether cost model was loaded successfully. | ||
| */ | ||
| using FLoad = runtime::TypedPackedFunc<bool(String)>; | ||
| /*! | ||
| * \brief Save the cost model to given file location. | ||
| * \param file_location The file location. | ||
| * \return Whether cost model was saved successfully. | ||
| */ | ||
| using FSave = runtime::TypedPackedFunc<bool(String)>; | ||
| /*! | ||
| * \brief Update the cost model given running results. | ||
| * \param tune_context The tuning context. | ||
| * \param candidates The measure candidates. | ||
| * \param results The running results of the measure candidates. | ||
| * \return Whether cost model was updated successfully. | ||
| */ | ||
| using FUpdate = runtime::TypedPackedFunc<void(const TuneContext&, const Array<MeasureCandidate>&, | ||
| const Array<RunnerResult>&)>; | ||
| /*! | ||
| * \brief Predict the running results of given measure candidates. | ||
| * \param tune_context The tuning context. | ||
| * \param candidates The measure candidates. | ||
| * \param p_addr The address to save the the estimated running results. | ||
| */ | ||
| using FPredict = runtime::TypedPackedFunc<void(const TuneContext&, const Array<MeasureCandidate>&, | ||
| void* p_addr)>; | ||
| /*! | ||
| * \brief Get the cost model as string with name. | ||
| * \return The string representation of the cost model. | ||
| */ | ||
| using FAsString = runtime::TypedPackedFunc<String()>; | ||
|
|
||
| /*! \brief The packed function to the `Load` function. */ | ||
| FLoad f_load; | ||
| /*! \brief The packed function to the `Save` function. */ | ||
| FSave f_save; | ||
| /*! \brief The packed function to the `Update` function. */ | ||
| FUpdate f_update; | ||
| /*! \brief The packed function to the `Predict` function. */ | ||
| FPredict f_predict; | ||
| /*! \brief The packed function to the `AsString` function. */ | ||
| FAsString f_as_string; | ||
|
|
||
| void VisitAttrs(tvm::AttrVisitor* v) { | ||
| // `f_load` is not visited | ||
| // `f_save` is not visited | ||
| // `f_update` is not visited | ||
| // `f_predict` is not visited | ||
| // `f_as_string` is not visited | ||
| } | ||
|
|
||
| bool Load(const String& file_location) { | ||
| ICHECK(f_load != nullptr) << "PyCostModel's Load method not implemented!"; | ||
| return f_load(file_location); | ||
| } | ||
|
|
||
| bool Save(const String& file_location) { | ||
| ICHECK(f_save != nullptr) << "PyCostModel's Save method not implemented!"; | ||
| return f_save(file_location); | ||
| } | ||
|
|
||
| void Update(const TuneContext& tune_context, const Array<MeasureCandidate>& candidates, | ||
| const Array<RunnerResult>& results) { | ||
| ICHECK(f_update != nullptr) << "PyCostModel's Update method not implemented!"; | ||
| f_update(tune_context, candidates, results); | ||
| } | ||
|
|
||
| std::vector<double> Predict(const TuneContext& tune_context, | ||
| const Array<MeasureCandidate>& candidates) { | ||
| ICHECK(f_predict != nullptr) << "PyCostModel's Predict method not implemented!"; | ||
| std::vector<double> result(candidates.size(), 0.0); | ||
| f_predict(tune_context, candidates, result.data()); | ||
| return result; | ||
| } | ||
|
|
||
| static constexpr const char* _type_key = "meta_schedule.PyCostModel"; | ||
| TVM_DECLARE_FINAL_OBJECT_INFO(PyCostModelNode, CostModelNode); | ||
| }; | ||
|
|
||
| /*! | ||
| * \brief Managed reference to CostModelNode | ||
| * \sa CostModelNode | ||
| */ | ||
| class CostModel : public runtime::ObjectRef { | ||
| public: | ||
| /*! | ||
| * \brief Create a feature extractor with customized methods on the python-side. | ||
| * \param f_load The packed function of `Load`. | ||
| * \param f_save The packed function of `Save`. | ||
| * \param f_update The packed function of `Update`. | ||
| * \param f_predict The packed function of `Predict`. | ||
| * \param f_as_string The packed function of `AsString`. | ||
| * \return The feature extractor created. | ||
| */ | ||
| TVM_DLL static CostModel PyCostModel(PyCostModelNode::FLoad f_load, // | ||
| PyCostModelNode::FSave f_save, // | ||
| PyCostModelNode::FUpdate f_update, // | ||
| PyCostModelNode::FPredict f_predict, // | ||
| PyCostModelNode::FAsString f_as_string); | ||
| TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(CostModel, ObjectRef, CostModelNode); | ||
| }; | ||
|
|
||
| } // namespace meta_schedule | ||
| } // namespace tvm | ||
|
|
||
| #endif // TVM_META_SCHEDULE_COST_MODEL_H_ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,109 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, | ||
| * software distributed under the License is distributed on an | ||
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| * KIND, either express or implied. See the License for the | ||
| * specific language governing permissions and limitations | ||
| * under the License. | ||
| */ | ||
|
|
||
| #ifndef TVM_META_SCHEDULE_FEATURE_EXTRACTOR_H_ | ||
| #define TVM_META_SCHEDULE_FEATURE_EXTRACTOR_H_ | ||
|
|
||
| #include <tvm/meta_schedule/search_strategy.h> | ||
|
|
||
| namespace tvm { | ||
| namespace meta_schedule { | ||
|
|
||
| class TuneContext; | ||
|
|
||
| /*! \brief Extractor for features from measure candidates for use in cost model. */ | ||
| class FeatureExtractorNode : public runtime::Object { | ||
| public: | ||
| /*! \brief Virtual destructor. */ | ||
| virtual ~FeatureExtractorNode() = default; | ||
|
|
||
| void VisitAttrs(tvm::AttrVisitor* v) {} | ||
|
|
||
| /*! | ||
| * \brief Extract features from the given measure candidate. | ||
| * \param tune_context The tuning context for feature extraction. | ||
| * \param candidates The measure candidates to extract features from. | ||
| * \return The feature ndarray extracted. | ||
| */ | ||
| virtual Array<tvm::runtime::NDArray> ExtractFrom(const TuneContext& tune_context, | ||
| const Array<MeasureCandidate>& candidates) = 0; | ||
|
|
||
| static constexpr const char* _type_key = "meta_schedule.FeatureExtractor"; | ||
| TVM_DECLARE_BASE_OBJECT_INFO(FeatureExtractorNode, Object); | ||
| }; | ||
|
|
||
| /*! \brief The feature extractor with customized methods on the python-side. */ | ||
| class PyFeatureExtractorNode : public FeatureExtractorNode { | ||
| public: | ||
| /*! | ||
| * \brief Extract features from the given measure candidate. | ||
| * \param tune_context The tuning context for feature extraction. | ||
| * \param candidates The measure candidates to extract features from. | ||
| * \return The feature ndarray extracted. | ||
| */ | ||
| using FExtractFrom = runtime::TypedPackedFunc<Array<tvm::runtime::NDArray>( | ||
| const TuneContext& tune_context, const Array<MeasureCandidate>& candidates)>; | ||
| /*! | ||
| * \brief Get the feature extractor as string with name. | ||
| * \return The string of the feature extractor. | ||
| */ | ||
| using FAsString = runtime::TypedPackedFunc<String()>; | ||
|
|
||
| /*! \brief The packed function to the `ExtractFrom` function. */ | ||
| FExtractFrom f_extract_from; | ||
| /*! \brief The packed function to the `AsString` function. */ | ||
| FAsString f_as_string; | ||
|
|
||
| void VisitAttrs(tvm::AttrVisitor* v) { | ||
| // `f_extract_from` is not visited | ||
| // `f_as_string` is not visited | ||
| } | ||
|
|
||
| Array<tvm::runtime::NDArray> ExtractFrom(const TuneContext& tune_context, | ||
| const Array<MeasureCandidate>& candidates) { | ||
| ICHECK(f_extract_from != nullptr) << "PyFeatureExtractor's ExtractFrom method not implemented!"; | ||
| return f_extract_from(tune_context, candidates); | ||
| } | ||
|
|
||
| static constexpr const char* _type_key = "meta_schedule.PyFeatureExtractor"; | ||
| TVM_DECLARE_FINAL_OBJECT_INFO(PyFeatureExtractorNode, FeatureExtractorNode); | ||
| }; | ||
|
|
||
| /*! | ||
| * \brief Managed reference to FeatureExtractorNode | ||
| * \sa FeatureExtractorNode | ||
| */ | ||
| class FeatureExtractor : public runtime::ObjectRef { | ||
| public: | ||
| /*! | ||
| * \brief Create a feature extractor with customized methods on the python-side. | ||
| * \param f_extract_from The packed function of `ExtractFrom`. | ||
| * \param f_as_string The packed function of `AsString`. | ||
| * \return The feature extractor created. | ||
| */ | ||
| TVM_DLL static FeatureExtractor PyFeatureExtractor( | ||
| PyFeatureExtractorNode::FExtractFrom f_extract_from, // | ||
| PyFeatureExtractorNode::FAsString f_as_string); | ||
| TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(FeatureExtractor, ObjectRef, FeatureExtractorNode); | ||
| }; | ||
|
|
||
| } // namespace meta_schedule | ||
| } // namespace tvm | ||
|
|
||
| #endif // TVM_META_SCHEDULE_FEATURE_EXTRACTOR_H_ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.