From: Lianmin Zheng Date: Thu, 6 Aug 2020 10:46:12 +0000 (-0700) Subject: [Ansor][AutoTVM v2.0] Phase 1: The base class for cost models (#6187) X-Git-Tag: upstream/0.7.0~299 X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=7ef89ad2dec184a033ed135dbef9fa599f55808b;p=platform%2Fupstream%2Ftvm.git [Ansor][AutoTVM v2.0] Phase 1: The base class for cost models (#6187) * add the base class for cost models * address comments * Update tests/python/unittest/test_auto_scheduler_cost_model.py Disable test if user doesn't have llvm Co-authored-by: Zhao Wu --- diff --git a/include/tvm/auto_scheduler/cost_model.h b/include/tvm/auto_scheduler/cost_model.h new file mode 100644 index 0000000..a6da93f --- /dev/null +++ b/include/tvm/auto_scheduler/cost_model.h @@ -0,0 +1,160 @@ +/* + * 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. + */ + +/*! + * \file tvm/auto_scheduler/cost_model.h + * \brief Cost models that estimate the performance of programs + */ + +#ifndef TVM_AUTO_SCHEDULER_COST_MODEL_H_ +#define TVM_AUTO_SCHEDULER_COST_MODEL_H_ + +#include +#include +#include +#include + +#include + +namespace tvm { +namespace auto_scheduler { + +using runtime::PackedFunc; +using runtime::TypedPackedFunc; + +/*! \brief The base class for cost model */ +class CostModelNode : public Object { + public: + /*! + * \brief Update the cost model according to new measurement results (training data). + * \param inputs The measure inputs + * \param results The measure results + */ + virtual void Update(const Array& inputs, const Array& results) = 0; + + /*! + * \brief Predict the scores of states + * \param task The search task of states + * \param states The input states + * \param scores The predicted scores for all states + */ + virtual void Predict(const SearchTask& task, const std::vector& states, + std::vector* scores) = 0; + + /*! + * \brief Predict the scores of all stages in states. This is the breakdown version of `Predict` + * \param task The search task + * \param states The input states + * \param state_scores The predicted scores for all states + * \param stage_scores The predicted scores for all stages in all stages + */ + virtual void PredictStages(const SearchTask& task, const std::vector& states, + std::vector* state_scores, + std::vector>* stage_scores) { + LOG(FATAL) << "Not implemented"; + } + + static constexpr const char* _type_key = "auto_scheduler.CostModel"; + TVM_DECLARE_BASE_OBJECT_INFO(CostModelNode, Object); +}; + +/*! + * \brief Managed reference to CostModelNode. + * \sa CostModelNode + */ +class CostModel : public ObjectRef { + public: + TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(CostModel, ObjectRef, CostModelNode); +}; + +/*! \brief The cost model returning random value for all predictions */ +class RandomModelNode : public CostModelNode { + public: + /*! \brief Pointer to a random number generator function */ + const TypedPackedFunc* random_number_func; + + void Update(const Array& inputs, const Array& results) final; + + void Predict(const SearchTask& task, const std::vector& states, + std::vector* scores) final; + + static constexpr const char* _type_key = "auto_scheduler.RandomModel"; + TVM_DECLARE_FINAL_OBJECT_INFO(RandomModelNode, CostModelNode); +}; + +/*! + * \brief Managed reference to RandomModelNode. + * \sa RandomModelNode + */ +class RandomModel : public CostModel { + public: + RandomModel(); + explicit RandomModel(::tvm::runtime::ObjectPtr<::tvm::runtime::Object> n) : CostModel(n) {} + + RandomModelNode* operator->() const { return static_cast(data_.get()); } + + TVM_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(RandomModel); + using ContainerType = RandomModelNode; +}; + +/*! \brief A wrapper for cost model defined by python code + * This class will call functions defined in the python */ +class PythonBasedModelNode : public CostModelNode { + public: + /*! \brief Pointer to the update funcion in python */ + PackedFunc update_func; + /*! \brief Pointer to the predict funcion in python */ + PackedFunc predict_func; + /*! \brief Pointer to the predict funcion in python */ + PackedFunc predict_stage_func; + + void Update(const Array& inputs, const Array& results) final; + + void Predict(const SearchTask& task, const std::vector& states, + std::vector* scores) final; + + void PredictStages(const SearchTask& task, const std::vector& states, + std::vector* state_scores, + std::vector>* stage_scores) final; + + static constexpr const char* _type_key = "auto_scheduler.PythonBasedModel"; + TVM_DECLARE_FINAL_OBJECT_INFO(PythonBasedModelNode, CostModelNode); +}; + +/*! + * \brief Managed reference to PythonBasedModelNode. + * \sa PythonBasedModelNode + */ +class PythonBasedModel : public CostModel { + public: + /*! + * \brief The constructor. + * \param update_func The pointer to the update function defined in python + * \param predict_func The pointer to the prediction function defined in python + * \param predict_stage_func The pointer to the prediction function defined in python + */ + PythonBasedModel(PackedFunc update_func, PackedFunc predict_func, PackedFunc predict_stage_func); + + TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PythonBasedModel, CostModel, PythonBasedModelNode); +}; + +} // namespace auto_scheduler +} // namespace tvm + +#endif // TVM_AUTO_SCHEDULER_COST_MODEL_H_ diff --git a/python/tvm/auto_scheduler/__init__.py b/python/tvm/auto_scheduler/__init__.py index c3a3712..62ebce0 100644 --- a/python/tvm/auto_scheduler/__init__.py +++ b/python/tvm/auto_scheduler/__init__.py @@ -25,9 +25,10 @@ from . import utils from . import workload_registry # Shortcut -from .compute_dag import ComputeDAG from .auto_schedule import SearchTask, TuningOptions, HardwareParams, \ auto_schedule, EmptyPolicy +from .compute_dag import ComputeDAG +from .cost_model import RandomModel from .measure import MeasureInput, LocalBuilder, LocalRunner, RPCRunner, \ LocalRPCMeasureContext from .measure_record import RecordToFile, RecordReader, load_best, \ diff --git a/python/tvm/auto_scheduler/cost_model/__init__.py b/python/tvm/auto_scheduler/cost_model/__init__.py new file mode 100644 index 0000000..fc3821c --- /dev/null +++ b/python/tvm/auto_scheduler/cost_model/__init__.py @@ -0,0 +1,20 @@ +# 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. +# pylint: disable=unused-import, redefined-builtin +""" Cost model that estimates the performance of programs """ + +from .cost_model import RandomModel diff --git a/python/tvm/auto_scheduler/cost_model/cost_model.py b/python/tvm/auto_scheduler/cost_model/cost_model.py new file mode 100644 index 0000000..699719d --- /dev/null +++ b/python/tvm/auto_scheduler/cost_model/cost_model.py @@ -0,0 +1,150 @@ +# 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. + +""" Cost model that estimates the performance of programs """ +import ctypes +import numpy as np + +import tvm._ffi +from tvm.runtime import Object +from .. import _ffi_api + + +@tvm._ffi.register_object("auto_scheduler.CostModel") +class CostModel(Object): + """The base class for cost model""" + +@tvm._ffi.register_object("auto_scheduler.RandomModel") +class RandomModel(CostModel): + """A model returns random estimation for all inputs""" + def __init__(self): + self.__init_handle_by_constructor__(_ffi_api.RandomModel) + + def update(self, inputs, results): + """Update the cost model according to new measurement results (training data). + + Parameters + ---------- + inputs : List[MeasureInput] + The measurement inputs + results : List[MeasureResult] + The measurement results + """ + _ffi_api.CostModelUpdate(self, inputs, results) + + def predict(self, search_task, states): + """Predict the scores of states + + Parameters + ---------- + search_task : SearchTask + The search task of states + statse : List[State] + The input states + + Returns + ------- + scores: List[float] + The predicted scores for all states + """ + return [x.value for x in _ffi_api.CostModelPredict(self, search_task, states)] + + +@tvm._ffi.register_func("auto_scheduler.cost_model.random_fill_float") +def random_fill_float(size, return_ptr): + """Fills a c++ float array with random numbers in [0, 1] + + Parameters + ---------- + size: int + The size of the array + return_ptr: + A pointer to a c++ float array + """ + if size == 0: + return + return_ptr = ctypes.cast(return_ptr, ctypes.POINTER(ctypes.c_float)) + array_wrapper = np.ctypeslib.as_array(return_ptr, shape=(size,)) + array_wrapper[:] = np.random.uniform(0, 1, (size,)) + + +@tvm._ffi.register_object("auto_scheduler.PythonBasedModel") +class PythonBasedModel(CostModel): + """Base class for cost models implemented in python""" + def __init__(self): + def update_func(inputs, results): + self.update(inputs, results) + + def predict_func(task, states, return_ptr): + return_ptr = ctypes.cast(return_ptr, ctypes.POINTER(ctypes.c_float)) + array_wrapper = np.ctypeslib.as_array(return_ptr, shape=(len(states),)) + array_wrapper[:] = self.predict(task, states) + + def predict_stage_func(task, states, return_ptr): + ret = self.predict_stages(task, states) + return_ptr = ctypes.cast(return_ptr, ctypes.POINTER(ctypes.c_float)) + array_wrapper = np.ctypeslib.as_array(return_ptr, shape=ret.shape) + array_wrapper[:] = ret + + self.__init_handle_by_constructor__(_ffi_api.PythonBasedModel, update_func, + predict_func, predict_stage_func) + + def update(self, inputs, results): + """Update the cost model according to new measurement results (training data). + + Parameters + ---------- + inputs : List[MeasureInput] + The measurement inputs + results : List[MeasureResult] + The measurement results + """ + raise NotImplementedError + + def predict(self, task, states): + """Predict the scores of states + + Parameters + ---------- + search_task : SearchTask + The search task of states + statse : List[State] + The input states + + Returns + ------- + scores: List[float] + The predicted scores for all states + """ + raise NotImplementedError + + def predict_stages(self, task, states): + """Predict the scores of all stages in states. This is the breakdown version of `predict`. + + Parameters + ---------- + search_task : SearchTask + The search task of states + statse : List[State] + The input states + + Returns + ------- + scores: List[float] + The predicted scores for all stages in all states in the packed format + """ + raise NotImplementedError diff --git a/src/auto_scheduler/cost_model.cc b/src/auto_scheduler/cost_model.cc new file mode 100644 index 0000000..2f1c309 --- /dev/null +++ b/src/auto_scheduler/cost_model.cc @@ -0,0 +1,156 @@ +/* + * 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. + */ + +/*! + * \file auto_scheduler/cost_model.cc + * \brief Cost models that estimate the performance of programs + */ + +#include + +namespace tvm { +namespace auto_scheduler { + +TVM_REGISTER_OBJECT_TYPE(CostModelNode); +TVM_REGISTER_OBJECT_TYPE(RandomModelNode); +TVM_REGISTER_OBJECT_TYPE(PythonBasedModelNode); + +RandomModel::RandomModel() { + ObjectPtr node = make_object(); + const auto* f = runtime::Registry::Get("auto_scheduler.cost_model.random_fill_float"); + CHECK(f != nullptr); + node->random_number_func = reinterpret_cast*>(f); + data_ = std::move(node); +} + +void RandomModelNode::Update(const Array& inputs, + const Array& results) {} + +void RandomModelNode::Predict(const SearchTask& task, const std::vector& states, + std::vector* scores) { + scores->resize(states.size()); + (*random_number_func)(states.size(), static_cast(scores->data())); +} + +PythonBasedModel::PythonBasedModel(PackedFunc update_func, PackedFunc predict_func, + PackedFunc predict_stage_func) { + auto node = make_object(); + node->update_func = std::move(update_func); + node->predict_func = std::move(predict_func); + node->predict_stage_func = std::move(predict_stage_func); + data_ = std::move(node); +} + +void PythonBasedModelNode::Update(const Array& inputs, + const Array& results) { + update_func(inputs, results); +} + +void PythonBasedModelNode::Predict(const SearchTask& task, const std::vector& states, + std::vector* scores) { + scores->resize(states.size()); + predict_func(task, Array(states.begin(), states.end()), + static_cast(scores->data())); +} + +void PythonBasedModelNode::PredictStages(const SearchTask& task, const std::vector& states, + std::vector* state_scores, + std::vector>* stage_scores) { + size_t n_states = states.size(); + size_t n_stages = task->compute_dag->init_state->stages.size(); + std::vector flatten_scores; + // Allocate sufficient spaces. + flatten_scores.resize(n_states * n_stages * 2); + predict_stage_func(task, Array(states.begin(), states.end()), + static_cast(flatten_scores.data())); + + // Unpack flatten scores. + state_scores->clear(); + stage_scores->clear(); + + // Score of each states. + for (size_t i = 0; i < n_states; ++i) { + state_scores->push_back(flatten_scores[i]); + } + + // Score of each stage in each states. + size_t idx = n_states; + for (size_t i = 0; i < n_states; ++i) { + CHECK_LE(idx, flatten_scores.size()); + + // Number of scored stages of this state. + int s_length = static_cast(flatten_scores[idx++]); + + if (s_length > 0) { + std::vector scores; + int offset = 0; + + if ((*state_scores)[i] > -INFINITY) { + // If the score is valid. Copy scored stages and assign 0 to placeholder + // and inlined stages. If the score is 0, meaning this state failed to + // be lowered. Just bypass to update offset. + for (const Stage& stage : states[i]->stages) { + if (stage->op_type == StageKind::kPlaceholder) { + scores.push_back(0); + continue; + } + if (stage->compute_at == ComputeAtKind::kInlined) { + scores.push_back(0); + continue; + } + scores.push_back(flatten_scores[idx + offset]); + offset++; + } + CHECK_EQ(offset, s_length); + stage_scores->push_back(std::move(scores)); + } + idx += s_length; + } else { + // Cost model does not provide any stage score details. + stage_scores->push_back({}); + } + } +} + +TVM_REGISTER_GLOBAL("auto_scheduler.RandomModel").set_body_typed([]() { return RandomModel(); }); + +TVM_REGISTER_GLOBAL("auto_scheduler.PythonBasedModel") + .set_body_typed([](PackedFunc update_func, PackedFunc predict_func, + PackedFunc predict_stage_func) { + return PythonBasedModel(update_func, predict_func, predict_stage_func); + }); + +TVM_REGISTER_GLOBAL("auto_scheduler.CostModelUpdate") + .set_body_typed([](CostModel model, Array inputs, Array results) { + model->Update(inputs, results); + }); + +TVM_REGISTER_GLOBAL("auto_scheduler.CostModelPredict") + .set_body_typed([](CostModel model, SearchTask task, Array states) { + std::vector scores; + model->Predict(task, std::vector(states.begin(), states.end()), &scores); + Array ret; + for (auto x : scores) { + ret.push_back(FloatImm(DataType::Float(32), x)); + } + return ret; + }); + +} // namespace auto_scheduler +} // namespace tvm diff --git a/tests/python/unittest/test_auto_scheduler_cost_model.py b/tests/python/unittest/test_auto_scheduler_cost_model.py new file mode 100644 index 0000000..3d152a4 --- /dev/null +++ b/tests/python/unittest/test_auto_scheduler_cost_model.py @@ -0,0 +1,42 @@ +# 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. + +"""Test cost models""" + +import tvm +from tvm import auto_scheduler + +from test_auto_scheduler_common import matmul_auto_scheduler_test + + +def test_random_model(): + if not tvm.runtime.enabled("llvm"): + return + N = 128 + workload_key = auto_scheduler.make_workload_key(matmul_auto_scheduler_test, (N, N, N)) + dag = auto_scheduler.ComputeDAG(workload_key) + target = tvm.target.create('llvm') + task = auto_scheduler.SearchTask(dag, workload_key, target) + + model = auto_scheduler.RandomModel() + model.update([], []) + scores = model.predict(task, [dag.init_state, dag.init_state]) + assert len(scores) == 2 + + +if __name__ == "__main__": + test_random_model()