[ir] Generate yaml documentation for statement classes (#2192)

misc/generate_ir_design_doc.py

@@ -0,0 +1,72 @@
+import taichi as ti
+import os
+import yaml
+
+
+def extract_doc(doc_filename=None):
+    repo_dir = ti.get_repo_directory()
+    statements_fn = os.path.join(repo_dir, 'taichi/ir/statements.h')
+    with open(statements_fn, 'r') as f:
+        statements = f.readlines()
+
+    class_doc = {}
+
+    for i in range(len(statements)):
+        line = statements[i]
+        start_pos = line.find('/**')
+        if start_pos == -1:
+            continue
+        current_doc = line[start_pos + 3:].strip()
+        doc_ends_at_line = 0
+        for j in range(i + 1, len(statements)):
+            next_line = statements[j]
+            end_pos = next_line.find('*/')
+            if end_pos != -1:
+                doc_ends_at_line = j
+                break
+            next_line = next_line.strip()
+            if next_line.startswith('*'):
+                next_line = next_line[1:].strip()
+            if next_line == '':  # an empty line
+                current_doc += '\n'
+            else:
+                current_doc += ' ' + next_line
+        current_doc = current_doc.strip()
+
+        line = statements[doc_ends_at_line + 1]
+        start_pos = line.find('class')
+        if start_pos == -1:
+            print('We only support doc for classes now. '
+                  f'The following doc at line {i}-{doc_ends_at_line} '
+                  'cannot be recognized:\n'
+                  f'{current_doc}')
+            continue
+        class_name = line[start_pos + 5:].strip().split()[0]
+        class_doc[class_name] = current_doc
+
+    if doc_filename is None:
+        doc_filename = 'ir_design_doc.yml'
+    with open(doc_filename, 'w') as f:
+        yaml.dump(class_doc, f, Dumper=yaml.SafeDumper)
+
+
+def yml_to_md(yml_filename=None, md_filename=None):
+    if yml_filename is None:
+        yml_filename = 'ir_design_doc.yml'
+    if md_filename is None:
+        md_filename = 'ir_design_doc.md'
+    with open(yml_filename, 'r') as f:
+        doc_yml = yaml.load(f, Loader=yaml.SafeLoader)
+    doc_md = ''
+    for (class_name, class_doc) in doc_yml.items():
+        doc_md += f'### {class_name}\n'
+        doc_md += class_doc
+        doc_md += '\n\n'
+    with open(md_filename, 'w') as f:
+        f.write(doc_md)
+
+
+if __name__ == '__main__':
+    extract_doc()
+    yml_to_md()
+    print('Done!')

taichi/ir/statements.h

@@ -6,6 +6,9 @@
 
 TLANG_NAMESPACE_BEGIN
 
+/**
+ * Allocate a local variable with initial value 0.
+ */
 class AllocaStmt : public Stmt {
  public:
   AllocaStmt(DataType type) {
@@ -30,7 +33,9 @@ class AllocaStmt : public Stmt {
   TI_DEFINE_ACCEPT_AND_CLONE
 };
 
-// updates mask, break if no active
+/**
+ * Updates mask, break if all bits of the mask are 0.
+ */
 class WhileControlStmt : public Stmt {
  public:
   Stmt *mask;
@@ -43,6 +48,9 @@ class WhileControlStmt : public Stmt {
   TI_DEFINE_ACCEPT_AND_CLONE;
 };
 
+/**
+ * Jump to the next loop iteration, i.e., `continue` in C++.
+ */
 class ContinueStmt : public Stmt {
  public:
   // This is the loop on which this continue stmt has effects. It can be either