Skip to content

Commit 9f4bf38

Browse files
yeliteyelite
authored
[TVMScript] Doc Base Class & DocPrinter Scaffolding (#11971)
This PR addes: - Doc base class - DocPrinter base class - PythonDocPrinter - LiteralDoc and its support in DocPrinter Tracking issue: #11912
1 parent 1392e64 commit 9f4bf38

File tree

14 files changed

+725
-0
lines changed

14 files changed

+725
-0
lines changed

‎CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -281,6 +281,7 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
281281
src/parser/*.cc
282282
src/printer/*.cc
283283
src/support/*.cc
284+
src/script/*.cc
284285
)
285286

286287
tvm_file_glob(GLOB CODEGEN_SRCS

‎include/tvm/script/printer/doc.h‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
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+
namespace tvm {
27+
namespace script {
28+
namespace printer {
29+
30+
/*!
31+
* \brief The base class of all Doc.
32+
*
33+
* Doc is an intermediate representation between IR from TVM
34+
* and the TVMScript code.
35+
* During printing, IR graph is first translated into Doc tree,
36+
* then the Doc tree is translated to the target language in
37+
* text format.
38+
*
39+
* \sa Doc
40+
*/
41+
class DocNode : public Object {
42+
public:
43+
void VisitAttrs(AttrVisitor* v) {}
44+
45+
static constexpr const char* _type_key = "script.printer.Doc";
46+
TVM_DECLARE_BASE_OBJECT_INFO(DocNode, Object);
47+
48+
public:
49+
virtual ~DocNode() = default;
50+
};
51+
52+
/*!
53+
* \brief Reference type of DocNode.
54+
*
55+
* \sa DocNode
56+
*/
57+
class Doc : public ObjectRef {
58+
protected:
59+
Doc() = default;
60+
61+
public:
62+
virtual ~Doc() = default;
63+
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Doc, ObjectRef, DocNode);
64+
};
65+
66+
/*!
67+
* \brief The base class of expression doc.
68+
*
69+
* \sa ExprDoc
70+
*/
71+
class ExprDocNode : public DocNode {
72+
public:
73+
void VisitAttrs(AttrVisitor* v) { DocNode::VisitAttrs(v); }
74+
75+
static constexpr const char* _type_key = "script.printer.ExprDoc";
76+
TVM_DECLARE_BASE_OBJECT_INFO(ExprDocNode, DocNode);
77+
};
78+
79+
/*!
80+
* \brief Reference type of ExprDocNode.
81+
*
82+
* \sa ExprDocNode
83+
*/
84+
class ExprDoc : public Doc {
85+
protected:
86+
ExprDoc() = default;
87+
88+
public:
89+
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ExprDoc, Doc, ExprDocNode);
90+
};
91+
92+
/*!
93+
* \brief Doc that represents literal value.
94+
*
95+
* \sa LiteralDoc
96+
*/
97+
class LiteralDocNode : public ExprDocNode {
98+
public:
99+
/*!
100+
* \brief the internal representation of the literal value.
101+
*
102+
* Possible actual types:
103+
* - IntImm (integer or boolean)
104+
* - FloatImm
105+
* - String
106+
* - null
107+
*/
108+
ObjectRef value;
109+
110+
void VisitAttrs(AttrVisitor* v) {
111+
ExprDocNode::VisitAttrs(v);
112+
v->Visit("value", &value);
113+
}
114+
115+
static constexpr const char* _type_key = "script.printer.LiteralDoc";
116+
TVM_DECLARE_FINAL_OBJECT_INFO(LiteralDocNode, ExprDocNode);
117+
};
118+
119+
/*!
120+
* \brief Reference type of LiteralDocNode.
121+
*
122+
* \sa LiteralDocNode
123+
*/
124+
class LiteralDoc : public ExprDoc {
125+
protected:
126+
explicit LiteralDoc(ObjectRef value);
127+
128+
public:
129+
/*!
130+
* \brief Create a LiteralDoc to represent None/null/empty value.
131+
*/
132+
static LiteralDoc None() { return LiteralDoc(ObjectRef(nullptr)); }
133+
134+
/*!
135+
* \brief Create a LiteralDoc to represent integer.
136+
* \param v The integer value.
137+
*/
138+
static LiteralDoc Int(int v) { return LiteralDoc(IntImm(DataType::Int(64), v)); }
139+
140+
/*!
141+
* \brief Create a LiteralDoc to represent boolean.
142+
* \param v The boolean value.
143+
*/
144+
static LiteralDoc Boolean(bool v) { return LiteralDoc(IntImm(DataType::Bool(), v)); }
145+
146+
/*!
147+
* \brief Create a LiteralDoc to represent float.
148+
* \param v The float value.
149+
*/
150+
static LiteralDoc Float(double v) { return LiteralDoc(FloatImm(DataType::Float(64), v)); }
151+
152+
/*!
153+
* \brief Create a LiteralDoc to represent string.
154+
* \param v The string value.
155+
*/
156+
static LiteralDoc Str(const String& v) { return LiteralDoc(v); }
157+
158+
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LiteralDoc, ExprDoc, LiteralDocNode);
159+
};
160+
161+
} // namespace printer
162+
} // namespace script
163+
} // namespace tvm
164+
165+
#endif // TVM_SCRIPT_PRINTER_DOC_H_

