Exporting build_model api, need to change hard paths, and allow for selecting models without export script

Author: BujSet

extension/embedded/export_add.py

@@ -0,0 +1,25 @@
+import torch
+from torch.export import export
+from executorch.exir import to_edge
+
+# Start with a PyTorch model that adds two input tensors (matrices)
+class Add(torch.nn.Module):
+  def __init__(self):
+    super(Add, self).__init__()
+
+  def forward(self, x: torch.Tensor, y: torch.Tensor):
+      return x + y
+
+# 1. torch.export: Defines the program with the ATen operator set.
+aten_dialect = export(Add(), (torch.ones(1), torch.ones(1)))
+
+# 2. to_edge: Make optimizations for Edge devices
+edge_program = to_edge(aten_dialect)
+
+# 3. to_executorch: Convert the graph to an ExecuTorch program
+executorch_program = edge_program.to_executorch()
+
+# 4. Save the compiled .pte program
+with open("add.pte", "wb") as file:
+    file.write(executorch_program.buffer)
+

extension/embedded/pte_to_header.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+"""
+Generates .pte model, operator definitions, and header files
+"""
+
+import os
+import sys
+import subprocess
+import argparse
+from pathlib import Path
+
+def run_command(cmd, cwd=None, description=""):
+    """Run a command and handle errors"""
+    print(f"Running: {description}")
+    print(f"Command: {' '.join(cmd)}")
+    
+    try:
+        result = subprocess.run(cmd, cwd=cwd, check=True, capture_output=True, text=True)
+        print(f"✓ {description} completed successfully")
+        if result.stdout:
+            print(f"Output: {result.stdout}")
+        return result
+    except subprocess.CalledProcessError as e:
+        print(f"✗ {description} failed")
+        print(f"Error: {e.stderr}")
+        sys.exit(1)
+
+def main():
+    parser = argparse.ArgumentParser(description="Build ExecuTorch ARM Hello World model")
+    parser.add_argument("--executorch-root", default="~/optional/modules/lib/executorch", 
+                       help="Path to ExecuTorch root directory")
+    parser.add_argument("--model-name", default="add", 
+                       help="Name of the model (default: add)")
+    parser.add_argument("--clean", action="store_true", 
+                       help="Clean generated files before building")
+    
+    args = parser.parse_args()
+    
+    # Paths
+    script_dir = Path(__file__).parent
+    project_root = script_dir.parent.parent.parent.parent.parent.parent  # Go up to petriok root
+    executorch_root = project_root / args.executorch_root
+    example_files_dir = "/home/zephyruser/zephyr/samples/modules/executorch/arm/hello_world/example_files"
+    src_dir = script_dir / "src"
+    
+    model_name = args.model_name
+    pte_file = f"{model_name}.pte"
+    ops_def_file = "gen_ops_def.yml"
+    header_file = "model_pte.h"
+    
+    print(f"Building ExecuTorch model: {model_name}")
+    print(f"ExecuTorch root: {executorch_root}")
+    print(f"Working directory: {script_dir}")
+    
+    # Clean previous build if requested
+    if args.clean:
+        files_to_clean = [pte_file, ops_def_file, src_dir / header_file]
+        for file_path in files_to_clean:
+            if Path(file_path).exists():
+                Path(file_path).unlink()
+                print(f"Cleaned: {file_path}")
+    
+    # Step 1: Generate the .pte model file
+    export_script = os.path.join(example_files_dir, f"export_{model_name}.py")
+    if not os.path.exists(export_script):
+        print(f"Error: Export script not found: {export_script}")
+        sys.exit(1)
+    
+    try:
+        run_command(
+            [sys.executable, str(export_script)],
+            cwd=script_dir,
+            description="Generating .pte model file"
+        )
+    except SystemExit:
+        print(f"\n❌ Model generation failed. This is likely because PyTorch/ExecuTorch is not installed.")
+        print(f"For now, using dummy model_pte.h for compilation testing.")
+        print(f"To generate a real model, install PyTorch and ExecuTorch:")
+        print(f"  pip install torch")
+        print(f"  # Install ExecuTorch according to documentation")
+        print(f"  python build_model.py")
+        return
+    
+    if not Path(script_dir / pte_file).exists():
+        print(f"Error: Model file {pte_file} was not generated")
+        sys.exit(1)
+    
+    # Step 2: Generate operator definitions
+
+    gen_ops_script = "/home/zephyruser/optional/modules/lib/executorch/codegen/tools/gen_ops_def.py"
+    if not os.path.exists(gen_ops_script):
+        print(f"Error: gen_ops_def.py not found at {gen_ops_script}")
+        sys.exit(1)
+    
+    run_command(
+        [sys.executable, str(gen_ops_script), 
+         "--output_path", ops_def_file,
+         "--model_file_path", pte_file],
+        cwd=script_dir,
+        description="Generating operator definitions"
+    )
+    
+    # Step 3: Convert .pte to header file
+    #pte_to_header_script = executorch_root / "examples" / "arm" / "executor_runner" / "pte_to_header.py"
+    pte_to_header_script = "/home/zephyruser/optional/modules/lib/executorch/examples/arm/executor_runner/pte_to_header.py"
+    if not os.path.exists(pte_to_header_script):
+        print(f"Error: pte_to_header.py not found at {pte_to_header_script}")
+        sys.exit(1)
+    
+    run_command(
+        [sys.executable, str(pte_to_header_script),
+         "--pte", pte_file,
+         "--outdir", "src"],
+        cwd=script_dir,
+        description="Converting .pte to header file"
+    )
+    
+    # Step 4: Make the generated array const and remove section attribute
+    header_path = src_dir / header_file
+    if header_path.exists():
+        content = header_path.read_text()
+        
+        # Remove section attribute and replace with Zephyr alignment macro
+        import re
+        # Replace section+aligned pattern with Zephyr __ALIGN macro
+        content = re.sub(r'__attribute__\s*\(\s*\(\s*section\s*\([^)]*\)\s*,\s*aligned\s*\(([^)]*)\)\s*\)\s*\)\s*', r'__ALIGN(\1) ', content)
+        # Remove any remaining section-only attributes  
+        content = re.sub(r'__attribute__\s*\(\s*\(\s*section\s*\([^)]*\)\s*\)\s*\)\s*', '', content)
+        # Also replace any standalone __attribute__((aligned(n))) with __ALIGN(n)
+        content = re.sub(r'__attribute__\s*\(\s*\(\s*aligned\s*\(([^)]*)\)\s*\)\s*\)\s*', r'__ALIGN(\1) ', content)
+        
+        # Replace 'char model_pte_data[]' with 'const char model_pte_data[]'
+        content = content.replace('char model_pte_data[]', 'const char model_pte_data[]')
+        # Also handle 'char model_pte[]' variant
+        content = content.replace('char model_pte[]', 'const char model_pte[]')
+        
+        header_path.write_text(content)
+        print(f"✓ Made model data const and removed section attributes in {header_file}")
+    else:
+        print(f"Warning: Header file {header_file} not found")
+    
+    print("\n=== Build Summary ===")
+    print(f"✓ Generated: {pte_file}")
+    print(f"✓ Generated: {ops_def_file}")
+    print(f"✓ Generated: src/{header_file}")
+    print("\nNext steps:")
+    print("1. Review gen_ops_def.yml and customize if needed")
+    print("2. Build the Zephyr application with west build")
+
+if __name__ == "__main__":
+    main()