add ROCm binary wheels (#3575)

* add ROCm binary wheels

* fix lint error
Co-authored-by: Eli Uriegas <1700823+seemethere@users.noreply.github.com>
Co-authored-by: Francisco Massa <fvsmassa@gmail.com>

.circleci/config.yml

@@ -829,6 +829,11 @@ workflows:
           name: binary_linux_wheel_py3.6_cu111
           python_version: '3.6'
           wheel_docker_image: pytorch/manylinux-cuda111
+      - binary_linux_wheel:
+          cu_version: rocm4.0.1
+          name: binary_linux_wheel_py3.6_rocm4.0.1
+          python_version: '3.6'
+          wheel_docker_image: pytorch/manylinux-rocm:4.0.1
       - binary_linux_wheel:
           conda_docker_image: pytorch/conda-builder:cpu
           cu_version: cpu
@@ -853,6 +858,11 @@ workflows:
           name: binary_linux_wheel_py3.7_cu111
           python_version: '3.7'
           wheel_docker_image: pytorch/manylinux-cuda111
+      - binary_linux_wheel:
+          cu_version: rocm4.0.1
+          name: binary_linux_wheel_py3.7_rocm4.0.1
+          python_version: '3.7'
+          wheel_docker_image: pytorch/manylinux-rocm:4.0.1
       - binary_linux_wheel:
           conda_docker_image: pytorch/conda-builder:cpu
           cu_version: cpu
@@ -877,6 +887,11 @@ workflows:
           name: binary_linux_wheel_py3.8_cu111
           python_version: '3.8'
           wheel_docker_image: pytorch/manylinux-cuda111
+      - binary_linux_wheel:
+          cu_version: rocm4.0.1
+          name: binary_linux_wheel_py3.8_rocm4.0.1
+          python_version: '3.8'
+          wheel_docker_image: pytorch/manylinux-rocm:4.0.1
       - binary_linux_wheel:
           conda_docker_image: pytorch/conda-builder:cpu
           cu_version: cpu
@@ -901,6 +916,11 @@ workflows:
           name: binary_linux_wheel_py3.9_cu111
           python_version: '3.9'
           wheel_docker_image: pytorch/manylinux-cuda111
+      - binary_linux_wheel:
+          cu_version: rocm4.0.1
+          name: binary_linux_wheel_py3.9_rocm4.0.1
+          python_version: '3.9'
+          wheel_docker_image: pytorch/manylinux-rocm:4.0.1
       - binary_macos_wheel:
           conda_docker_image: pytorch/conda-builder:cpu
           cu_version: cpu
@@ -1607,6 +1627,36 @@ workflows:
           python_version: '3.6'
           requires:
           - nightly_binary_linux_wheel_py3.6_cu111_upload
+      - binary_linux_wheel:
+          cu_version: rocm4.0.1
+          filters:
+            branches:
+              only: nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.6_rocm4.0.1
+          python_version: '3.6'
+          wheel_docker_image: pytorch/manylinux-rocm:4.0.1
+      - binary_wheel_upload:
+          context: org-member
+          filters:
+            branches:
+              only: nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.6_rocm4.0.1_upload
+          requires:
+          - nightly_binary_linux_wheel_py3.6_rocm4.0.1
+          subfolder: rocm4.0.1/
+      - smoke_test_linux_pip:
+          filters:
+            branches:
+              only:
+              - nightly
+          name: nightly_binary_linux_wheel_py3.6_rocm4.0.1_smoke_test_pip
+          python_version: '3.6'
+          requires:
+          - nightly_binary_linux_wheel_py3.6_rocm4.0.1_upload
       - binary_linux_wheel:
           conda_docker_image: pytorch/conda-builder:cpu
           cu_version: cpu
@@ -1731,6 +1781,36 @@ workflows:
           python_version: '3.7'
           requires:
           - nightly_binary_linux_wheel_py3.7_cu111_upload
+      - binary_linux_wheel:
+          cu_version: rocm4.0.1
+          filters:
+            branches:
+              only: nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.7_rocm4.0.1
+          python_version: '3.7'
+          wheel_docker_image: pytorch/manylinux-rocm:4.0.1
+      - binary_wheel_upload:
+          context: org-member
+          filters:
+            branches:
+              only: nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.7_rocm4.0.1_upload
+          requires:
+          - nightly_binary_linux_wheel_py3.7_rocm4.0.1
+          subfolder: rocm4.0.1/
+      - smoke_test_linux_pip:
+          filters:
+            branches:
+              only:
+              - nightly
+          name: nightly_binary_linux_wheel_py3.7_rocm4.0.1_smoke_test_pip
+          python_version: '3.7'
+          requires:
+          - nightly_binary_linux_wheel_py3.7_rocm4.0.1_upload
       - binary_linux_wheel:
           conda_docker_image: pytorch/conda-builder:cpu
           cu_version: cpu
@@ -1855,6 +1935,36 @@ workflows:
           python_version: '3.8'
           requires:
           - nightly_binary_linux_wheel_py3.8_cu111_upload
+      - binary_linux_wheel:
+          cu_version: rocm4.0.1
+          filters:
+            branches:
+              only: nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.8_rocm4.0.1
+          python_version: '3.8'
+          wheel_docker_image: pytorch/manylinux-rocm:4.0.1
+      - binary_wheel_upload:
+          context: org-member
+          filters:
+            branches:
+              only: nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.8_rocm4.0.1_upload
+          requires:
+          - nightly_binary_linux_wheel_py3.8_rocm4.0.1
+          subfolder: rocm4.0.1/
+      - smoke_test_linux_pip:
+          filters:
+            branches:
+              only:
+              - nightly
+          name: nightly_binary_linux_wheel_py3.8_rocm4.0.1_smoke_test_pip
+          python_version: '3.8'
+          requires:
+          - nightly_binary_linux_wheel_py3.8_rocm4.0.1_upload
       - binary_linux_wheel:
           conda_docker_image: pytorch/conda-builder:cpu
           cu_version: cpu
@@ -1979,6 +2089,36 @@ workflows:
           python_version: '3.9'
           requires:
           - nightly_binary_linux_wheel_py3.9_cu111_upload
+      - binary_linux_wheel:
+          cu_version: rocm4.0.1
+          filters:
+            branches:
+              only: nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.9_rocm4.0.1
+          python_version: '3.9'
+          wheel_docker_image: pytorch/manylinux-rocm:4.0.1
+      - binary_wheel_upload:
+          context: org-member
+          filters:
+            branches:
+              only: nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.9_rocm4.0.1_upload
+          requires:
+          - nightly_binary_linux_wheel_py3.9_rocm4.0.1
+          subfolder: rocm4.0.1/
+      - smoke_test_linux_pip:
+          filters:
+            branches:
+              only:
+              - nightly
+          name: nightly_binary_linux_wheel_py3.9_rocm4.0.1_smoke_test_pip
+          python_version: '3.9'
+          requires:
+          - nightly_binary_linux_wheel_py3.9_rocm4.0.1_upload
       - binary_macos_wheel:
           conda_docker_image: pytorch/conda-builder:cpu
           cu_version: cpu

.circleci/regenerate.py

@@ -29,12 +29,15 @@ def build_workflows(prefix='', filter_branch=None, upload=False, indentation=6,
     for btype in ["wheel", "conda"]:
         for os_type in ["linux", "macos", "win"]:
             python_versions = PYTHON_VERSIONS
-            cu_versions_dict = {"linux": ["cpu", "cu101", "cu102", "cu111"],
+            cu_versions_dict = {"linux": ["cpu", "cu101", "cu102", "cu111", "rocm4.0.1"],
                                 "win": ["cpu", "cu101", "cu102", "cu111"],
                                 "macos": ["cpu"]}
             cu_versions = cu_versions_dict[os_type]
             for python_version in python_versions:
                 for cu_version in cu_versions:
+                    # ROCm conda packages not yet supported
+                    if cu_version.startswith('rocm') and btype == "conda":
+                        continue
                     for unicode in ([False, True] if btype == "wheel" and python_version == "2.7" else [False]):
                         fb = filter_branch
                         if windows_latest_only and os_type == "win" and filter_branch is None and \
@@ -108,18 +111,22 @@ manylinux_images = {
 
 
 def get_manylinux_image(cu_version):
-    cu_suffix = "102"
-    if cu_version.startswith('cu'):
+    if cu_version == "cpu":
+        return "pytorch/manylinux-cuda102"
+    elif cu_version.startswith('cu'):
         cu_suffix = cu_version[len('cu'):]
-    return f"pytorch/manylinux-cuda{cu_suffix}"
+        return f"pytorch/manylinux-cuda{cu_suffix}"
+    elif cu_version.startswith('rocm'):
+        rocm_suffix = cu_version[len('rocm'):]
+        return f"pytorch/manylinux-rocm:{rocm_suffix}"
 
 
 def get_conda_image(cu_version):
     if cu_version == "cpu":
         return "pytorch/conda-builder:cpu"
-    if cu_version.startswith('cu'):
+    elif cu_version.startswith('cu'):
         cu_suffix = cu_version[len('cu'):]
-    return f"pytorch/conda-builder:cuda{cu_suffix}"
+        return f"pytorch/conda-builder:cuda{cu_suffix}"
 
 
 def generate_base_workflow(base_workflow_name, python_version, cu_version,
@@ -136,7 +143,9 @@ def generate_base_workflow(base_workflow_name, python_version, cu_version,
 
     if os_type != "win":
         d["wheel_docker_image"] = get_manylinux_image(cu_version)
-        d["conda_docker_image"] = get_conda_image(cu_version)
+        # ROCm conda packages not yet supported
+        if "rocm" not in cu_version:
+            d["conda_docker_image"] = get_conda_image(cu_version)
 
     if filter_branch is not None:
         d["filters"] = {

packaging/pkg_helpers.bash

@@ -126,6 +126,9 @@ setup_cuda() {
       ;;
     cpu)
       ;;
+    rocm*)
+      export FORCE_CUDA=1
+      ;;
     *)
       echo "Unrecognized CU_VERSION=$CU_VERSION"
       exit 1