Repository navigation
Commit 7586c16
authored
CUDA backend: per-program weights file names and antialiased resize (#24314)
This change fixes several problems in the CUDA backend. Each fix is its
own commit.
## Name the weights file of each program
The backend saves every program's weights to `aoti_cuda_blob.ptd`.
Exporting several programs into one folder overwrites the weights, and
the earlier programs then fail to load. Renaming each file after export
works, but is easy to miss. A new option names the file:
```python
partitioner = CudaPartitioner([
CudaBackend.generate_method_name_compile_spec("forward"),
CudaBackend.generate_data_name_compile_spec("encoder"),
]) # writes encoder.ptd
```
Without the option nothing changes. The name is checked: up to 200 ASCII
letters, digits, `_`, `-` and `.`, not starting with `.` or `-`, not
ending with `.ptd`, and not a Windows device name such as `NUL`. All
methods of a program must share it. A backend that keeps its data inside
the `.pte`, such as Metal, refuses it.
## Lower antialiased resizes
`F.interpolate(x, mode="bilinear" or "bicubic", antialias=True)` failed
to lower with "missing fallback kernels". The backend now rewrites it
into plain tensor math, with PyTorch's window and weight formulas, for
dynamic sizes and float32, float16 and bfloat16. The results are not
bit-identical: in float32 they stay within about 2e-5 of PyTorch's
kernel for inputs between 0 and 1.
## Tests
Ran on an A100 GPU (Linux x86_64): `test_cuda_export.py`,
`test_antialiased_resize.py`, `test_missing_weights_blob.py`,
`test_merge_ptes.py` and `test_cuda_weight_metadata.py`. Each test added
for a fix fails when that fix is removed.1 parent fb34f59 commit 7586c16
10 files changed
Lines changed: 511 additions & 17 deletions
File tree
- backends
- aoti
- cuda
- tests
- docs/source/backends
- cuda
- tensorrt
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
28 | 28 | | |
29 | 29 | | |
30 | 30 | | |
| 31 | + | |
| 32 | + | |
| 33 | + | |
| 34 | + | |
| 35 | + | |
| 36 | + | |
| 37 | + | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
| 51 | + | |
| 52 | + | |
| 53 | + | |
31 | 54 | | |
32 | 55 | | |
33 | 56 | | |
| |||
368 | 391 | | |
369 | 392 | | |
370 | 393 | | |
371 | | - | |
| 394 | + | |
| 395 | + | |
| 396 | + | |
372 | 397 | | |
373 | 398 | | |
374 | 399 | | |
| |||
404 | 429 | | |
405 | 430 | | |
406 | 431 | | |
| 432 | + | |
| 433 | + | |
| 434 | + | |
| 435 | + | |
| 436 | + | |
| 437 | + | |
| 438 | + | |
| 439 | + | |
| 440 | + | |
| 441 | + | |
| 442 | + | |
| 443 | + | |
| 444 | + | |
| 445 | + | |
| 446 | + | |
| 447 | + | |
| 448 | + | |
| 449 | + | |
| 450 | + | |
| 451 | + | |
| 452 | + | |
| 453 | + | |
| 454 | + | |
| 455 | + | |
| 456 | + | |
| 457 | + | |
| 458 | + | |
407 | 459 | | |
408 | 460 | | |
409 | 461 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
9 | 9 | | |
10 | 10 | | |
11 | 11 | | |
| 12 | + | |
12 | 13 | | |
13 | 14 | | |
14 | 15 | | |
| |||
50 | 51 | | |
51 | 52 | | |
52 | 53 | | |
| 54 | + | |
53 | 55 | | |
| 56 | + | |
54 | 57 | | |
55 | 58 | | |
56 | 59 | | |
| |||
548 | 551 | | |
549 | 552 | | |
550 | 553 | | |
| 554 | + | |
| 555 | + | |
| 556 | + | |
| 557 | + | |
| 558 | + | |
| 559 | + | |
| 560 | + | |
| 561 | + | |
| 562 | + | |
| 563 | + | |
| 564 | + | |
| 565 | + | |
| 566 | + | |
| 567 | + | |
| 568 | + | |
| 569 | + | |
| 570 | + | |
| 571 | + | |
| 572 | + | |
| 573 | + | |
| 574 | + | |
| 575 | + | |
| 576 | + | |
| 577 | + | |
| 578 | + | |
| 579 | + | |
| 580 | + | |
| 581 | + | |
| 582 | + | |
| 583 | + | |
| 584 | + | |
| 585 | + | |
| 586 | + | |
| 587 | + | |
| 588 | + | |
| 589 | + | |
| 590 | + | |
| 591 | + | |
| 592 | + | |
| 593 | + | |
| 594 | + | |
| 595 | + | |
| 596 | + | |
| 597 | + | |
| 598 | + | |
| 599 | + | |
| 600 | + | |
| 601 | + | |
| 602 | + | |
| 603 | + | |
| 604 | + | |
| 605 | + | |
| 606 | + | |
| 607 | + | |
| 608 | + | |
| 609 | + | |
| 610 | + | |
| 611 | + | |
| 612 | + | |
| 613 | + | |
| 614 | + | |
| 615 | + | |
| 616 | + | |
| 617 | + | |
| 618 | + | |
| 619 | + | |
| 620 | + | |
| 621 | + | |
| 622 | + | |
| 623 | + | |
| 624 | + | |
| 625 | + | |
| 626 | + | |
| 627 | + | |
| 628 | + | |
| 629 | + | |
| 630 | + | |
| 631 | + | |
| 632 | + | |
| 633 | + | |
| 634 | + | |
| 635 | + | |
| 636 | + | |
| 637 | + | |
| 638 | + | |
| 639 | + | |
| 640 | + | |
| 641 | + | |
| 642 | + | |
| 643 | + | |
| 644 | + | |
| 645 | + | |
| 646 | + | |
| 647 | + | |
| 648 | + | |
| 649 | + | |
| 650 | + | |
| 651 | + | |
| 652 | + | |
| 653 | + | |
| 654 | + | |
| 655 | + | |
| 656 | + | |
| 657 | + | |
| 658 | + | |
| 659 | + | |
| 660 | + | |
| 661 | + | |
| 662 | + | |
| 663 | + | |
| 664 | + | |
551 | 665 | | |
552 | 666 | | |
553 | 667 | | |
| |||
704 | 818 | | |
705 | 819 | | |
706 | 820 | | |
| 821 | + | |
707 | 822 | | |
708 | 823 | | |
709 | 824 | | |
| |||
726 | 841 | | |
727 | 842 | | |
728 | 843 | | |
| 844 | + | |
| 845 | + | |
| 846 | + | |
| 847 | + | |
| 848 | + | |
| 849 | + | |
| 850 | + | |
| 851 | + | |
| 852 | + | |
| 853 | + | |
| 854 | + | |
729 | 855 | | |
730 | 856 | | |
731 | 857 | | |
| |||
852 | 978 | | |
853 | 979 | | |
854 | 980 | | |
| 981 | + | |
| 982 | + | |
| 983 | + | |
| 984 | + | |
| 985 | + | |
| 986 | + | |
| 987 | + | |
| 988 | + | |
| 989 | + | |
855 | 990 | | |
856 | 991 | | |
857 | 992 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
570 | 570 | | |
571 | 571 | | |
572 | 572 | | |
| 573 | + | |
573 | 574 | | |
574 | 575 | | |
575 | 576 | | |
| |||
589 | 590 | | |
590 | 591 | | |
591 | 592 | | |
592 | | - | |
| 593 | + | |
593 | 594 | | |
594 | 595 | | |
595 | 596 | | |
| |||
601 | 602 | | |
602 | 603 | | |
603 | 604 | | |
604 | | - | |
| 605 | + | |
605 | 606 | | |
606 | 607 | | |
607 | 608 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
208 | 208 | | |
209 | 209 | | |
210 | 210 | | |
| 211 | + | |
| 212 | + | |
| 213 | + | |
| 214 | + | |
| 215 | + | |
211 | 216 | | |
212 | 217 | | |
213 | 218 | | |
| |||
219 | 224 | | |
220 | 225 | | |
221 | 226 | | |
222 | | - | |
| 227 | + | |
223 | 228 | | |
224 | 229 | | |
225 | 230 | | |
| |||
534 | 539 | | |
535 | 540 | | |
536 | 541 | | |
537 | | - | |
| 542 | + | |
538 | 543 | | |
539 | 544 | | |
540 | 545 | | |
| |||
0 commit comments