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 | * Printer implementation for containers |
22 | * \file node/container_printint.cc |
23 | */ |
24 | #include <tvm/node/functor.h> |
25 | #include <tvm/node/repr_printer.h> |
26 | #include <tvm/runtime/registry.h> |
27 | |
28 | namespace tvm { |
29 | |
30 | // Container printer |
31 | TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) |
32 | .set_dispatch<ArrayNode>([](const ObjectRef& node, ReprPrinter* p) { |
33 | auto* op = static_cast<const ArrayNode*>(node.get()); |
34 | p->stream << '['; |
35 | for (size_t i = 0; i < op->size(); ++i) { |
36 | if (i != 0) { |
37 | p->stream << ", " ; |
38 | } |
39 | p->Print(op->at(i)); |
40 | } |
41 | p->stream << ']'; |
42 | }); |
43 | |
44 | TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) |
45 | .set_dispatch<MapNode>([](const ObjectRef& node, ReprPrinter* p) { |
46 | auto* op = static_cast<const MapNode*>(node.get()); |
47 | p->stream << '{'; |
48 | for (auto it = op->begin(); it != op->end(); ++it) { |
49 | if (it != op->begin()) { |
50 | p->stream << ", " ; |
51 | } |
52 | if (it->first->IsInstance<StringObj>()) { |
53 | p->stream << '\"' << Downcast<String>(it->first) << "\": " ; |
54 | } else { |
55 | p->Print(it->first); |
56 | p->stream << ": " ; |
57 | } |
58 | p->Print(it->second); |
59 | } |
60 | p->stream << '}'; |
61 | }); |
62 | |
63 | TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) |
64 | .set_dispatch<ShapeTupleObj>([](const ObjectRef& node, ReprPrinter* p) { |
65 | auto* op = static_cast<const ShapeTupleObj*>(node.get()); |
66 | p->stream << '['; |
67 | for (size_t i = 0; i < op->size; ++i) { |
68 | if (i != 0) { |
69 | p->stream << ", " ; |
70 | } |
71 | p->stream << op->data[i]; |
72 | } |
73 | p->stream << ']'; |
74 | }); |
75 | } // namespace tvm |
76 | |