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 | #include "./utils.h" |
20 | |
21 | namespace tvm { |
22 | namespace script { |
23 | namespace printer { |
24 | |
25 | TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) |
26 | .set_dispatch<String>("" , [](String s, ObjectPath p, IRDocsifier d) -> Doc { |
27 | if (HasMultipleLines(s)) { |
28 | return d->AddMetadata(s); |
29 | } |
30 | return LiteralDoc::Str(s, p); |
31 | }); |
32 | |
33 | TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) |
34 | .set_dispatch<Array<ObjectRef>>( // |
35 | "" , [](Array<ObjectRef> array, ObjectPath p, IRDocsifier d) -> Doc { |
36 | int n = array.size(); |
37 | Array<ExprDoc> results; |
38 | results.reserve(n); |
39 | for (int i = 0; i < n; ++i) { |
40 | results.push_back(d->AsDoc<ExprDoc>(array[i], p->ArrayIndex(i))); |
41 | } |
42 | return ListDoc(results); |
43 | }); |
44 | |
45 | TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) |
46 | .set_dispatch<Map<ObjectRef, ObjectRef>>( // |
47 | "" , [](Map<ObjectRef, ObjectRef> dict, ObjectPath p, IRDocsifier d) -> Doc { |
48 | using POO = std::pair<ObjectRef, ObjectRef>; |
49 | std::vector<POO> items{dict.begin(), dict.end()}; |
50 | bool is_str_map = true; |
51 | for (const auto& kv : items) { |
52 | if (!kv.first.as<runtime::StringObj>()) { |
53 | is_str_map = false; |
54 | break; |
55 | } |
56 | } |
57 | if (is_str_map) { |
58 | std::sort(items.begin(), items.end(), [](const POO& lhs, const POO& rhs) { |
59 | return Downcast<String>(lhs.first) < Downcast<String>(rhs.first); |
60 | }); |
61 | } else { |
62 | std::sort(items.begin(), items.end(), [](const POO& lhs, const POO& rhs) { |
63 | return lhs.first.get() < rhs.first.get(); |
64 | }); |
65 | } |
66 | int n = dict.size(); |
67 | Array<ExprDoc> ks; |
68 | Array<ExprDoc> vs; |
69 | ks.reserve(n); |
70 | vs.reserve(n); |
71 | for (int i = 0; i < n; ++i) { |
72 | ks.push_back(d->AsDoc<ExprDoc>(items[i].first, p->MissingMapEntry())); |
73 | vs.push_back(d->AsDoc<ExprDoc>(items[i].second, p->MapValue(items[i].first))); |
74 | } |
75 | return DictDoc(ks, vs); |
76 | }); |
77 | |
78 | } // namespace printer |
79 | } // namespace script |
80 | } // namespace tvm |
81 | |