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 | |
20 | /*! |
21 | * \file install_debug_spans.h |
22 | * \brief Interface of the InstallDebugSpans pass |
23 | */ |
24 | |
25 | #ifndef TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_H_ |
26 | #define TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_H_ |
27 | |
28 | #include <tvm/tir/expr.h> |
29 | #include <tvm/tir/expr_functor.h> |
30 | #include <tvm/tir/stmt.h> |
31 | #include <tvm/tir/stmt_functor.h> |
32 | |
33 | #include <string> |
34 | #include <unordered_map> |
35 | |
36 | #ifndef TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_OPS_H_ |
37 | #define TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_OPS_H_ |
38 | |
39 | #define TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_SUPPORTED_EXPRS \ |
40 | X(Call) \ |
41 | X(Add) \ |
42 | X(Sub) \ |
43 | X(Mul) \ |
44 | X(Div) \ |
45 | X(Mod) \ |
46 | X(FloorDiv) \ |
47 | X(FloorMod) \ |
48 | X(Min) \ |
49 | X(Max) \ |
50 | X(EQ) \ |
51 | X(NE) \ |
52 | X(LT) \ |
53 | X(LE) \ |
54 | X(GT) \ |
55 | X(GE) \ |
56 | X(And) \ |
57 | X(Or) \ |
58 | X(Reduce) \ |
59 | X(Cast) \ |
60 | X(Not) \ |
61 | X(Select) \ |
62 | X(Ramp) \ |
63 | X(Broadcast) \ |
64 | X(Shuffle) \ |
65 | X(IntImm) \ |
66 | X(FloatImm) \ |
67 | X(StringImm) |
68 | |
69 | #define TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_SUPPORTED_STMTS \ |
70 | X(AttrStmt) \ |
71 | X(IfThenElse) \ |
72 | X(LetStmt) \ |
73 | X(For) \ |
74 | X(While) \ |
75 | X(Allocate) \ |
76 | X(AllocateConst) \ |
77 | X(DeclBuffer) \ |
78 | X(Store) \ |
79 | X(BufferStore) \ |
80 | X(BufferRealize) \ |
81 | X(AssertStmt) \ |
82 | X(ProducerStore) \ |
83 | X(ProducerRealize) \ |
84 | X(Prefetch) \ |
85 | X(SeqStmt) \ |
86 | X(Evaluate) \ |
87 | X(BlockRealize) |
88 | |
89 | #endif // TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_OPS_H_ |
90 | |
91 | namespace tvm { |
92 | namespace tir { |
93 | |
94 | /*! |
95 | * \brief This Pass prints out the provided 'stmt' through the TIR debug printer |
96 | while recording the statements and expressions printed on each line. Running |
97 | this pass uses the per-line information to change the Spans attached to each |
98 | statement and expression to the source location in the printed TIR. This pass |
99 | also writes to a file called '<name>.tir' so the line information used is |
100 | saved to disk. |
101 | */ |
102 | class DebugInfoInstaller : public StmtExprMutator { |
103 | public: |
104 | static Stmt InstallInfo(const std::string& name, const Stmt& stmt); |
105 | |
106 | PrimExpr VisitExpr(const PrimExpr& expr) override; |
107 | Stmt VisitStmt(const Stmt& stmt) override; |
108 | |
109 | protected: |
110 | DebugInfoInstaller(const Stmt& stmt, const std::string& filename); |
111 | |
112 | #define X(TypeName) PrimExpr VisitExpr_(const TypeName##Node* op) override; |
113 | TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_SUPPORTED_EXPRS |
114 | #undef X |
115 | |
116 | #define X(TypeName) Stmt VisitStmt_(const TypeName##Node* op) override; |
117 | TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_SUPPORTED_STMTS |
118 | #undef X |
119 | |
120 | private: |
121 | std::unordered_map<const StmtNode*, size_t> stmt_lines_; |
122 | std::unordered_map<const PrimExprNode*, size_t> expr_lines_; |
123 | std::string filename_; |
124 | |
125 | Span MaybeSpan(const StmtNode* op); |
126 | Span MaybeSpan(const PrimExprNode* op); |
127 | }; |
128 | |
129 | } // namespace tir |
130 | } // namespace tvm |
131 | |
132 | #endif // TVM_TIR_TRANSFORMS_INSTALL_DEBUG_SPANS_H_ |
133 | |