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
24void PyMeasureCallbackNode::Apply(const TaskScheduler& task_scheduler, //
25 int task_id, //
26 const Array<MeasureCandidate>& measure_candidates, //
27 const Array<BuilderResult>& builds, //
28 const Array<RunnerResult>& results) {
29 ICHECK(f_apply != nullptr) << "PyMeasureCallback's Apply method not implemented!";
30 auto _ = Profiler::TimedScope("MeasureCallback/" + this->f_as_string());
31 return f_apply(task_scheduler, task_id, measure_candidates, builds, results);
32}
33
34MeasureCallback MeasureCallback::PyMeasureCallback(PyMeasureCallbackNode::FApply f_apply, //
35 PyMeasureCallbackNode::FAsString f_as_string) {
36 ObjectPtr<PyMeasureCallbackNode> n = make_object<PyMeasureCallbackNode>();
37 n->f_apply = std::move(f_apply);
38 n->f_as_string = std::move(f_as_string);
39 return MeasureCallback(n);
40}
41
42Array<MeasureCallback, void> MeasureCallback::Default() {
43 return {
44 MeasureCallback::AddToDatabase(),
45 MeasureCallback::RemoveBuildArtifact(),
46 MeasureCallback::UpdateCostModel(),
47 };
48}
49
50TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
51 .set_dispatch<PyMeasureCallbackNode>([](const ObjectRef& n, ReprPrinter* p) {
52 const auto* self = n.as<PyMeasureCallbackNode>();
53 ICHECK(self);
54 PyMeasureCallbackNode::FAsString f_as_string = (*self).f_as_string;
55 ICHECK(f_as_string != nullptr) << "PyMeasureCallback's AsString method not implemented!";
56 p->stream << f_as_string();
57 });
58
59TVM_REGISTER_OBJECT_TYPE(MeasureCallbackNode);
60TVM_REGISTER_NODE_TYPE(PyMeasureCallbackNode);
61
62TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackApply")
63 .set_body_method<MeasureCallback>(&MeasureCallbackNode::Apply);
64TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackPyMeasureCallback")
65 .set_body_typed(MeasureCallback::PyMeasureCallback);
66TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackDefault")
67 .set_body_typed(MeasureCallback::Default);
68
69} // namespace meta_schedule
70} // namespace tvm
71