pytorch

1
import argparse
2
import os
3
import pathlib
4
import re
5
from collections import Counter, defaultdict, namedtuple
6
from typing import Dict, List, Optional, Sequence, Set, Union
7

8
import yaml
9

10
import torchgen.api.dispatcher as dispatcher
11
import torchgen.dest as dest
12
from torchgen.api.types import DispatcherSignature
13
from torchgen.code_template import CodeTemplate
14
from torchgen.context import native_function_manager
15
from torchgen.gen import get_grouped_native_functions, parse_native_yaml
16
from torchgen.model import (
17
    BackendIndex,
18
    BackendMetadata,
19
    DispatchKey,
20
    NativeFunction,
21
    NativeFunctionsGroup,
22
    OperatorName,
23
)
24
from torchgen.selective_build.selector import SelectiveBuilder
25
from torchgen.utils import concatMap, context, FileManager, NamespaceHelper, Target
26
from torchgen.yaml_utils import YamlLoader
27

28

29
# Parses the external backend's yaml, and adds a new BackendIndex for the backend's dispatch key.
30
# Returns a Tuple of (backend_key, autograd_key, cpp_namespace, updated BackendIndex mapping)
31
ParsedExternalYaml = namedtuple(
32
    "ParsedExternalYaml",
33
    ["backend_key", "autograd_key", "class_name", "cpp_namespace", "backend_indices"],
34
)
35

36

37
def parse_backend_yaml(
38
    backend_yaml_path: str,
39
    grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]],
40
    backend_indices: Dict[DispatchKey, BackendIndex],
41
) -> ParsedExternalYaml:
42
    native_functions_map: Dict[OperatorName, NativeFunction] = {
43
        f.func.name: f
44
        for f in concatMap(
45
            lambda f: [f] if isinstance(f, NativeFunction) else list(f.functions()),
46
            grouped_native_functions,
47
        )
48
    }
49

50
    with open(backend_yaml_path) as f:
51
        yaml_values = yaml.load(f, Loader=YamlLoader)
52
    assert isinstance(yaml_values, dict)
53

54
    valid_keys = [
55
        "backend",
56
        "class_name",
57
        "cpp_namespace",
58
        "extra_headers",
59
        "supported",
60
        "autograd",
61
        "full_codegen",
62
        "non_native",
63
        "ir_gen",
64
        "symint",
65
    ]
66

67
    backend = yaml_values.pop("backend", None)
68
    assert backend is not None, 'You must provide a value for "backend"'
69

70
    class_name = yaml_values.pop("class_name", None)
71

72
    cpp_namespace = yaml_values.pop("cpp_namespace", None)
73
    assert cpp_namespace is not None, 'You must provide a value for "cpp_namespace"'
74

75
    # Mostly just defaulting to false to stick with LazyTensor convention.
76
    use_out_as_primary = yaml_values.pop("use_out_as_primary", False)
77
    assert isinstance(
78
        use_out_as_primary, bool
79
    ), f"You must provide either True or False for use_out_as_primary. Provided: {use_out_as_primary}"
80

81
    use_device_guard = yaml_values.pop("device_guard", False)
82
    assert isinstance(
83
        use_device_guard, bool
84
    ), f"You must provide either True or False for device_guard. Provided: {use_device_guard}"
85

86
    supported = yaml_values.pop("supported", [])
87
    if supported is None:
88
        supported = []  # Allow an empty list of supported ops
89
    assert isinstance(
90
        supported, list
91
    ), f'expected "supported" to be a list, but got: {supported} (of type {type(supported)})'
92

93
    symint = yaml_values.pop("symint", [])
94
    if symint is None:
95
        symint = []  # Allow an empty list of symint ops
96
    assert isinstance(
97
        symint, list
98
    ), f'expected "symint" to be a list, but got: {supported} (of type {type(supported)})'
99
    symint_set = set(symint)
100

101
    supported_autograd = yaml_values.pop("autograd", [])
102
    assert isinstance(
103
        supported_autograd, list
104
    ), f'expected "autograd" to be a list, but got: {supported_autograd}'
105

106
    # full_codegen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py
107
    full_codegen = yaml_values.pop("full_codegen", [])
108
    supported.extend(full_codegen)
109

110
    # non_native is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py
111
    non_native = yaml_values.pop("non_native", {})
112

113
    # ir_gen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py
114
    _ = yaml_values.pop("ir_gen", {})
115

116
    assert (
117
        len(yaml_values.keys()) == 0
118
    ), f'{backend_yaml_path} contains unexpected keys: {", ".join(yaml_values.keys())}. \
119
Only the following keys are supported: {", ".join(valid_keys)}'
120

