| /* |
| * 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 <tvm/auto_scheduler/cost_model.h> |
| |
| namespace tvm { |
| namespace auto_scheduler { |
| |
| TVM_REGISTER_OBJECT_TYPE(CostModelNode); |
| TVM_REGISTER_OBJECT_TYPE(RandomModelNode); |
| TVM_REGISTER_OBJECT_TYPE(PythonBasedModelNode); |
| |
| RandomModel::RandomModel() { |
| ObjectPtr<RandomModelNode> node = make_object<RandomModelNode>(); |
| const auto* f = runtime::Registry::Get("auto_scheduler.cost_model.random_fill_float"); |
| CHECK(f != nullptr); |
| node->random_number_func = reinterpret_cast<const TypedPackedFunc<void(size_t, void*)>*>(f); |
| data_ = std::move(node); |
| } |
| |
| void RandomModelNode::Update(const Array<MeasureInput>& inputs, |
| const Array<MeasureResult>& results) {} |
| |
| void RandomModelNode::Predict(const SearchTask& task, const Array<State>& states, |
| std::vector<float>* scores) { |
| scores->resize(states.size()); |
| (*random_number_func)(states.size(), static_cast<void*>(scores->data())); |
| } |
| |
| PythonBasedModel::PythonBasedModel(PackedFunc update_func, PackedFunc predict_func, |
| PackedFunc predict_stage_func) { |
| auto node = make_object<PythonBasedModelNode>(); |
| 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<MeasureInput>& inputs, |
| const Array<MeasureResult>& results) { |
| update_func(inputs, results); |
| } |
| |
| void PythonBasedModelNode::Predict(const SearchTask& task, const Array<State>& states, |
| std::vector<float>* scores) { |
| scores->resize(states.size()); |
| predict_func(task, states, static_cast<void*>(scores->data())); |
| } |
| |
| void PythonBasedModelNode::PredictStages(const SearchTask& task, const Array<State>& states, |
| std::vector<float>* state_scores, |
| std::vector<std::vector<float>>* stage_scores) { |
| size_t n_states = states.size(); |
| size_t n_stages = task->compute_dag->init_state->stages.size(); |
| std::vector<float> flatten_scores; |
| // Allocate sufficient spaces. |
| flatten_scores.resize(n_states * n_stages * 2); |
| predict_stage_func(task, states, static_cast<void*>(flatten_scores.data())); |
| |
| /* For faster data copy between c++ and python, the python part returns scores in a |
| * single flatten array using a packed format. The c++ part then unpacks the flatten array. |
| * |
| * The packed format is: |
| * { |
| * float scores[N]; // scores[i] is the score for states[i]. |
| * int n_stage_0; // the number of stages in states[0] |
| * float stage_scores_0[[n_stage_0] // the scores for all stages in states[0] |
| * int n_stage_1; // the number of stages in states[1] |
| * float stage_scores_1[n_stage_1]; // the scores for all stages in states[1] |
| * ... |
| * int n_stage_i; // the number of stages in states[i] |
| * float stage_scores_1[n_stage_i]; // the scores for all stages in states[i] |
| * ... // until i == N - 1 |
| * } |
| * To implement this format, we also store int as float, so we can store all numbers |
| * into a single float array. |
| */ |
| |
| // 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<int>(flatten_scores[idx++]); |
| |
| if (s_length > 0) { |
| std::vector<float> 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<MeasureInput> inputs, Array<MeasureResult> results) { |
| model->Update(inputs, results); |
| }); |
| |
| TVM_REGISTER_GLOBAL("auto_scheduler.CostModelPredict") |
| .set_body_typed([](CostModel model, SearchTask task, Array<State> states) { |
| std::vector<float> scores; |
| model->Predict(task, states, &scores); |
| Array<FloatImm> ret; |
| for (auto x : scores) { |
| ret.push_back(FloatImm(DataType::Float(32), x)); |
| } |
| return ret; |
| }); |
| |
| } // namespace auto_scheduler |
| } // namespace tvm |