Skip to content

Commit a20cb0d

Browse files
tom-armtom-arm
authored
Arm backend: Add documentation generation (pytorch#16699)
* Add script that auto-generates documentation from docstrings in source Signed-off-by: Tom Allsop <tom.allsop@arm.com> Co-authored-by: Agrima Khare <agrima.khare@arm.com> Co-authored-by: Adrian Lundell <adrian.lundell@arm.com> Co-authored-by: Erik Lundell <erik.lundell@arm.com> cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai
1 parent 466e5e3 commit a20cb0d

9 files changed

Lines changed: 848 additions & 0 deletions

‎backends/arm/scripts/docgen/docgen.py‎

Lines changed: 269 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,269 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
import inspect
7+
import json
8+
9+
from dataclasses import dataclass
10+
from pathlib import Path
11+
12+
from executorch.backends.arm.ethosu import EthosUCompileSpec, EthosUPartitioner
13+
from executorch.backends.arm.quantizer import EthosUQuantizer, VgfQuantizer
14+
from executorch.backends.arm.vgf.partitioner import VgfCompileSpec, VgfPartitioner
15+
16+
17+
@dataclass
18+
class DocumentationJob:
19+
template_path: Path
20+
placeholder: str
21+
replacement_text: str
22+
output_path: Path
23+
24+
25+
def get_docstring(obj) -> str:
26+
"""
27+
Returns the docstring of an object, formatted in markdown to resemble pytorch-style
28+
docs, e.g. an argument list like:
29+
30+
Args:
31+
arg1: description
32+
33+
is converted to
34+
35+
Args:
36+
- **arg1**: description
37+
- **arg2**: description
38+
"""
39+
40+
docstring = inspect.getdoc(obj)
41+
if docstring is None:
42+
print(
43+
f"WARNING: Docstring not found for object '{obj.__name__}'. Please document all user facing components of the Arm Backend properly."
44+
)
45+
docstring = ""
46+
47+
lines = docstring.split("\n")
48+
for line in lines:
49+
if ":" in line and line.startswith(" "):
50+
new_line = line.strip()
51+
pos = new_line.index(":")
52+
new_line = f"- **{new_line[:pos]}**" + new_line[pos:]
53+
docstring = docstring.replace(line, new_line)
54+
55+
return docstring
56+
57+
58+
def get_function_docstring(cls, func) -> str:
59+
"""
60+
Returns a function's signature and docstring formatted in markdown.
61+
"""
62+
return f"```python\ndef {cls.__name__}.{func.__name__}{inspect.signature(func)}:\n```\n{get_docstring(func)}\n\n"
63+
64+
65+
def get_class_docstring(cls, filter_funcs=()) -> str:
66+
"""
67+
Returns a class signature and docstring, as well as documentation for all its public
68+
methods with names not matching strings listed in filter_funcs.
69+
"""
70+
header = f"```python\nclass {cls.__name__}{inspect.signature(cls)}\n```\n{get_docstring(cls)}\n\n"
71+
72+
class_functions = [
73+
getattr(cls, name)
74+
for name in dir(cls)
75+
if callable(getattr(cls, name))
76+
and not name.startswith("_")
77+
and not any(f in name for f in filter_funcs)
78+
]
79+
80+
function_docstrings = [
81+
get_function_docstring(cls, func) for func in class_functions
82+
]
83+
84+
return header + "".join(function_docstrings)
85+
86+
87+
def get_jupyter_code(path, get_bash, which_cells: list[int] | None = None) -> str:
88+
"""
89+
Returns all code cells from the jupyter notebook at 'path'. If get_bash is True,
90+
only bash cells are returned, otherwise only python cells are returned.
91+
which_cells lets you supply a list of cell indicies to return.
92+
"""
93+
output = f"```{'bash' if get_bash else 'python'}\n"
94+
with open(path, "r") as f:
95+
j = json.load(f)
96+
i = -1
97+
for cell in j["cells"]:
98+
is_code = cell["cell_type"] == "code"
99+
if len(cell["source"]) == 0:
100+
continue
101+
is_bash = "bash" in cell["source"][0]
102+
is_copyright = "Copyright" in cell["source"][0]
103+
if is_code and is_bash == get_bash and not is_copyright:
104+
i += 1
105+
if which_cells is not None:
106+
if i not in which_cells:
107+
continue
108+
for line in cell["source"]:
109+
is_print = "print_readable" in line
110+
is_bash_line = "bash" in line
111+
is_setup_path = "setup_path.sh" in line
112+
if not (is_print or is_bash_line or is_setup_path):
113+
output += line
114+
output += "\n"
115+
output += "```\n"
116+
return output
117+
118+
119+
def generate_document(job: DocumentationJob):
120+
"""Generates a markdown document based on a DocumentationJob."""
121+
with open(job.template_path, "r") as f:
122+
content = f.read()
123+
124+
content = content.replace(job.placeholder, job.replacement_text)
125+
126+
# Remove multiple new lines at end of document if it exists
127+
if content.endswith("\n\n"):
128+
content = content.removesuffix("\n")
129+
130+
with open(job.output_path, "w") as f:
131+
f.write(content)
132+
133+
134+
def generate_ethos_u_docs():
135+
"""Generates documentation for the Ethos-U components in the backend."""
136+
compilespec_string = get_class_docstring(
137+
EthosUCompileSpec,
138+
("DebugMode", "to_list", "from_list", "from_list_hook", "validate"),
139+
)
140+
partitioner_string = get_class_docstring(EthosUPartitioner)
141+
quantizer_string = get_class_docstring(
142+
EthosUQuantizer, ("prepare_obs_or_fq_callback", "annotate", "validate")
143+
)
144+
example_string = get_jupyter_code(
145+
"./examples/arm/ethos_u_minimal_example.ipynb", get_bash=False
146+
)
147+
148+
documentation_jobs = [
149+
DocumentationJob(
150+
Path(
151+
"backends/arm/scripts/docgen/ethos-u/backends-arm-ethos-u-overview.md.in"
152+
),
153+
"$COMPILE_SPEC",
154+
compilespec_string,
155+
Path("docs/source/backends/arm-ethos-u/arm-ethos-u-overview.md"),
156+
),
157+
DocumentationJob(
158+
Path(
159+
"backends/arm/scripts/docgen/ethos-u/backends-arm-ethos-u-partitioner.md.in"
160+
),
161+
"$PARTITIONER",
162+
partitioner_string,
163+
Path("docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md"),
164+
),
165+
DocumentationJob(
166+
Path(
167+
"backends/arm/scripts/docgen/ethos-u/backends-arm-ethos-u-quantization.md.in"
168+
),
169+
"$QUANTIZER",
170+
quantizer_string,
171+
Path("docs/source/backends/arm-ethos-u/arm-ethos-u-quantization.md"),
172+
),
173+
DocumentationJob(
174+
Path(
175+
"backends/arm/scripts/docgen/ethos-u/ethos-u-getting-started-tutorial.md.in"
176+
),
177+
"$MINIMAL_EXAMPLE",
178+
example_string,
179+
Path(
180+
"docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md"
181+
),
182+
),
183+
]
184+
185+
for job in documentation_jobs:
186+
generate_document(job)
187+
188+
189+
def generate_vgf_docs():
190+
"""Generates documentation for the VGF components in the backend."""
191+
compilespec_string = get_class_docstring(
192+
VgfCompileSpec,
193+
("DebugMode", "to_list", "from_list", "from_list_hook", "validate"),
194+
)
195+
partitioner_string = get_class_docstring(VgfPartitioner)
196+
quantizer_string = get_class_docstring(
197+
VgfQuantizer, ("prepare_obs_or_fq_callback", "annotate", "validate")
198+
)
199+
example_string = get_jupyter_code(
200+
"./examples/arm/vgf_minimal_example.ipynb",
201+
get_bash=False,
202+
which_cells=[0, 2, 3],
203+
)
204+
205+
documentation_jobs = [
206+
DocumentationJob(
207+
Path("backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in"),
208+
"$COMPILE_SPEC",
209+
compilespec_string,
210+
Path("docs/source/backends/arm-vgf/arm-vgf-overview.md"),
211+
),
212+
DocumentationJob(
213+
Path("backends/arm/scripts/docgen/vgf/backends-arm-vgf-partitioner.md.in"),
214+
"$PARTITIONER",
215+
partitioner_string,
216+
Path("docs/source/backends/arm-vgf/arm-vgf-partitioner.md"),
217+
),
218+
DocumentationJob(
219+
Path("backends/arm/scripts/docgen/vgf/backends-arm-vgf-quantization.md.in"),
220+
"$QUANTIZER",
221+
quantizer_string,
222+
Path("docs/source/backends/arm-vgf/arm-vgf-quantization.md"),
223+
),
224+
DocumentationJob(
225+
Path("backends/arm/scripts/docgen/vgf/vgf-getting-started-tutorial.md.in"),
226+
"$MINIMAL_EXAMPLE",
227+
example_string,
228+
Path("docs/source/backends/arm-vgf/tutorials/vgf-getting-started.md"),
229+
),
230+
]
231+
232+
for job in documentation_jobs:
233+
generate_document(job)
234+
235+
236+
def generate_ethosu_tutorial():
237+
"""Generates the tutorial for the Ethos-U minimal example."""
238+
ethosu_example = get_jupyter_code(
239+
"./examples/arm/ethos_u_minimal_example.ipynb", get_bash=False
240+
)
241+
doc = "tutorial-arm-ethos-u.md"
242+
with open(f"backends/arm/scripts/docgen/{doc}.in", "r") as f:
243+
content = f.read()
244+
content = content.replace("$MINIMAL_EXAMPLE", ethosu_example)
245+
246+
with open(f"docs/source/{doc}", "w") as f:
247+
f.write(content)
248+
249+
250+
def generate_vgf_tutorial():
251+
"""Generates the tutorial for the VGF minimal example."""
252+
vgf_example = get_jupyter_code(
253+
"./examples/arm/vgf_minimal_example.ipynb",
254+
get_bash=False,
255+
which_cells=[0, 2, 3],
256+
)
257+
258+
doc = "tutorial-arm-vgf.md"
259+
with open(f"backends/arm/scripts/docgen/{doc}.in", "r") as f:
260+
content = f.read()
261+
content = content.replace("$MINIMAL_EXAMPLE", vgf_example)
262+
263+
with open(f"docs/source/{doc}", "w") as f:
264+
f.write(content)
265+
266+
267+
if __name__ == "__main__":
268+
generate_ethos_u_docs()
269+
generate_vgf_docs()

0 commit comments

Comments
 (0)