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#ifndef TVM_META_SCHEDULE_PROFILER_H_
20#define TVM_META_SCHEDULE_PROFILER_H_
21
22#include <tvm/ir/module.h>
23#include <tvm/node/reflection.h>
24#include <tvm/runtime/container/array.h>
25#include <tvm/runtime/container/optional.h>
26#include <tvm/runtime/container/string.h>
27#include <tvm/runtime/object.h>
28#include <tvm/runtime/packed_func.h>
29#include <tvm/target/target.h>
30
31#include <string>
32#include <unordered_map>
33#include <utility>
34#include <vector>
35
36namespace tvm {
37namespace meta_schedule {
38
39class ScopedTimer {
40 public:
41 ~ScopedTimer() {
42 if (deferred_ != nullptr) {
43 deferred_();
44 }
45 }
46
47 private:
48 friend class Profiler;
49
50 explicit ScopedTimer(runtime::TypedPackedFunc<void()> deferred) : deferred_(deferred) {}
51 runtime::TypedPackedFunc<void()> deferred_;
52};
53
54/*! \brief A generic profiler */
55class ProfilerNode : public runtime::Object {
56 public:
57 /*! \brief The segments that are already profiled */
58 std::unordered_map<std::string, double> stats_sec;
59 /*! \brief Counter for the total time used */
60 runtime::PackedFunc total_timer;
61
62 void VisitAttrs(tvm::AttrVisitor* v) {
63 // `stats_sec` is not visited.
64 // `total_timer` is not visited.
65 }
66
67 static constexpr const char* _type_key = "meta_schedule.Profiler";
68 TVM_DECLARE_FINAL_OBJECT_INFO(ProfilerNode, runtime::Object);
69
70 public:
71 /*! \brief Get the internal stats of the running time */
72 Map<String, FloatImm> Get() const;
73 /*! \brief Return a summary of profiling results as table format */
74 String Table() const;
75};
76
77/*!
78 * \brief Managed reference to ProfilerNode
79 * \sa ProfilerNode
80 */
81class Profiler : public runtime::ObjectRef {
82 public:
83 Profiler();
84 TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Profiler, runtime::ObjectRef, ProfilerNode);
85
86 /*! \brief Entering the scope of the context manager */
87 void EnterWithScope();
88 /*! \brief Exiting the scope of the context manager */
89 void ExitWithScope();
90 /*! \brief Returns the current profiler */
91 static Optional<Profiler> Current();
92 /*!
93 * \brief Profile the time usage in the given scope in the given name.
94 * \param name Name for the scope.
95 * \return A scope timer for time profiling.
96 */
97 static ScopedTimer TimedScope(String name);
98};
99
100} // namespace meta_schedule
101} // namespace tvm
102
103#endif // TVM_META_SCHEDULE_PROFILER_H_
104