mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
36420b5e8c
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/76275 In preparation for addressing https://github.com/pytorch/pytorch/issues/73212 Diff was generated with: ``` git mv tools/codegen torchgen git grep -l 'tools.codegen' | xargs sed -i 's/tools.codegen/torchgen/g' sed -i "s/\${TOOLS_PATH}\/codegen/\${TORCH_ROOT}\/torchgen/g" caffe2/CMakeLists.txt ``` and a manual edits to: * tools/test/test_gen_backend_stubs.py * torchgen/build.bzl * torchgen/gen_backend_stubs.py aka this diff: ``` diff --git a/tools/test/test_gen_backend_stubs.py b/tools/test/test_gen_backend_stubs.py index 3dc26c6d2d..104054575e 100644 --- a/tools/test/test_gen_backend_stubs.py +++ b/tools/test/test_gen_backend_stubs.py @@ -9,7 +9,7 @@ from torchgen.gen_backend_stubs import run from torchgen.gen import _GLOBAL_PARSE_NATIVE_YAML_CACHE # noqa: F401 path = os.path.dirname(os.path.realpath(__file__)) -gen_backend_stubs_path = os.path.join(path, '../torchgen/gen_backend_stubs.py') +gen_backend_stubs_path = os.path.join(path, '../../torchgen/gen_backend_stubs.py') # gen_backend_stubs.py is an integration point that is called directly by external backends. # The tests here are to confirm that badly formed inputs result in reasonable error messages. diff --git a/torchgen/build.bzl b/torchgen/build.bzl index ed04e35a43..d00078a3cf 100644 --- a/torchgen/build.bzl +++ b/torchgen/build.bzl @@ -1,6 +1,6 @@ def define_targets(rules): rules.py_library( - name = "codegen", + name = "torchgen", srcs = rules.glob(["**/*.py"]), deps = [ rules.requirement("PyYAML"), @@ -11,6 +11,6 @@ def define_targets(rules): rules.py_binary( name = "gen", - srcs = [":codegen"], + srcs = [":torchgen"], visibility = ["//visibility:public"], ) diff --git a/torchgen/gen_backend_stubs.py b/torchgen/gen_backend_stubs.py index c1a672a655..beee7a15e0 100644 --- a/torchgen/gen_backend_stubs.py +++ b/torchgen/gen_backend_stubs.py @@ -474,7 +474,7 @@ def run( ) -> None: # Assumes that this file lives at PYTORCH_ROOT/torchgen/gen_backend_stubs.py - pytorch_root = pathlib.Path(__file__).parent.parent.parent.absolute() + pytorch_root = pathlib.Path(__file__).parent.parent.absolute() template_dir = os.path.join(pytorch_root, "aten/src/ATen/templates") def make_file_manager(install_dir: str) -> FileManager: ``` run_all_fbandroid_tests Test Plan: sandcastle Reviewed By: albanD, ngimel Differential Revision: D35770317 fbshipit-source-id: 153ac4a7fef15b1e750812a90bfafdbc8f1ebcdf (cherry picked from commit c6d485d1d4648fa1c8a4c14c5bf3d8e899b9b4dd)
95 lines
2.3 KiB
Python
95 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from torch.jit._decompositions import decomposition_table
|
|
|
|
# from torchgen.code_template import CodeTemplate
|
|
|
|
DECOMP_HEADER = r"""
|
|
/**
|
|
* @generated
|
|
* This is an auto-generated file. Please do not modify it by hand.
|
|
* To re-generate, please run:
|
|
* cd ~/pytorch && python torchgen/decompositions/gen_jit_decompositions.py
|
|
*/
|
|
#include <torch/csrc/jit/jit_log.h>
|
|
#include <torch/csrc/jit/passes/inliner.h>
|
|
#include <torch/csrc/jit/runtime/operator.h>
|
|
#include <torch/csrc/jit/runtime/decomposition_registry_util.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
|
|
const std::string decomp_funcs =
|
|
R"("""
|
|
|
|
|
|
DECOMP_CENTER = r"""
|
|
)";
|
|
|
|
const std::string& GetSerializedDecompositions() {
|
|
return decomp_funcs;
|
|
}
|
|
|
|
const OperatorMap<std::string>& GetDecompositionMapping() {
|
|
// clang-format off
|
|
static const OperatorMap<std::string> decomposition_mapping {
|
|
"""
|
|
|
|
DECOMP_END = r"""
|
|
};
|
|
// clang-format on
|
|
|
|
return decomposition_mapping;
|
|
}
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|
|
"""
|
|
|
|
|
|
DECOMPOSITION_UTIL_FILE_NAME = "decomposition_registry_util.cpp"
|
|
|
|
|
|
def gen_serialized_decompisitions() -> str:
|
|
return "\n".join(
|
|
[scripted_func.code for scripted_func in decomposition_table.values()]
|
|
)
|
|
|
|
|
|
def gen_decomposition_mappings() -> str:
|
|
decomposition_mappings = []
|
|
for schema, scripted_func in decomposition_table.items():
|
|
decomposition_mappings.append(
|
|
' {"' + schema + '", "' + scripted_func.name + '"},'
|
|
)
|
|
return "\n".join(decomposition_mappings)
|
|
|
|
|
|
def write_decomposition_util_file(path: str) -> None:
|
|
decomposition_str = gen_serialized_decompisitions()
|
|
decomposition_mappings = gen_decomposition_mappings()
|
|
file_components = [
|
|
DECOMP_HEADER,
|
|
decomposition_str,
|
|
DECOMP_CENTER,
|
|
decomposition_mappings,
|
|
DECOMP_END,
|
|
]
|
|
print("writing file to : ", path + "/" + DECOMPOSITION_UTIL_FILE_NAME)
|
|
with open(os.path.join(path, DECOMPOSITION_UTIL_FILE_NAME), "wb") as out_file:
|
|
final_output = "".join(file_components)
|
|
out_file.write(final_output.encode("utf-8"))
|
|
|
|
|
|
def main() -> None:
|
|
pytorch_dir = Path(__file__).resolve().parents[3]
|
|
upgrader_path = pytorch_dir / "torch" / "csrc" / "jit" / "runtime"
|
|
write_decomposition_util_file(str(upgrader_path))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|