mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
947c7a8215
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/55047 Added namespaces to all of the `CTypes` printed in the codegen. This is pretty much required if we want to use codegen externally, since we can no longer assume that we're inside of the `at::` namespace. Important changes are in `types.py`. How do we add the notion of namespaces to C++ types without people having to write "at::Tensor" everywhere? Before this PR, `CType` held a raw string representing the type, i.e. `BaseCType("Tensor", binds)`. This PR introduces a set of singleton base C++ types in `types.py`, that know how to print their namespace. Instead, we'd write `BaseCType(tensorT, binds)`, where printing `tensorT` will properly print out "at::Tensor". This also means that you can't create arbitrary `CTypes`. If we need a new C++ type in the codegen, we need to add it to the list in `types.py`. One blip in the design: we don't want to change `RegistrationDeclarations.yaml`, since that'll break external backends that ingest it. I added separate functions to display types without the namespace that are used to create RegistrationDeclarations.yaml`. With an external codegen API though, we can eventually kill it :) I also didn't realize until this PR that `Declarations.yaml` is still directly in use, by some python/autograd codegen. Rather than keep that yaml byte-for-byte compatible, I just updated the callsites in the autograd codegen to work with namespaces. In the NEXT pr, I try to clean up some of the autograd codegen to stop using raw strings to match against C++ types. Test Plan: Imported from OSS Reviewed By: bhosmer Differential Revision: D27708349 Pulled By: bdhirsh fbshipit-source-id: 56a4f81fc101795bcb9ee1f722121480fb2356ad
111 lines
4.5 KiB
Python
111 lines
4.5 KiB
Python
from tools.codegen.model import (Argument, FunctionSchema, Return,
|
|
SelfArgument, TensorOptionsArguments, Type,
|
|
assert_never)
|
|
|
|
from tools.codegen.api.types import (ArgName, BaseCType, Binding,
|
|
ConstRefCType, CType, MutRefCType, ListCType,
|
|
OptionalCType, tensorT, scalarT, layoutT,
|
|
deviceT, boolT, scalarTypeT)
|
|
from tools.codegen.api import cpp
|
|
|
|
from typing import Union, Sequence, List, Optional
|
|
|
|
# This file describes the translation of JIT schema to the native functions API.
|
|
# This looks a lot like the C++ API (which makes historical sense, because the
|
|
# idea was you wrote native functions to implement functions in the C++ API),
|
|
# but over time we have evolved the C++ API without actually changing our
|
|
# native:: kernels. The intention is to make native API and dispatcher API
|
|
# line up as closely as possible, since this results in the least overhead
|
|
# (no translation is needed from dispatcher API to native API).
|
|
|
|
def name(func: FunctionSchema) -> str:
|
|
name = str(func.name.name)
|
|
# TODO: delete this!
|
|
if func.is_out_fn():
|
|
name += '_out'
|
|
if func.name.overload_name:
|
|
name += f'_{func.name.overload_name}'
|
|
return name
|
|
|
|
def argumenttype_type(t: Type, *, mutable: bool, binds: ArgName) -> CType:
|
|
if str(t) == 'Tensor?':
|
|
tensor_type: OptionalCType = OptionalCType(BaseCType(tensorT, binds))
|
|
if mutable:
|
|
return MutRefCType(tensor_type)
|
|
else:
|
|
return ConstRefCType(tensor_type)
|
|
elif str(t) == 'Tensor?[]':
|
|
return ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT, binds))))
|
|
elif str(t) == 'Scalar':
|
|
return ConstRefCType(BaseCType(scalarT, binds))
|
|
elif str(t) == 'Scalar?':
|
|
return ConstRefCType(OptionalCType(BaseCType(scalarT, binds)))
|
|
return cpp.argumenttype_type(t, mutable=mutable, binds=binds)
|
|
|
|
def returns_type(rs: Sequence[Return]) -> CType:
|
|
return cpp.returns_type(rs)
|
|
|
|
def argument_type(a: Argument, *, binds: ArgName) -> CType:
|
|
return argumenttype_type(a.type, mutable=a.is_write, binds=binds)
|
|
|
|
def argument(a: Union[Argument, SelfArgument, TensorOptionsArguments], *, is_out: bool) -> List[Binding]:
|
|
# Ideally, we NEVER default native functions. However, there are a number
|
|
# of functions that call native:: directly and rely on the defaulting
|
|
# existing. So for BC, we generate defaults for non-out variants (but not
|
|
# for out variants, where it is impossible to generate an appropriate
|
|
# default)
|
|
should_default = not is_out
|
|
if isinstance(a, Argument):
|
|
default: Optional[str] = None
|
|
if should_default and a.default is not None:
|
|
default = cpp.default_expr(a.default, a.type)
|
|
return [Binding(
|
|
ctype=argument_type(a, binds=a.name),
|
|
name=a.name,
|
|
default=default,
|
|
argument=a,
|
|
)]
|
|
elif isinstance(a, SelfArgument):
|
|
# Erase SelfArgument from the distinction
|
|
return argument(a.argument, is_out=is_out)
|
|
elif isinstance(a, TensorOptionsArguments):
|
|
default = None
|
|
if should_default:
|
|
default = '{}'
|
|
# TODO: Not sure why the arguments assigned here are for
|
|
# TensorOptionsArguments and not the constituent pieces. It seems
|
|
# to matter
|
|
return [
|
|
Binding(
|
|
ctype=OptionalCType(BaseCType(scalarTypeT, 'dtype')),
|
|
name='dtype',
|
|
default=default,
|
|
argument=a,
|
|
),
|
|
Binding(
|
|
ctype=OptionalCType(BaseCType(layoutT, 'layout')),
|
|
name='layout',
|
|
default=default,
|
|
argument=a,
|
|
),
|
|
Binding(
|
|
ctype=OptionalCType(BaseCType(deviceT, 'device')),
|
|
name='device',
|
|
default=default,
|
|
argument=a,
|
|
),
|
|
Binding(
|
|
ctype=OptionalCType(BaseCType(boolT, 'pin_memory')),
|
|
name='pin_memory',
|
|
default=default,
|
|
argument=a,
|
|
)]
|
|
else:
|
|
assert_never(a)
|
|
|
|
def arguments(func: FunctionSchema) -> List[Binding]:
|
|
args: List[Union[Argument, TensorOptionsArguments, SelfArgument]] = []
|
|
args.extend(func.arguments.non_out)
|
|
args.extend(func.arguments.out)
|
|
return [r for arg in args for r in argument(arg, is_out=func.is_out_fn())]
|