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 AddToDatabaseNode : 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 if (!task_scheduler->database_.defined()) {
31 return;
32 }
33 auto _ = Profiler::TimedScope("MeasureCallback/AddToDatabase");
34 TuneContext task = task_scheduler->tasks_[task_id]->ctx;
35 Database database = task_scheduler->database_.value();
36 Workload workload = database->CommitWorkload(task->mod.value());
37 Target target = task->target.value();
38 ICHECK_EQ(runner_results.size(), measure_candidates.size());
39 int n = runner_results.size();
40 for (int i = 0; i < n; ++i) {
41 RunnerResult result = runner_results[i];
42 MeasureCandidate candidate = measure_candidates[i];
43 Array<FloatImm> run_secs{nullptr};
44 if (result->run_secs.defined()) {
45 run_secs = result->run_secs.value();
46 } else {
47 run_secs = Array<FloatImm>{FloatImm(DataType::Float(32), 1e10)};
48 }
49 database->CommitTuningRecord(TuningRecord(
50 /*trace=*/candidate->sch->trace().value(),
51 /*workload=*/workload,
52 /*run_secs=*/run_secs,
53 /*target=*/target,
54 /*args_info=*/candidate->args_info));
55 }
56 }
57
58 static constexpr const char* _type_key = "meta_schedule.AddToDatabase";
59 TVM_DECLARE_FINAL_OBJECT_INFO(AddToDatabaseNode, MeasureCallbackNode);
60};
61
62MeasureCallback MeasureCallback::AddToDatabase() {
63 ObjectPtr<AddToDatabaseNode> n = make_object<AddToDatabaseNode>();
64 return MeasureCallback(n);
65}
66
67TVM_REGISTER_NODE_TYPE(AddToDatabaseNode);
68TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackAddToDatabase")
69 .set_body_typed(MeasureCallback::AddToDatabase);
70
71} // namespace meta_schedule
72} // namespace tvm
73