def compute_cubin_sizes()

in analytics/cubinsizes.py [0:0]

• 30 lines of code
• 11 McCabe index (conditional complexity)

def compute_cubin_sizes(file_name, section_name='.nv_fatbin', debug=False):
    with open(file_name, 'rb') as f:
        elf_file = ELFFile(f)
        nv_fatbin = elf_file.get_section_by_name(section_name)
        if nv_fatbin is None:
            return {}
        data = nv_fatbin.data()
        idx, offs = 0, 0
        elf_sizes = {}
        while offs < len(data):
            (magic, version, header_size, fatbin_size) = struct.unpack('IHHL', data[offs: offs + 16])
            if magic != 0xba55ed50 or version != 1:
                raise RuntimeError(f"Unexpected fatbin magic {hex(magic)} or version {version}")
            if debug:
                print(f"Found fatbin at {offs}  header_size={header_size} fatbin_size={fatbin_size}")
            offs += header_size
            fatbin_end = offs + fatbin_size
            while offs < fatbin_end:
                (kind, version, hdr_size, elf_size, empty, code_ver, sm_ver) = struct.unpack('HHILLIH', data[offs: offs + 30])
                if version != 0x0101 or kind not in [1, 2]:
                    raise RuntimeError(f"Unexpected cubin version {hex(version)} or kind {kind}")
                sm_ver = f'{"ptx" if kind == 1 else "sm"}_{sm_ver}'
                if debug:
                    print(f"    {idx}: elf_size={elf_size} code_ver={hex(code_ver)} sm={sm_ver}")
                if sm_ver not in elf_sizes:
                    elf_sizes[sm_ver] = 0
                elf_sizes[sm_ver] += elf_size
                idx, offs = idx + 1, offs + hdr_size + elf_size
            offs = fatbin_end
        return elf_sizes