Update on "nn.Linear: dispatch to bsr_dense_mm for half and bfloat16"

cc alexsamardzic pearu cpuhrsch amjames bhosmer

[ghstack-poisoned]

.ci/pytorch/macos-test.sh

@@ -25,6 +25,7 @@ setup_test_python() {
   # using the address associated with the loopback interface.
   export GLOO_SOCKET_IFNAME=lo0
   echo "Ninja version: $(ninja --version)"
+  echo "Python version: $(python --version)"
 
   # Increase default limit on open file handles from 256 to 1024
   ulimit -n 1024

.github/ci_commit_pins/vision.txt

@@ -1 +1 @@
-5b07d6c9c6c14cf88fc545415d63021456874744
+29757104250dd088386fef1ec3d70ed0b0c1be8a

.github/merge_rules.yaml

@@ -370,7 +370,7 @@
   - '**functorch**'
   - '**pytree**'
   approved_by:
-  - kshiteej12345
+  - kshitij12345
   - srossross
   - chillee
   - zou3519

.github/requirements/pip-requirements-macOS.txt

@@ -21,3 +21,4 @@ sympy==1.11.1
 unittest-xml-reporting<=3.2.0,>=2.0.0
 xdoctest==1.1.0
 filelock==3.6.0
+sympy==1.11.1

.github/scripts/on_cancel_merge.py

@@ -0,0 +1,31 @@
+import argparse
+from typing import Any
+
+from gitutils import get_git_remote_name, get_git_repo_dir, GitRepo
+from trymerge import GitHubPR, MERGE_IN_PROGRESS_LABEL
+
+
+def parse_args() -> Any:
+    parser = argparse.ArgumentParser(
+        description="Perform actions when a merge workflow is cancelled"
+    )
+    parser.add_argument(
+        "--pr-num",
+        type=int,
+        required=True,
+        help="The PR number to cancel the merge for",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    repo = GitRepo(get_git_repo_dir(), get_git_remote_name(), debug=True)
+    org, project = repo.gh_owner_and_name()
+    pr_num = args.pr_num
+
+    GitHubPR(org, project, pr_num).remove_label(MERGE_IN_PROGRESS_LABEL)
+
+
+if __name__ == "__main__":
+    main()

.github/workflows/_mac-build.yml

@@ -27,7 +27,7 @@ on:
         description: |
           If this is set, our linter will use this to make sure that every other
           job with the same `sync-tag` is identical.
-      python_version:
+      python-version:
         required: false
         type: string
         default: "3.8"
@@ -58,23 +58,12 @@ on:
         value: ${{ jobs.build.outputs.build-outcome }}
         description: The outcome of the build step. This is used to influence test filtering logic later on.
 
-    secrets:
-      MACOS_SCCACHE_S3_ACCESS_KEY_ID:
-        required: true
-        description: Access key for S3 bucket for macOS sccache.
-      MACOS_SCCACHE_S3_SECRET_ACCESS_KEY:
-        required: true
-        description: Secret for S3 bucket for macOS sccache.
-
 jobs:
   build:
     # Don't run on forked repos.
     if: github.repository_owner == 'pytorch'
     runs-on: ${{ inputs.runner-type }}
     env:
-      # For sccache access (only on non-forked PRs)
-      AWS_ACCESS_KEY_ID: ${{ secrets.MACOS_SCCACHE_S3_ACCESS_KEY_ID }}
-      AWS_SECRET_ACCESS_KEY: ${{ secrets.MACOS_SCCACHE_S3_SECRET_ACCESS_KEY }}
       BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
       SCCACHE_USE_GHA: ${{ inputs.sccache-use-gha }}  # this is placed here instead of the sccache step to appease actionlint
     outputs:
@@ -99,7 +88,7 @@ jobs:
         if: inputs.environment-file == ''
         uses: pytorch/test-infra/.github/actions/setup-miniconda@main
         with:
-          python-version: ${{ inputs.python_version }}
+          python-version: ${{ inputs.python-version }}
           environment-file: .github/requirements/conda-env-${{ runner.os }}-${{ runner.arch }}
           pip-requirements-file: .github/requirements/pip-requirements-${{ runner.os }}.txt
 
@@ -109,21 +98,10 @@ jobs:
         if: inputs.environment-file != ''
         uses: pytorch/test-infra/.github/actions/setup-miniconda@main
         with:
-          python-version: ${{ inputs.python_version }}
+          python-version: ${{ inputs.python-version }}
           environment-file: ${{ inputs.environment-file }}
           pip-requirements-file: .github/requirements/pip-requirements-${{ runner.os }}.txt
 
-      - name: Install macOS homebrew dependencies
-        uses: nick-fields/retry@v2.8.2
-        with:
-          timeout_minutes: 5
-          max_attempts: 3
-          retry_wait_seconds: 90
-          command: |
-            # Install dependencies
-            brew install libomp
-            brew link --force libomp
-
       - name: Install sccache (only for non-forked PRs, and pushes to trunk)
         uses: nick-fields/retry@v2.8.2
         if: ${{ github.event_name == 'push' || github.event.pull_request.head.repo.full_name == github.repository }}
@@ -132,17 +110,40 @@ jobs:
           max_attempts: 3
           retry_wait_seconds: 90
           command: |
+            set -ex
+
+            DOWNLOAD_SCCACHE=0
+            SCCACHE_VERSION="0.4.1"
+            LOCAL_PATH="/usr/local/bin"
+
+            if [ ! -f "${LOCAL_PATH}/sccache" ]; then
+              DOWNLOAD_SCCACHE=1
+            else
+              LOCAL_VERSION=$("${LOCAL_PATH}/sccache" --version | cut -d" " -f2)
+
+              if [ "${LOCAL_VERSION}" != "${SCCACHE_VERSION}" ]; then
+                DOWNLOAD_SCCACHE=1
+              fi
+            fi
+
+            if [ "${DOWNLOAD_SCCACHE}" == "1" ]; then
+              sudo curl --retry 3 --retry-all-errors "https://s3.amazonaws.com/ossci-macos/sccache/sccache-v0.4.1-${RUNNER_ARCH}" --output "${LOCAL_PATH}/sccache"
+              sudo chmod +x "${LOCAL_PATH}/sccache"
+            fi
+
             if [[ "${SCCACHE_USE_GHA}" == "true" ]]; then
-              sudo curl --retry 3 --retry-all-errors https://s3.amazonaws.com/ossci-macos/sccache_v0.4.0-pre.7 --output /usr/local/bin/sccache
               echo "ACTIONS_CACHE_URL=${ACTIONS_CACHE_URL}" >> "${GITHUB_ENV}"
               echo "ACTIONS_RUNTIME_TOKEN=${ACTIONS_RUNTIME_TOKEN}" >> "${GITHUB_ENV}"
               echo "SCCACHE_GHA_ENABLED=on" >> "${GITHUB_ENV}"
             else
-              sudo curl --retry 3 --retry-all-errors https://s3.amazonaws.com/ossci-macos/sccache_v2.15 --output /usr/local/bin/sccache
+              # The runner has access to the S3 bucket via IAM profile without the need
+              # for any credential
               echo "SCCACHE_BUCKET=ossci-compiler-cache-circleci-v2" >> "${GITHUB_ENV}"
               echo "SCCACHE_S3_KEY_PREFIX=${GITHUB_WORKFLOW}" >> "${GITHUB_ENV}"
             fi
-            sudo chmod +x /usr/local/bin/sccache
+
+            # This is needed so that later build script could find sccache (which sccache)
+            echo "${LOCAL_PATH}" >> $GITHUB_PATH
 
       - name: Get workflow job id
         id: get-job-id

.github/workflows/_mac-test-mps.yml

@@ -14,6 +14,12 @@ on:
         description: |
           If this is set, our linter will use this to make sure that every other
           job with the same `sync-tag` is identical.
+      python-version:
+        required: false
+        type: string
+        default: "3.8"
+        description: |
+          The python version to be used. Will be 3.8 by default
       test-matrix:
         required: true
         type: string
@@ -77,24 +83,10 @@ jobs:
           name: ${{ inputs.build-environment }}
           use-gha: true
 
-      # This is copied from the main macos test workflow. It was missed in the earlier fix because macos M1
-      # runners are shared and not ephemeral, so the issue wasn't manifested if the runners with the fix were
-      # used
-      - name: Install macOS homebrew dependencies
-        uses: nick-fields/retry@v2.8.2
-        with:
-          timeout_minutes: 5
-          max_attempts: 3
-          retry_wait_seconds: 90
-          command: |
-            # Install dependencies
-            brew install libomp
-            brew link --force libomp
-
       - name: Setup miniconda
         uses: pytorch/test-infra/.github/actions/setup-miniconda@main
         with:
-          python-version: 3.9
+          python-version: ${{ inputs.python-version }}
           environment-file: .github/requirements/conda-env-${{ runner.os }}-${{ runner.arch }}
           pip-requirements-file: .github/requirements/pip-requirements-${{ runner.os }}.txt
 
@@ -117,6 +109,14 @@ jobs:
             export PATH="$CONDA_ENV/bin":$PATH
           fi
 
+          # Print out some information about the test environment
+          which conda
+          conda --version
+          ${CONDA_RUN} which python3
+          ${CONDA_RUN} python3 --version
+          ${CONDA_RUN} which python
+          ${CONDA_RUN} python --version
+
           # As wheels are cross-compiled they are reported as x86_64 ones
           ORIG_WHLNAME=$(ls -1 dist/*.whl); ARM_WHLNAME=${ORIG_WHLNAME/x86_64/arm64}; mv ${ORIG_WHLNAME} ${ARM_WHLNAME}
           ${CONDA_RUN} python3 -mpip install --no-index --no-deps dist/*.whl

.github/workflows/_mac-test.yml

@@ -18,6 +18,12 @@ on:
         description: |
           If this is set, our linter will use this to make sure that every other
           job with the same `sync-tag` is identical.
+      python-version:
+        required: false
+        type: string
+        default: "3.8"
+        description: |
+          The python version to be used. Will be 3.8 by default
       arch:
         required: true
         type: string
@@ -29,13 +35,6 @@ on:
         default: 270
         description: |
           Set the maximum (in minutes) how long the workflow should take to finish
-    secrets:
-      AWS_OSSCI_METRICS_V2_ACCESS_KEY_ID:
-        required: true
-        description: access key id for test stats upload
-      AWS_OSSCI_METRICS_V2_SECRET_ACCESS_KEY:
-        required: true
-        description: secret acess key for test stats upload
 
 jobs:
   # This needs to be run right before the test starts so that it can gather the
@@ -106,19 +105,10 @@ jobs:
           name: ${{ inputs.build-environment }}
           use-gha: true
 
-      - name: Setup miniconda (x86, py3.8)
-        if: ${{ runner.arch == 'X64' }}
-        uses: pytorch/test-infra/.github/actions/setup-miniconda@main
-        with:
-          python-version: 3.8
-          environment-file: .github/requirements/conda-env-${{ runner.os }}-${{ runner.arch }}
-          pip-requirements-file: .github/requirements/pip-requirements-${{ runner.os }}.txt
-
-      - name: Setup miniconda (arm64, py3.9)
-        if: ${{ runner.arch == 'ARM64' }}
+      - name: Setup miniconda
         uses: pytorch/test-infra/.github/actions/setup-miniconda@main
         with:
-          python-version: 3.9
+          python-version: ${{ inputs.python-version }}
           environment-file: .github/requirements/conda-env-${{ runner.os }}-${{ runner.arch }}
           pip-requirements-file: .github/requirements/pip-requirements-${{ runner.os }}.txt
 
@@ -129,17 +119,6 @@ jobs:
           ${CONDA_RUN} python3 -m tools.stats.monitor > usage_log.txt 2>&1 &
           echo "monitor-script-pid=${!}" >> "${GITHUB_OUTPUT}"
 
-      - name: Install macOS homebrew dependencies
-        uses: nick-fields/retry@v2.8.2
-        with:
-          timeout_minutes: 5
-          max_attempts: 3
-          retry_wait_seconds: 90
-          command: |
-            # Install dependencies
-            brew install libomp
-            brew link --force libomp
-
       - name: Parse ref
         id: parse-ref
         run: .github/scripts/parse_ref.py
@@ -165,6 +144,9 @@ jobs:
           PYTORCH_TEST_CUDA_MEM_LEAK_CHECK: ${{ matrix.mem_leak_check && '1' || '0' }}
           PYTORCH_TEST_RERUN_DISABLED_TESTS: ${{ matrix.rerun_disabled_tests && '1' || '0' }}
         run: |
+          # shellcheck disable=SC1090
+          set -ex
+
           COMMIT_MESSAGES=$(git cherry -v "origin/${GIT_DEFAULT_BRANCH:-master}")
 
           # sanitize the input commit message and PR body here:
@@ -185,6 +167,14 @@ jobs:
             export PATH="$CONDA_ENV/bin":$PATH
           fi
 
+          # Print out some information about the test environment
+          which conda
+          conda --version
+          ${CONDA_RUN} which python3
+          ${CONDA_RUN} python3 --version
+          ${CONDA_RUN} which python
+          ${CONDA_RUN} python --version
+
           ${CONDA_RUN} python3 -mpip install --no-index --no-deps $(echo dist/*.whl)
           ${CONDA_RUN} .ci/pytorch/macos-test.sh
 

.github/workflows/mac-mps.yml

@@ -20,7 +20,7 @@ jobs:
       runner-type: macos-m1-12
       build-generates-artifacts: true
       # To match the one pre-installed in the m1 runners
-      python_version: 3.9.12
+      python-version: 3.9.12
       # We need to set the environment file here instead of trying to detect it automatically because
       # MacOS arm64 is cross-compiled from x86-64. Specifically, it means that arm64 conda environment
       # is needed when building PyTorch MacOS arm64 from x86-64
@@ -29,9 +29,6 @@ jobs:
         { include: [
           { config: "default", shard: 1, num_shards: 1 },
         ]}
-    secrets:
-      MACOS_SCCACHE_S3_ACCESS_KEY_ID: ${{ secrets.MACOS_SCCACHE_S3_ACCESS_KEY_ID }}
-      MACOS_SCCACHE_S3_SECRET_ACCESS_KEY: ${{ secrets.MACOS_SCCACHE_S3_SECRET_ACCESS_KEY }}
 
   macos-12-py3-arm64-mps-test:
     name: macos-12-py3-arm64-mps
@@ -40,6 +37,8 @@ jobs:
     with:
       sync-tag: macos-12-py3-arm64-mps-test
       build-environment: macos-12-py3-arm64
+      # Same as the build job
+      python-version: 3.9.12
       test-matrix: ${{ needs.macos-12-py3-arm64-build.outputs.test-matrix }}
 
   macos-13-py3-arm64-mps-test:
@@ -48,5 +47,7 @@ jobs:
     needs: macos-12-py3-arm64-build
     with:
       build-environment: macos-12-py3-arm64
+      # Same as the build job
+      python-version: 3.9.12
       test-matrix: ${{ needs.macos-12-py3-arm64-build.outputs.test-matrix }}
       runs-on: macos-m1-13

.github/workflows/periodic.yml

@@ -223,9 +223,6 @@ jobs:
           { config: "default", shard: 3, num_shards: 4, runner: "macos-12" },
           { config: "default", shard: 4, num_shards: 4, runner: "macos-12" },
         ]}
-    secrets:
-      MACOS_SCCACHE_S3_ACCESS_KEY_ID: ${{ secrets.MACOS_SCCACHE_S3_ACCESS_KEY_ID }}
-      MACOS_SCCACHE_S3_SECRET_ACCESS_KEY: ${{ secrets.MACOS_SCCACHE_S3_SECRET_ACCESS_KEY }}
 
   macos-12-py3-x86-64-test:
     name: macos-12-py3-x86-64
@@ -235,9 +232,6 @@ jobs:
       build-environment: macos-12-py3-x86-64
       test-matrix: ${{ needs.macos-12-py3-x86-64-build.outputs.test-matrix }}
       arch: x86_64
-    secrets:
-      AWS_OSSCI_METRICS_V2_ACCESS_KEY_ID: ${{ secrets.AWS_OSSCI_METRICS_V2_ACCESS_KEY_ID }}
-      AWS_OSSCI_METRICS_V2_SECRET_ACCESS_KEY: ${{ secrets.AWS_OSSCI_METRICS_V2_SECRET_ACCESS_KEY }}
 
   macos-12-py3-x86-64-lite-interpreter-build-test:
     name: macos-12-py3-x86-64-lite-interpreter
@@ -252,9 +246,6 @@ jobs:
         { include: [
           { config: "default", shard: 1, num_shards: 1, runner: "macos-12" },
         ]}
-    secrets:
-      MACOS_SCCACHE_S3_ACCESS_KEY_ID: ${{ secrets.MACOS_SCCACHE_S3_ACCESS_KEY_ID }}
-      MACOS_SCCACHE_S3_SECRET_ACCESS_KEY: ${{ secrets.MACOS_SCCACHE_S3_SECRET_ACCESS_KEY }}
 
   android-emulator-build-test:
     name: android-emulator-build-test