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
91namespace tvm {
92namespace 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 */
102class 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