main() — pytorch Function Reference

Architecture documentation for the main() function in generate_aiter_embedded_hsa.py from the pytorch codebase.

Function python ComputeKernels ActivationOps calls 1 called by 1

Entity Profile

ComputeKernels→ ActivationOps→ main() — pytorch Function Reference

Dependency Diagram

graph TD
  be927a11_dfb0_c7e8_9987_812acb08f1cc["main()"]
  207b5d39_71a9_eb57_c653_b70c55fa39a3["sys()"]
  207b5d39_71a9_eb57_c653_b70c55fa39a3 -->|calls| be927a11_dfb0_c7e8_9987_812acb08f1cc
  45a8c845_25fb_a7fb_78e2_dec6d2d767fd["generate_embedded_hsa_header()"]
  be927a11_dfb0_c7e8_9987_812acb08f1cc -->|calls| 45a8c845_25fb_a7fb_78e2_dec6d2d767fd
  style be927a11_dfb0_c7e8_9987_812acb08f1cc fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

aten/src/ATen/native/transformers/hip/flash_attn/ck/fav_v3/generate_aiter_embedded_hsa.py lines 119–144

def main():
    parser = argparse.ArgumentParser(
        description="Generate aiter_embedded_hsa.h with embedded binary .co files"
    )
    parser.add_argument(
        "--hsa-dir", required=True, type=Path, help="Path to the aiter hsa directory"
    )
    parser.add_argument(
        "--output", required=True, type=Path, help="Path to the output header file"
    )
    parser.add_argument(
        "--subdirs",
        nargs="+",
        default=["gfx942/fmha_v3_bwd", "gfx950/fmha_v3_bwd"],
        help="Subdirectories to scan for .co files",
    )

    args = parser.parse_args()

    if not args.hsa_dir.exists():
        print(f"Error: HSA directory does not exist: {args.hsa_dir}", file=sys.stderr)
        return 1

    count = generate_embedded_hsa_header(args.hsa_dir, args.output, args.subdirs)
    print(f"Generated {args.output} with {count} embedded .co files")
    return 0