121
    def create_backend_index(
122
        backend_ops: List[str],
123
        symint_ops: Set[str],
124
        dispatch_key: DispatchKey,
125
        *,
126
        use_out_as_primary: bool,
127
        use_device_guard: bool,
128
    ) -> BackendIndex:
129
        metadata: Dict[OperatorName, BackendMetadata] = {}
130
        for op in backend_ops:
131
            op_name = OperatorName.parse(op)
132
            assert (
133
                op_name in native_functions_map
134
            ), f"Found an invalid operator name: {op_name}"
135
            # See Note [External Backends Follow Dispatcher API]
136
            kernel_name = dispatcher.name(native_functions_map[op_name].func)
137
            if op in symint_ops:
138
                kernel_name += "_symint"
139
            # TODO: allow structured external backends later.
140
            m = BackendMetadata(
141
                kernel=kernel_name, structured=False, cpp_namespace=cpp_namespace
142
            )
143
            metadata[op_name] = m
144
        return BackendIndex(
145
            dispatch_key=dispatch_key,
146
            use_out_as_primary=use_out_as_primary,
147
            external=True,
148
            device_guard=use_device_guard,
149
            index=metadata,
150
        )
151

152
    backend_key: Optional[DispatchKey] = None
153
    if len(supported) > 0:
154
        with context(
155
            lambda: f'The provided value for "backend" must be a valid DispatchKey, but got {backend}.'
156
        ):
157
            backend_key = DispatchKey.parse(backend)
158

159
        backend_idx = create_backend_index(
160
            supported,
161
            symint_set,
162
            backend_key,
163
            use_out_as_primary=use_out_as_primary,
164
            use_device_guard=use_device_guard,
165
        )
166
        assert backend_key not in backend_indices
167
        backend_indices[backend_key] = backend_idx
168

169
    autograd_key: Optional[DispatchKey] = None
170
    if len(supported_autograd) > 0:
171
        with context(
172
            lambda: f'The "autograd" key was specified, which indicates that you would like to override \
173
the behavior of autograd for some operators on your backend. However "Autograd{backend}" is not a valid DispatchKey.'
174
        ):
175
            autograd_key = DispatchKey.parse(f"Autograd{backend}")
176

177
        autograd_idx = create_backend_index(
178
            supported_autograd,
179
            symint_set,
180
            autograd_key,
181
            use_out_as_primary=use_out_as_primary,
182
            use_device_guard=use_device_guard,
183
        )
184
        assert autograd_key not in backend_indices
185
        backend_indices[autograd_key] = autograd_idx
186

187
    for g in grouped_native_functions:
188
        if isinstance(g, NativeFunction):
189
            forward_kernels = (
190
                []
191
                if backend_key is None
192
                else [
193
                    m
194
                    for m in [backend_indices[backend_key].get_kernel(g)]
195
                    if m is not None
196
                ]
197
            )
198
            backward_kernels = (
199
                []
200
                if autograd_key is None
201
                else [
202
                    m
203
                    for m in [backend_indices[autograd_key].get_kernel(g)]
204
                    if m is not None
205
                ]
206
            )
207
        else:
208
            forward_kernels = (
209
                []
210
                if backend_key is None
211
                else [
212
                    m
213
                    for m in [
214
                        backend_indices[backend_key].get_kernel(f)
215
                        for f in g.functions()
216
                    ]
217
                    if m is not None
218
                ]
219
            )
220
            backward_kernels = (
221
                []
222
                if autograd_key is None
223
                else [
224
                    m
225
                    for m in [
226
                        backend_indices[autograd_key].get_kernel(f)
227
                        for f in g.functions()
228
                    ]
229
                    if m is not None
230
                ]
231
            )
232

233
        forward_kernels = [f for f in forward_kernels if f is not None]
234
        backward_kernels = [f for f in backward_kernels if f is not None]
235
        assert (
236
            len(forward_kernels) == 0 or len(backward_kernels) == 0
237
        ), f'Currently, all variants of an op must either be registered to a backend key, or to a backend\'s \
238
autograd key. They cannot be mix and matched. If this is something you need, feel free to create an issue! \
239
{forward_kernels[0].kernel} is listed under "supported", but {backward_kernels[0].kernel} is listed under "autograd".'
240

241
    return ParsedExternalYaml(
242
        backend_key, autograd_key, class_name, cpp_namespace, backend_indices
243
    )
244

245

