xref: /aosp_15_r20/external/pytorch/tools/setup_helpers/generate_linker_script.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1import subprocess
2
3
4def gen_linker_script(
5    filein: str = "cmake/prioritized_text.txt", fout: str = "cmake/linker_script.ld"
6) -> None:
7    with open(filein) as f:
8        prioritized_text = f.readlines()
9        prioritized_text = [
10            line.replace("\n", "") for line in prioritized_text if line != "\n"
11        ]
12
13    linker_script_lines = subprocess.check_output(["ld", "-verbose"], text=True).split(
14        "\n"
15    )
16
17    indices = [
18        i
19        for i, x in enumerate(linker_script_lines)
20        if x == "=================================================="
21    ]
22    linker_script_lines = linker_script_lines[indices[0] + 1 : indices[1]]
23
24    text_line_start = [
25        i for i, line in enumerate(linker_script_lines) if ".text           :" in line
26    ]
27    assert len(text_line_start) == 1, "The linker script has multiple text sections!"
28    text_line_start = text_line_start[0]
29
30    with open(fout, "w") as f:
31        for lineid, line in enumerate(linker_script_lines):
32            if lineid == text_line_start + 2:
33                f.write("    *(\n")
34                for plines in prioritized_text:
35                    f.write(f"      .text.{plines}\n")
36                f.write("    )\n")
37            f.write(f"{line}\n")
38