Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,635 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
from collections import Counter, defaultdict, namedtuple
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import yaml
|
||||
|
||||
import torchgen.api.dispatcher as dispatcher
|
||||
import torchgen.dest as dest
|
||||
from torchgen.api.types import DispatcherSignature
|
||||
from torchgen.code_template import CodeTemplate
|
||||
from torchgen.context import native_function_manager
|
||||
from torchgen.gen import get_grouped_native_functions, parse_native_yaml
|
||||
from torchgen.model import (
|
||||
BackendIndex,
|
||||
BackendMetadata,
|
||||
DispatchKey,
|
||||
NativeFunction,
|
||||
NativeFunctionsGroup,
|
||||
OperatorName,
|
||||
)
|
||||
from torchgen.selective_build.selector import SelectiveBuilder
|
||||
from torchgen.utils import concatMap, context, FileManager, NamespaceHelper, Target
|
||||
from torchgen.yaml_utils import YamlLoader
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
|
||||
# Parses the external backend's yaml, and adds a new BackendIndex for the backend's dispatch key.
|
||||
# Returns a Tuple of (backend_key, autograd_key, cpp_namespace, updated BackendIndex mapping)
|
||||
ParsedExternalYaml = namedtuple(
|
||||
"ParsedExternalYaml",
|
||||
["backend_key", "autograd_key", "class_name", "cpp_namespace", "backend_indices"],
|
||||
)
|
||||
|
||||
|
||||
def parse_backend_yaml(
|
||||
backend_yaml_path: str,
|
||||
grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
) -> ParsedExternalYaml:
|
||||
native_functions_map: dict[OperatorName, NativeFunction] = {
|
||||
f.func.name: f
|
||||
for f in concatMap(
|
||||
lambda f: [f] if isinstance(f, NativeFunction) else list(f.functions()),
|
||||
grouped_native_functions,
|
||||
)
|
||||
}
|
||||
|
||||
with open(backend_yaml_path) as f:
|
||||
yaml_values = yaml.load(f, Loader=YamlLoader)
|
||||
if not isinstance(yaml_values, dict):
|
||||
raise AssertionError(
|
||||
f"Expected yaml_values to be a dict, got {type(yaml_values)}"
|
||||
)
|
||||
|
||||
valid_keys = [
|
||||
"backend",
|
||||
"class_name",
|
||||
"cpp_namespace",
|
||||
"extra_headers",
|
||||
"supported",
|
||||
"autograd",
|
||||
"full_codegen",
|
||||
"non_native",
|
||||
"ir_gen",
|
||||
"symint",
|
||||
]
|
||||
|
||||
backend = yaml_values.pop("backend", None)
|
||||
if backend is None:
|
||||
raise AssertionError('You must provide a value for "backend"')
|
||||
|
||||
class_name = yaml_values.pop("class_name", None)
|
||||
|
||||
cpp_namespace = yaml_values.pop("cpp_namespace", None)
|
||||
if cpp_namespace is None:
|
||||
raise AssertionError('You must provide a value for "cpp_namespace"')
|
||||
|
||||
# Mostly just defaulting to false to stick with LazyTensor convention.
|
||||
use_out_as_primary = yaml_values.pop("use_out_as_primary", False)
|
||||
if not isinstance(use_out_as_primary, bool):
|
||||
raise AssertionError(
|
||||
f"You must provide either True or False for use_out_as_primary. Provided: {use_out_as_primary}"
|
||||
)
|
||||
|
||||
use_device_guard = yaml_values.pop("device_guard", False)
|
||||
if not isinstance(use_device_guard, bool):
|
||||
raise AssertionError(
|
||||
f"You must provide either True or False for device_guard. Provided: {use_device_guard}"
|
||||
)
|
||||
|
||||
supported = yaml_values.pop("supported", [])
|
||||
if supported is None:
|
||||
supported = [] # Allow an empty list of supported ops
|
||||
if not isinstance(supported, list):
|
||||
raise AssertionError(
|
||||
f'expected "supported" to be a list, but got: {supported} (of type {type(supported)})'
|
||||
)
|
||||
|
||||
symint = yaml_values.pop("symint", [])
|
||||
if symint is None:
|
||||
symint = [] # Allow an empty list of symint ops
|
||||
if not isinstance(symint, list):
|
||||
raise AssertionError(
|
||||
f'expected "symint" to be a list, but got: {symint} (of type {type(symint)})'
|
||||
)
|
||||
symint_set = set(symint)
|
||||
|
||||
supported_autograd = yaml_values.pop("autograd", [])
|
||||
if not isinstance(supported_autograd, list):
|
||||
raise AssertionError(
|
||||
f'expected "autograd" to be a list, but got: {supported_autograd}'
|
||||
)
|
||||
|
||||
# full_codegen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py
|
||||
full_codegen = yaml_values.pop("full_codegen", [])
|
||||
supported.extend(full_codegen)
|
||||
|
||||
# non_native is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py
|
||||
yaml_values.pop("non_native", {})
|
||||
|
||||
# ir_gen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py
|
||||
yaml_values.pop("ir_gen", {})
|
||||
|
||||
if len(yaml_values.keys()) != 0:
|
||||
raise AssertionError(
|
||||
f"{backend_yaml_path} contains unexpected keys: {', '.join(yaml_values.keys())}. "
|
||||
f"Only the following keys are supported: {', '.join(valid_keys)}"
|
||||
)
|
||||
|
||||
def create_backend_index(
|
||||
backend_ops: list[str],
|
||||
symint_ops: set[str],
|
||||
dispatch_key: DispatchKey,
|
||||
*,
|
||||
use_out_as_primary: bool,
|
||||
use_device_guard: bool,
|
||||
) -> BackendIndex:
|
||||
metadata: dict[OperatorName, BackendMetadata] = {}
|
||||
for op in backend_ops:
|
||||
op_name = OperatorName.parse(op)
|
||||
if op_name not in native_functions_map:
|
||||
raise AssertionError(f"Found an invalid operator name: {op_name}")
|
||||
# See Note [External Backends Follow Dispatcher API]
|
||||
kernel_name = dispatcher.name(native_functions_map[op_name].func)
|
||||
if op in symint_ops:
|
||||
kernel_name += "_symint"
|
||||
# TODO: allow structured external backends later.
|
||||
m = BackendMetadata(
|
||||
kernel=kernel_name, structured=False, cpp_namespace=cpp_namespace
|
||||
)
|
||||
metadata[op_name] = m
|
||||
return BackendIndex(
|
||||
dispatch_key=dispatch_key,
|
||||
use_out_as_primary=use_out_as_primary,
|
||||
external=True,
|
||||
device_guard=use_device_guard,
|
||||
index=metadata,
|
||||
)
|
||||
|
||||
backend_key: DispatchKey | None = None
|
||||
if len(supported) > 0:
|
||||
with context(
|
||||
lambda: f'The provided value for "backend" must be a valid DispatchKey, but got {backend}.'
|
||||
):
|
||||
backend_key = DispatchKey.parse(backend)
|
||||
|
||||
backend_idx = create_backend_index(
|
||||
supported,
|
||||
symint_set,
|
||||
backend_key,
|
||||
use_out_as_primary=use_out_as_primary,
|
||||
use_device_guard=use_device_guard,
|
||||
)
|
||||
if backend_key in backend_indices:
|
||||
raise AssertionError(f"Duplicate backend key: {backend_key}")
|
||||
backend_indices[backend_key] = backend_idx
|
||||
|
||||
autograd_key: DispatchKey | None = None
|
||||
if len(supported_autograd) > 0:
|
||||
with context(
|
||||
lambda: f'The "autograd" key was specified, which indicates that you would like to override \
|
||||
the behavior of autograd for some operators on your backend. However "Autograd{backend}" is not a valid DispatchKey.'
|
||||
):
|
||||
autograd_key = DispatchKey.parse(f"Autograd{backend}")
|
||||
|
||||
autograd_idx = create_backend_index(
|
||||
supported_autograd,
|
||||
symint_set,
|
||||
autograd_key,
|
||||
use_out_as_primary=use_out_as_primary,
|
||||
use_device_guard=use_device_guard,
|
||||
)
|
||||
if autograd_key in backend_indices:
|
||||
raise AssertionError(f"Duplicate autograd key: {autograd_key}")
|
||||
backend_indices[autograd_key] = autograd_idx
|
||||
|
||||
for g in grouped_native_functions:
|
||||
if isinstance(g, NativeFunction):
|
||||
forward_kernels = (
|
||||
[]
|
||||
if backend_key is None
|
||||
else [
|
||||
m
|
||||
for m in [backend_indices[backend_key].get_kernel(g)]
|
||||
if m is not None
|
||||
]
|
||||
)
|
||||
backward_kernels = (
|
||||
[]
|
||||
if autograd_key is None
|
||||
else [
|
||||
m
|
||||
for m in [backend_indices[autograd_key].get_kernel(g)]
|
||||
if m is not None
|
||||
]
|
||||
)
|
||||
else:
|
||||
forward_kernels = (
|
||||
[]
|
||||
if backend_key is None
|
||||
else [
|
||||
m
|
||||
for m in [
|
||||
backend_indices[backend_key].get_kernel(f)
|
||||
for f in g.functions()
|
||||
]
|
||||
if m is not None
|
||||
]
|
||||
)
|
||||
backward_kernels = (
|
||||
[]
|
||||
if autograd_key is None
|
||||
else [
|
||||
m
|
||||
for m in [
|
||||
backend_indices[autograd_key].get_kernel(f)
|
||||
for f in g.functions()
|
||||
]
|
||||
if m is not None
|
||||
]
|
||||
)
|
||||
|
||||
forward_kernels = [f for f in forward_kernels if f is not None]
|
||||
backward_kernels = [f for f in backward_kernels if f is not None]
|
||||
if not (len(forward_kernels) == 0 or len(backward_kernels) == 0):
|
||||
raise AssertionError(
|
||||
f"Currently, all variants of an op must either be registered to a backend key, "
|
||||
f"or to a backend's autograd key. They cannot be mix and matched. "
|
||||
f"If this is something you need, feel free to create an issue! "
|
||||
f'{forward_kernels[0].kernel} is listed under "supported", '
|
||||
f'but {backward_kernels[0].kernel} is listed under "autograd".'
|
||||
)
|
||||
|
||||
return ParsedExternalYaml(
|
||||
backend_key, autograd_key, class_name, cpp_namespace, backend_indices
|
||||
)
|
||||
|
||||
|
||||
def error_on_missing_kernels(
|
||||
native_functions: Sequence[NativeFunction],
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
backend_key: DispatchKey,
|
||||
autograd_key: DispatchKey | None,
|
||||
class_name: str,
|
||||
kernel_defn_file_path: str,
|
||||
full_codegen: list[OperatorName] | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
with open(kernel_defn_file_path) as f:
|
||||
backend_defns = f.read()
|
||||
except OSError as e:
|
||||
raise AssertionError(
|
||||
f"Unable to read from the specified impl_path file: {kernel_defn_file_path}"
|
||||
) from e
|
||||
|
||||
if full_codegen is None:
|
||||
full_codegen = []
|
||||
|
||||
indices = [backend_indices[backend_key].index] + (
|
||||
[] if autograd_key is None else [backend_indices[autograd_key].index]
|
||||
)
|
||||
# Quick mapping from each OperatorName used by the external backend
|
||||
# to its backend kernel name
|
||||
expected_backend_op_names: dict[OperatorName, str] = dict(
|
||||
list(
|
||||
concatMap(
|
||||
lambda index: [
|
||||
(op_name, metadata.kernel) for op_name, metadata in index.items()
|
||||
],
|
||||
indices,
|
||||
)
|
||||
)
|
||||
)
|
||||
expected_backend_native_funcs: list[NativeFunction] = [
|
||||
f
|
||||
for f in native_functions
|
||||
if f.func.name in expected_backend_op_names and f.func.name not in full_codegen
|
||||
]
|
||||
expected_backend_kernel_name_counts: dict[str, list[NativeFunction]] = defaultdict(
|
||||
list
|
||||
)
|
||||
for native_f in expected_backend_native_funcs:
|
||||
expected_backend_kernel_name_counts[
|
||||
expected_backend_op_names[native_f.func.name]
|
||||
].append(native_f)
|
||||
|
||||
# This just looks for lines containing "foo(", and assumes that the kernel foo has been implemented.
|
||||
# It might cause false negatives (we won't catch all cases), but that's ok - if we catch a missing kernel
|
||||
# here, then we get a nicer error message. If we miss it, you get a linker error.
|
||||
kernel_defn_regex = rf"(.*){class_name}::\s*([\w\d]*)\("
|
||||
actual_backend_kernel_name_counts = Counter(
|
||||
# A bit unwieldy (this could probably be moved into regex),
|
||||
# but we don't want to include kernel names that come from function calls,
|
||||
# like "return torch_xla::XLANativeFunctions::empty_strided_symint(...)".
|
||||
# Easy check is to ignore any lines with colons before the class name.
|
||||
[
|
||||
y
|
||||
for (x, y) in re.findall(kernel_defn_regex, backend_defns)
|
||||
if not x.endswith(":")
|
||||
]
|
||||
)
|
||||
|
||||
missing_kernels_err_msg = ""
|
||||
for expected_name, funcs in expected_backend_kernel_name_counts.items():
|
||||
expected_overload_count = len(funcs)
|
||||
actual_overload_count = actual_backend_kernel_name_counts[expected_name]
|
||||
if expected_overload_count != actual_overload_count:
|
||||
|
||||
def create_decl(f: NativeFunction) -> str:
|
||||
with native_function_manager(f):
|
||||
return DispatcherSignature.from_schema(f.func).decl()
|
||||
|
||||
expected_schemas_str = "\n".join([create_decl(f) for f in funcs])
|
||||
missing_kernels_err_msg += f"""
|
||||
{class_name} is missing a kernel definition for {expected_name}. We found {actual_overload_count} kernel(s) with that name,
|
||||
but expected {expected_overload_count} kernel(s). The expected function schemas for the missing operator are:
|
||||
{expected_schemas_str}
|
||||
|
||||
"""
|
||||
if missing_kernels_err_msg != "":
|
||||
raise AssertionError(missing_kernels_err_msg)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Generate backend stub files")
|
||||
parser.add_argument(
|
||||
"-s",
|
||||
"--source-yaml",
|
||||
"--source_yaml",
|
||||
help="path to source yaml file containing operator external definitions",
|
||||
)
|
||||
parser.add_argument("-o", "--output-dir", "--output_dir", help="output directory")
|
||||
parser.add_argument(
|
||||
"--dry-run", "--dry_run", type=bool, default=False, help="output directory"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--impl-path",
|
||||
"--impl_path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="path to the source C++ file containing kernel definitions",
|
||||
)
|
||||
options = parser.parse_args()
|
||||
|
||||
run(options.source_yaml, options.output_dir, options.dry_run, options.impl_path)
|
||||
|
||||
|
||||
def gen_dispatchkey_nativefunc_headers(
|
||||
fm: FileManager,
|
||||
class_name: str,
|
||||
cpp_namespace: str,
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],
|
||||
backend_dispatch_key: DispatchKey,
|
||||
autograd_dispatch_key: DispatchKey | None,
|
||||
backend_name: str = "",
|
||||
) -> None:
|
||||
if class_name is None:
|
||||
raise AssertionError("class_name must not be None")
|
||||
generated_comment = (
|
||||
"Autogenerated file by gen_backend_stubs.py. Do not edit directly!"
|
||||
)
|
||||
|
||||
# Convert to a set first to remove duplicate kernel names.
|
||||
# Backends are allowed to repeat kernel names; only generate the declaration once!
|
||||
# Sort for deterministic output.
|
||||
backend_declarations = sorted(
|
||||
set(
|
||||
concatMap(
|
||||
lambda f: dest.compute_native_function_declaration(
|
||||
f, backend_indices[backend_dispatch_key]
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
)
|
||||
)
|
||||
autograd_declarations = sorted(
|
||||
set(
|
||||
concatMap(
|
||||
lambda f: []
|
||||
if autograd_dispatch_key is None
|
||||
else dest.compute_native_function_declaration(
|
||||
f, backend_indices[autograd_dispatch_key]
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
ns_helper = NamespaceHelper(cpp_namespace)
|
||||
fm.write_with_template(
|
||||
f"{backend_dispatch_key}NativeFunctions.h",
|
||||
"DispatchKeyNativeFunctions.h",
|
||||
lambda: {
|
||||
"generated_comment": generated_comment,
|
||||
"namespace_prologue": ns_helper.prologue,
|
||||
"class_name": class_name,
|
||||
"namespace_epilogue": ns_helper.epilogue,
|
||||
"dispatch_declarations": backend_declarations + autograd_declarations,
|
||||
"BackendName": backend_name,
|
||||
"DispatchKey": backend_dispatch_key,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def gen_dispatcher_registrations(
|
||||
fm: FileManager,
|
||||
output_dir: str,
|
||||
class_name: str,
|
||||
backend_indices: dict[DispatchKey, BackendIndex],
|
||||
grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],
|
||||
backend_dispatch_key: DispatchKey,
|
||||
dispatch_key: DispatchKey,
|
||||
selector: SelectiveBuilder,
|
||||
# build_in_tree is true for lazy TS backend and affects include paths, not used for external backends
|
||||
build_in_tree: bool = False,
|
||||
per_operator_headers: bool = False,
|
||||
backend_name: str = "",
|
||||
eager_registration: bool = True,
|
||||
) -> None:
|
||||
headers = [
|
||||
f"{output_dir}/{backend_dispatch_key}NativeFunctions.h",
|
||||
]
|
||||
if build_in_tree:
|
||||
external_backend_headers_str = "\n".join(f"#include <{h}>" for h in headers)
|
||||
else:
|
||||
external_backend_headers_str = "\n".join(f'#include "{h}"' for h in headers)
|
||||
|
||||
if class_name is None:
|
||||
raise AssertionError("class_name must not be None")
|
||||
backend_index = backend_indices[dispatch_key]
|
||||
|
||||
dispatch_registrations_body = list(
|
||||
concatMap(
|
||||
dest.RegisterDispatchKey(
|
||||
backend_index,
|
||||
Target.REGISTRATION,
|
||||
selector,
|
||||
rocm=False,
|
||||
symint=True,
|
||||
class_method_name=f"{class_name}",
|
||||
skip_dispatcher_op_registration=False,
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
)
|
||||
newline = "\n"
|
||||
ns_helper = NamespaceHelper(namespace_str="at")
|
||||
deferred_dispatch_registrations = ""
|
||||
static_init_dispatch_registrations = ""
|
||||
if eager_registration:
|
||||
static_template = CodeTemplate(
|
||||
"""\
|
||||
TORCH_LIBRARY_IMPL(aten, $dispatch_key, m) {
|
||||
$dispatch_registrations_body
|
||||
}"""
|
||||
)
|
||||
static_init_dispatch_registrations = static_template.substitute(
|
||||
dispatch_key=dispatch_key,
|
||||
dispatch_registrations_body=dispatch_registrations_body,
|
||||
)
|
||||
else:
|
||||
deferred_template = CodeTemplate(
|
||||
"""\
|
||||
TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions();
|
||||
TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() {
|
||||
static auto m = MAKE_TORCH_LIBRARY_IMPL(aten, $dispatch_key);
|
||||
$dispatch_registrations_body
|
||||
}"""
|
||||
)
|
||||
deferred_dispatch_registrations = deferred_template.substitute(
|
||||
backend_name=backend_name,
|
||||
dispatch_key=dispatch_key,
|
||||
dispatch_registrations_body=dispatch_registrations_body,
|
||||
)
|
||||
|
||||
fm.write_with_template(
|
||||
f"Register{dispatch_key}.cpp",
|
||||
"RegisterDispatchKey.cpp",
|
||||
lambda: {
|
||||
"extra_cuda_headers": "",
|
||||
"external_backend_headers": external_backend_headers_str,
|
||||
"ops_headers": "#include <ATen/Functions.h>"
|
||||
if not per_operator_headers
|
||||
else "",
|
||||
"DispatchKey": dispatch_key,
|
||||
"dispatch_namespace": dispatch_key.lower(),
|
||||
"dispatch_headers": dest.gen_registration_headers(
|
||||
backend_index, per_operator_headers=per_operator_headers, rocm=False
|
||||
),
|
||||
"dispatch_helpers": dest.gen_registration_helpers(backend_index),
|
||||
"dispatch_definitions": fm.substitute_with_template(
|
||||
"RegisterDispatchDefinitions.ini",
|
||||
lambda: {
|
||||
"ns_prologue": ns_helper.prologue,
|
||||
"ns_epilogue": ns_helper.epilogue,
|
||||
"static_init_dispatch_registrations": static_init_dispatch_registrations,
|
||||
"deferred_dispatch_registrations": deferred_dispatch_registrations,
|
||||
"dispatch_namespace": dispatch_key.lower(),
|
||||
"dispatch_namespaced_definitions": "",
|
||||
"dispatch_anonymous_definitions": list(
|
||||
concatMap(
|
||||
dest.RegisterDispatchKey(
|
||||
backend_index,
|
||||
Target.ANONYMOUS_DEFINITION,
|
||||
selector,
|
||||
rocm=False,
|
||||
symint=True,
|
||||
class_method_name=f"{class_name}",
|
||||
skip_dispatcher_op_registration=False,
|
||||
),
|
||||
grouped_native_functions,
|
||||
)
|
||||
),
|
||||
},
|
||||
).split(newline),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def run(
|
||||
source_yaml: str, output_dir: str, dry_run: bool, impl_path: str | None = None
|
||||
) -> None:
|
||||
# Assumes that this file lives at torchgen/gen_backend_stubs.py
|
||||
root = Path(__file__).absolute().parent.parent
|
||||
common_dir = os.path.join(root, "aten/src") # Assumes root is pytorch_root
|
||||
if not os.path.exists(common_dir): # This file is out-of-tree.
|
||||
common_dir = os.path.join(root, "torchgen/packaged")
|
||||
|
||||
template_dir = os.path.join(common_dir, "ATen/templates")
|
||||
|
||||
def make_file_manager(install_dir: str) -> FileManager:
|
||||
return FileManager(
|
||||
install_dir=install_dir, template_dir=template_dir, dry_run=dry_run
|
||||
)
|
||||
|
||||
fm = make_file_manager(output_dir)
|
||||
|
||||
native_yaml_path = os.path.join(common_dir, "ATen/native/native_functions.yaml")
|
||||
tags_yaml_path = os.path.join(common_dir, "ATen/native/tags.yaml")
|
||||
parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)
|
||||
native_functions, backend_indices = (
|
||||
parsed_yaml.native_functions,
|
||||
parsed_yaml.backend_indices,
|
||||
)
|
||||
grouped_native_functions = get_grouped_native_functions(native_functions)
|
||||
parsed_backend_yaml = parse_backend_yaml(
|
||||
source_yaml, grouped_native_functions, backend_indices
|
||||
)
|
||||
backend_key = parsed_backend_yaml.backend_key
|
||||
autograd_key = parsed_backend_yaml.autograd_key
|
||||
cpp_namespace = parsed_backend_yaml.cpp_namespace
|
||||
class_name = parsed_backend_yaml.class_name
|
||||
backend_indices = parsed_backend_yaml.backend_indices
|
||||
|
||||
selector = SelectiveBuilder.get_nop_selector()
|
||||
|
||||
if backend_key is None:
|
||||
# This could be useful if a backend wants to quickly set up a noop yaml file but doesn't have any kernels ready yet.
|
||||
return
|
||||
|
||||
if class_name is None:
|
||||
# class_name is an optional argument to backend yaml file.
|
||||
# if specified it allows an external backend to override
|
||||
# the name of the class that all generated kernel definitions live under.
|
||||
# if not specified, its value is given as native_function_class_name.
|
||||
class_name = backend_indices[backend_key].native_function_class_name()
|
||||
if class_name is None:
|
||||
raise AssertionError("class_name must not be None")
|
||||
|
||||
if impl_path is not None:
|
||||
error_on_missing_kernels(
|
||||
native_functions,
|
||||
backend_indices,
|
||||
backend_key,
|
||||
autograd_key,
|
||||
class_name,
|
||||
impl_path,
|
||||
)
|
||||
|
||||
gen_dispatchkey_nativefunc_headers(
|
||||
fm,
|
||||
class_name,
|
||||
cpp_namespace,
|
||||
backend_indices,
|
||||
grouped_native_functions,
|
||||
backend_key,
|
||||
autograd_key,
|
||||
)
|
||||
|
||||
for dispatch_key in (
|
||||
[backend_key] if autograd_key is None else [backend_key, autograd_key]
|
||||
):
|
||||
gen_dispatcher_registrations(
|
||||
fm,
|
||||
output_dir,
|
||||
class_name,
|
||||
backend_indices,
|
||||
grouped_native_functions,
|
||||
backend_key,
|
||||
dispatch_key,
|
||||
selector,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user