246
def error_on_missing_kernels(
247
    native_functions: Sequence[NativeFunction],
248
    backend_indices: Dict[DispatchKey, BackendIndex],
249
    backend_key: DispatchKey,
250
    autograd_key: Optional[DispatchKey],
251
    class_name: str,
252
    kernel_defn_file_path: str,
253
    full_codegen: Optional[List[OperatorName]] = None,
254
) -> None:
255
    try:
256
        with open(kernel_defn_file_path) as f:
257
            backend_defns = f.read()
258
    except OSError as e:
259
        raise AssertionError(
260
            f"Unable to read from the specified impl_path file: {kernel_defn_file_path}"
261
        ) from e
262

263
    if full_codegen is None:
264
        full_codegen = []
265

266
    indices = [backend_indices[backend_key].index] + (
267
        [] if autograd_key is None else [backend_indices[autograd_key].index]
268
    )
269
    # Quick mapping from each OperatorName used by the external backend
270
    # to its backend kernel name
271
    expected_backend_op_names: Dict[OperatorName, str] = dict(
272
        list(
273
            concatMap(
274
                lambda index: [
275
                    (op_name, metadata.kernel) for op_name, metadata in index.items()
276
                ],
277
                indices,
278
            )
279
        )
280
    )
281
    expected_backend_native_funcs: List[NativeFunction] = [
282
        f
283
        for f in native_functions
284
        if f.func.name in expected_backend_op_names.keys()
285
        and f.func.name not in full_codegen
286
    ]
287
    expected_backend_kernel_name_counts: Dict[str, List[NativeFunction]] = defaultdict(
288
        list
289
    )
290
    for native_f in expected_backend_native_funcs:
291
        expected_backend_kernel_name_counts[
292
            expected_backend_op_names[native_f.func.name]
293
        ].append(native_f)
294

295
    # This just looks for lines containing "foo(", and assumes that the kernel foo has been implemented.
296
    # It might cause false negatives (we won't catch all cases), but that's ok - if we catch a missing kernel
297
    # here, then we get a nicer error message. If we miss it, you get a linker error.
298
    kernel_defn_regex = rf"(.*){class_name}::\s*([\w\d]*)\("
299
    actual_backend_kernel_name_counts = Counter(
300
        # A bit unwieldy (this could probably be moved into regex),
301
        # but we don't want to include kernel names that come from function calls,
302
        # like "return torch_xla::XLANativeFunctions::empty_strided_symint(...)".
303
        # Easy check is to ignore any lines with colons before the class name.
304
        [
305
            y
306
            for (x, y) in re.findall(kernel_defn_regex, backend_defns)
307
            if not x.endswith(":")
308
        ]
309
    )
310

311
    missing_kernels_err_msg = ""
312
    for expected_name, funcs in expected_backend_kernel_name_counts.items():
313
        expected_overload_count = len(funcs)
314
        actual_overload_count = actual_backend_kernel_name_counts[expected_name]
315
        if expected_overload_count != actual_overload_count:
316

317
            def create_decl(f: NativeFunction) -> str:
318
                with native_function_manager(f):
319
                    return DispatcherSignature.from_schema(f.func).decl()
320

321
            expected_schemas_str = "\n".join([create_decl(f) for f in funcs])
322
            missing_kernels_err_msg += f"""
323
{class_name} is missing a kernel definition for {expected_name}. We found {actual_overload_count} kernel(s) with that name,
324
but expected {expected_overload_count} kernel(s). The expected function schemas for the missing operator are:
325
{expected_schemas_str}
326

327
"""
328
    assert missing_kernels_err_msg == "", missing_kernels_err_msg
329

330

331
def main() -> None:
332
    parser = argparse.ArgumentParser(description="Generate backend stub files")
333
    parser.add_argument(
334
        "-s",
335
        "--source-yaml",
336
        "--source_yaml",
337
        help="path to source yaml file containing operator external definitions",
338
    )
339
    parser.add_argument("-o", "--output-dir", "--output_dir", help="output directory")
340
    parser.add_argument(
341
        "--dry-run", "--dry_run", type=bool, default=False, help="output directory"
342
    )
343
    parser.add_argument(
344
        "--impl-path",
345
        "--impl_path",
346
        type=str,
347
        default=None,
348
        help="path to the source C++ file containing kernel definitions",
349
    )
350
    options = parser.parse_args()
351

352
    run(options.source_yaml, options.output_dir, options.dry_run, options.impl_path)
353

354

355
def gen_dispatchkey_nativefunc_headers(
356
    fm: FileManager,
357
    class_name: str,
358
    cpp_namespace: str,
359
    backend_indices: Dict[DispatchKey, BackendIndex],
360
    grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]],
361
    backend_dispatch_key: DispatchKey,
362
    autograd_dispatch_key: Optional[DispatchKey],