‎include/tvm/script/printer/doc_printer.h‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
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_PRINTER_H_
20+
#define TVM_SCRIPT_PRINTER_DOC_PRINTER_H_
21+
22+
#include <tvm/script/printer/doc.h>
23+
24+
namespace tvm {
25+
namespace script {
26+
namespace printer {
27+
28+
/*!
29+
* \brief Convert Doc into Python script.
30+
*
31+
* This function unpacks the DocPrinterOptions into function arguments
32+
* to be FFI friendly.
33+
*
34+
* \param doc the doc to be converted
35+
* \param indent_spaces the number of spaces used for indention
36+
*/
37+
String DocToPythonScript(Doc doc, int indent_spaces = 4);
38+
39+
} // namespace printer
40+
} // namespace script
41+
} // namespace tvm
42+
43+
#endif // TVM_SCRIPT_PRINTER_DOC_PRINTER_H_

‎python/tvm/script/printer/__init__.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""
18+
TVMScript Unified Printer
19+
20+
This package provides a set of APIs to print supported TVM IR into TVMScript
21+
in a roundtrippable way.
22+
23+
https://github.com/apache/tvm-rfcs/blob/main/rfcs/0074-tvmscript-unified-printer.md
24+
"""
25+
26+
from . import _ffi_api

‎python/tvm/script/printer/_ffi_api.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""FFI APIs for tvm.script.printer"""
18+
import tvm._ffi
19+
20+
tvm._ffi._init_api("script.printer", __name__)

‎python/tvm/script/printer/doc.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""Doc types for TVMScript Unified Printer"""
18+
19+
import tvm._ffi
20+
from tvm.runtime import Object
21+
22+
from . import _ffi_api
23+
24+
25+
class Doc(Object):
26+
"""Base class of all Docs"""
27+
28+
29+
class ExprDoc(Object):
30+
"""Base class of all expression Docs"""
31+
32+
33+
@tvm._ffi.register_object("script.printer.LiteralDoc")
34+
class LiteralDoc(ExprDoc):
35+
"""Doc that represents literal value"""
36+
37+
def __init__(self, value):
38+
if value is None:
39+
self.__init_handle_by_constructor__(_ffi_api.LiteralDocNone) # type: ignore
40+
elif isinstance(value, str):
41+
self.__init_handle_by_constructor__(_ffi_api.LiteralDocStr, value) # type: ignore
42+
elif isinstance(value, float):
43+
self.__init_handle_by_constructor__(_ffi_api.LiteralDocFloat, value) # type: ignore
44+
elif isinstance(value, bool):
45+
self.__init_handle_by_constructor__(_ffi_api.LiteralDocBoolean, value) # type: ignore
46+
elif isinstance(value, int):
47+
self.__init_handle_by_constructor__(_ffi_api.LiteralDocInt, value) # type: ignore
48+
else:
49+
raise TypeError(f"Unsupported type {type(value)} for LiteralDoc")

‎python/tvm/script/printer/doc_printer.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""Functions to print doc into text format"""
18+
19+
from . import _ffi_api
20+
from .doc import Doc
21+
22+
23+
def to_python_script(doc: Doc, indent_spaces: int = 4) -> str:
24+
"""
25+
Convert Doc into Python script.
26+
27+
Parameters
28+
----------
29+
doc : Doc
30+
The doc to convert into Python script
31+
indent_spaces : int
32+
The number of indent spaces to use in the output
33+
34+
Returns
35+
-------
36+
script : str
37+
The text representation of Doc in Python syntax
38+
"""
39+
return _ffi_api.DocToPythonScript(doc, indent_spaces) # type: ignore

0 commit comments

Comments
 (0)