mirror of
https://github.com/pytorch/pytorch.git
synced 2025-11-02 14:34:54 +08: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)
103 lines
3.0 KiB
Python
103 lines
3.0 KiB
Python
from torchgen.model import (
|
|
Argument,
|
|
FunctionSchema,
|
|
Return,
|
|
SelfArgument,
|
|
TensorOptionsArguments,
|
|
Type,
|
|
)
|
|
|
|
from torchgen.api.types import ArgName, Binding, NamedCType, CType
|
|
from torchgen.api import cpp
|
|
from torchgen.utils import concatMap, assert_never
|
|
|
|
import itertools
|
|
from typing import Sequence, List, Union
|
|
|
|
# This file describes the translation of JIT schema to the dispatcher
|
|
# API, the *unboxed* calling convention by which invocations through
|
|
# the dispatcher are made. Historically, the dispatcher API matched
|
|
# the C++ API, but with the establishment of the boxed API, we've
|
|
# made changes to the dispatcher API to so that the unboxed API
|
|
# better aligns with the boxed API. The dispatcher API hooks heavily
|
|
# into our template based boxing/unboxing machinery, so changes
|
|
# to this convention will usually need template updates too.
|
|
#
|
|
# Prominent characteristics of the dispatcher API:
|
|
#
|
|
# - dtype, layout, device and pin_memory are represented as separate
|
|
# arguments.
|
|
#
|
|
|
|
|
|
def name(func: FunctionSchema) -> str:
|
|
return cpp.name(func)
|
|
|
|
|
|
def argumenttype_type(
|
|
t: Type, *, mutable: bool, binds: ArgName, remove_non_owning_ref_types: bool = False
|
|
) -> NamedCType:
|
|
# This is a faux amis. If it makes sense in the future to add
|
|
# more special cases here, or invert things so cpp.argument_type
|
|
# calls this, or just completely inline the function, please do
|
|
# it.
|
|
return cpp.argumenttype_type(
|
|
t,
|
|
mutable=mutable,
|
|
binds=binds,
|
|
remove_non_owning_ref_types=remove_non_owning_ref_types,
|
|
)
|
|
|
|
|
|
def argument_type(
|
|
a: Argument, *, binds: ArgName, remove_non_owning_ref_types: bool = False
|
|
) -> NamedCType:
|
|
return argumenttype_type(
|
|
a.type,
|
|
mutable=a.is_write,
|
|
binds=binds,
|
|
remove_non_owning_ref_types=remove_non_owning_ref_types,
|
|
)
|
|
|
|
|
|
def returns_type(rs: Sequence[Return]) -> CType:
|
|
# At present, there is no difference. But there could be!
|
|
return cpp.returns_type(rs)
|
|
|
|
|
|
def jit_arguments(func: FunctionSchema) -> List[Argument]:
|
|
def to_argument(
|
|
a: Union[Argument, TensorOptionsArguments, SelfArgument]
|
|
) -> List[Argument]:
|
|
if isinstance(a, Argument):
|
|
return [a]
|
|
elif isinstance(a, SelfArgument):
|
|
return [a.argument]
|
|
elif isinstance(a, TensorOptionsArguments):
|
|
return [a.dtype, a.layout, a.device, a.pin_memory]
|
|
else:
|
|
assert_never(a)
|
|
|
|
return list(
|
|
concatMap(
|
|
to_argument,
|
|
itertools.chain(
|
|
func.arguments.positional, func.arguments.kwarg_only, func.arguments.out
|
|
),
|
|
)
|
|
)
|
|
|
|
|
|
def argument(a: Argument, *, remove_non_owning_ref_types: bool = False) -> Binding:
|
|
return Binding(
|
|
nctype=argument_type(
|
|
a, binds=a.name, remove_non_owning_ref_types=remove_non_owning_ref_types
|
|
),
|
|
name=a.name,
|
|
argument=a,
|
|
)
|
|
|
|
|
|
def arguments(func: FunctionSchema) -> List[Binding]:
|
|
return [argument(a) for a in jit_arguments(func)]
|