| 1 | This patch fixes some scripts for generating source files. For
|
| 2 | gen_jit_decompositions.py, gen_mobile_upgraders.py and
|
| 3 | gen_jit_shape_functions.py, which depend on the compiled PyTorch library, the
|
| 4 | option to generate "dummy" source files is added for the initial build, which
|
| 5 | is later corrected. codegen_external.py is patched to avoid duplicate
|
| 6 | functions and add the static keyword as in the existing generated file.
|
| 7 |
|
| 8 | diff --git a/tools/gen_flatbuffers.sh b/tools/gen_flatbuffers.sh
|
| 9 | index cc0263dbb..ac34e84b8 100644
|
| 10 | |
| 11 | |
| 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"
|
| 34 | diff --git a/torch/csrc/jit/tensorexpr/codegen_external.py b/torch/csrc/jit/tensorexpr/codegen_external.py
|
| 35 | index 5dcf1b284..0e20b0c10 100644
|
| 36 | |
| 37 | |
| 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,
|
| 63 | diff --git a/torchgen/decompositions/gen_jit_decompositions.py b/torchgen/decompositions/gen_jit_decompositions.py
|
| 64 | index b42948045..e1cfc73a5 100644
|
| 65 | |
| 66 | |
| 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 |
|
| 90 | diff --git a/torchgen/operator_versions/gen_mobile_upgraders.py b/torchgen/operator_versions/gen_mobile_upgraders.py
|
| 91 | index 845034cb7..a1c5767c2 100644
|
| 92 | |
| 93 | |
| 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 | @@ -260,7 +263,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 | @@ -372,7 +378,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)))
|
| 134 | diff --git a/torchgen/shape_functions/gen_jit_shape_functions.py b/torchgen/shape_functions/gen_jit_shape_functions.py
|
| 135 | index 56a3d8bf0..ffd0785fd 100644
|
| 136 | |
| 137 | |
| 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 | /**
|