363
    backend_name: str = "",
364
) -> None:
365
    assert class_name is not None
366
    generated_comment = (
367
        "Autogenerated file by gen_backend_stubs.py. Do not edit directly!"
368
    )
369

370
    # Convert to a set first to remove duplicate kernel names.
371
    # Backends are allowed to repeat kernel names; only generate the declaration once!
372
    # Sort for deterministic output.
373
    backend_declarations = sorted(
374
        set(
375
            concatMap(
376
                lambda f: dest.compute_native_function_declaration(
377
                    f, backend_indices[backend_dispatch_key]
378
                ),
379
                grouped_native_functions,
380
            )
381
        )
382
    )
383
    autograd_declarations = sorted(
384
        set(
385
            concatMap(
386
                lambda f: []
387
                if autograd_dispatch_key is None
388
                else dest.compute_native_function_declaration(
389
                    f, backend_indices[autograd_dispatch_key]
390
                ),
391
                grouped_native_functions,
392
            )
393
        )
394
    )
395

396
    ns_helper = NamespaceHelper(cpp_namespace)
397
    fm.write_with_template(
398
        f"{backend_dispatch_key}NativeFunctions.h",
399
        "DispatchKeyNativeFunctions.h",
400
        lambda: {
401
            "generated_comment": generated_comment,
402
            "namespace_prologue": ns_helper.prologue,
403
            "class_name": class_name,
404
            "namespace_epilogue": ns_helper.epilogue,
405
            "dispatch_declarations": backend_declarations + autograd_declarations,
406
            "BackendName": backend_name,
407
            "DispatchKey": backend_dispatch_key,
408
        },
409
    )
410

411

412
def gen_dispatcher_registrations(
413
    fm: FileManager,
414
    output_dir: str,
415
    class_name: str,
416
    backend_indices: Dict[DispatchKey, BackendIndex],
417
    grouped_native_functions: Sequence[Union[NativeFunction, NativeFunctionsGroup]],
418
    backend_dispatch_key: DispatchKey,
419
    dispatch_key: DispatchKey,
420
    selector: "SelectiveBuilder",
421
    # build_in_tree is true for lazy TS backend and affects include paths, not used for external backends
422
    build_in_tree: bool = False,
423
    per_operator_headers: bool = False,
424
    backend_name: str = "",
425
    eager_registration: bool = True,
426
) -> None:
427
    headers = [
428
        f"{output_dir}/{backend_dispatch_key}NativeFunctions.h",
429
    ]
430
    if build_in_tree:
431
        external_backend_headers_str = "\n".join(f"#include <{h}>" for h in headers)
432
    else:
433
        external_backend_headers_str = "\n".join(f'#include "{h}"' for h in headers)
434

435
    assert class_name is not None
436
    backend_index = backend_indices[dispatch_key]
437

438
    dispatch_registrations_body = list(
439
        concatMap(
440
            dest.RegisterDispatchKey(
441
                backend_index,
442
                Target.REGISTRATION,
443
                selector,
444
                rocm=False,
445
                symint=True,
446
                class_method_name=f"{class_name}",
447
                skip_dispatcher_op_registration=False,
448
            ),
449
            grouped_native_functions,
450
        )
451
    )
452
    newline = "\n"
453
    ns_helper = NamespaceHelper(namespace_str="at")
454
    deferred_dispatch_registrations = ""
455
    static_init_dispatch_registrations = ""
456
    if eager_registration:
457
        static_template = CodeTemplate(
458
            """\
459
TORCH_LIBRARY_IMPL(aten, $dispatch_key, m) {
460
    $dispatch_registrations_body
461
};"""
462
        )
463
        static_init_dispatch_registrations = static_template.substitute(
464
            dispatch_key=dispatch_key,
465
            dispatch_registrations_body=dispatch_registrations_body,
466
        )
467
    else:
468
        deferred_template = CodeTemplate(
469
            """\
470
TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions();
471
TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() {
472
    static auto m = MAKE_TORCH_LIBRARY_IMPL(aten, $dispatch_key);
473
    $dispatch_registrations_body
474
}"""
475
        )
476
        deferred_dispatch_registrations = deferred_template.substitute(
477
            backend_name=backend_name,
478
            dispatch_key=dispatch_key,
479
            dispatch_registrations_body=dispatch_registrations_body,
480
        )
481

