[ROCm] add rocm wheels (#1477)

* add rocm binary use the term cu_version update pkg_helpers add docker wheel option update config yml add force_cuda option * update config.yml * remove conda builds
pytorch · May 28, 2021 · 3bb3287 · 3bb3287
1 parent c8239c6
commit 3bb3287
Show file tree

Hide file tree

Showing 3 changed files with 184 additions and 3 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -650,6 +650,13 @@ workflows:
           requires:
           - download_third_parties_nix
           wheel_docker_image: pytorch/manylinux-cuda111
+      - binary_linux_wheel:
+          cuda_version: rocm4.1
+          name: binary_linux_wheel_py3.6_rocm4.1
+          python_version: '3.6'
+          requires:
+          - download_third_parties_nix
+          wheel_docker_image: pytorch/manylinux-rocm:4.1
       - binary_linux_wheel:
           cuda_version: cpu
           name: binary_linux_wheel_py3.7_cpu
@@ -670,6 +677,13 @@ workflows:
           requires:
           - download_third_parties_nix
           wheel_docker_image: pytorch/manylinux-cuda111
+      - binary_linux_wheel:
+          cuda_version: rocm4.1
+          name: binary_linux_wheel_py3.7_rocm4.1
+          python_version: '3.7'
+          requires:
+          - download_third_parties_nix
+          wheel_docker_image: pytorch/manylinux-rocm:4.1
       - binary_linux_wheel:
           cuda_version: cpu
           name: binary_linux_wheel_py3.8_cpu
@@ -690,6 +704,13 @@ workflows:
           requires:
           - download_third_parties_nix
           wheel_docker_image: pytorch/manylinux-cuda111
+      - binary_linux_wheel:
+          cuda_version: rocm4.1
+          name: binary_linux_wheel_py3.8_rocm4.1
+          python_version: '3.8'
+          requires:
+          - download_third_parties_nix
+          wheel_docker_image: pytorch/manylinux-rocm:4.1
       - binary_linux_wheel:
           cuda_version: cpu
           name: binary_linux_wheel_py3.9_cpu
@@ -710,6 +731,13 @@ workflows:
           requires:
           - download_third_parties_nix
           wheel_docker_image: pytorch/manylinux-cuda111
+      - binary_linux_wheel:
+          cuda_version: rocm4.1
+          name: binary_linux_wheel_py3.9_rocm4.1
+          python_version: '3.9'
+          requires:
+          - download_third_parties_nix
+          wheel_docker_image: pytorch/manylinux-rocm:4.1
       - binary_macos_wheel:
           cuda_version: cpu
           name: binary_macos_wheel_py3.6_cpu
@@ -1217,6 +1245,42 @@ workflows:
           python_version: '3.6'
           requires:
           - nightly_binary_linux_wheel_py3.6_cu111_upload
+      - binary_linux_wheel:
+          cuda_version: rocm4.1
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.6_rocm4.1
+          python_version: '3.6'
+          requires:
+          - download_third_parties_nix
+          wheel_docker_image: pytorch/manylinux-rocm:4.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.1_upload
+          requires:
+          - nightly_binary_linux_wheel_py3.6_rocm4.1
+      - smoke_test_linux_pip:
+          cuda_version: rocm4.1
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.6_rocm4.1_smoke_test_pip
+          python_version: '3.6'
+          requires:
+          - nightly_binary_linux_wheel_py3.6_rocm4.1_upload
       - binary_linux_wheel:
           cuda_version: cpu
           filters:
@@ -1324,6 +1388,42 @@ workflows:
           python_version: '3.7'
           requires:
           - nightly_binary_linux_wheel_py3.7_cu111_upload
+      - binary_linux_wheel:
+          cuda_version: rocm4.1
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.7_rocm4.1
+          python_version: '3.7'
+          requires:
+          - download_third_parties_nix
+          wheel_docker_image: pytorch/manylinux-rocm:4.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.1_upload
+          requires:
+          - nightly_binary_linux_wheel_py3.7_rocm4.1
+      - smoke_test_linux_pip:
+          cuda_version: rocm4.1
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.7_rocm4.1_smoke_test_pip
+          python_version: '3.7'
+          requires:
+          - nightly_binary_linux_wheel_py3.7_rocm4.1_upload
       - binary_linux_wheel:
           cuda_version: cpu
           filters:
@@ -1431,6 +1531,42 @@ workflows:
           python_version: '3.8'
           requires:
           - nightly_binary_linux_wheel_py3.8_cu111_upload
+      - binary_linux_wheel:
+          cuda_version: rocm4.1
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.8_rocm4.1
+          python_version: '3.8'
+          requires:
+          - download_third_parties_nix
+          wheel_docker_image: pytorch/manylinux-rocm:4.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.1_upload
+          requires:
+          - nightly_binary_linux_wheel_py3.8_rocm4.1
+      - smoke_test_linux_pip:
+          cuda_version: rocm4.1
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.8_rocm4.1_smoke_test_pip
+          python_version: '3.8'
+          requires:
+          - nightly_binary_linux_wheel_py3.8_rocm4.1_upload
       - binary_linux_wheel:
           cuda_version: cpu
           filters:
@@ -1538,6 +1674,42 @@ workflows:
           python_version: '3.9'
           requires:
           - nightly_binary_linux_wheel_py3.9_cu111_upload
+      - binary_linux_wheel:
+          cuda_version: rocm4.1
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.9_rocm4.1
+          python_version: '3.9'
+          requires:
+          - download_third_parties_nix
+          wheel_docker_image: pytorch/manylinux-rocm:4.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.1_upload
+          requires:
+          - nightly_binary_linux_wheel_py3.9_rocm4.1
+      - smoke_test_linux_pip:
+          cuda_version: rocm4.1
+          filters:
+            branches:
+              only:
+              - nightly
+            tags:
+              only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
+          name: nightly_binary_linux_wheel_py3.9_rocm4.1_smoke_test_pip
+          python_version: '3.9'
+          requires:
+          - nightly_binary_linux_wheel_py3.9_rocm4.1_upload
       - binary_macos_wheel:
           cuda_version: cpu
           filters:

diff --git a/.circleci/regenerate.py b/.circleci/regenerate.py
@@ -21,7 +21,7 @@
 
 
 PYTHON_VERSIONS = ["3.6", "3.7", "3.8", "3.9"]
-CU_VERSIONS_DICT = {"linux": ["cpu", "cu102", "cu111"],
+CU_VERSIONS_DICT = {"linux": ["cpu", "cu102", "cu111","rocm4.1"],
                     "windows": ["cpu", "cu102", "cu111"],
                     "macos": ["cpu"]}
 
@@ -36,7 +36,10 @@ def build_workflows(prefix='', upload=False, filter_branch=None, indentation=6):
         for os_type in ["linux", "macos", "windows"]:
             for python_version in PYTHON_VERSIONS:
                 for cu_version in CU_VERSIONS_DICT[os_type]:
-                    w += build_workflow_pair(btype, os_type, python_version, cu_version, filter_branch, prefix, upload)
+                    if cu_version.startswith("rocm") and btype=="conda":
+                        pass
+                    else:
+                        w += build_workflow_pair(btype, os_type, python_version, cu_version, filter_branch, prefix, upload)
 
     if not filter_branch:
         # Build on every pull request, but upload only on nightly and tags
@@ -124,8 +127,10 @@ def generate_base_workflow(base_workflow_name, python_version, cu_version, filte
         d['requires'] = ['download_third_parties_nix']
     if btype == 'conda':
         d['conda_docker_image'] = f'pytorch/conda-builder:{cu_version.replace("cu1","cuda1")}'
-    elif cu_version != 'cpu':
+    elif cu_version.startswith('cu'):
         d['wheel_docker_image'] = f'pytorch/manylinux-{cu_version.replace("cu1","cuda1")}'
+    elif cu_version.startswith('rocm'):
+        d["wheel_docker_image"] = f"pytorch/manylinux-rocm:{cu_version[len('rocm'):]}"
 
     if filter_branch:
         d["filters"] = gen_filter_branch_tree(filter_branch)

diff --git a/packaging/pkg_helpers.bash b/packaging/pkg_helpers.bash
@@ -107,6 +107,10 @@ setup_cuda() {
       export FORCE_CUDA=1
       export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0"
       ;;
+    rocm*)
+      export FORCE_CUDA=1
+      export USE_ROCM=1
+      ;;
     cpu)
       ;;
     *)