projects/guix·blob kernel-updates · refs

178 lines · 6.7 KiB

1This patch fixes some scripts for generating source files. For
2gen_jit_decompositions.py, gen_mobile_upgraders.py and
3gen_jit_shape_functions.py, which depend on the compiled PyTorch library, the
4option to generate "dummy" source files is added for the initial build, which
5is later corrected. codegen_external.py is patched to avoid duplicate
6functions and add the static keyword as in the existing generated file.
7
8diff --git a/tools/gen_flatbuffers.sh b/tools/gen_flatbuffers.sh
9index cc0263dbb..ac34e84b8 100644
10--- a/tools/gen_flatbuffers.sh
11+++ b/tools/gen_flatbuffers.sh
12@@ -1,13 +1,13 @@
13 #!/bin/bash
14 ROOT=$(pwd)
15-FF_LOCATION="$ROOT/third_party/flatbuffers"
16-cd "$FF_LOCATION" || exit
17-mkdir build
18-cd build || exit
19-cmake ..
20-cmake --build . --target flatc
21-mkdir -p "$ROOT/build/torch/csrc/jit/serialization"
22-./flatc --cpp --gen-mutable --scoped-enums \
23+#FF_LOCATION="$ROOT/third_party/flatbuffers"
24+#cd "$FF_LOCATION" || exit
25+#mkdir build
26+#cd build || exit
27+#cmake ..
28+#cmake --build . --target flatc
29+#mkdir -p "$ROOT/build/torch/csrc/jit/serialization"
30+flatc --cpp --gen-mutable --scoped-enums \
31 -o "$ROOT/torch/csrc/jit/serialization" \
32 -c "$ROOT/torch/csrc/jit/serialization/mobile_bytecode.fbs"
33 echo '// @generated' >> "$ROOT/torch/csrc/jit/serialization/mobile_bytecode_generated.h"
34diff --git a/torch/csrc/jit/tensorexpr/codegen_external.py b/torch/csrc/jit/tensorexpr/codegen_external.py
35index 5dcf1b284..0e20b0c10 100644
36--- a/torch/csrc/jit/tensorexpr/codegen_external.py
37+++ b/torch/csrc/jit/tensorexpr/codegen_external.py
38@@ -21,9 +21,14 @@ def gen_external(native_functions_path, tags_path, external_path):
39 native_functions = parse_native_yaml(native_functions_path, tags_path)
40 func_decls = []
41 func_registrations = []
42- for func in native_functions:
43+ done_names = set()
44+ for func in native_functions[0]:
45 schema = func.func
46 name = schema.name.name.base
47+ if name in done_names:
48+ continue
49+ else:
50+ done_names.add(name)
51 args = schema.arguments
52 # Only supports extern calls for functions with out variants
53 if not schema.is_out_fn():
54@@ -63,7 +68,7 @@ def gen_external(native_functions_path, tags_path, external_path):
55
56 # print(tensor_decls, name, arg_names)
57 func_decl = f"""\
58-void nnc_aten_{name}(
59+static void nnc_aten_{name}(
60 int64_t bufs_num,
61 void** buf_data,
62 int64_t* buf_ranks,
63diff --git a/torchgen/decompositions/gen_jit_decompositions.py b/torchgen/decompositions/gen_jit_decompositions.py
64index b42948045..e1cfc73a5 100644
65--- a/torchgen/decompositions/gen_jit_decompositions.py
66+++ b/torchgen/decompositions/gen_jit_decompositions.py
67@@ -1,8 +1,12 @@
68 #!/usr/bin/env python3
69 import os
70 from pathlib import Path
71+import sys
72
73-from torch.jit._decompositions import decomposition_table
74+if len(sys.argv) < 2 or sys.argv[1] != "dummy":
75+ from torch.jit._decompositions import decomposition_table
76+else:
77+ decomposition_table = {}
78
79
80 # from torchgen.code_template import CodeTemplate
81@@ -86,7 +90,7 @@ def write_decomposition_util_file(path: str) -> None:
82
83
84 def main() -> None:
85- pytorch_dir = Path(__file__).resolve().parents[3]
86+ pytorch_dir = Path(__file__).resolve().parents[2]
87 upgrader_path = pytorch_dir / "torch" / "csrc" / "jit" / "runtime"
88 write_decomposition_util_file(str(upgrader_path))
89
90diff --git a/torchgen/operator_versions/gen_mobile_upgraders.py b/torchgen/operator_versions/gen_mobile_upgraders.py
91index 845034cb7..a1c5767c2 100644
92--- a/torchgen/operator_versions/gen_mobile_upgraders.py
93+++ b/torchgen/operator_versions/gen_mobile_upgraders.py
94@@ -6,10 +6,13 @@ import os
95 from enum import Enum
96 from operator import itemgetter
97 from pathlib import Path
98+import sys
99 from typing import Any
100
101-import torch
102-from torch.jit.generate_bytecode import generate_upgraders_bytecode
103+if len(sys.argv) < 2 or sys.argv[1] != "dummy":
104+ import torch
105+ from torch.jit.generate_bytecode import generate_upgraders_bytecode
106+
107 from torchgen.code_template import CodeTemplate
108 from torchgen.operator_versions.gen_mobile_upgraders_constant import (
109 MOBILE_UPGRADERS_HEADER_DESCRIPTION,
110@@ -263,7 +266,10 @@ def construct_register_size(register_size_from_yaml: int) -> str:
111 def construct_version_maps(
112 upgrader_bytecode_function_to_index_map: dict[str, Any],
113 ) -> str:
114- version_map = torch._C._get_operator_version_map()
115+ if len(sys.argv) < 2 or sys.argv[1] != "dummy":
116+ version_map = torch._C._get_operator_version_map()
117+ else:
118+ version_map = {}
119 sorted_version_map_ = sorted(version_map.items(), key=itemgetter(0)) # type: ignore[no-any-return]
120 sorted_version_map = dict(sorted_version_map_)
121
122@@ -375,7 +381,10 @@ def sort_upgrader(upgrader_list: list[dict[str, Any]]) -> list[dict[str, Any]]:
123
124
125 def main() -> None:
126- upgrader_list = generate_upgraders_bytecode()
127+ if len(sys.argv) < 2 or sys.argv[1] != "dummy":
128+ upgrader_list = generate_upgraders_bytecode()
129+ else:
130+ upgrader_list = []
131 sorted_upgrader_list = sort_upgrader(upgrader_list)
132 for up in sorted_upgrader_list:
133 print("after sort upgrader : ", next(iter(up)))
134diff --git a/torchgen/shape_functions/gen_jit_shape_functions.py b/torchgen/shape_functions/gen_jit_shape_functions.py
135index 56a3d8bf0..ffd0785fd 100644
136--- a/torchgen/shape_functions/gen_jit_shape_functions.py
137+++ b/torchgen/shape_functions/gen_jit_shape_functions.py
138@@ -1,6 +1,7 @@
139 #!/usr/bin/env python3
140 import os
141 import sys
142+import importlib
143 from importlib.util import module_from_spec, spec_from_file_location
144 from itertools import chain
145 from pathlib import Path
146@@ -18,17 +19,21 @@ you are in the root directory of the Pytorch git repo"""
147 if not file_path.exists():
148 raise Exception(err_msg) # noqa: TRY002
149
150-spec = spec_from_file_location(module_name, file_path)
151-assert spec is not None
152-module = module_from_spec(spec)
153-sys.modules[module_name] = module
154-assert spec.loader is not None
155-assert module is not None
156-spec.loader.exec_module(module)
157-
158-bounded_compute_graph_mapping = module.bounded_compute_graph_mapping
159-shape_compute_graph_mapping = module.shape_compute_graph_mapping
160-
161+if len(sys.argv) < 2 or sys.argv[1] != "dummy":
162+ spec = importlib.util.spec_from_file_location(module_name, file_path)
163+ assert spec is not None
164+ module = importlib.util.module_from_spec(spec)
165+ sys.modules[module_name] = module
166+ assert spec.loader is not None
167+ assert module is not None
168+ spec.loader.exec_module(module)
169+
170+ bounded_compute_graph_mapping = module.bounded_compute_graph_mapping
171+ shape_compute_graph_mapping = module.shape_compute_graph_mapping
172+
173+else:
174+ bounded_compute_graph_mapping = {}
175+ shape_compute_graph_mapping = {}
176
177 SHAPE_HEADER = r"""
178 /**