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#include "../utils.h"
20
21namespace tvm {
22namespace meta_schedule {
23
24class UpdateCostModelNode : public MeasureCallbackNode {
25 public:
26 void Apply(const TaskScheduler& task_scheduler, int task_id,
27 const Array<MeasureCandidate>& measure_candidates,
28 const Array<BuilderResult>& builder_results,
29 const Array<RunnerResult>& runner_results) final {
30 auto _ = Profiler::TimedScope("MeasureCallback/UpdateCostModel");
31 const TaskRecord& task = task_scheduler->tasks_[task_id];
32 if (!task_scheduler->cost_model_.defined()) {
33 return;
34 }
35 CostModel cost_model = task_scheduler->cost_model_.value();
36 ICHECK(task->measure_candidates.defined()) << "Task's measure candidates must be present!";
37 ICHECK_EQ(measure_candidates.size(), builder_results.size());
38 ICHECK_EQ(runner_results.size(), builder_results.size());
39 int n = builder_results.size();
40 Array<MeasureCandidate> pruned_candidate;
41 Array<RunnerResult> pruned_runner_result;
42 pruned_candidate.reserve(n);
43 pruned_runner_result.reserve(n);
44 for (int i = 0; i < n; i++) {
45 if (!builder_results[i]->error_msg.defined() && //
46 (runner_results[i]->error_msg.defined() || //
47 (runner_results[i]->run_secs.defined() &&
48 Sum(runner_results[i]->run_secs.value()) > 0))) {
49 pruned_candidate.push_back(measure_candidates[i]);
50 pruned_runner_result.push_back(runner_results[i]);
51 }
52 }
53 cost_model->Update(task->ctx, pruned_candidate, pruned_runner_result);
54 }
55
56 static constexpr const char* _type_key = "meta_schedule.UpdateCostModel";
57 TVM_DECLARE_FINAL_OBJECT_INFO(UpdateCostModelNode, MeasureCallbackNode);
58};
59
60MeasureCallback MeasureCallback::UpdateCostModel() {
61 ObjectPtr<UpdateCostModelNode> n = make_object<UpdateCostModelNode>();
62 return MeasureCallback(n);
63}
64
65TVM_REGISTER_NODE_TYPE(UpdateCostModelNode);
66TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackUpdateCostModel")
67 .set_body_typed(MeasureCallback::UpdateCostModel);
68
69} // namespace meta_schedule
70} // namespace tvm
71