tvm
trace.h
Go to the documentation of this file.
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_TIR_SCHEDULE_TRACE_H_
20 #define TVM_TIR_SCHEDULE_TRACE_H_
21 
23 
24 namespace tvm {
25 namespace tir {
26 
27 // Forward declaration
28 class Trace;
29 
40  ffi::TypedFunction<Any(const Instruction& inst, const Array<Any>& inputs,
41  const Array<Any>& attrs, const Any& decision)>;
42 
58 class TraceNode : public runtime::Object {
59  public:
61  Array<Instruction> insts;
63  Map<Instruction, Any> decisions;
64 
65  static void RegisterReflection() {
66  namespace refl = tvm::ffi::reflection;
67  refl::ObjectDef<TraceNode>()
68  .def_ro("insts", &TraceNode::insts)
69  .def_ro("decisions", &TraceNode::decisions);
70  }
71 
72  static constexpr const char* _type_key = "tir.Trace";
74 
75  public:
81  Any GetDecision(const Instruction& inst) const;
86  void Append(Instruction inst);
94  void Append(Instruction inst, Any decision);
99  Optional<Instruction> Pop();
108  void ApplyToSchedule(Schedule sch, bool remove_postproc,
109  FTraceDecisionProvider decision_provider = nullptr) const;
115  ObjectRef AsJSON(bool remove_postproc) const;
121  Array<String> AsPython(bool remove_postproc) const;
130  Trace WithDecision(Instruction inst, Any decision, bool remove_postproc) const;
136  Trace Simplified(bool remove_postproc) const;
137 };
138 
143 class Trace : public runtime::ObjectRef {
144  public:
146  Trace();
152  explicit Trace(Array<Instruction> insts, Map<Instruction, Any> decisions);
158  static void ApplyJSONToSchedule(ObjectRef json, Schedule sch);
159 
161 };
162 
163 } // namespace tir
164 } // namespace tvm
165 
166 #endif // TVM_TIR_SCHEDULE_TRACE_H_
Managed reference to InstructionNode.
Definition: instruction.h:191
Managed reference to ScheduleNode.
Definition: schedule.h:880
An execution trace of a scheduling program.
Definition: trace.h:58
static void RegisterReflection()
Definition: trace.h:65
Map< Instruction, Any > decisions
The random decisions made upon those instructions.
Definition: trace.h:63
Trace WithDecision(Instruction inst, Any decision, bool remove_postproc) const
Create a new trace with an instruction whose decision is changed, assuming this instruction exists in...
ObjectRef AsJSON(bool remove_postproc) const
Serialize the trace as a JSON-style object.
void ApplyToSchedule(Schedule sch, bool remove_postproc, FTraceDecisionProvider decision_provider=nullptr) const
Apply the trace to a TensorIR schedule.
void Append(Instruction inst, Any decision)
Append a new instruction with a random decision to the trace.
TVM_DECLARE_FINAL_OBJECT_INFO(TraceNode, runtime::Object)
Optional< Instruction > Pop()
Remove the last instruction, along with the decision made on that instruction, if any.
Array< String > AsPython(bool remove_postproc) const
Serialize the trace as a sequence of python statements.
Trace Simplified(bool remove_postproc) const
Simplify the trace with dead-code elimination.
static constexpr const char * _type_key
Definition: trace.h:72
Array< Instruction > insts
The instructions invoked so far in the program execution.
Definition: trace.h:61
Any GetDecision(const Instruction &inst) const
Retrieve the decision made on a specific instruction.
void Append(Instruction inst)
Append a new instruction to the trace.
Managed reference to TraceNode.
Definition: trace.h:143
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Trace, runtime::ObjectRef, TraceNode)
Trace()
Default constructor. Creating an empty trace.
Trace(Array< Instruction > insts, Map< Instruction, Any > decisions)
Constructor. Creating a trace from existing instructions and their decisions.
static void ApplyJSONToSchedule(ObjectRef json, Schedule sch)
Apply a JSON-serialized trace to a TensorIR schedule.
Definition: repr_printer.h:91
ffi::TypedFunction< Any(const Instruction &inst, const Array< Any > &inputs, const Array< Any > &attrs, const Any &decision)> FTraceDecisionProvider
A callback that allows users to mutate decisions on the fly when applying instructions....
Definition: trace.h:41
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37