feat: generate docs from build file and source

README.md

@@ -29,6 +29,7 @@ the activation folder.
 - [Building kernels with Nix](./docs/nix.md)
 - [Local kernel development](docs/local-dev.md) (IDE integration)
 - [Why Nix?](./docs/why-nix.md)
+- [Documentation Generator](./docs/documentation-generator.md)
 
 ## Credits
 

docs/documentation-generator.md

@@ -0,0 +1,74 @@
+# Documentation Generator
+
+This tool helps you generate consistent documentation for your CUDA kernel projects built with kernel-builder. It analyzes source files and build configuration to create structured Markdown documentation.
+
+## Features
+
+- Extracts documentation from CUDA, C++, and header files
+- Parses build configuration from build.toml
+- Identifies kernels and function definitions with their parameters
+- Generates Markdown documentation with a table of contents
+- Includes build settings and project information
+
+## Usage
+
+```bash
+# Generate documentation for a kernel project
+python scripts/gen-docs.py /path/to/your/kernel/project
+
+# Specify custom output directory (default is "docs")
+python scripts/gen-docs.py /path/to/your/kernel/project --output custom-docs
+```
+
+## Comment Format
+
+For best results, document your kernels and functions using the following format:
+
+```cpp
+/**
+ * This is a description of the kernel or function.
+ * 
+ * Any additional details about the implementation or usage can go here.
+ */
+__global__ void my_kernel(float* input, float* output, int size) {
+    // kernel implementation
+}
+```
+
+The generator will extract this documentation and include it in the generated files.
+
+## Output Structure
+
+The generated documentation includes:
+
+1. **Table of Contents** - Navigation for all documentation sections
+2. **Project Overview** - Basic information about the project (from build.toml if available)
+3. **Build Configuration** - Settings from build.toml including:
+   - Kernel definitions
+   - CUDA capabilities
+   - Source files
+   - Dependencies
+4. **API Documentation** - Documentation for each source file:
+   - Function and kernel signatures
+   - Parameter tables with types
+   - Documentation comments
+
+## Example
+
+Given a project with the following structure:
+
+```
+my_kernel/
+├── build.toml
+├── kernel.cu
+└── torch-ext/
+    └── ...
+```
+
+Running the documentation generator:
+
+```bash
+python scripts/gen-docs.py my_kernel
+```
+
+Will produce `my_kernel.md` in the `my_kernel/docs` directory.

scripts/gen-docs.py

@@ -0,0 +1,143 @@
+# /// script
+# dependencies = ["toml", "minijinja"]
+# ///
+import toml
+import sys
+import os
+import glob
+import re
+from minijinja import Environment
+
+
+def extract_ops(file_path):
+    ops = []
+    with open(file_path, "r") as f:
+        content = f.read()
+        # Find TORCH_LIBRARY_EXPAND blocks
+        lib_blocks = re.findall(
+            r"TORCH_LIBRARY_EXPAND\([^)]+\)([^}]+)}", content, re.DOTALL
+        )
+        for block in lib_blocks:
+            # Extract ops.def lines
+            op_defs = re.findall(r"ops\.def\(\"([^\"]+)\"", block)
+            ops.extend(op_defs)
+    return ops
+
+
+def main():
+    if len(sys.argv) < 2:
+        print("Usage: script.py <project_directory>")
+        return
+
+    project_dir = sys.argv[1]
+
+    # Read build.toml or return if not found
+    build_config_path = os.path.join(project_dir, "build.toml")
+    if not os.path.exists(build_config_path):
+        print(f"Error: build.toml not found at {build_config_path}")
+        return
+    try:
+        build_config = toml.load(build_config_path)
+        print(f"Successfully parsed build configuration from {build_config_path}")
+
+    except Exception as e:
+        print(f"Error parsing build.toml: {e}")
+        return
+
+    if not build_config:
+        return
+
+    # Get all kernels from the config
+    config_kernels = build_config.get("kernel", {})
+
+    kernels = []
+    for kernel_name, kernel_info in config_kernels.items():
+        kernels.append(
+            {
+                "name": kernel_name,
+                "cuda-capabilities": kernel_info.get("cuda-capabilities", []),
+                "src": kernel_info.get("src", []),
+                "dependencies": kernel_info.get("depends", []),
+            }
+        )
+
+    # Find torch-ext directory and extract all ops
+    ops = []
+    torch_ext_dirs = glob.glob(
+        os.path.join(project_dir, "**/torch-ext"), recursive=True
+    )
+    for torch_ext_dir in torch_ext_dirs:
+        for file in glob.glob(os.path.join(torch_ext_dir, "**/*.cpp"), recursive=True):
+            for op in extract_ops(file):
+                relative_file = "../" + os.path.relpath(file, project_dir)
+                ops.append(dict(op=op, file=relative_file))
+
+    # Prepare template data
+    template_data = {
+        "project_name": os.path.basename(project_dir),
+        "build_config": {"kernels": kernels},
+        "ops": ops,
+    }
+
+    # Render template
+    env = Environment(
+        templates={
+            "doc_template": """
+# `{{ project_name }}` Documentation
+
+> __Generated on 2025-03-25__
+
+## Table of Contents
+
+- [Project Overview](#project-overview)
+- [Build Configuration](#build-configuration)
+  - [Kernels](#kernels){% for kernel_info in build_config.get("kernels", {}) %}\n    - [{{ kernel_info.get("name") }}](#{{ kernel_info.get("name") }}){% endfor %}
+  - [Operations](#operations){% for op in ops %}\n    - [{{ op.get("op") }}]({{ op.get("file") }}){% endfor %}
+
+## Project Overview
+{{ project_name }} is a CUDA kernel project.
+
+## Build Configuration
+
+### Kernels
+{% for kernel_info in build_config.get("kernels", {}) %}
+#### {{ kernel_info.get("name") }}
+
+**CUDA Capabilities:**
+- `{{ kernel_info.get("cuda-capabilities", []) }}`
+
+**Source Files:**
+{% for source in kernel_info.get("src", []) %}
+- `{{ source }}`
+{% endfor %}
+
+**Dependencies:**
+{% for dep in kernel_info.get("dependencies", []) %}
+- `{{ dep }}`
+{% endfor %}
+{% endfor %}
+
+### operations
+{% for op in ops %}
+```cpp
+{{ op.get("op") }}
+```
+[defined]({{ op.get("file") }})
+
+{% endfor %}
+"""
+        }
+    )
+    output = env.render_template("doc_template", **template_data)
+
+    # Write output file
+    project_name = os.path.basename(os.path.abspath(project_dir))
+    output_dir = os.path.join(project_dir, "docs")
+    os.makedirs(output_dir, exist_ok=True)
+    output_path = os.path.join(output_dir, f"{project_name}.md")
+    with open(output_path, "w", encoding="utf-8") as f:
+        f.write(output)
+
+
+if __name__ == "__main__":
+    main()

scripts/init-kernel.py

@@ -163,6 +163,11 @@ def main():
         f"  {Colors.YELLOW}{1}.{Colors.ENDC} {Colors.BOLD}pytest -vv tests/{Colors.ENDC}"
     )
 
+    print(f"\n{Colors.CYAN}{Colors.BOLD}Generate documentation{Colors.ENDC}")
+    print(
+        f"  {Colors.YELLOW}{1}.{Colors.ENDC} {Colors.BOLD}uv run scripts/gen-docs.py ./{Colors.ENDC}"
+    )
+
     print("")