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 tir_text_printer.cc
22 * \brief Printer to print out the IR text format
23 * that can be parsed by a parser.
24 */
25
26#include "tir_text_printer_debug.h"
27
28#include <optional>
29#include <string>
30
31namespace tvm {
32namespace relay {
33
34std::optional<std::string> span_text(const Span& span) {
35 if (!span.defined()) {
36 return std::nullopt;
37 }
38
39 std::string source("main.tir");
40 if (span->source_name.defined() && span->source_name->name.get()) {
41 source = span->source_name->name;
42 }
43 return source + ":" + std::to_string(span->line) + ":" + std::to_string(span->column);
44}
45
46template <typename ObjectPtr>
47void add_all_relevant_lines(const std::vector<std::tuple<const ObjectPtr*, size_t>>& data,
48 size_t current_line, Doc* output) {
49 ICHECK(output) << "output must be a valid Doc";
50 for (const auto& item : data) {
51 if (std::get<1>(item) != current_line - 1) {
52 // Item is not relevant for this line, skip it
53 continue;
54 }
55
56 // Print out the item's span info if present
57 auto text = span_text(std::get<0>(item)->span);
58 if (text.has_value()) {
59 *output << *text;
60 } else {
61 *output << "missing";
62 }
63 *output << ", ";
64 }
65}
66
67Doc TIRTextPrinterDebug::NewLine() {
68 current_line_ += 1;
69
70 if (!show_spans_) {
71 return TIRTextPrinter::NewLine();
72 }
73
74 Doc output;
75
76 output << " [";
77
78 add_all_relevant_lines(exprs_by_line_, current_line_, &output);
79 add_all_relevant_lines(stmts_by_line_, current_line_, &output);
80
81 output << "]" << TIRTextPrinter::NewLine();
82
83 return output;
84}
85
86Doc TIRTextPrinterDebug::VisitStmt(const tvm::tir::Stmt& n) {
87 stmts_by_line_.push_back(std::make_tuple(n.get(), current_line_));
88 return TIRTextPrinter::VisitStmt(n);
89}
90
91Doc TIRTextPrinterDebug::VisitExpr(const PrimExpr& e) {
92 exprs_by_line_.push_back(std::make_tuple(e.get(), current_line_));
93 return TIRTextPrinter::VisitExpr(e);
94}
95
96} // namespace relay
97} // namespace tvm
98