mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
02c4d877b4
Add codegen infrastructure to generate IR nodes for non-native ops.
The proposed change is to add a `non_native` key to the `{backend}_native_functions.yaml` file that contains schema definitions similar to what is found in `native_functions.yaml`. e.g.
```
non_native:
...
- func: expand(Tensor input, int[] size, bool is_scalar_expand) -> Tensor
...
```
these definitions are parsed into a `LazyIrSchema` that can be used for generating IR nodes using `GenLazyIR`.
Fixes #74628
CC: @wconstab @desertfire @henrytwo
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76535
Approved by: https://github.com/wconstab
20 lines
801 B
Python
20 lines
801 B
Python
from .lazy_ir import GenLazyIR as GenLazyIR
|
|
from .lazy_ir import GenLazyShapeInferenceDefinition as GenLazyShapeInferenceDefinition
|
|
from .lazy_ir import GenLazyNativeFuncDefinition as GenLazyNativeFuncDefinition
|
|
from .lazy_ir import (
|
|
generate_non_native_lazy_ir_nodes as generate_non_native_lazy_ir_nodes,
|
|
)
|
|
from .register_dispatch_key import (
|
|
RegisterDispatchKey as RegisterDispatchKey,
|
|
gen_registration_helpers as gen_registration_helpers,
|
|
gen_registration_headers as gen_registration_headers,
|
|
)
|
|
from .native_functions import (
|
|
compute_native_function_declaration as compute_native_function_declaration,
|
|
)
|
|
from .ufunc import (
|
|
compute_ufunc_cuda as compute_ufunc_cuda,
|
|
compute_ufunc_cpu as compute_ufunc_cpu,
|
|
compute_ufunc_cpu_kernel as compute_ufunc_cpu_kernel,
|
|
)
|