Support for Building Without CPU

.github/workflows/generate_binary_build_matrix.yml

@@ -28,7 +28,11 @@ on:
         default: "enable"
         type: string
       with-rocm:
-        description: "Build with Cuda?"
+        description: "Build with Rocm?"
+        default: "enable"
+        type: string
+      with-cpu:
+        description: "Build with CPU?"
         default: "enable"
         type: string
     outputs:
@@ -59,6 +63,7 @@ jobs:
           CHANNEL: ${{ inputs.channel }}
           WITH_CUDA: ${{ inputs.with-cuda }}
           WITH_ROCM: ${{ inputs.with-rocm }}
+          WITH_CPU: ${{ inputs.with-cpu }}
           # limit pull request builds to one version of python unless ciflow/binaries/all is applied to the workflow
           # should not affect builds that are from events that are not the pull_request event
           LIMIT_PR_BUILDS: ${{ github.event_name == 'pull_request' && !contains( github.event.pull_request.labels.*.name, 'ciflow/binaries/all') }}

tools/scripts/generate_binary_build_matrix.py

@@ -203,13 +203,26 @@ def get_wheel_install_command(os: str, channel: str, gpu_arch_type: str, gpu_arc
         whl_install_command = f"{WHL_INSTALL_BASE} --pre {PACKAGES_TO_INSTALL_WHL}" if channel == "nightly" else f"{WHL_INSTALL_BASE} {PACKAGES_TO_INSTALL_WHL}"
         return f"{whl_install_command} --index-url {get_base_download_url_for_repo('whl', channel, gpu_arch_type, desired_cuda)}"
 
-def generate_conda_matrix(os: str, channel: str, with_cuda: str, with_rocm: str, limit_pr_builds: bool) -> List[Dict[str, str]]:
+def generate_conda_matrix(
+    os: str,
+    channel: str,
+    with_cuda: str,
+    with_rocm: str,
+    with_cpu: str,
+    limit_pr_builds: bool
+) -> List[Dict[str, str]]:
+
     ret: List[Dict[str, str]] = []
-    arches = ["cpu"]
     python_versions = list(mod.PYTHON_ARCHES)
 
-    if with_cuda == ENABLE and (os == "linux" or os == "windows"):
-        arches += mod.CUDA_ARCHES
+    arches = []
+
+    if with_cpu == ENABLE:
+        arches += ["cpu"]
+
+    if with_cuda == ENABLE:
+        if os == "linux" or os == "windows":
+            arches += mod.CUDA_ARCHES
 
     if limit_pr_builds:
         python_versions = [ python_versions[0] ]
@@ -248,6 +261,7 @@ def generate_libtorch_matrix(
     channel: str,
     with_cuda: str,
     with_rocm: str,
+    with_cpu: str,
     limit_pr_builds: str,
     abi_versions: Optional[List[str]] = None,
     arches: Optional[List[str]] = None,
@@ -261,12 +275,13 @@ def generate_libtorch_matrix(
         return ret
 
     if arches is None:
-        arches = ["cpu"]
+        arches = []
+
+        if with_cpu == ENABLE:
+            arches += ["cpu"]
 
         if with_cuda == ENABLE:
-            if os == "linux":
-                arches += mod.CUDA_ARCHES
-            elif os == "windows":
+            if os == "linux" or os == "windows":
                 arches += mod.CUDA_ARCHES
 
         if with_rocm == ENABLE:
@@ -342,6 +357,7 @@ def generate_wheels_matrix(
     channel: str,
     with_cuda: str,
     with_rocm: str,
+    with_cpu: str,
     limit_pr_builds: bool,
     arches: Optional[List[str]] = None,
     python_versions: Optional[List[str]] = None,
@@ -359,19 +375,21 @@ def generate_wheels_matrix(
     upload_to_base_bucket = "yes"
     if arches is None:
         # Define default compute architectures
-        arches = ["cpu"]
+        arches = []
+
+        if with_cpu == ENABLE:
+            arches += ["cpu"]
 
         if with_cuda == ENABLE:
             upload_to_base_bucket = "no"
-            if os == "linux":
-                arches += mod.CUDA_ARCHES
-            elif os == "windows":
+            if os == "linux" or os == "windows":
                 arches += mod.CUDA_ARCHES
 
         if with_rocm == ENABLE:
             if os == "linux":
                 arches += mod.ROCM_ARCHES
 
+
     if limit_pr_builds:
         python_versions = [ python_versions[0] ]
 
@@ -415,6 +433,7 @@ def generate_build_matrix(
         channel: str,
         with_cuda: str,
         with_rocm: str,
+        with_cpu: str,
         limit_pr_builds: str) -> Dict[str, List[Dict[str, str]]]:
     includes = []
 
@@ -432,6 +451,7 @@ def generate_build_matrix(
                                                             channel,
                                                             with_cuda,
                                                             with_rocm,
+                                                            with_cpu,
                                                             limit_pr_builds == "true")
                 )
 
@@ -473,6 +493,13 @@ def main(args) -> None:
         choices=[ENABLE, DISABLE],
         default=os.getenv("WITH_ROCM", ENABLE),
     )
+    parser.add_argument(
+        "--with-cpu",
+        help="Build with CPU?",
+        type=str,
+        choices=[ENABLE, DISABLE],
+        default=os.getenv("WITH_CPU", ENABLE),
+    )
     # By default this is false for this script but expectation is that the caller
     # workflow will default this to be true most of the time, where a pull
     # request is synchronized and does not contain the label "ciflow/binaries/all"
@@ -487,12 +514,15 @@ def main(args) -> None:
 
     options = parser.parse_args(args)
 
+    assert (options.with_cuda or options.with_rocm or options.with_cpu), "Must build with either CUDA, ROCM, or CPU support."
+
     build_matrix = generate_build_matrix(
             options.package_type,
             options.operating_system,
             options.channel,
             options.with_cuda,
             options.with_rocm,
+            options.with_cpu,
             options.limit_pr_builds)
 
     print(json.dumps(build_matrix))

tools/tests/assets/build_matrix_linux_wheel_nocpu.json

@@ -0,0 +1 @@
+{"include": [{"python_version": "3.8", "gpu_arch_type": "cuda", "gpu_arch_version": "11.8", "desired_cuda": "cu118", "container_image": "pytorch/manylinux-builder:cuda11.8", "package_type": "manywheel", "build_name": "manywheel-py3_8-cuda11_8", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.8", "gpu_arch_type": "cuda", "gpu_arch_version": "12.1", "desired_cuda": "cu121", "container_image": "pytorch/manylinux-builder:cuda12.1", "package_type": "manywheel", "build_name": "manywheel-py3_8-cuda12_1", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.8", "gpu_arch_type": "rocm", "gpu_arch_version": "5.5", "desired_cuda": "rocm5.5", "container_image": "pytorch/manylinux-builder:rocm5.5", "package_type": "manywheel", "build_name": "manywheel-py3_8-rocm5_5", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.5", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.8", "gpu_arch_type": "rocm", "gpu_arch_version": "5.6", "desired_cuda": "rocm5.6", "container_image": "pytorch/manylinux-builder:rocm5.6", "package_type": "manywheel", "build_name": "manywheel-py3_8-rocm5_6", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.6", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.9", "gpu_arch_type": "cuda", "gpu_arch_version": "11.8", "desired_cuda": "cu118", "container_image": "pytorch/manylinux-builder:cuda11.8", "package_type": "manywheel", "build_name": "manywheel-py3_9-cuda11_8", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.9", "gpu_arch_type": "cuda", "gpu_arch_version": "12.1", "desired_cuda": "cu121", "container_image": "pytorch/manylinux-builder:cuda12.1", "package_type": "manywheel", "build_name": "manywheel-py3_9-cuda12_1", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.9", "gpu_arch_type": "rocm", "gpu_arch_version": "5.5", "desired_cuda": "rocm5.5", "container_image": "pytorch/manylinux-builder:rocm5.5", "package_type": "manywheel", "build_name": "manywheel-py3_9-rocm5_5", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.5", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.9", "gpu_arch_type": "rocm", "gpu_arch_version": "5.6", "desired_cuda": "rocm5.6", "container_image": "pytorch/manylinux-builder:rocm5.6", "package_type": "manywheel", "build_name": "manywheel-py3_9-rocm5_6", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.6", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.10", "gpu_arch_type": "cuda", "gpu_arch_version": "11.8", "desired_cuda": "cu118", "container_image": "pytorch/manylinux-builder:cuda11.8", "package_type": "manywheel", "build_name": "manywheel-py3_10-cuda11_8", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.10", "gpu_arch_type": "cuda", "gpu_arch_version": "12.1", "desired_cuda": "cu121", "container_image": "pytorch/manylinux-builder:cuda12.1", "package_type": "manywheel", "build_name": "manywheel-py3_10-cuda12_1", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.10", "gpu_arch_type": "rocm", "gpu_arch_version": "5.5", "desired_cuda": "rocm5.5", "container_image": "pytorch/manylinux-builder:rocm5.5", "package_type": "manywheel", "build_name": "manywheel-py3_10-rocm5_5", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.5", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.10", "gpu_arch_type": "rocm", "gpu_arch_version": "5.6", "desired_cuda": "rocm5.6", "container_image": "pytorch/manylinux-builder:rocm5.6", "package_type": "manywheel", "build_name": "manywheel-py3_10-rocm5_6", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.6", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.11", "gpu_arch_type": "cuda", "gpu_arch_version": "11.8", "desired_cuda": "cu118", "container_image": "pytorch/manylinux-builder:cuda11.8", "package_type": "manywheel", "build_name": "manywheel-py3_11-cuda11_8", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.11", "gpu_arch_type": "cuda", "gpu_arch_version": "12.1", "desired_cuda": "cu121", "container_image": "pytorch/manylinux-builder:cuda12.1", "package_type": "manywheel", "build_name": "manywheel-py3_11-cuda12_1", "validation_runner": "linux.g5.4xlarge.nvidia.gpu", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.11", "gpu_arch_type": "rocm", "gpu_arch_version": "5.5", "desired_cuda": "rocm5.5", "container_image": "pytorch/manylinux-builder:rocm5.5", "package_type": "manywheel", "build_name": "manywheel-py3_11-rocm5_5", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.5", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}, {"python_version": "3.11", "gpu_arch_type": "rocm", "gpu_arch_version": "5.6", "desired_cuda": "rocm5.6", "container_image": "pytorch/manylinux-builder:rocm5.6", "package_type": "manywheel", "build_name": "manywheel-py3_11-rocm5_6", "validation_runner": "linux.2xlarge", "installation": "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.6", "channel": "nightly", "upload_to_base_bucket": "no", "stable_version": "2.0.1"}]}

tools/tests/test_generate_binary_build_matrix.py

@@ -10,14 +10,21 @@
 
 class GenerateBuildMatrixTest(TestCase):
     def matrix_compare_helper(
-        self, package_type, operating_system, cuda, rocm, reference_output_file
-    ):
+        self,
+        package_type: str,
+        operating_system: str,
+        cuda: bool,
+        rocm: bool,
+        cpu: bool,
+        reference_output_file: str,
+    ) -> None:
         out = generate_build_matrix(
             package_type,
             operating_system,
             "nightly",
             "enable" if cuda else "disable",
             "enable" if rocm else "disable",
+            "enable" if cpu else "disable",
             "false",
         )
 
@@ -35,6 +42,7 @@ def test_linux_wheel_cuda(self):
             operating_system="linux",
             cuda=True,
             rocm=True,
+            cpu=True,
             reference_output_file="build_matrix_linux_wheel_cuda.json",
         )
 
@@ -44,6 +52,7 @@ def test_linux_conda_cuda(self):
             operating_system="linux",
             cuda=True,
             rocm=True,
+            cpu=True,
             reference_output_file="build_matrix_linux_conda_cuda.json",
         )
 
@@ -53,6 +62,7 @@ def test_macos_wheel(self):
             operating_system="macos",
             cuda=False,
             rocm=False,
+            cpu=True,
             reference_output_file="build_matrix_macos_wheel.json",
         )
 
@@ -62,6 +72,7 @@ def test_macos_conda(self):
             operating_system="macos",
             cuda=False,
             rocm=False,
+            cpu=True,
             reference_output_file="build_matrix_macos_conda.json",
         )
 
@@ -71,6 +82,7 @@ def test_windows_wheel_cuda(self):
             operating_system="windows",
             cuda=True,
             rocm=True,
+            cpu=True,
             reference_output_file="build_matrix_windows_wheel_cuda.json",
         )
 
@@ -80,6 +92,7 @@ def test_windows_conda_cuda(self):
             operating_system="windows",
             cuda=True,
             rocm=True,
+            cpu=True,
             reference_output_file="build_matrix_windows_conda_cuda.json",
         )
 
@@ -89,9 +102,30 @@ def test_linux_wheel_cuda_norocm(self):
             operating_system="linux",
             cuda=True,
             rocm=False,
+            cpu=True,
             reference_output_file="build_matrix_linux_wheel_cuda_norocm.json",
         )
 
+    def test_linux_wheel_cuda_nocpu(self):
+        self.matrix_compare_helper(
+            package_type="wheel",
+            operating_system="linux",
+            cuda=True,
+            rocm=True,
+            cpu=False,
+            reference_output_file="build_matrix_linux_wheel_cuda_nocpu.json",
+        )
+
+    def test_linux_wheel_cuda_rocm_nocpu(self):
+        self.matrix_compare_helper(
+            package_type="wheel",
+            operating_system="linux",
+            cuda=True,
+            rocm=True,
+            cpu=False,
+            reference_output_file="build_matrix_linux_wheel_nocpu.json",
+        )
+
 
 if __name__ == "__main__":
     main()