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#ifndef TVM_SCRIPT_PRINTER_DOC_H_
20#define TVM_SCRIPT_PRINTER_DOC_H_
21
22#include <tvm/ir/expr.h>
23#include <tvm/node/node.h>
24#include <tvm/runtime/data_type.h>
25
26#include <string>
27
28namespace tvm {
29namespace script {
30namespace printer {
31
32// Forward declaration
33class Doc;
34
35/*!
36 * \brief Convert Doc into Python script.
37 * \param doc Doc to be converted
38 * \param cfg The configuration of the printer
39 */
40String DocToPythonScript(Doc doc, const PrinterConfig& cfg);
41
42/*!
43 * \brief The base class of all Doc.
44 *
45 * Doc is an intermediate representation between IR from TVM
46 * and the TVMScript code.
47 * During printing, IR graph is first translated into Doc tree,
48 * then the Doc tree is translated to the target language in
49 * text format.
50 *
51 * \sa Doc
52 */
53class DocNode : public Object {
54 public:
55 /*!
56 * \brief The list of object paths of the source IR node.
57 *
58 * This is used to trace back to the IR node position where
59 * this Doc is generated, in order to position the diagnostic
60 * message.
61 */
62 mutable Array<ObjectPath> source_paths;
63
64 void VisitAttrs(AttrVisitor* v) { v->Visit("source_paths", &source_paths); }
65
66 static constexpr const char* _type_key = "script.printer.Doc";
67 TVM_DECLARE_BASE_OBJECT_INFO(DocNode, Object);
68
69 public:
70 virtual ~DocNode() = default;
71};
72
73/*!
74 * \brief Reference type of DocNode.
75 *
76 * \sa DocNode
77 */
78class Doc : public ObjectRef {
79 protected:
80 Doc() = default;
81
82 public:
83 virtual ~Doc() = default;
84 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Doc, ObjectRef, DocNode);
85};
86
87class ExprDoc;
88
89/*!
90 * \brief The base class of expression doc.
91 *
92 * \sa ExprDoc
93 */
94class ExprDocNode : public DocNode {
95 public:
96 /*!
97 * \brief Create a doc representing attribute access on the current ExprDoc
98 * \param attr The attribute to access.
99 */
100 ExprDoc Attr(String attr) const;
101
102 /*!
103 * \brief Create a doc representing index access on the current ExprDoc
104 * \param indices The indices to access.
105 */
106 ExprDoc operator[](Array<Doc> indices) const;
107
108 /*!
109 * \brief Create a doc representing calling the current ExprDoc
110 * \param args The positional arguments of the function call.
111 */
112 ExprDoc Call(Array<ExprDoc, void> args) const;
113
114 /*!
115 * \brief Create a doc representing attribute access on the current ExprDoc
116 * \param args The positional arguments of the function call.
117 * \param kwargs_keys Keys of keywords arguments of the function call.
118 * \param kwargs_values Values of keywords arguments of the function call.
119 */
120 ExprDoc Call(Array<ExprDoc, void> args, //
121 Array<String> kwargs_keys, //
122 Array<ExprDoc, void> kwargs_values) const;
123
124 void VisitAttrs(AttrVisitor* v) { DocNode::VisitAttrs(v); }
125
126 static constexpr const char* _type_key = "script.printer.ExprDoc";
127 TVM_DECLARE_BASE_OBJECT_INFO(ExprDocNode, DocNode);
128};
129
130/*!
131 * \brief Reference type of ExprDocNode.
132 *
133 * \sa ExprDocNode
134 */
135class ExprDoc : public Doc {
136 protected:
137 ExprDoc() = default;
138
139 public:
140 /*!
141 * \brief Create a doc representing index access on the current ExprDoc
142 * \param indices The indices to access.
143 */
144 ExprDoc operator[](Array<Doc> indices) const;
145
146 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ExprDoc, Doc, ExprDocNode);
147};
148
149/*!
150 * \brief The base class of statement doc.
151 *
152 * \sa StmtDoc
153 */
154class StmtDocNode : public DocNode {
155 public:
156 /*!
157 * \brief The comment of this doc.
158 *
159 * The actual position of the comment depends on the type of Doc
160 * and also the DocPrinter implementation. It could be on the same
161 * line as the statement, or the line above, or inside the statement
162 * if it spans over multiple lines.
163 * */
164 mutable Optional<String> comment{NullOpt};
165
166 void VisitAttrs(AttrVisitor* v) {
167 DocNode::VisitAttrs(v);
168 v->Visit("comment", &comment);
169 }
170
171 static constexpr const char* _type_key = "script.printer.StmtDoc";
172 TVM_DECLARE_BASE_OBJECT_INFO(StmtDocNode, DocNode);
173};
174
175/*!
176 * \brief Reference type of StmtDocNode.
177 *
178 * \sa StmtDocNode
179 */
180class StmtDoc : public Doc {
181 protected:
182 StmtDoc() = default;
183
184 public:
185 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(StmtDoc, Doc, StmtDocNode);
186};
187
188/*!
189 * \brief The container doc that holds a list of StmtDoc.
190 * \note `StmtBlockDoc` is never used in the IR, but a temporary container that allows holding a
191 * list of StmtDoc.
192 * \sa StmtBlockDoc
193 */
194class StmtBlockDocNode : public DocNode {
195 public:
196 /*! \brief The list of statements. */
197 Array<StmtDoc> stmts;
198
199 void VisitAttrs(AttrVisitor* v) {
200 DocNode::VisitAttrs(v);
201 v->Visit("stmts", &stmts);
202 }
203
204 static constexpr const char* _type_key = "script.printer.StmtBlockDoc";
205 TVM_DECLARE_FINAL_OBJECT_INFO(StmtBlockDocNode, DocNode);
206};
207
208/*!
209 * \brief Reference type of StmtBlockDocNode.
210 * \sa StmtBlockDocNode
211 */
212class StmtBlockDoc : public Doc {
213 public:
214 /*!
215 * \brief Constructor of StmtBlockDoc.
216 * \param stmts The list of statements.
217 */
218 explicit StmtBlockDoc(Array<StmtDoc> stmts);
219 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(StmtBlockDoc, Doc, StmtBlockDocNode);
220};
221
222/*!
223 * \brief Doc that represents literal value.
224 *
225 * \sa LiteralDoc
226 */
227class LiteralDocNode : public ExprDocNode {
228 public:
229 /*!
230 * \brief the internal representation of the literal value.
231 *
232 * Possible actual types:
233 * - IntImm (integer or boolean)
234 * - FloatImm
235 * - String
236 * - null
237 */
238 ObjectRef value;
239
240 void VisitAttrs(AttrVisitor* v) {
241 ExprDocNode::VisitAttrs(v);
242 v->Visit("value", &value);
243 }
244
245 static constexpr const char* _type_key = "script.printer.LiteralDoc";
246 TVM_DECLARE_FINAL_OBJECT_INFO(LiteralDocNode, ExprDocNode);
247};
248
249/*!
250 * \brief Reference type of LiteralDocNode.
251 *
252 * \sa LiteralDocNode
253 */
254class LiteralDoc : public ExprDoc {
255 protected:
256 explicit LiteralDoc(ObjectRef value, const Optional<ObjectPath>& object_path);
257
258 public:
259 /*!
260 * \brief Create a LiteralDoc to represent None/null/empty value.
261 * \param p The object path
262 */
263 static LiteralDoc None(const Optional<ObjectPath>& p) {
264 return LiteralDoc(ObjectRef(nullptr), p);
265 }
266 /*!
267 * \brief Create a LiteralDoc to represent integer.
268 * \param v The integer value.
269 * \param p The object path
270 */
271 static LiteralDoc Int(int64_t v, const Optional<ObjectPath>& p) {
272 return LiteralDoc(IntImm(DataType::Int(64), v), p);
273 }
274 /*!
275 * \brief Create a LiteralDoc to represent boolean.
276 * \param v The boolean value.
277 * \param p The object path
278 */
279 static LiteralDoc Boolean(bool v, const Optional<ObjectPath>& p) {
280 return LiteralDoc(IntImm(DataType::Bool(), v), p);
281 }
282 /*!
283 * \brief Create a LiteralDoc to represent float.
284 * \param v The float value.
285 * \param p The object path
286 */
287 static LiteralDoc Float(double v, const Optional<ObjectPath>& p) {
288 return LiteralDoc(FloatImm(DataType::Float(64), v), p);
289 }
290 /*!
291 * \brief Create a LiteralDoc to represent string.
292 * \param v The string value.
293 * \param p The object path
294 */
295 static LiteralDoc Str(const String& v, const Optional<ObjectPath>& p) { return LiteralDoc(v, p); }
296 /*!
297 * \brief Create a LiteralDoc to represent string.
298 * \param v The string value.
299 * \param p The object path
300 */
301 static LiteralDoc DataType(const runtime::DataType& v, const Optional<ObjectPath>& p) {
302 std::string dtype = v.is_void() ? "void" : runtime::DLDataType2String(v);
303 return LiteralDoc::Str(dtype, p);
304 }
305
306 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LiteralDoc, ExprDoc, LiteralDocNode);
307};
308
309/*!
310 * \brief Doc that represents identifier.
311 *
312 * \sa IdDoc
313 */
314class IdDocNode : public ExprDocNode {
315 public:
316 /*! \brief The name of the identifier */
317 String name;
318
319 void VisitAttrs(AttrVisitor* v) {
320 ExprDocNode::VisitAttrs(v);
321 v->Visit("name", &name);
322 }
323
324 static constexpr const char* _type_key = "script.printer.IdDoc";
325 TVM_DECLARE_FINAL_OBJECT_INFO(IdDocNode, ExprDocNode);
326};
327
328/*!
329 * \brief Reference type of IdDocNode.
330 *
331 * \sa IdDocNode
332 */
333class IdDoc : public ExprDoc {
334 public:
335 /*!
336 * \brief Constructor of IdDoc.
337 * \param name The name of identifier.
338 */
339 explicit IdDoc(String name);
340 explicit IdDoc(std::nullptr_t) : ExprDoc(nullptr) {}
341 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IdDoc, ExprDoc, IdDocNode);
342};
343
344/*!
345 * \brief Doc that represents attribute access on another expression.
346 *
347 * \sa AttrAccessDoc
348 */
349class AttrAccessDocNode : public ExprDocNode {
350 public:
351 /*! \brief The target expression to be accessed */
352 ExprDoc value{nullptr};
353 /*! \brief The attribute to be accessed */
354 String name;
355
356 void VisitAttrs(AttrVisitor* v) {
357 ExprDocNode::VisitAttrs(v);
358 v->Visit("value", &value);
359 v->Visit("name", &name);
360 }
361
362 static constexpr const char* _type_key = "script.printer.AttrAccessDoc";
363 TVM_DECLARE_FINAL_OBJECT_INFO(AttrAccessDocNode, ExprDocNode);
364};
365
366/*!
367 * \brief Reference type of AttrAccessDocNode.
368 *
369 * \sa AttrAccessDocNode
370 */
371class AttrAccessDoc : public ExprDoc {
372 public:
373 /*!
374 * \brief Constructor of AttrAccessDoc
375 * \param value The target expression of attribute access.
376 * \param name The name of attribute to access.
377 */
378 explicit AttrAccessDoc(ExprDoc value, String name);
379 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AttrAccessDoc, ExprDoc, AttrAccessDocNode);
380};
381
382/*!
383 * \brief Doc that represents index access on another expression.
384 *
385 * \sa IndexDoc
386 */
387class IndexDocNode : public ExprDocNode {
388 public:
389 /*! \brief The container value to be accessed */
390 ExprDoc value{nullptr};
391 /*!
392 * \brief The indices to access
393 *
394 * Possible actual types:
395 * - ExprDoc (single point access like a[1, 2])
396 * - SliceDoc (slice access like a[1:5, 2])
397 */
398 Array<Doc> indices; // Each element is union of: Slice / ExprDoc
399
400 void VisitAttrs(AttrVisitor* v) {
401 ExprDocNode::VisitAttrs(v);
402 v->Visit("value", &value);
403 v->Visit("indices", &indices);
404 }
405
406 static constexpr const char* _type_key = "script.printer.IndexDoc";
407 TVM_DECLARE_FINAL_OBJECT_INFO(IndexDocNode, ExprDocNode);
408};
409
410/*!
411 * \brief Reference type of IndexDocNode.
412 *
413 * \sa IndexDocNode
414 */
415class IndexDoc : public ExprDoc {
416 public:
417 /*!
418 * \brief Constructor of IndexDoc
419 * \param value The target expression of index access.
420 * \param indices The indices to access.
421 */
422 explicit IndexDoc(ExprDoc value, Array<Doc> indices);
423 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IndexDoc, ExprDoc, IndexDocNode);
424};
425
426/*!
427 * \brief Doc that represents function call.
428 *
429 * \sa CallDoc
430 */
431class CallDocNode : public ExprDocNode {
432 public:
433 /*! \brief The callee of this function call */
434 ExprDoc callee{nullptr};
435 /*! \brief The positional arguments */
436 Array<ExprDoc> args;
437 /*! \brief The keys of keyword arguments */
438 Array<String> kwargs_keys;
439 /*!
440 * \brief The values of keyword arguments.
441 *
442 * The i-th element is the value of the i-th key in `kwargs_keys`.
443 * It must have the same length as `kwargs_keys`.
444 */
445 Array<ExprDoc> kwargs_values;
446
447 void VisitAttrs(AttrVisitor* v) {
448 ExprDocNode::VisitAttrs(v);
449 v->Visit("callee", &callee);
450 v->Visit("args", &args);
451 v->Visit("kwargs_keys", &kwargs_keys);
452 v->Visit("kwargs_values", &kwargs_values);
453 }
454
455 static constexpr const char* _type_key = "script.printer.CallDoc";
456 TVM_DECLARE_FINAL_OBJECT_INFO(CallDocNode, ExprDocNode);
457};
458
459/*!
460 * \brief Reference type of CallDocNode.
461 *
462 * \sa CallDocNode
463 */
464class CallDoc : public ExprDoc {
465 public:
466 /*!
467 * \brief Constructor of CallDoc
468 * \param callee The callee of this function call.
469 * \param args The positional arguments.
470 * \param kwargs_keys Keys of keyword arguments.
471 * \param kwargs_values Values of keyword arguments, must have the same length as `kwargs_keys.
472 */
473 CallDoc(ExprDoc callee, Array<ExprDoc> args, Array<String> kwargs_keys,
474 Array<ExprDoc> kwargs_values);
475 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(CallDoc, ExprDoc, CallDocNode);
476};
477
478/*!
479 * \brief Doc that represents operation.
480 *
481 * It can be unary, binary and other special operators (for example,
482 * the if-then-else expression).
483 *
484 * \sa OperationDoc
485 */
486class OperationDocNode : public ExprDocNode {
487 public:
488 enum class Kind : int32_t {
489 // Unary operators
490 kUnaryStart = 0,
491 kUSub = 1, // -x
492 kInvert = 2, // ~x
493 kNot = 3, // not x
494 kUnaryEnd = 4,
495
496 // Binary operators
497 kBinaryStart = 5,
498 kAdd = 6, // +
499 kSub = 7, // -
500 kMult = 8, // *
501 kDiv = 9, // /
502 kFloorDiv = 10, // // in Python
503 kMod = 11, // % in Python
504 kPow = 12, // ** in Python
505 kLShift = 13, // <<
506 kRShift = 14, // >>
507 kBitAnd = 15, // &
508 kBitOr = 16, // |
509 kBitXor = 17, // ^
510 kLt = 18, // <
511 kLtE = 19, // <=
512 kEq = 20, // ==
513 kNotEq = 21, // !=
514 kGt = 22, // >
515 kGtE = 23, // >=
516 kAnd = 24, // and
517 kOr = 25, // or
518 kBinaryEnd = 26,
519
520 // Special
521 kSpecialStart = 27,
522 kIfThenElse = 28, // <operands[1]> if <operands[0]> else <operands[2]>
523 kSpecialEnd = 29
524 };
525
526 /*! \brief The kind of operation (operator) */
527 Kind kind;
528 /*! \brief Operands of this expression */
529 Array<ExprDoc> operands;
530
531 void VisitAttrs(AttrVisitor* v) {
532 ExprDocNode::VisitAttrs(v);
533 v->Visit("kind", &kind);
534 v->Visit("operands", &operands);
535 }
536
537 static constexpr const char* _type_key = "script.printer.OperationDoc";
538 TVM_DECLARE_FINAL_OBJECT_INFO(OperationDocNode, ExprDocNode);
539};
540
541/*!
542 * \brief Reference type of OperationDocNode.
543 *
544 * \sa OperationDocNode
545 */
546class OperationDoc : public ExprDoc {
547 public:
548 /*!
549 * \brief Constructor of OperationDoc
550 * \param kind The kind of operation.
551 * \param operands Operands of this expression.
552 */
553 explicit OperationDoc(OperationDocNode::Kind kind, Array<ExprDoc> operands);
554 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(OperationDoc, ExprDoc, OperationDocNode);
555};
556
557/*!
558 * \brief Doc that represents anonymous function.
559 *
560 * LambdaDoc can only have positional arguments without type annotation,
561 * and a single expression as body.
562 *
563 * \sa LambdaDoc
564 */
565class LambdaDocNode : public ExprDocNode {
566 public:
567 /*! \brief The arguments of this anonymous function */
568 Array<IdDoc> args;
569 /*! \brief The body of this anonymous function */
570 ExprDoc body{nullptr};
571
572 void VisitAttrs(AttrVisitor* v) {
573 ExprDocNode::VisitAttrs(v);
574 v->Visit("args", &args);
575 v->Visit("body", &body);
576 }
577
578 static constexpr const char* _type_key = "script.printer.LambdaDoc";
579 TVM_DECLARE_FINAL_OBJECT_INFO(LambdaDocNode, ExprDocNode);
580};
581
582/*!
583 * \brief Reference type of LambdaDocNode.
584 *
585 * \sa LambdaDocNode
586 */
587class LambdaDoc : public ExprDoc {
588 public:
589 /*!
590 * \brief Constructor of LambdaDoc
591 * \param args Arguments of this function.
592 * \param body Body expression of this function.
593 */
594 explicit LambdaDoc(Array<IdDoc> args, ExprDoc body);
595 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LambdaDoc, ExprDoc, LambdaDocNode);
596};
597
598/*!
599 * \brief Doc that represents tuple literal.
600 *
601 * \sa TupleDoc
602 */
603class TupleDocNode : public ExprDocNode {
604 public:
605 /*! \brief Elements of tuple */
606 Array<ExprDoc> elements;
607
608 void VisitAttrs(AttrVisitor* v) {
609 ExprDocNode::VisitAttrs(v);
610 v->Visit("elements", &elements);
611 }
612
613 static constexpr const char* _type_key = "script.printer.TupleDoc";
614 TVM_DECLARE_FINAL_OBJECT_INFO(TupleDocNode, ExprDocNode);
615};
616
617/*!
618 * \brief Reference type of TupleDocNode.
619 *
620 * \sa TupleDocNode
621 */
622class TupleDoc : public ExprDoc {
623 public:
624 /*!
625 * \brief Create an empty TupleDoc
626 */
627 TupleDoc() : TupleDoc(runtime::make_object<TupleDocNode>()) {}
628 /*!
629 * \brief Constructor of TupleDoc
630 * \param elements Elements of tuple.
631 */
632 explicit TupleDoc(Array<ExprDoc> elements);
633 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TupleDoc, ExprDoc, TupleDocNode);
634};
635
636/*!
637 * \brief Doc that represents list literal.
638 *
639 * \sa AttrAccessDoc
640 */
641class ListDocNode : public ExprDocNode {
642 public:
643 /*! \brief Elements of list */
644 Array<ExprDoc> elements;
645
646 void VisitAttrs(AttrVisitor* v) {
647 ExprDocNode::VisitAttrs(v);
648 v->Visit("elements", &elements);
649 }
650
651 static constexpr const char* _type_key = "script.printer.ListDoc";
652 TVM_DECLARE_FINAL_OBJECT_INFO(ListDocNode, ExprDocNode);
653};
654
655/*!
656 * \brief Reference type of ListDocNode.
657 *
658 * \sa ListDocNode
659 */
660class ListDoc : public ExprDoc {
661 public:
662 /*!
663 * \brief Create an empty ListDoc
664 */
665 ListDoc() : ListDoc(runtime::make_object<ListDocNode>()) {}
666 /*!
667 * \brief Constructor of ListDoc
668 * \param elements Elements of list.
669 */
670 explicit ListDoc(Array<ExprDoc> elements);
671 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ListDoc, ExprDoc, ListDocNode);
672};
673
674/*!
675 * \brief Doc that represents dictionary literal.
676 *
677 * \sa AttrAccessDoc
678 */
679class DictDocNode : public ExprDocNode {
680 public:
681 /*! \brief keys of dictionary */
682 Array<ExprDoc> keys;
683 /*!
684 * \brief Values of dictionary
685 *
686 * The i-th element is the value of the i-th element of `keys`.
687 * It must have the same length as `keys`.
688 */
689 Array<ExprDoc> values;
690
691 void VisitAttrs(AttrVisitor* v) {
692 ExprDocNode::VisitAttrs(v);
693 v->Visit("keys", &keys);
694 v->Visit("values", &values);
695 }
696
697 static constexpr const char* _type_key = "script.printer.DictDoc";
698 TVM_DECLARE_FINAL_OBJECT_INFO(DictDocNode, ExprDocNode);
699};
700
701/*!
702 * \brief Reference type of DictDocNode.
703 *
704 * \sa DictDocNode
705 */
706class DictDoc : public ExprDoc {
707 public:
708 /*!
709 * \brief Create an empty dictionary
710 */
711 DictDoc() : DictDoc(runtime::make_object<DictDocNode>()) {}
712 /*!
713 * \brief Constructor of DictDoc
714 * \param keys Keys of dictionary.
715 * \param values Values of dictionary, must have same length as `keys`.
716 */
717 explicit DictDoc(Array<ExprDoc> keys, Array<ExprDoc> values);
718 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(DictDoc, ExprDoc, DictDocNode);
719};
720
721/*!
722 * \brief Doc that represents slice in Index expression.
723 *
724 * This doc can only appear in IndexDoc::indices.
725 *
726 * \sa AttrAccessDoc
727 */
728class SliceDocNode : public DocNode {
729 public:
730 /*! \brief The start of slice */
731 Optional<ExprDoc> start;
732 /*! \brief The exclusive end of slice */
733 Optional<ExprDoc> stop;
734 /*! \brief The step of slice */
735 Optional<ExprDoc> step;
736
737 void VisitAttrs(AttrVisitor* v) {
738 DocNode::VisitAttrs(v);
739 v->Visit("start", &start);
740 v->Visit("stop", &stop);
741 v->Visit("step", &step);
742 }
743
744 static constexpr const char* _type_key = "script.printer.SliceDoc";
745 TVM_DECLARE_FINAL_OBJECT_INFO(SliceDocNode, DocNode);
746};
747
748/*!
749 * \brief Reference type of SliceDocNode.
750 *
751 * \sa SliceDocNode
752 */
753class SliceDoc : public Doc {
754 public:
755 /*!
756 * \brief Constructor of SliceDoc
757 * \param start The start of slice.
758 * \param stop The exclusive end of slice.
759 * \param step The step of slice.
760 */
761 explicit SliceDoc(Optional<ExprDoc> start, Optional<ExprDoc> stop, Optional<ExprDoc> step);
762 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(SliceDoc, Doc, SliceDocNode);
763};
764
765/*!
766 * \brief Doc that represents assign statement.
767 *
768 * \sa AssignDoc
769 */
770class AssignDocNode : public StmtDocNode {
771 public:
772 /*! \brief The left hand side of the assignment */
773 ExprDoc lhs{nullptr};
774 /*!
775 * \brief The right hand side of the assignment.
776 *
777 * If null, this doc represents declaration, e.g. `A: T.Buffer((1,2))`
778 * */
779 Optional<ExprDoc> rhs;
780 /*! \brief The type annotation of this assignment. */
781 Optional<ExprDoc> annotation;
782
783 void VisitAttrs(AttrVisitor* v) {
784 StmtDocNode::VisitAttrs(v);
785 v->Visit("lhs", &lhs);
786 v->Visit("rhs", &rhs);
787 v->Visit("annotation", &annotation);
788 }
789
790 static constexpr const char* _type_key = "script.printer.AssignDoc";
791 TVM_DECLARE_FINAL_OBJECT_INFO(AssignDocNode, StmtDocNode);
792};
793
794/*!
795 * \brief Reference type of AssignDocNode.
796 *
797 * \sa AssignDoc
798 */
799class AssignDoc : public StmtDoc {
800 public:
801 /*!
802 * \brief Constructor of AssignDoc.
803 * \param lhs The left hand side of the assignment.
804 * \param rhs The right hand side of the assignment.
805 * \param annotation The type annotation of this assignment.
806 */
807 explicit AssignDoc(ExprDoc lhs, Optional<ExprDoc> rhs, Optional<ExprDoc> annotation);
808 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AssignDoc, StmtDoc, AssignDocNode);
809};
810
811/*!
812 * \brief Doc that represent if-then-else statement.
813 *
814 * \sa IfDoc
815 */
816class IfDocNode : public StmtDocNode {
817 public:
818 /*! \brief The predicate of the if-then-else statement. */
819 ExprDoc predicate{nullptr};
820 /*! \brief The then branch of the if-then-else statement. */
821 Array<StmtDoc> then_branch;
822 /*! \brief The else branch of the if-then-else statement. */
823 Array<StmtDoc> else_branch;
824
825 void VisitAttrs(AttrVisitor* v) {
826 StmtDocNode::VisitAttrs(v);
827 v->Visit("predicate", &predicate);
828 v->Visit("then_branch", &then_branch);
829 v->Visit("else_branch", &else_branch);
830 }
831
832 static constexpr const char* _type_key = "script.printer.IfDoc";
833 TVM_DECLARE_FINAL_OBJECT_INFO(IfDocNode, StmtDocNode);
834};
835
836/*!
837 * \brief Reference type of IfDocNode.
838 *
839 * \sa IfDocNode
840 */
841class IfDoc : public StmtDoc {
842 public:
843 /*!
844 * \brief Constructor of IfDoc.
845 * \param predicate The predicate of the if-then-else statement.
846 * \param then_branch The then branch of the if-then-else statement.
847 * \param else_branch The else branch of the if-then-else statement.
848 */
849 explicit IfDoc(ExprDoc predicate, Array<StmtDoc> then_branch, Array<StmtDoc> else_branch);
850 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IfDoc, StmtDoc, IfDocNode);
851};
852
853/*!
854 * \brief Doc that represents while statement.
855 *
856 * \sa WhileDoc
857 */
858class WhileDocNode : public StmtDocNode {
859 public:
860 /*! \brief The predicate of the while statement. */
861 ExprDoc predicate{nullptr};
862 /*! \brief The body of the while statement. */
863 Array<StmtDoc> body;
864
865 void VisitAttrs(AttrVisitor* v) {
866 StmtDocNode::VisitAttrs(v);
867 v->Visit("predicate", &predicate);
868 v->Visit("body", &body);
869 }
870
871 static constexpr const char* _type_key = "script.printer.WhileDoc";
872 TVM_DECLARE_FINAL_OBJECT_INFO(WhileDocNode, StmtDocNode);
873};
874
875/*!
876 * \brief Reference type of WhileDocNode.
877 *
878 * \sa WhileDocNode
879 */
880class WhileDoc : public StmtDoc {
881 public:
882 /*!
883 * \brief Constructor of WhileDoc.
884 * \param predicate The predicate of the while statement.
885 * \param body The body of the while statement.
886 */
887 explicit WhileDoc(ExprDoc predicate, Array<StmtDoc> body);
888 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(WhileDoc, StmtDoc, WhileDocNode);
889};
890
891/*!
892 * \brief Doc that represents for statement.
893 *
894 * Example:
895 * for 'lhs' in 'rhs':
896 * 'body...'
897 *
898 * \sa ForDoc
899 */
900class ForDocNode : public StmtDocNode {
901 public:
902 /*! \brief The left hand side of the assignment of iterating variable. */
903 ExprDoc lhs{nullptr};
904 /*! \brief The right hand side of the assignment of iterating variable. */
905 ExprDoc rhs{nullptr};
906 /*! \brief The body of the for statement. */
907 Array<StmtDoc> body;
908
909 void VisitAttrs(AttrVisitor* v) {
910 StmtDocNode::VisitAttrs(v);
911 v->Visit("lhs", &lhs);
912 v->Visit("rhs", &rhs);
913 v->Visit("body", &body);
914 }
915
916 static constexpr const char* _type_key = "script.printer.ForDoc";
917 TVM_DECLARE_FINAL_OBJECT_INFO(ForDocNode, StmtDocNode);
918};
919
920/*!
921 * \brief Reference type of ForDocNode.
922 *
923 * \sa ForDocNode
924 */
925class ForDoc : public StmtDoc {
926 public:
927 /*!
928 * \brief Constructor of ForDoc.
929 * \param lhs The left hand side of the assignment of iterating variable.
930 * \param rhs The right hand side of the assignment of iterating variable.
931 * \param body The body of the for statement.
932 */
933 explicit ForDoc(ExprDoc lhs, ExprDoc rhs, Array<StmtDoc> body);
934 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ForDoc, StmtDoc, ForDocNode);
935};
936
937/*!
938 * \brief Doc that represents special scopes.
939 *
940 * Specifically, this means the with statement in Python:
941 *
942 * with 'rhs' as 'lhs':
943 * 'body...'
944 *
945 * \sa ScopeDoc
946 */
947class ScopeDocNode : public StmtDocNode {
948 public:
949 /*! \brief The name of the scoped variable. */
950 Optional<ExprDoc> lhs{NullOpt};
951 /*! \brief The value of the scoped variable. */
952 ExprDoc rhs{nullptr};
953 /*! \brief The body of the scope doc. */
954 Array<StmtDoc> body;
955
956 void VisitAttrs(AttrVisitor* v) {
957 StmtDocNode::VisitAttrs(v);
958 v->Visit("lhs", &lhs);
959 v->Visit("rhs", &rhs);
960 v->Visit("body", &body);
961 }
962
963 static constexpr const char* _type_key = "script.printer.ScopeDoc";
964 TVM_DECLARE_FINAL_OBJECT_INFO(ScopeDocNode, StmtDocNode);
965};
966
967/*!
968 * \brief Reference type of ScopeDocNode.
969 *
970 * \sa ScopeDocNode
971 */
972class ScopeDoc : public StmtDoc {
973 public:
974 /*!
975 * \brief Constructor of ScopeDoc.
976 * \param lhs The name of the scoped variable.
977 * \param rhs The value of the scoped variable.
978 * \param body The body of the scope doc.
979 */
980 explicit ScopeDoc(Optional<ExprDoc> lhs, ExprDoc rhs, Array<StmtDoc> body);
981
982 /*!
983 * \brief Constructor of ScopeDoc.
984 * \param rhs The value of the scoped variable.
985 * \param body The body of the scope doc.
986 */
987 explicit ScopeDoc(ExprDoc rhs, Array<StmtDoc> body);
988
989 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ScopeDoc, StmtDoc, ScopeDocNode);
990};
991
992/*!
993 * \brief Doc that represents an expression as statement.
994 *
995 * \sa ExprStmtDoc
996 */
997class ExprStmtDocNode : public StmtDocNode {
998 public:
999 /*! \brief The expression represented by this doc. */
1000 ExprDoc expr{nullptr};
1001
1002 void VisitAttrs(AttrVisitor* v) {
1003 StmtDocNode::VisitAttrs(v);
1004 v->Visit("expr", &expr);
1005 }
1006
1007 static constexpr const char* _type_key = "script.printer.ExprStmtDoc";
1008 TVM_DECLARE_FINAL_OBJECT_INFO(ExprStmtDocNode, StmtDocNode);
1009};
1010
1011/*!
1012 * \brief Reference type of ExprStmtDocNode.
1013 *
1014 * \sa ExprStmtDocNode
1015 */
1016class ExprStmtDoc : public StmtDoc {
1017 public:
1018 /*!
1019 * \brief Constructor of ExprStmtDoc.
1020 * \param expr The expression represented by this doc.
1021 */
1022 explicit ExprStmtDoc(ExprDoc expr);
1023 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ExprStmtDoc, StmtDoc, ExprStmtDocNode);
1024};
1025
1026/*!
1027 * \brief Doc that represents assert statement.
1028 *
1029 * \sa AssertDoc
1030 */
1031class AssertDocNode : public StmtDocNode {
1032 public:
1033 /*! \brief The expression to test. */
1034 ExprDoc test{nullptr};
1035 /*! \brief The optional error message when assertion failed. */
1036 Optional<ExprDoc> msg{NullOpt};
1037
1038 void VisitAttrs(AttrVisitor* v) {
1039 StmtDocNode::VisitAttrs(v);
1040 v->Visit("test", &test);
1041 v->Visit("msg", &msg);
1042 }
1043
1044 static constexpr const char* _type_key = "script.printer.AssertDoc";
1045 TVM_DECLARE_FINAL_OBJECT_INFO(AssertDocNode, StmtDocNode);
1046};
1047
1048/*!
1049 * \brief Reference type of AssertDocNode.
1050 *
1051 * \sa AssertDocNode
1052 */
1053class AssertDoc : public StmtDoc {
1054 public:
1055 /*!
1056 * \brief Constructor of AssertDoc.
1057 * \param test The expression to test.
1058 * \param msg The optional error message when assertion failed.
1059 */
1060 explicit AssertDoc(ExprDoc test, Optional<ExprDoc> msg = NullOpt);
1061 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AssertDoc, StmtDoc, AssertDocNode);
1062};
1063
1064/*!
1065 * \brief Doc that represents return statement.
1066 *
1067 * \sa ReturnDoc
1068 */
1069class ReturnDocNode : public StmtDocNode {
1070 public:
1071 /*! \brief The value to return. */
1072 ExprDoc value{nullptr};
1073
1074 void VisitAttrs(AttrVisitor* v) {
1075 StmtDocNode::VisitAttrs(v);
1076 v->Visit("value", &value);
1077 }
1078
1079 static constexpr const char* _type_key = "script.printer.ReturnDoc";
1080 TVM_DECLARE_FINAL_OBJECT_INFO(ReturnDocNode, StmtDocNode);
1081};
1082
1083/*!
1084 * \brief Reference type of ReturnDocNode.
1085 *
1086 * \sa ReturnDocNode
1087 */
1088class ReturnDoc : public StmtDoc {
1089 public:
1090 /*!
1091 * \brief Constructor of ReturnDoc.
1092 * \param value The value to return.
1093 */
1094 explicit ReturnDoc(ExprDoc value);
1095 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ReturnDoc, StmtDoc, ReturnDocNode);
1096};
1097
1098/*!
1099 * \brief Doc that represents function definition.
1100 *
1101 * \sa FunctionDoc
1102 */
1103class FunctionDocNode : public StmtDocNode {
1104 public:
1105 /*! \brief The name of function. */
1106 IdDoc name{nullptr};
1107 /*!
1108 * \brief The arguments of function.
1109 *
1110 * The `lhs` means argument name,
1111 * `annotation` means argument type,
1112 * and `rhs` means default value.
1113 */
1114 Array<AssignDoc> args;
1115 /*! \brief Decorators of function. */
1116 Array<ExprDoc> decorators;
1117 /*! \brief The return type of function. */
1118 Optional<ExprDoc> return_type{NullOpt};
1119 /*! \brief The body of function. */
1120 Array<StmtDoc> body;
1121
1122 void VisitAttrs(AttrVisitor* v) {
1123 StmtDocNode::VisitAttrs(v);
1124 v->Visit("name", &name);
1125 v->Visit("args", &args);
1126 v->Visit("decorators", &decorators);
1127 v->Visit("return_type", &return_type);
1128 v->Visit("body", &body);
1129 }
1130
1131 static constexpr const char* _type_key = "script.printer.FunctionDoc";
1132 TVM_DECLARE_FINAL_OBJECT_INFO(FunctionDocNode, StmtDocNode);
1133};
1134
1135/*!
1136 * \brief Reference type of FunctionDocNode.
1137 *
1138 * \sa FunctionDocNode
1139 */
1140class FunctionDoc : public StmtDoc {
1141 public:
1142 /*!
1143 * \brief Constructor of FunctionDoc.
1144 * \param name The name of function..
1145 * \param args The arguments of function.
1146 * \param decorators The decorator of function.
1147 * \param return_type The return type of function.
1148 * \param body The body of function.
1149 */
1150 explicit FunctionDoc(IdDoc name, Array<AssignDoc> args, Array<ExprDoc> decorators,
1151 Optional<ExprDoc> return_type, Array<StmtDoc> body);
1152 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(FunctionDoc, StmtDoc, FunctionDocNode);
1153};
1154
1155/*!
1156 * \brief Doc that represents class definition.
1157 *
1158 * \sa ClassDoc
1159 */
1160class ClassDocNode : public StmtDocNode {
1161 public:
1162 /*! \brief The name of class. */
1163 IdDoc name{nullptr};
1164 /*! \brief Decorators of class. */
1165 Array<ExprDoc> decorators;
1166 /*! \brief The body of class. */
1167 Array<StmtDoc> body;
1168
1169 void VisitAttrs(AttrVisitor* v) {
1170 StmtDocNode::VisitAttrs(v);
1171 v->Visit("name", &name);
1172 v->Visit("decorators", &decorators);
1173 v->Visit("body", &body);
1174 }
1175
1176 static constexpr const char* _type_key = "script.printer.ClassDoc";
1177 TVM_DECLARE_FINAL_OBJECT_INFO(ClassDocNode, StmtDocNode);
1178};
1179
1180/*!
1181 * \brief Reference type of ClassDocNode.
1182 *
1183 * \sa ClassDocNode
1184 */
1185class ClassDoc : public StmtDoc {
1186 public:
1187 /*!
1188 * \brief Constructor of ClassDoc.
1189 * \param name The name of class.
1190 * \param decorators The decorator of class.
1191 * \param body The body of class.
1192 */
1193 explicit ClassDoc(IdDoc name, Array<ExprDoc> decorators, Array<StmtDoc> body);
1194 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ClassDoc, StmtDoc, ClassDocNode);
1195};
1196
1197/*!
1198 * \brief Doc that represents comment.
1199 *
1200 * \sa CommentDoc
1201 */
1202class CommentDocNode : public StmtDocNode {
1203 public:
1204 static constexpr const char* _type_key = "script.printer.CommentDoc";
1205 TVM_DECLARE_FINAL_OBJECT_INFO(CommentDocNode, StmtDocNode);
1206};
1207
1208/*!
1209 * \brief Reference type of CommentDocNode.
1210 *
1211 * \sa CommentDocNode
1212 */
1213class CommentDoc : public StmtDoc {
1214 public:
1215 explicit CommentDoc(String comment);
1216 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(CommentDoc, StmtDoc, CommentDocNode);
1217};
1218
1219/*!
1220 * \brief Doc that represents docstring.
1221 *
1222 * \sa DocStringDoc
1223 */
1224class DocStringDocNode : public StmtDocNode {
1225 public:
1226 static constexpr const char* _type_key = "script.printer.DocStringDoc";
1227 TVM_DECLARE_FINAL_OBJECT_INFO(DocStringDocNode, StmtDocNode);
1228};
1229
1230/*!
1231 * \brief Reference type of DocStringDocNode.
1232 *
1233 * \sa DocStringDocNode
1234 */
1235class DocStringDoc : public StmtDoc {
1236 public:
1237 explicit DocStringDoc(String docs);
1238 TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(DocStringDoc, StmtDoc, DocStringDocNode);
1239};
1240
1241} // namespace printer
1242} // namespace script
1243} // namespace tvm
1244
1245#endif // TVM_SCRIPT_PRINTER_DOC_H_
1246