482
    fm.write_with_template(
483
        f"Register{dispatch_key}.cpp",
484
        "RegisterDispatchKey.cpp",
485
        lambda: {
486
            "extra_cuda_headers": "",
487
            "external_backend_headers": external_backend_headers_str,
488
            "ops_headers": "#include <ATen/Functions.h>"
489
            if not per_operator_headers
490
            else "",
491
            "DispatchKey": dispatch_key,
492
            "dispatch_namespace": dispatch_key.lower(),
493
            "dispatch_headers": dest.gen_registration_headers(
494
                backend_index, per_operator_headers=per_operator_headers, rocm=False
495
            ),
496
            "dispatch_definitions": fm.substitute_with_template(
497
                "RegisterDispatchDefinitions.ini",
498
                lambda: {
499
                    "ns_prologue": ns_helper.prologue,
500
                    "ns_epilogue": ns_helper.epilogue,
501
                    "static_init_dispatch_registrations": static_init_dispatch_registrations,
502
                    "deferred_dispatch_registrations": deferred_dispatch_registrations,
503
                    "dispatch_helpers": dest.gen_registration_helpers(backend_index),
504
                    "dispatch_namespace": dispatch_key.lower(),
505
                    "dispatch_namespaced_definitions": "",
506
                    "dispatch_anonymous_definitions": list(
507
                        concatMap(
508
                            dest.RegisterDispatchKey(
509
                                backend_index,
510
                                Target.ANONYMOUS_DEFINITION,
511
                                selector,
512
                                rocm=False,
513
                                symint=True,
514
                                class_method_name=f"{class_name}",
515
                                skip_dispatcher_op_registration=False,
516
                            ),
517
                            grouped_native_functions,
518
                        )
519
                    ),
520
                },
521
            ).split(newline),
522
        },
523
    )
524

525

526
def run(
527
    source_yaml: str, output_dir: str, dry_run: bool, impl_path: Optional[str] = None
528
) -> None:
529
    # Assumes that this file lives at PYTORCH_ROOT/torchgen/gen_backend_stubs.py
530
    pytorch_root = pathlib.Path(__file__).parent.parent.absolute()
531
    template_dir = os.path.join(pytorch_root, "aten/src/ATen/templates")
532

533
    def make_file_manager(install_dir: str) -> FileManager:
534
        return FileManager(
535
            install_dir=install_dir, template_dir=template_dir, dry_run=dry_run
536
        )
537

538
    fm = make_file_manager(output_dir)
539

540
    native_yaml_path = os.path.join(
541
        pytorch_root, "aten/src/ATen/native/native_functions.yaml"
542
    )
543
    tags_yaml_path = os.path.join(pytorch_root, "aten/src/ATen/native/tags.yaml")
544
    parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)
545
    native_functions, backend_indices = (
546
        parsed_yaml.native_functions,
547
        parsed_yaml.backend_indices,
548
    )
549
    grouped_native_functions = get_grouped_native_functions(native_functions)
550
    parsed_backend_yaml = parse_backend_yaml(
551
        source_yaml, grouped_native_functions, backend_indices
552
    )
553
    backend_key = parsed_backend_yaml.backend_key
554
    autograd_key = parsed_backend_yaml.autograd_key
555
    cpp_namespace = parsed_backend_yaml.cpp_namespace
556
    class_name = parsed_backend_yaml.class_name
557
    backend_indices = parsed_backend_yaml.backend_indices
558

559
    selector = SelectiveBuilder.get_nop_selector()
560

561
    if backend_key is None:
562
        # This could be useful if a backend wants to quickly set up a noop yaml file but doesn't have any kernels ready yet.
563
        return
564

565
    if class_name is None:
566
        # class_name is an optional argument to backend yaml file.
567
        # if specified it allows an external backend to override
568
        # the name of the class that all generated kernel definitions live under.
569
        # if not specified, its value is given as native_function_class_name.
570
        class_name = backend_indices[backend_key].native_function_class_name()
571
    assert class_name is not None
572

573
    if impl_path is not None:
574
        error_on_missing_kernels(
575
            native_functions,
576
            backend_indices,
577
            backend_key,
578
            autograd_key,
579
            class_name,
580
            impl_path,
581
        )
582

583
    gen_dispatchkey_nativefunc_headers(
584
        fm,
585
        class_name,
586
        cpp_namespace,
587
        backend_indices,
588
        grouped_native_functions,
589
        backend_key,
590
        autograd_key,
591
    )
592

593
    for dispatch_key in (
594
        [backend_key] if autograd_key is None else [backend_key, autograd_key]
595
    ):
596
        gen_dispatcher_registrations(
597
            fm,
598
            output_dir,
599
            class_name,
600
            backend_indices,
601
            grouped_native_functions,
602
            backend_key,
603
            dispatch_key,
604
            selector,
605
        )
606

607

608
if __name__ == "__main__":
609
    main()
610