def main()

in analytics/cubinsizes.py [0:0]

• 31 lines of code
• 12 McCabe index (conditional complexity)

def main():
    if sys.platform != 'linux':
        print('This script only works with Linux ELF files')
        return
    if len(sys.argv) < 2:
        print(f"{sys.argv[0]} invoked without any arguments trying to infer location of libtorch_cuda")
        import torch
        fname = os.path.join(os.path.dirname(torch.__file__), 'lib', 'libtorch_cuda.so')
    else:
        fname = sys.argv[1]

    if not os.path.exists(fname):
        print(f"Can't find {fname}")
        sys.exit(-1)

    section_names = ['.nv_fatbin', '__nv_relfatbin']
    results = {name: {} for name in section_names}
    print(f"Analyzing {fname}")
    if os.path.splitext(fname)[1] == '.a':
        with ArFileCtx(fname):
            for fname in os.listdir("."):
                if not fname.endswith(".o"): continue
                for section_name in section_names:
                    elf_sizes = compute_cubin_sizes(fname, section_name)
                    dict_add(results[section_name], elf_sizes)
    else:
        for section_name in ['.nv_fatbin', '__nv_relfatbin']:
            dict_add(results[section_name], compute_cubin_sizes(fname, section_name))

    for section_name in section_names:
        elf_sizes = results[section_name]
        print(f"{section_name} size {sizeof_fmt(sum(elf_sizes.values()))}")
        for (sm_ver, total_size) in elf_sizes.items():
            print(f"  {sm_ver}: {sizeof_fmt(total_size)}")