|
592 | 592 | "target": "cu132", |
593 | 593 | }, |
594 | 594 | ], |
| 595 | + "rocm": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 596 | + "rocm-sdk-core": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 597 | + "rocm-sdk-libraries": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 598 | + "rocm-sdk-device-gfx1010": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 599 | + "rocm-sdk-device-gfx1011": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 600 | + "rocm-sdk-device-gfx1012": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 601 | + "rocm-sdk-device-gfx1030": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 602 | + "rocm-sdk-device-gfx1031": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 603 | + "rocm-sdk-device-gfx1032": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 604 | + "rocm-sdk-device-gfx1033": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 605 | + "rocm-sdk-device-gfx1034": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 606 | + "rocm-sdk-device-gfx1035": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 607 | + "rocm-sdk-device-gfx1036": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 608 | + "rocm-sdk-device-gfx1100": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 609 | + "rocm-sdk-device-gfx1101": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 610 | + "rocm-sdk-device-gfx1102": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 611 | + "rocm-sdk-device-gfx1103": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 612 | + "rocm-sdk-device-gfx1150": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 613 | + "rocm-sdk-device-gfx1151": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 614 | + "rocm-sdk-device-gfx1152": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 615 | + "rocm-sdk-device-gfx1153": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 616 | + "rocm-sdk-device-gfx1200": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 617 | + "rocm-sdk-device-gfx1201": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 618 | + "rocm-sdk-device-gfx1250": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 619 | + "rocm-sdk-device-gfx908": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 620 | + "rocm-sdk-device-gfx90a": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 621 | + "rocm-sdk-device-gfx942": [{"project": "torch_rocm", "target": "rocm7.14"}], |
| 622 | + "rocm-sdk-device-gfx950": [{"project": "torch_rocm", "target": "rocm7.14"}], |
595 | 623 | "arpeggio": [{"project": "triton"}], |
596 | 624 | "caliper-reader": [{"project": "triton"}], |
597 | 625 | "contourpy": [{"project": "triton"}], |
@@ -786,12 +814,19 @@ def is_nvidia_package(pkg_name: str) -> bool: |
786 | 814 | return pkg_name.startswith("nvidia-") or pkg_name.startswith("cuda-") |
787 | 815 |
|
788 | 816 |
|
| 817 | +def is_amd_package(pkg_name: str) -> bool: |
| 818 | + """Check if a package is from AMD and should use repo.amd.com""" |
| 819 | + name = pkg_name.lower() |
| 820 | + return "rocm" in name or "amd" in name |
| 821 | + |
| 822 | + |
789 | 823 | def get_package_source_url(pkg_name: str) -> str: |
790 | 824 | """Get the source URL for a package based on its type""" |
791 | 825 | if is_nvidia_package(pkg_name): |
792 | 826 | return f"https://pypi.nvidia.com/{pkg_name}/" |
793 | | - else: |
794 | | - return f"https://pypi.org/simple/{pkg_name}/" |
| 827 | + if is_amd_package(pkg_name): |
| 828 | + return f"https://repo.amd.com/rocm/whl-multi-arch/{pkg_name}/" |
| 829 | + return f"https://pypi.org/simple/{pkg_name}/" |
795 | 830 |
|
796 | 831 |
|
797 | 832 | def download(url: str) -> bytes: |
@@ -989,11 +1024,14 @@ def upload_package_using_simple_index( |
989 | 1024 | Works for both NVIDIA and non-NVIDIA packages. |
990 | 1025 | """ |
991 | 1026 | source_url = get_package_source_url(pkg_name) |
992 | | - is_nvidia = is_nvidia_package(pkg_name) |
| 1027 | + if is_nvidia_package(pkg_name): |
| 1028 | + source_label = "NVIDIA" |
| 1029 | + elif is_amd_package(pkg_name): |
| 1030 | + source_label = "AMD" |
| 1031 | + else: |
| 1032 | + source_label = "PyPI" |
993 | 1033 |
|
994 | | - print( |
995 | | - f"Processing {pkg_name} using {'NVIDIA' if is_nvidia else 'PyPI'} Simple Index: {source_url}" |
996 | | - ) |
| 1034 | + print(f"Processing {pkg_name} using {source_label} Simple Index: {source_url}") |
997 | 1035 |
|
998 | 1036 | # Parse the index and get raw HTML |
999 | 1037 | try: |
@@ -1025,20 +1063,26 @@ def get_packages_for_target(target: str) -> List[str]: |
1025 | 1063 | Get packages from PACKAGES_PER_PROJECT that should be initialized for a target. |
1026 | 1064 |
|
1027 | 1065 | Returns packages where: |
1028 | | - - project is "torch" AND |
| 1066 | + - project is "torch" (or "torch_rocm" for ROCm targets) AND |
1029 | 1067 | - either no target is specified (universal packages like filelock, numpy) |
1030 | 1068 | - or the target matches the specified target |
1031 | 1069 | - nvidia/cuda packages are only included for CUDA targets (cu*) |
| 1070 | + - amd/rocm packages are only included for ROCm targets (rocm*) |
1032 | 1071 | """ |
1033 | 1072 | is_cuda_target = target.startswith("cu") |
| 1073 | + is_rocm_target = target.startswith("rocm") |
| 1074 | + allowed_projects = ("torch", "torch_rocm") if is_rocm_target else ("torch",) |
1034 | 1075 | packages = [] |
1035 | 1076 | for pkg_name, pkg_configs in PACKAGES_PER_PROJECT.items(): |
1036 | 1077 | # Skip nvidia/cuda packages for non-CUDA targets |
1037 | 1078 | if not is_cuda_target and is_nvidia_package(pkg_name): |
1038 | 1079 | continue |
| 1080 | + # Skip amd/rocm packages for non-ROCm targets |
| 1081 | + if not is_rocm_target and is_amd_package(pkg_name): |
| 1082 | + continue |
1039 | 1083 |
|
1040 | 1084 | for config in pkg_configs: |
1041 | | - if config.get("project") != "torch": |
| 1085 | + if config.get("project") not in allowed_projects: |
1042 | 1086 | continue |
1043 | 1087 | pkg_target = config.get("target", "") |
1044 | 1088 | # Include if no target specified (universal) or target matches |
|
0 commit comments