def gen_unstructured(f: NativeFunction, backend_index: BackendIndex) -> Optional[str]: sig = kernel_signature(f, backend_index) metadata = backend_index.get_kernel(f) if metadata is None: return None if "legacy::" in metadata.kernel: return None else: prefix = '' if backend_index.external else 'TORCH_API ' return f"{prefix}{sig.decl(name=metadata.kernel)};"
def gen_structured(g: NativeFunctionsGroup, backend_index: BackendIndex) -> List[str]: meta_name = meta.name(g) out_args = structured.impl_arguments(g) metadata = backend_index.get_kernel(g) if metadata is None: return [] prefix = '' if backend_index.external else 'TORCH_API ' return [ f"""\ struct {prefix}structured_{metadata.kernel} : public at::meta::{meta_name} {{ void impl({', '.join(a.decl() for a in out_args)}); }}; """ ]
def compute_native_function_declaration( g: Union[NativeFunctionsGroup, NativeFunction], backend_index: BackendIndex) -> List[str]: metadata = backend_index.get_kernel(g) if isinstance(g, NativeFunctionsGroup): if metadata is not None and metadata.structured: if backend_index.external: # Structured hasn't been tested with external backends yet. raise AssertionError( "Structured external backend functions are not implemented yet." ) else: return gen_structured(g, backend_index) else: return list( mapMaybe(lambda f: gen_unstructured(f, backend_index), g.functions())) else: x = gen_unstructured(g, backend_index) return [] if x is None else [x]