main() — pytorch Function Reference

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

Function python CoreTensor Dispatcher calls 2 called by 1

Entity Profile

CoreTensor→ Dispatcher→ main() — pytorch Function Reference

Dependency Diagram

graph TD
  aad87944_33f6_d3c8_e945_0c0de37e4fe0["main()"]
  2a7b91f6_0c36_b743_0ba0_d6b276e5dc23["parser()"]
  2a7b91f6_0c36_b743_0ba0_d6b276e5dc23 -->|calls| aad87944_33f6_d3c8_e945_0c0de37e4fe0
  39541fcd_017c_75bf_8b33_705eeef8b2bc["write_decl_impl()"]
  aad87944_33f6_d3c8_e945_0c0de37e4fe0 -->|calls| 39541fcd_017c_75bf_8b33_705eeef8b2bc
  94e5a4b8_623e_3785_583d_a91baa59116d["get_all()"]
  aad87944_33f6_d3c8_e945_0c0de37e4fe0 -->|calls| 94e5a4b8_623e_3785_583d_a91baa59116d
  style aad87944_33f6_d3c8_e945_0c0de37e4fe0 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,
    )