main() — pytorch Function Reference

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

Function python ComputeKernels QuantizedCPU calls 2 called by 1

Entity Profile

ComputeKernels→ QuantizedCPU→ main() — pytorch Function Reference

Dependency Diagram

graph TD
  20303473_0948_19b1_a212_748886b73572["main()"]
  7fedab6d_8981_08c3_ca38_2f0eddd64dcc["parser()"]
  7fedab6d_8981_08c3_ca38_2f0eddd64dcc -->|calls| 20303473_0948_19b1_a212_748886b73572
  1ec91f05_2749_ffa8_d3f5_5f37a18418dc["write_decl_impl()"]
  20303473_0948_19b1_a212_748886b73572 -->|calls| 1ec91f05_2749_ffa8_d3f5_5f37a18418dc
  0e771794_8e99_3f41_d4a0_5d896b90ba41["get_all()"]
  20303473_0948_19b1_a212_748886b73572 -->|calls| 0e771794_8e99_3f41_d4a0_5d896b90ba41
  style 20303473_0948_19b1_a212_748886b73572 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernels/generate_kernels.py lines 381–397

def main(output_dir: Optional[str]) -> None:
    if output_dir is None:
        output_dir = Path(__file__).parent
    else:
        output_dir = Path(output_dir)
    write_decl_impl(
        FwdKernel.get_all(),
        "cutlassF",
        impl_file="<ATen/native/transformers/cuda/mem_eff_attention/kernel_forward.h>",
        autogen_dir=output_dir,
    )
    write_decl_impl(
        BwdKernel.get_all(),
        "cutlassB",
        impl_file="<ATen/native/transformers/cuda/mem_eff_attention/kernel_backward.h>",
        autogen_dir=output_dir,
    )