allow providing full fr trace path

Summary:
- allow users to specify the full path instead of fr suffixing the rank id
- this will be used by torchft to provide the global rank id accross all replicas
- we can't just prefix the replica id because analysis tool expects the file name to provide a unique integer

.ci/docker/build.sh

@@ -113,7 +113,6 @@ case "$tag" in
     UCX_COMMIT=${_UCX_COMMIT}
     UCC_COMMIT=${_UCC_COMMIT}
     TRITON=yes
-    INSTALL_MINGW=yes
     ;;
   pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11)
     CUDA_VERSION=13.0.0
@@ -362,7 +361,6 @@ docker build \
        --build-arg "OPENBLAS=${OPENBLAS:-}" \
        --build-arg "SKIP_SCCACHE_INSTALL=${SKIP_SCCACHE_INSTALL:-}" \
        --build-arg "SKIP_LLVM_SRC_BUILD_INSTALL=${SKIP_LLVM_SRC_BUILD_INSTALL:-}" \
-       --build-arg "INSTALL_MINGW=${INSTALL_MINGW:-}" \
        -f $(dirname ${DOCKERFILE})/Dockerfile \
        -t "$tmp_tag" \
        "$@" \

.ci/docker/common/install_mingw.sh

@@ -1,10 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-# Install MinGW-w64 for Windows cross-compilation
-apt-get update
-apt-get install -y g++-mingw-w64-x86-64-posix
-
-echo "MinGW-w64 installed successfully"
-x86_64-w64-mingw32-g++ --version

.ci/docker/common/install_onnx.sh

@@ -20,7 +20,7 @@ pip_install \
 
 pip_install coloredlogs packaging
 pip_install onnxruntime==1.23.0
-pip_install onnxscript==0.5.4
+pip_install onnxscript==0.5.3
 
 # Cache the transformers model to be used later by ONNX tests. We need to run the transformers
 # package to download the model. By default, the model is cached at ~/.cache/huggingface/hub/

.ci/docker/libtorch/build.sh

@@ -39,13 +39,9 @@ case ${DOCKER_TAG_PREFIX} in
         DOCKER_GPU_BUILD_ARG=""
         ;;
     rocm*)
-        # we want the patch version of 7.0 instead
-        if [[ "$GPU_ARCH_VERSION" == *"7.0"* ]]; then
-            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2"
-        fi
         # we want the patch version of 6.4 instead
         if [[ "$GPU_ARCH_VERSION" == *"6.4"* ]]; then
-            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.4"
+            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2"
         fi
         BASE_TARGET=rocm
         GPU_IMAGE=rocm/dev-ubuntu-22.04:${GPU_ARCH_VERSION}-complete

.ci/docker/manywheel/build.sh

@@ -75,13 +75,9 @@ case ${image} in
         DOCKERFILE_SUFFIX="_cuda_aarch64"
         ;;
     manylinux2_28-builder:rocm*)
-        # we want the patch version of 7.0 instead
-        if [[ "$GPU_ARCH_VERSION" == *"7.0"* ]]; then
-            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2"
-        fi
         # we want the patch version of 6.4 instead
         if [[ "$GPU_ARCH_VERSION" == *"6.4"* ]]; then
-            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.4"
+            GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.2"
         fi
         TARGET=rocm_final
         MANY_LINUX_VERSION="2_28"

.ci/docker/ubuntu/Dockerfile

@@ -103,11 +103,6 @@ COPY ci_commit_pins/torchbench.txt torchbench.txt
 RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_deps.sh; fi
 RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface-requirements.txt torchbench.txt
 
-ARG INSTALL_MINGW
-COPY ./common/install_mingw.sh install_mingw.sh
-RUN if [ -n "${INSTALL_MINGW}" ]; then bash ./install_mingw.sh; fi
-RUN rm install_mingw.sh
-
 ARG TRITON
 ARG TRITON_CPU
 

.ci/pytorch/test.sh

@@ -485,22 +485,6 @@ test_inductor_aoti() {
   /usr/bin/env "${TEST_ENVS[@]}" python test/run_test.py --cpp --verbose -i cpp/test_aoti_abi_check cpp/test_aoti_inference cpp/test_vec_half_AVX2 -dist=loadfile
 }
 
-test_inductor_aoti_cross_compile_for_windows() {
-
-  TEST_REPORTS_DIR=$(pwd)/test/test-reports
-  mkdir -p "$TEST_REPORTS_DIR"
-
-  # Set WINDOWS_CUDA_HOME environment variable
-  WINDOWS_CUDA_HOME="$(pwd)/win-torch-wheel-extracted"
-  export WINDOWS_CUDA_HOME
-
-  echo "WINDOWS_CUDA_HOME is set to: $WINDOWS_CUDA_HOME"
-  echo "Contents:"
-  ls -lah "$(pwd)/win-torch-wheel-extracted/lib/x64/" || true
-
-  python test/inductor/test_aoti_cross_compile_windows.py -k compile --package-dir "$TEST_REPORTS_DIR" --win-torch-lib-dir "$(pwd)/win-torch-wheel-extracted/torch/lib"
-}
-
 test_inductor_cpp_wrapper_shard() {
   if [[ -z "$NUM_TEST_SHARDS" ]]; then
     echo "NUM_TEST_SHARDS must be defined to run a Python test shard"
@@ -916,7 +900,7 @@ test_inductor_set_cpu_affinity(){
   export LD_PRELOAD="$JEMALLOC_LIB":"$LD_PRELOAD"
   export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms:-1,muzzy_decay_ms:-1"
 
-  if [[ "$(uname -m)" != "aarch64" ]]; then
+  if [[ "${TEST_CONFIG}" != *aarch64* ]]; then
     # Use Intel OpenMP for x86
     IOMP_LIB="$(dirname "$(which python)")/../lib/libiomp5.so"
     export LD_PRELOAD="$IOMP_LIB":"$LD_PRELOAD"
@@ -930,7 +914,7 @@ test_inductor_set_cpu_affinity(){
   cores=$((cpus / thread_per_core))
 
   # Set number of cores to 16 on aarch64 for performance runs
-  if [[ "$(uname -m)" == "aarch64" && $cores -gt 16 ]]; then
+  if [[ "${TEST_CONFIG}" == *aarch64* && $cores -gt 16 ]]; then
     cores=16
   fi
   export OMP_NUM_THREADS=$cores
@@ -1683,7 +1667,7 @@ if [[ "${TEST_CONFIG}" == *numpy_2* ]]; then
     python -m pip install --pre numpy==2.0.2 scipy==1.13.1 numba==0.60.0
   fi
   python test/run_test.py --include dynamo/test_functions.py dynamo/test_unspec.py test_binary_ufuncs.py test_fake_tensor.py test_linalg.py test_numpy_interop.py test_tensor_creation_ops.py test_torch.py torch_np/test_basic.py
-elif [[ "${BUILD_ENVIRONMENT}" == *aarch64* && "${TEST_CONFIG}" == 'default' ]]; then
+elif [[ "${BUILD_ENVIRONMENT}" == *aarch64* && "${TEST_CONFIG}" != *perf_cpu_aarch64* ]]; then
   test_linux_aarch64
 elif [[ "${TEST_CONFIG}" == *backward* ]]; then
   test_forward_backward_compatibility
@@ -1734,8 +1718,6 @@ elif [[ "${TEST_CONFIG}" == *inductor-triton-cpu* ]]; then
   test_inductor_triton_cpu
 elif [[ "${TEST_CONFIG}" == *inductor-micro-benchmark* ]]; then
   test_inductor_micro_benchmark
-elif [[ "${TEST_CONFIG}" == *aoti_cross_compile_for_windows* ]]; then
-  test_inductor_aoti_cross_compile_for_windows
 elif [[ "${TEST_CONFIG}" == *huggingface* ]]; then
   install_torchvision
   id=$((SHARD_NUMBER-1))

.flake8

@@ -13,6 +13,8 @@ ignore =
     EXE001,
     # these ignores are from flake8-bugbear; please fix!
     B007,B008,B017,B019,B023,B028,B903,B905,B906,B907,B908,B910
+    # these ignores are from flake8-comprehensions; please fix!
+    C407,
     # these ignores are from flake8-logging-format; please fix!
     G100,G101,G200
     # these ignores are from flake8-simplify. please fix or ignore with commented reason

.github/pytorch-probot.yml

@@ -3,7 +3,6 @@ ciflow_tracking_issue: 64124
 ciflow_push_tags:
 - ciflow/b200
 - ciflow/b200-symm-mem
-- ciflow/b200-distributed
 - ciflow/binaries
 - ciflow/binaries_libtorch
 - ciflow/binaries_wheel

.github/scripts/generate_binary_build_matrix.py

@@ -241,11 +241,7 @@ def generate_libtorch_matrix(
             arches += CUDA_ARCHES
             arches += ROCM_ARCHES
         elif os == "windows":
-            # TODO (huydhn): Only build CUDA 12.9 for Linux. This logic is to be cleaned up
-            # in 2.10
-            windows_cuda_arches = CUDA_ARCHES.copy()
-            windows_cuda_arches.remove("12.9")
-            arches += windows_cuda_arches
+            arches += CUDA_ARCHES
     if libtorch_variants is None:
         libtorch_variants = [
             "shared-with-deps",
@@ -309,11 +305,7 @@ def generate_wheels_matrix(
         if os == "linux":
             arches += CUDA_ARCHES + ROCM_ARCHES + XPU_ARCHES
         elif os == "windows":
-            # TODO (huydhn): Only build CUDA 12.9 for Linux. This logic is to be cleaned up
-            # in 2.10
-            windows_cuda_arches = CUDA_ARCHES.copy()
-            windows_cuda_arches.remove("12.9")
-            arches += windows_cuda_arches + XPU_ARCHES
+            arches += CUDA_ARCHES + XPU_ARCHES
         elif os == "linux-aarch64":
             # Separate new if as the CPU type is different and
             # uses different build/test scripts

.github/scripts/trymerge.py

@@ -1092,7 +1092,7 @@ class GitHubPR:
         editor = node["editor"]
         return GitHubComment(
             body_text=node["bodyText"],
-            created_at=node.get("createdAt", ""),
+            created_at=node["createdAt"] if "createdAt" in node else "",
             author_login=node["author"]["login"],
             author_url=node["author"].get("url", None),
             author_association=node["authorAssociation"],

.github/workflows/_linux-build.yml

@@ -37,7 +37,7 @@ on:
       runner:
         required: false
         type: string
-        default: "linux.c7i.2xlarge"
+        default: "linux.2xlarge"
         description: |
           Label of the runner this job should run on.
       test-matrix:

.github/workflows/_linux-test.yml

@@ -224,46 +224,6 @@ jobs:
         continue-on-error: true
         uses: ./.github/actions/download-td-artifacts
 
-      - name: Download Windows torch wheel for cross-compilation
-        if: matrix.win_torch_wheel_artifact != ''
-        uses: seemethere/download-artifact-s3@1da556a7aa0a088e3153970611f6c432d58e80e6 # v4.2.0
-        with:
-          name: ${{ matrix.win_torch_wheel_artifact }}
-          path: win-torch-wheel
-
-      - name: Extract Windows wheel and setup CUDA libraries
-        if: matrix.win_torch_wheel_artifact != ''
-        shell: bash
-        run: |
-          set -x
-
-          # Find the wheel file
-          WHEEL_FILE=$(find win-torch-wheel -name "*.whl" -type f | head -n 1)
-          if [ -z "$WHEEL_FILE" ]; then
-            echo "Error: No wheel file found in win-torch-wheel directory"
-            exit 1
-          fi
-          echo "Found wheel file: $WHEEL_FILE"
-
-          # Unzip the wheel file
-          unzip -q "$WHEEL_FILE" -d win-torch-wheel-extracted
-          echo "Extracted wheel contents"
-
-          # Setup CUDA libraries (cuda.lib and cudart.lib) directory
-          mkdir -p win-torch-wheel-extracted/lib/x64
-          if [ -f "win-torch-wheel/cuda.lib" ]; then
-            mv win-torch-wheel/cuda.lib win-torch-wheel-extracted/lib/x64/
-            echo "Moved cuda.lib to win-torch-wheel-extracted/lib/x64/"
-          fi
-          if [ -f "win-torch-wheel/cudart.lib" ]; then
-            mv win-torch-wheel/cudart.lib win-torch-wheel-extracted/lib/x64/
-            echo "Moved cudart.lib to win-torch-wheel-extracted/lib/x64/"
-          fi
-
-          # Verify CUDA libraries are present
-          echo "CUDA libraries:"
-          ls -la win-torch-wheel-extracted/lib/x64/ || echo "No CUDA libraries found"
-
       - name: Parse ref
         id: parse-ref
         run: .github/scripts/parse_ref.py

.github/workflows/_win-build.yml

@@ -168,31 +168,6 @@ jobs:
         run: |
           .ci/pytorch/win-build.sh
 
-      # Collect Windows torch libs and CUDA libs for cross-compilation
-      - name: Collect Windows CUDA libs for cross-compilation
-        if: steps.build.outcome != 'skipped' && inputs.cuda-version != 'cpu'
-        shell: bash
-        run: |
-          set -ex
-
-          # Create directory structure if does not exist
-          mkdir -p /c/${{ github.run_id }}/build-results
-
-          # Copy CUDA libs
-          CUDA_PATH="/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v${{ inputs.cuda-version }}"
-
-          if [ -f "${CUDA_PATH}/lib/x64/cuda.lib" ]; then
-            cp "${CUDA_PATH}/lib/x64/cuda.lib" /c/${{ github.run_id }}/build-results/
-          fi
-
-          if [ -f "${CUDA_PATH}/lib/x64/cudart.lib" ]; then
-            cp "${CUDA_PATH}/lib/x64/cudart.lib" /c/${{ github.run_id }}/build-results/
-          fi
-
-          # List collected files
-          echo "Collected CUDA libs:"
-          ls -lah /c/${{ github.run_id }}/build-results/*.lib
-
       # Upload to github so that people can click and download artifacts
       - name: Upload artifacts to s3
         if: steps.build.outcome != 'skipped'

.github/workflows/b200-distributed.yml

@@ -1,62 +0,0 @@
-name: CI for distributed tests on B200
-
-on:
-  pull_request:
-    paths:
-      - .github/workflows/b200-distributed.yml
-  workflow_dispatch:
-  push:
-    tags:
-      - ciflow/b200-distributed/*
-  schedule:
-    - cron: 46 8 * * *  # about 1:46am PDT
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }}
-  cancel-in-progress: true
-
-permissions:
-  id-token: write
-  contents: read
-
-jobs:
-
-  get-label-type:
-    if: github.repository_owner == 'pytorch'
-    name: get-label-type
-    uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main
-    with:
-      triggering_actor: ${{ github.triggering_actor }}
-      issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
-      curr_branch: ${{ github.head_ref || github.ref_name }}
-      curr_ref_type: ${{ github.ref_type }}
-
-  linux-jammy-cuda12_8-py3_10-gcc11-build-distributed-b200:
-    name: linux-jammy-cuda12.8-py3.10-gcc11-build-distributed-b200
-    uses: ./.github/workflows/_linux-build.yml
-    needs: get-label-type
-    with:
-      runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
-      runner: linux.12xlarge.memory
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-distributed-b200
-      docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11
-      cuda-arch-list: '10.0'
-      test-matrix: |
-        { include: [
-          { config: "distributed", shard: 1, num_shards: 2, runner: "linux.dgx.b200.8" },
-          { config: "distributed", shard: 2, num_shards: 2, runner: "linux.dgx.b200.8" },
-        ]}
-    secrets: inherit
-
-  linux-jammy-cuda12_8-py3_10-gcc11-test-distributed-b200:
-    name: linux-jammy-cuda12.8-py3.10-gcc11-test-b200
-    uses: ./.github/workflows/_linux-test.yml
-    needs:
-      - linux-jammy-cuda12_8-py3_10-gcc11-build-distributed-b200
-    with:
-      timeout-minutes: 1200
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11-distributed-b200
-      docker-image: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build-distributed-b200.outputs.docker-image }}
-      test-matrix: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build-distributed-b200.outputs.test-matrix }}
-      aws-role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only
-    secrets: inherit

.github/workflows/generated-windows-binary-libtorch-debug-nightly.yml

@@ -788,6 +788,256 @@ jobs:
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
     uses: ./.github/workflows/_binary-upload.yml
+  libtorch-cuda12_9-shared-with-deps-debug-build:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs: get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
+    timeout-minutes: 360
+    env:
+      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PACKAGE_TYPE: libtorch
+      # TODO: This is a legacy variable that we eventually want to get rid of in
+      #       favor of GPU_ARCH_VERSION
+      DESIRED_CUDA: cu129
+      GPU_ARCH_VERSION: "12.9"
+      GPU_ARCH_TYPE: cuda
+      SKIP_ALL_TESTS: 1
+      LIBTORCH_CONFIG: debug
+      LIBTORCH_VARIANT: shared-with-deps
+      # This is a dummy value for libtorch to work correctly with our batch scripts
+      # without this value pip does not get installed for some reason
+      DESIRED_PYTHON: "3.10"
+    steps:
+      # NOTE: These environment variables are put here so that they can be applied on every job equally
+      #       They are also here because setting them at a workflow level doesn't give us access to the
+      #       runner.temp variable, which we need.
+      - name: Populate binary env
+        shell: bash
+        run: |
+          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
+          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
+          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
+      - name: Display EC2 information
+        shell: bash
+        run: |
+          set -euo pipefail
+          function get_ec2_metadata() {
+            # Pulled from instance metadata endpoint for EC2
+            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
+            category=$1
+            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
+          }
+          echo "ami-id: $(get_ec2_metadata ami-id)"
+          echo "instance-id: $(get_ec2_metadata instance-id)"
+          echo "instance-type: $(get_ec2_metadata instance-type)"
+          echo "system info $(uname -a)"
+      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
+        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        continue-on-error: true
+        with:
+          github-secret: ${{ secrets.GITHUB_TOKEN }}
+      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
+        shell: bash
+        run: |
+          git config --global core.longpaths true
+          git config --global core.symlinks true
+
+          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
+          # the directory on Windows and prevent GHA from checking out as reported
+          # in https://github.com/actions/checkout/issues/1018
+          git config --global core.fsmonitor false
+      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
+      - name: Enable long paths on Windows
+        shell: powershell
+        run: |
+          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
+      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
+      # removed once Windows Defender is removed from the AMI
+      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
+        continue-on-error: true
+        shell: powershell
+        run: |
+          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
+          # Let's both exclude the path and disable Windows Defender completely just to be sure
+          # that it doesn't interfere
+          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          submodules: recursive
+          path: pytorch
+          show-progress: false
+      - name: Clean PyTorch checkout
+        run: |
+          # Remove any artifacts from the previous checkouts
+          git clean -fxd
+        working-directory: pytorch
+      - name: Populate binary env
+        shell: bash
+        run: |
+          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+      - name: Build PyTorch binary
+        shell: bash
+        run: |
+          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cuda12_9-shared-with-deps-debug
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
+      - name: Wait until all sessions have drained
+        shell: powershell
+        working-directory: pytorch
+        if: always()
+        timeout-minutes: 120
+        run: |
+          .github\scripts\wait_for_ssh_to_drain.ps1
+      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
+        shell: powershell
+        working-directory: pytorch
+        if: always()
+        run: |
+          .github\scripts\kill_active_ssh_sessions.ps1
+
+  libtorch-cuda12_9-shared-with-deps-debug-test:  # Testing
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - libtorch-cuda12_9-shared-with-deps-debug-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
+    timeout-minutes: 360
+    env:
+      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PACKAGE_TYPE: libtorch
+      # TODO: This is a legacy variable that we eventually want to get rid of in
+      #       favor of GPU_ARCH_VERSION
+      DESIRED_CUDA: cu129
+      GPU_ARCH_VERSION: "12.9"
+      GPU_ARCH_TYPE: cuda
+      SKIP_ALL_TESTS: 1
+      LIBTORCH_CONFIG: debug
+      LIBTORCH_VARIANT: shared-with-deps
+      # This is a dummy value for libtorch to work correctly with our batch scripts
+      # without this value pip does not get installed for some reason
+      DESIRED_PYTHON: "3.10"
+    steps:
+      - name: Display EC2 information
+        shell: bash
+        run: |
+          set -euo pipefail
+          function get_ec2_metadata() {
+            # Pulled from instance metadata endpoint for EC2
+            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
+            category=$1
+            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
+          }
+          echo "ami-id: $(get_ec2_metadata ami-id)"
+          echo "instance-id: $(get_ec2_metadata instance-id)"
+          echo "instance-type: $(get_ec2_metadata instance-type)"
+          echo "system info $(uname -a)"
+      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
+        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        continue-on-error: true
+        with:
+          github-secret: ${{ secrets.GITHUB_TOKEN }}
+      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
+        shell: bash
+        run: |
+          git config --global core.longpaths true
+          git config --global core.symlinks true
+
+          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
+          # the directory on Windows and prevent GHA from checking out as reported
+          # in https://github.com/actions/checkout/issues/1018
+          git config --global core.fsmonitor false
+      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
+      - name: Enable long paths on Windows
+        shell: powershell
+        run: |
+          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
+      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
+      # removed once Windows Defender is removed from the AMI
+      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
+        continue-on-error: true
+        shell: powershell
+        run: |
+          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
+          # Let's both exclude the path and disable Windows Defender completely just to be sure
+          # that it doesn't interfere
+          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          submodules: recursive
+          path: pytorch
+          show-progress: false
+      - name: Clean PyTorch checkout
+        run: |
+          # Remove any artifacts from the previous checkouts
+          git clean -fxd
+        working-directory: pytorch
+      # NOTE: These environment variables are put here so that they can be applied on every job equally
+      #       They are also here because setting them at a workflow level doesn't give us access to the
+      #       runner.temp variable, which we need.
+      - name: Populate binary env
+        shell: bash
+        run: |
+          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
+          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
+          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: libtorch-cuda12_9-shared-with-deps-debug
+          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
+      - name: Populate binary env
+        shell: bash
+        run: |
+          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+      - name: Test PyTorch binary
+        shell: bash
+        run: |
+          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+      - name: Wait until all sessions have drained
+        shell: powershell
+        working-directory: pytorch
+        if: always()
+        timeout-minutes: 120
+        run: |
+          .github\scripts\wait_for_ssh_to_drain.ps1
+      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
+        shell: powershell
+        working-directory: pytorch
+        if: always()
+        run: |
+          .github\scripts\kill_active_ssh_sessions.ps1
+  libtorch-cuda12_9-shared-with-deps-debug-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cuda12_9-shared-with-deps-debug-test
+    with:
+      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PACKAGE_TYPE: libtorch
+      # TODO: This is a legacy variable that we eventually want to get rid of in
+      #       favor of GPU_ARCH_VERSION
+      DESIRED_CUDA: cu129
+      GPU_ARCH_VERSION: "12.9"
+      GPU_ARCH_TYPE: cuda
+      LIBTORCH_CONFIG: debug
+      LIBTORCH_VARIANT: shared-with-deps
+      # This is a dummy value for libtorch to work correctly with our batch scripts
+      # without this value pip does not get installed for some reason
+      DESIRED_PYTHON: "3.10"
+      build_name: libtorch-cuda12_9-shared-with-deps-debug
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+    uses: ./.github/workflows/_binary-upload.yml
   libtorch-cuda13_0-shared-with-deps-debug-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type

.github/workflows/generated-windows-binary-libtorch-release-nightly.yml

@@ -788,6 +788,256 @@ jobs:
     secrets:
       github-token: ${{ secrets.GITHUB_TOKEN }}
     uses: ./.github/workflows/_binary-upload.yml
+  libtorch-cuda12_9-shared-with-deps-release-build:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs: get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.4xlarge"
+    timeout-minutes: 360
+    env:
+      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PACKAGE_TYPE: libtorch
+      # TODO: This is a legacy variable that we eventually want to get rid of in
+      #       favor of GPU_ARCH_VERSION
+      DESIRED_CUDA: cu129
+      GPU_ARCH_VERSION: "12.9"
+      GPU_ARCH_TYPE: cuda
+      SKIP_ALL_TESTS: 1
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      # This is a dummy value for libtorch to work correctly with our batch scripts
+      # without this value pip does not get installed for some reason
+      DESIRED_PYTHON: "3.10"
+    steps:
+      # NOTE: These environment variables are put here so that they can be applied on every job equally
+      #       They are also here because setting them at a workflow level doesn't give us access to the
+      #       runner.temp variable, which we need.
+      - name: Populate binary env
+        shell: bash
+        run: |
+          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
+          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
+          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
+      - name: Display EC2 information
+        shell: bash
+        run: |
+          set -euo pipefail
+          function get_ec2_metadata() {
+            # Pulled from instance metadata endpoint for EC2
+            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
+            category=$1
+            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
+          }
+          echo "ami-id: $(get_ec2_metadata ami-id)"
+          echo "instance-id: $(get_ec2_metadata instance-id)"
+          echo "instance-type: $(get_ec2_metadata instance-type)"
+          echo "system info $(uname -a)"
+      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
+        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        continue-on-error: true
+        with:
+          github-secret: ${{ secrets.GITHUB_TOKEN }}
+      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
+        shell: bash
+        run: |
+          git config --global core.longpaths true
+          git config --global core.symlinks true
+
+          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
+          # the directory on Windows and prevent GHA from checking out as reported
+          # in https://github.com/actions/checkout/issues/1018
+          git config --global core.fsmonitor false
+      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
+      - name: Enable long paths on Windows
+        shell: powershell
+        run: |
+          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
+      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
+      # removed once Windows Defender is removed from the AMI
+      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
+        continue-on-error: true
+        shell: powershell
+        run: |
+          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
+          # Let's both exclude the path and disable Windows Defender completely just to be sure
+          # that it doesn't interfere
+          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          submodules: recursive
+          path: pytorch
+          show-progress: false
+      - name: Clean PyTorch checkout
+        run: |
+          # Remove any artifacts from the previous checkouts
+          git clean -fxd
+        working-directory: pytorch
+      - name: Populate binary env
+        shell: bash
+        run: |
+          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+      - name: Build PyTorch binary
+        shell: bash
+        run: |
+          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_build.sh"
+      - uses: actions/upload-artifact@v4.4.0
+        if: always()
+        with:
+          name: libtorch-cuda12_9-shared-with-deps-release
+          retention-days: 14
+          if-no-files-found: error
+          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
+      - name: Wait until all sessions have drained
+        shell: powershell
+        working-directory: pytorch
+        if: always()
+        timeout-minutes: 120
+        run: |
+          .github\scripts\wait_for_ssh_to_drain.ps1
+      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
+        shell: powershell
+        working-directory: pytorch
+        if: always()
+        run: |
+          .github\scripts\kill_active_ssh_sessions.ps1
+
+  libtorch-cuda12_9-shared-with-deps-release-test:  # Testing
+    if: ${{ github.repository_owner == 'pytorch' }}
+    needs:
+      - libtorch-cuda12_9-shared-with-deps-release-build
+      - get-label-type
+    runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge"
+    timeout-minutes: 360
+    env:
+      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PACKAGE_TYPE: libtorch
+      # TODO: This is a legacy variable that we eventually want to get rid of in
+      #       favor of GPU_ARCH_VERSION
+      DESIRED_CUDA: cu129
+      GPU_ARCH_VERSION: "12.9"
+      GPU_ARCH_TYPE: cuda
+      SKIP_ALL_TESTS: 1
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      # This is a dummy value for libtorch to work correctly with our batch scripts
+      # without this value pip does not get installed for some reason
+      DESIRED_PYTHON: "3.10"
+    steps:
+      - name: Display EC2 information
+        shell: bash
+        run: |
+          set -euo pipefail
+          function get_ec2_metadata() {
+            # Pulled from instance metadata endpoint for EC2
+            # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
+            category=$1
+            curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
+          }
+          echo "ami-id: $(get_ec2_metadata ami-id)"
+          echo "instance-id: $(get_ec2_metadata instance-id)"
+          echo "instance-type: $(get_ec2_metadata instance-type)"
+          echo "system info $(uname -a)"
+      - name: "[FB EMPLOYEES] Enable SSH (Click me for login details)"
+        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        continue-on-error: true
+        with:
+          github-secret: ${{ secrets.GITHUB_TOKEN }}
+      - name: Enable git long paths and symlinks on Windows and disable fsmonitor daemon
+        shell: bash
+        run: |
+          git config --global core.longpaths true
+          git config --global core.symlinks true
+
+          # https://git-scm.com/docs/git-fsmonitor--daemon.  The daemon could lock
+          # the directory on Windows and prevent GHA from checking out as reported
+          # in https://github.com/actions/checkout/issues/1018
+          git config --global core.fsmonitor false
+      # Needed for binary builds, see: https://github.com/pytorch/pytorch/issues/73339#issuecomment-1058981560
+      - name: Enable long paths on Windows
+        shell: powershell
+        run: |
+          Set-ItemProperty -Path "HKLM:\\SYSTEM\CurrentControlSet\Control\FileSystem" -Name "LongPathsEnabled" -Value 1
+      # Since it's just a defensive command, the workflow should continue even the command fails. This step can be
+      # removed once Windows Defender is removed from the AMI
+      - name: Disables Windows Defender scheduled and real-time scanning for files in directories used by PyTorch
+        continue-on-error: true
+        shell: powershell
+        run: |
+          Add-MpPreference -ExclusionPath $(Get-Location).tostring(),$Env:TEMP -ErrorAction Ignore
+          # Let's both exclude the path and disable Windows Defender completely just to be sure
+          # that it doesn't interfere
+          Set-MpPreference -DisableRealtimeMonitoring $True -ErrorAction Ignore
+      - name: Checkout PyTorch
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+          submodules: recursive
+          path: pytorch
+          show-progress: false
+      - name: Clean PyTorch checkout
+        run: |
+          # Remove any artifacts from the previous checkouts
+          git clean -fxd
+        working-directory: pytorch
+      # NOTE: These environment variables are put here so that they can be applied on every job equally
+      #       They are also here because setting them at a workflow level doesn't give us access to the
+      #       runner.temp variable, which we need.
+      - name: Populate binary env
+        shell: bash
+        run: |
+          echo "BINARY_ENV_FILE=${RUNNER_TEMP}/env" >> "${GITHUB_ENV}"
+          echo "PYTORCH_FINAL_PACKAGE_DIR=${RUNNER_TEMP}/artifacts" >> "${GITHUB_ENV}"
+          echo "WIN_PACKAGE_WORK_DIR=${RUNNER_TEMP}"
+      - uses: actions/download-artifact@v4.1.7
+        name: Download Build Artifacts
+        with:
+          name: libtorch-cuda12_9-shared-with-deps-release
+          path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
+      - name: Populate binary env
+        shell: bash
+        run: |
+          "${PYTORCH_ROOT}/.circleci/scripts/binary_populate_env.sh"
+      - name: Test PyTorch binary
+        shell: bash
+        run: |
+          "${PYTORCH_ROOT}/.circleci/scripts/binary_windows_test.sh"
+      - name: Wait until all sessions have drained
+        shell: powershell
+        working-directory: pytorch
+        if: always()
+        timeout-minutes: 120
+        run: |
+          .github\scripts\wait_for_ssh_to_drain.ps1
+      - name: Kill active ssh sessions if still around (Useful if workflow was cancelled)
+        shell: powershell
+        working-directory: pytorch
+        if: always()
+        run: |
+          .github\scripts\kill_active_ssh_sessions.ps1
+  libtorch-cuda12_9-shared-with-deps-release-upload:  # Uploading
+    if: ${{ github.repository_owner == 'pytorch' }}
+    permissions:
+      id-token: write
+      contents: read
+    needs: libtorch-cuda12_9-shared-with-deps-release-test
+    with:
+      PYTORCH_ROOT: ${{ github.workspace }}/pytorch
+      PACKAGE_TYPE: libtorch
+      # TODO: This is a legacy variable that we eventually want to get rid of in
+      #       favor of GPU_ARCH_VERSION
+      DESIRED_CUDA: cu129
+      GPU_ARCH_VERSION: "12.9"
+      GPU_ARCH_TYPE: cuda
+      LIBTORCH_CONFIG: release
+      LIBTORCH_VARIANT: shared-with-deps
+      # This is a dummy value for libtorch to work correctly with our batch scripts
+      # without this value pip does not get installed for some reason
+      DESIRED_PYTHON: "3.10"
+      build_name: libtorch-cuda12_9-shared-with-deps-release
+    secrets:
+      github-token: ${{ secrets.GITHUB_TOKEN }}
+    uses: ./.github/workflows/_binary-upload.yml
   libtorch-cuda13_0-shared-with-deps-release-build:
     if: ${{ github.repository_owner == 'pytorch' }}
     needs: get-label-type

.github/workflows/inductor-perf-test-nightly-rocm-mi355.yml

@@ -88,27 +88,27 @@ jobs:
       docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3-benchmarks
       test-matrix: |
         { include: [
-          { config: "inductor_huggingface_perf_rocm_mi355", shard: 1, num_shards: 5, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_huggingface_perf_rocm_mi355", shard: 2, num_shards: 5, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_huggingface_perf_rocm_mi355", shard: 3, num_shards: 5, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_huggingface_perf_rocm_mi355", shard: 4, num_shards: 5, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_huggingface_perf_rocm_mi355", shard: 5, num_shards: 5, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_timm_perf_rocm_mi355", shard: 1, num_shards: 7, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_timm_perf_rocm_mi355", shard: 2, num_shards: 7, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_timm_perf_rocm_mi355", shard: 3, num_shards: 7, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_timm_perf_rocm_mi355", shard: 4, num_shards: 7, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_timm_perf_rocm_mi355", shard: 5, num_shards: 7, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_timm_perf_rocm_mi355", shard: 6, num_shards: 7, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_timm_perf_rocm_mi355", shard: 7, num_shards: 7, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_torchbench_perf_rocm_mi355", shard: 1, num_shards: 9, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_torchbench_perf_rocm_mi355", shard: 2, num_shards: 9, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_torchbench_perf_rocm_mi355", shard: 3, num_shards: 9, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_torchbench_perf_rocm_mi355", shard: 4, num_shards: 9, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_torchbench_perf_rocm_mi355", shard: 5, num_shards: 9, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_torchbench_perf_rocm_mi355", shard: 6, num_shards: 9, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_torchbench_perf_rocm_mi355", shard: 7, num_shards: 9, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_torchbench_perf_rocm_mi355", shard: 8, num_shards: 9, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "inductor_torchbench_perf_rocm_mi355", shard: 9, num_shards: 9, runner: "linux.rocm.gpu.mi355.1" },
+          { config: "inductor_huggingface_perf_rocm_mi355", shard: 1, num_shards: 5, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_huggingface_perf_rocm_mi355", shard: 2, num_shards: 5, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_huggingface_perf_rocm_mi355", shard: 3, num_shards: 5, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_huggingface_perf_rocm_mi355", shard: 4, num_shards: 5, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_huggingface_perf_rocm_mi355", shard: 5, num_shards: 5, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_timm_perf_rocm_mi355", shard: 1, num_shards: 7, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_timm_perf_rocm_mi355", shard: 2, num_shards: 7, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_timm_perf_rocm_mi355", shard: 3, num_shards: 7, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_timm_perf_rocm_mi355", shard: 4, num_shards: 7, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_timm_perf_rocm_mi355", shard: 5, num_shards: 7, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_timm_perf_rocm_mi355", shard: 6, num_shards: 7, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_timm_perf_rocm_mi355", shard: 7, num_shards: 7, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_torchbench_perf_rocm_mi355", shard: 1, num_shards: 9, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_torchbench_perf_rocm_mi355", shard: 2, num_shards: 9, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_torchbench_perf_rocm_mi355", shard: 3, num_shards: 9, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_torchbench_perf_rocm_mi355", shard: 4, num_shards: 9, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_torchbench_perf_rocm_mi355", shard: 5, num_shards: 9, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_torchbench_perf_rocm_mi355", shard: 6, num_shards: 9, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_torchbench_perf_rocm_mi355", shard: 7, num_shards: 9, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_torchbench_perf_rocm_mi355", shard: 8, num_shards: 9, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "inductor_torchbench_perf_rocm_mi355", shard: 9, num_shards: 9, runner: "linux.rocm.gpu.mi355.2" },
         ]}
     secrets: inherit
 

.github/workflows/lint.yml

@@ -118,9 +118,9 @@ jobs:
         CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}"
         echo "Running all other linters"
         if [ "$CHANGED_FILES" = '*' ]; then
-          ADDITIONAL_LINTRUNNER_ARGS="--skip CLANGTIDY,CLANGFORMAT,MYPY,MYPYSTRICT,PYREFLY --all-files" .github/scripts/lintrunner.sh
+          ADDITIONAL_LINTRUNNER_ARGS="--skip CLANGTIDY,CLANGFORMAT,MYPY,MYPYSTRICT --all-files" .github/scripts/lintrunner.sh
         else
-          ADDITIONAL_LINTRUNNER_ARGS="--skip CLANGTIDY,CLANGFORMAT,MYPY,MYPYSTRICT,PYREFLY ${CHANGED_FILES}" .github/scripts/lintrunner.sh
+          ADDITIONAL_LINTRUNNER_ARGS="--skip CLANGTIDY,CLANGFORMAT,MYPY,MYPYSTRICT ${CHANGED_FILES}" .github/scripts/lintrunner.sh
         fi
 
   quick-checks:

.github/workflows/operator_benchmark.yml

@@ -52,27 +52,3 @@ jobs:
       docker-image: ${{ needs.x86-opbenchmark-build.outputs.docker-image }}
       test-matrix: ${{ needs.x86-opbenchmark-build.outputs.test-matrix }}
     secrets: inherit
-
-  aarch64-opbenchmark-build:
-    if: github.repository_owner == 'pytorch'
-    name: aarch64-opbenchmark-build
-    uses: ./.github/workflows/_linux-build.yml
-    with:
-      build-environment: linux-jammy-aarch64-py3.10
-      runner: linux.arm64.m7g.4xlarge
-      docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc11
-      test-matrix: |
-        { include: [
-          { config: "cpu_operator_benchmark_short", shard: 1, num_shards: 1, runner: "linux.arm64.m8g.4xlarge" },
-        ]}
-    secrets: inherit
-
-  aarch64-opbenchmark-test:
-    name: aarch64-opbenchmark-test
-    uses: ./.github/workflows/_linux-test.yml
-    needs: aarch64-opbenchmark-build
-    with:
-      build-environment: linux-jammy-aarch64-py3.10
-      docker-image: ${{ needs.aarch64-opbenchmark-build.outputs.docker-image }}
-      test-matrix: ${{ needs.aarch64-opbenchmark-build.outputs.test-matrix }}
-    secrets: inherit

.github/workflows/rocm-mi355.yml

@@ -45,12 +45,12 @@ jobs:
       sync-tag: rocm-build
       test-matrix: |
         { include: [
-          { config: "default", shard: 1, num_shards: 6, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "default", shard: 2, num_shards: 6, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "default", shard: 3, num_shards: 6, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "default", shard: 4, num_shards: 6, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "default", shard: 5, num_shards: 6, runner: "linux.rocm.gpu.mi355.1" },
-          { config: "default", shard: 6, num_shards: 6, runner: "linux.rocm.gpu.mi355.1" },
+          { config: "default", shard: 1, num_shards: 6, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "default", shard: 2, num_shards: 6, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "default", shard: 3, num_shards: 6, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "default", shard: 4, num_shards: 6, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "default", shard: 5, num_shards: 6, runner: "linux.rocm.gpu.mi355.2" },
+          { config: "default", shard: 6, num_shards: 6, runner: "linux.rocm.gpu.mi355.2" },
         ]}
     secrets: inherit
 

.github/workflows/trunk.yml

@@ -200,23 +200,6 @@ jobs:
       cuda-arch-list: '8.0'
     secrets: inherit
 
-  # Test cross-compiled models with Windows libs extracted from wheel
-  cross-compile-linux-test:
-    name: cross-compile-linux-test
-    uses: ./.github/workflows/_linux-test.yml
-    needs:
-      - linux-jammy-cuda12_8-py3_10-gcc11-build
-      - get-label-type
-      - win-vs2022-cuda12_8-py3-build
-    with:
-      build-environment: linux-jammy-cuda12.8-py3.10-gcc11
-      docker-image: ${{ needs.linux-jammy-cuda12_8-py3_10-gcc11-build.outputs.docker-image }}
-      test-matrix: |
-        { include: [
-          { config: "aoti_cross_compile_for_windows", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g6.4xlarge.experimental.nvidia.gpu", win_torch_wheel_artifact: "win-vs2022-cuda12.8-py3" },
-        ]}
-    secrets: inherit
-
   verify-cachebench-cpu-build:
     name: verify-cachebench-cpu-build
     uses: ./.github/workflows/_linux-build.yml

.lintrunner.toml

@@ -209,46 +209,6 @@ command = [
     '@{{PATHSFILE}}'
 ]
 
-
-[[linter]]
-code = 'PYREFLY'
-include_patterns = [
-    'torch/**/*.py',
-    'torch/**/*.pyi',
-    'torchgen/**/*.py',
-    'torchgen/**/*.pyi',
-    'functorch/**/*.py',
-    'functorch/**/*.pyi',
-]
-exclude_patterns = []
-command = [
-    'python3',
-    'tools/linter/adapters/pyrefly_linter.py',
-    '--config=pyrefly.toml',
-]
-init_command = [
-    'python3',
-    'tools/linter/adapters/pip_init.py',
-    '--dry-run={{DRYRUN}}',
-    'numpy==2.1.0 ; python_version >= "3.12"',
-    'expecttest==0.3.0',
-    'pyrefly==0.36.2',
-    'sympy==1.13.3',
-    'types-requests==2.27.25',
-    'types-pyyaml==6.0.2',
-    'types-tabulate==0.8.8',
-    'types-protobuf==5.29.1.20250403',
-    'types-setuptools==79.0.0.20250422',
-    'types-jinja2==2.11.9',
-    'types-colorama==0.4.6',
-    'filelock==3.18.0',
-    'junitparser==2.1.1',
-    'rich==14.1.0',
-    'optree==0.17.0',
-    'types-openpyxl==3.1.5.20250919',
-    'types-python-dateutil==2.9.0.20251008'
-]
-
 [[linter]]
 code = 'CLANGTIDY'
 include_patterns = [

aten/src/ATen/core/PhiloxRNGEngine.h

@@ -229,10 +229,10 @@ private:
   }
 
 
-  static constexpr uint32_t kPhilox10A = 0x9E3779B9;
-  static constexpr uint32_t kPhilox10B = 0xBB67AE85;
-  static constexpr uint32_t kPhiloxSA = 0xD2511F53;
-  static constexpr uint32_t kPhiloxSB = 0xCD9E8D57;
+  static const uint32_t kPhilox10A = 0x9E3779B9;
+  static const uint32_t kPhilox10B = 0xBB67AE85;
+  static const uint32_t kPhiloxSA = 0xD2511F53;
+  static const uint32_t kPhiloxSB = 0xCD9E8D57;
 };
 
 typedef philox_engine Philox4_32;

aten/src/ATen/cpu/vec/vec128/vec128.h

@@ -8,7 +8,6 @@
 #include <ATen/cpu/vec/vec128/vec128_bfloat16_neon.h>
 #include <ATen/cpu/vec/vec128/vec128_float_neon.h>
 #include <ATen/cpu/vec/vec128/vec128_half_neon.h>
-#include <ATen/cpu/vec/vec128/vec128_int_aarch64.h>
 #endif
 
 #include <ATen/cpu/vec/vec128/vec128_convert.h>

aten/src/ATen/cpu/vec/vec128/vec128_int_aarch64.h

@@ -1,794 +0,0 @@
-#pragma once
-
-#include <ATen/cpu/vec/intrinsics.h>
-#include <ATen/cpu/vec/vec_base.h>
-#include <c10/macros/Macros.h>
-#include <c10/util/irange.h>
-
-namespace at::vec {
-// Note [CPU_CAPABILITY namespace]
-// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-// This header, and all of its subheaders, will be compiled with
-// different architecture flags for each supported set of vector
-// intrinsics. So we need to make sure they aren't inadvertently
-// linked together. We do this by declaring objects in an `inline
-// namespace` which changes the name mangling, but can still be
-// accessed as `at::vec`.
-inline namespace CPU_CAPABILITY {
-
-#define VEC_INT_NEON_TEMPLATE(vl, bit)                                        \
-  template <>                                                                 \
-  struct is_vec_specialized_for<int##bit##_t> : std::bool_constant<true> {};  \
-                                                                              \
-  template <>                                                                 \
-  class Vectorized<int##bit##_t> {                                            \
-    using neon_type = int##bit##x##vl##_t;                                    \
-                                                                              \
-   private:                                                                   \
-    neon_type values;                                                         \
-                                                                              \
-   public:                                                                    \
-    using value_type = int##bit##_t;                                          \
-    using size_type = int;                                                    \
-    static constexpr size_type size() {                                       \
-      return vl;                                                              \
-    }                                                                         \
-    Vectorized() {                                                            \
-      values = vdupq_n_s##bit(0);                                             \
-    }                                                                         \
-    Vectorized(neon_type v) : values(v) {}                                    \
-    Vectorized(int##bit##_t val);                                             \
-    template <                                                                \
-        typename... Args,                                                     \
-        typename = std::enable_if_t<(sizeof...(Args) == size())>>             \
-    Vectorized(Args... vals) {                                                \
-      __at_align__ int##bit##_t buffer[size()] = {vals...};                   \
-      values = vld1q_s##bit(buffer);                                          \
-    }                                                                         \
-    operator neon_type() const {                                              \
-      return values;                                                          \
-    }                                                                         \
-    static Vectorized<int##bit##_t> loadu(                                    \
-        const void* ptr,                                                      \
-        int64_t count = size());                                              \
-    void store(void* ptr, int64_t count = size()) const;                      \
-    template <int64_t mask>                                                   \
-    static Vectorized<int##bit##_t> blend(                                    \
-        const Vectorized<int##bit##_t>& a,                                    \
-        const Vectorized<int##bit##_t>& b);                                   \
-    static Vectorized<int##bit##_t> blendv(                                   \
-        const Vectorized<int##bit##_t>& a,                                    \
-        const Vectorized<int##bit##_t>& b,                                    \
-        const Vectorized<int##bit##_t>& mask_) {                              \
-      return vbslq_s##bit(vreinterpretq_u##bit##_s##bit(mask_.values), b, a); \
-    }                                                                         \
-    template <typename step_t>                                                \
-    static Vectorized<int##bit##_t> arange(                                   \
-        value_type base = 0,                                                  \
-        step_t step = static_cast<step_t>(1));                                \
-    static Vectorized<int##bit##_t> set(                                      \
-        const Vectorized<int##bit##_t>& a,                                    \
-        const Vectorized<int##bit##_t>& b,                                    \
-        int64_t count = size());                                              \
-    const int##bit##_t& operator[](int idx) const = delete;                   \
-    int##bit##_t& operator[](int idx) = delete;                               \
-    Vectorized<int##bit##_t> abs() const {                                    \
-      return vabsq_s##bit(values);                                            \
-    }                                                                         \
-    Vectorized<int##bit##_t> real() const {                                   \
-      return values;                                                          \
-    }                                                                         \
-    Vectorized<int##bit##_t> imag() const {                                   \
-      return vdupq_n_s##bit(0);                                               \
-    }                                                                         \
-    Vectorized<int##bit##_t> conj() const {                                   \
-      return values;                                                          \
-    }                                                                         \
-    Vectorized<int##bit##_t> neg() const {                                    \
-      return vnegq_s##bit(values);                                            \
-    }                                                                         \
-    int##bit##_t reduce_add() const {                                         \
-      return vaddvq_s##bit(values);                                           \
-    }                                                                         \
-    int##bit##_t reduce_max() const;                                          \
-    Vectorized<int##bit##_t> operator==(                                      \
-        const Vectorized<int##bit##_t>& other) const {                        \
-      return Vectorized<value_type>(                                          \
-          vreinterpretq_s##bit##_u##bit(vceqq_s##bit(values, other.values))); \
-    }                                                                         \
-    Vectorized<int##bit##_t> operator!=(                                      \
-        const Vectorized<int##bit##_t>& other) const;                         \
-    Vectorized<int##bit##_t> operator<(                                       \
-        const Vectorized<int##bit##_t>& other) const {                        \
-      return Vectorized<value_type>(                                          \
-          vreinterpretq_s##bit##_u##bit(vcltq_s##bit(values, other.values))); \
-    }                                                                         \
-    Vectorized<int##bit##_t> operator<=(                                      \
-        const Vectorized<int##bit##_t>& other) const {                        \
-      return Vectorized<value_type>(                                          \
-          vreinterpretq_s##bit##_u##bit(vcleq_s##bit(values, other.values))); \
-    }                                                                         \
-    Vectorized<int##bit##_t> operator>(                                       \
-        const Vectorized<int##bit##_t>& other) const {                        \
-      return Vectorized<value_type>(                                          \
-          vreinterpretq_s##bit##_u##bit(vcgtq_s##bit(values, other.values))); \
-    }                                                                         \
-    Vectorized<int##bit##_t> operator>=(                                      \
-        const Vectorized<int##bit##_t>& other) const {                        \
-      return Vectorized<value_type>(                                          \
-          vreinterpretq_s##bit##_u##bit(vcgeq_s##bit(values, other.values))); \
-    }                                                                         \
-    Vectorized<int##bit##_t> eq(const Vectorized<int##bit##_t>& other) const; \
-    Vectorized<int##bit##_t> ne(const Vectorized<int##bit##_t>& other) const; \
-    Vectorized<int##bit##_t> gt(const Vectorized<int##bit##_t>& other) const; \
-    Vectorized<int##bit##_t> ge(const Vectorized<int##bit##_t>& other) const; \
-    Vectorized<int##bit##_t> lt(const Vectorized<int##bit##_t>& other) const; \
-    Vectorized<int##bit##_t> le(const Vectorized<int##bit##_t>& other) const; \
-  };                                                                          \
-  template <>                                                                 \
-  Vectorized<int##bit##_t> inline operator+(                                  \
-      const Vectorized<int##bit##_t>& a, const Vectorized<int##bit##_t>& b) { \
-    return vaddq_s##bit(a, b);                                                \
-  }                                                                           \
-  template <>                                                                 \
-  Vectorized<int##bit##_t> inline operator-(                                  \
-      const Vectorized<int##bit##_t>& a, const Vectorized<int##bit##_t>& b) { \
-    return vsubq_s##bit(a, b);                                                \
-  }                                                                           \
-  template <>                                                                 \
-  Vectorized<int##bit##_t> inline operator&(                                  \
-      const Vectorized<int##bit##_t>& a, const Vectorized<int##bit##_t>& b) { \
-    return vandq_s##bit(a, b);                                                \
-  }                                                                           \
-  template <>                                                                 \
-  Vectorized<int##bit##_t> inline operator|(                                  \
-      const Vectorized<int##bit##_t>& a, const Vectorized<int##bit##_t>& b) { \
-    return vorrq_s##bit(a, b);                                                \
-  }                                                                           \
-  template <>                                                                 \
-  Vectorized<int##bit##_t> inline operator^(                                  \
-      const Vectorized<int##bit##_t>& a, const Vectorized<int##bit##_t>& b) { \
-    return veorq_s##bit(a, b);                                                \
-  }                                                                           \
-  Vectorized<int##bit##_t> inline Vectorized<int##bit##_t>::eq(               \
-      const Vectorized<int##bit##_t>& other) const {                          \
-    return (*this == other) & Vectorized<int##bit##_t>(1);                    \
-  }                                                                           \
-  Vectorized<int##bit##_t> inline Vectorized<int##bit##_t>::ne(               \
-      const Vectorized<int##bit##_t>& other) const {                          \
-    return (*this != other) & Vectorized<int##bit##_t>(1);                    \
-  }                                                                           \
-  Vectorized<int##bit##_t> inline Vectorized<int##bit##_t>::gt(               \
-      const Vectorized<int##bit##_t>& other) const {                          \
-    return (*this > other) & Vectorized<int##bit##_t>(1);                     \
-  }                                                                           \
-  Vectorized<int##bit##_t> inline Vectorized<int##bit##_t>::ge(               \
-      const Vectorized<int##bit##_t>& other) const {                          \
-    return (*this >= other) & Vectorized<int##bit##_t>(1);                    \
-  }                                                                           \
-  Vectorized<int##bit##_t> inline Vectorized<int##bit##_t>::lt(               \
-      const Vectorized<int##bit##_t>& other) const {                          \
-    return (*this < other) & Vectorized<int##bit##_t>(1);                     \
-  }                                                                           \
-  Vectorized<int##bit##_t> inline Vectorized<int##bit##_t>::le(               \
-      const Vectorized<int##bit##_t>& other) const {                          \
-    return (*this <= other) & Vectorized<int##bit##_t>(1);                    \
-  }
-
-VEC_INT_NEON_TEMPLATE(2, 64)
-VEC_INT_NEON_TEMPLATE(4, 32)
-VEC_INT_NEON_TEMPLATE(8, 16)
-VEC_INT_NEON_TEMPLATE(16, 8)
-
-inline int32_t Vectorized<int32_t>::reduce_max() const {
-  return vmaxvq_s32(values);
-}
-
-inline int16_t Vectorized<int16_t>::reduce_max() const {
-  return vmaxvq_s16(values);
-}
-
-inline int8_t Vectorized<int8_t>::reduce_max() const {
-  return vmaxvq_s8(values);
-}
-
-template <>
-Vectorized<int32_t> inline operator*(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& b) {
-  return vmulq_s32(a, b);
-}
-
-template <>
-Vectorized<int16_t> inline operator*(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& b) {
-  return vmulq_s16(a, b);
-}
-
-template <>
-Vectorized<int8_t> inline operator*(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& b) {
-  return vmulq_s8(a, b);
-}
-
-template <>
-inline Vectorized<int64_t> operator~(const Vectorized<int64_t>& a) {
-  int64x2_t val = a;
-  return ~val;
-}
-
-template <>
-inline Vectorized<int32_t> operator~(const Vectorized<int32_t>& a) {
-  return vmvnq_s32(a);
-}
-
-template <>
-inline Vectorized<int16_t> operator~(const Vectorized<int16_t>& a) {
-  return vmvnq_s16(a);
-}
-
-template <>
-inline Vectorized<int8_t> operator~(const Vectorized<int8_t>& a) {
-  return vmvnq_s8(a);
-}
-
-inline Vectorized<int64_t> Vectorized<int64_t>::operator!=(
-    const Vectorized<int64_t>& other) const {
-  return ~(*this == other);
-}
-
-inline Vectorized<int32_t> Vectorized<int32_t>::operator!=(
-    const Vectorized<int32_t>& other) const {
-  return ~(*this == other);
-}
-
-inline Vectorized<int16_t> Vectorized<int16_t>::operator!=(
-    const Vectorized<int16_t>& other) const {
-  return ~(*this == other);
-}
-
-inline Vectorized<int8_t> Vectorized<int8_t>::operator!=(
-    const Vectorized<int8_t>& other) const {
-  return ~(*this == other);
-}
-
-template <>
-Vectorized<int32_t> inline minimum(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& b) {
-  return vminq_s32(a, b);
-}
-
-template <>
-Vectorized<int16_t> inline minimum(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& b) {
-  return vminq_s16(a, b);
-}
-
-template <>
-Vectorized<int8_t> inline minimum(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& b) {
-  return vminq_s8(a, b);
-}
-
-template <>
-Vectorized<int32_t> inline maximum(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& b) {
-  return vmaxq_s32(a, b);
-}
-
-template <>
-Vectorized<int16_t> inline maximum(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& b) {
-  return vmaxq_s16(a, b);
-}
-
-template <>
-Vectorized<int8_t> inline maximum(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& b) {
-  return vmaxq_s8(a, b);
-}
-
-template <int64_t mask>
-Vectorized<int64_t> Vectorized<int64_t>::blend(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& b) {
-  // Build an array of flags: each bit of element is 1 if the corresponding bit
-  // in 'mask' is set, 0 otherwise.
-  uint64x2_t maskArray = {
-      (mask & 1LL) ? 0xFFFFFFFFFFFFFFFF : 0,
-      (mask & 2LL) ? 0xFFFFFFFFFFFFFFFF : 0};
-  // Use BSL to select elements from b where the mask is 1, else from a
-  return vbslq_s64(maskArray, b.values, a.values);
-}
-
-template <int64_t mask>
-Vectorized<int32_t> Vectorized<int32_t>::blend(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& b) {
-  // Build an array of flags: each bit of element is 1 if the corresponding bit
-  // in 'mask' is set, 0 otherwise.
-  uint32x4_t maskArray = {
-      (mask & 1LL) ? 0xFFFFFFFF : 0,
-      (mask & 2LL) ? 0xFFFFFFFF : 0,
-      (mask & 4LL) ? 0xFFFFFFFF : 0,
-      (mask & 8LL) ? 0xFFFFFFFF : 0};
-  // Use BSL to select elements from b where the mask is 1, else from a
-  return vbslq_s32(maskArray, b.values, a.values);
-}
-
-template <int64_t mask>
-Vectorized<int16_t> Vectorized<int16_t>::blend(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& b) {
-  // Build an array of flags: each bit of element is 1 if the corresponding bit
-  // in 'mask' is set, 0 otherwise.
-  uint16x8_t maskArray = {
-      (mask & 1LL) ? 0xFFFF : 0,
-      (mask & 2LL) ? 0xFFFF : 0,
-      (mask & 4LL) ? 0xFFFF : 0,
-      (mask & 8LL) ? 0xFFFF : 0,
-      (mask & 16LL) ? 0xFFFF : 0,
-      (mask & 32LL) ? 0xFFFF : 0,
-      (mask & 64LL) ? 0xFFFF : 0,
-      (mask & 128LL) ? 0xFFFF : 0};
-  // Use BSL to select elements from b where the mask is 1, else from a
-  return vbslq_s16(maskArray, b.values, a.values);
-}
-
-template <int64_t mask>
-Vectorized<int8_t> Vectorized<int8_t>::blend(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& b) {
-  // Build an array of flags: each bit of element is 1 if the corresponding bit
-  // in 'mask' is set, 0 otherwise.
-  uint8x16_t maskArray = {
-      (mask & 1LL) ? 0xFF : 0,
-      (mask & 2LL) ? 0xFF : 0,
-      (mask & 4LL) ? 0xFF : 0,
-      (mask & 8LL) ? 0xFF : 0,
-      (mask & 16LL) ? 0xFF : 0,
-      (mask & 32LL) ? 0xFF : 0,
-      (mask & 64LL) ? 0xFF : 0,
-      (mask & 128LL) ? 0xFF : 0,
-      (mask & 256LL) ? 0xFF : 0,
-      (mask & 512LL) ? 0xFF : 0,
-      (mask & 1024LL) ? 0xFF : 0,
-      (mask & 2048LL) ? 0xFF : 0,
-      (mask & 4096LL) ? 0xFF : 0,
-      (mask & 8192LL) ? 0xFF : 0,
-      (mask & 16384LL) ? 0xFF : 0,
-      (mask & 32768LL) ? 0xFF : 0};
-  // Use BSL to select elements from b where the mask is 1, else from a
-  return vbslq_s8(maskArray, b.values, a.values);
-}
-
-#define VEC_INT_NEON_OPS(vl, bit)                                             \
-  inline Vectorized<int##bit##_t>::Vectorized(int##bit##_t val) {             \
-    values = vdupq_n_s##bit(val);                                             \
-  }                                                                           \
-  inline Vectorized<int##bit##_t> Vectorized<int##bit##_t>::loadu(            \
-      const void* ptr, int64_t count) {                                       \
-    if (count == size()) {                                                    \
-      return vld1q_s##bit(reinterpret_cast<const int##bit##_t*>(ptr));        \
-    } else {                                                                  \
-      __at_align__ int##bit##_t tmp_values[size()];                           \
-      for (const auto i : c10::irange(size())) {                              \
-        tmp_values[i] = 0;                                                    \
-      }                                                                       \
-      std::memcpy(                                                            \
-          tmp_values,                                                         \
-          reinterpret_cast<const int##bit##_t*>(ptr),                         \
-          count * sizeof(int##bit##_t));                                      \
-      return vld1q_s##bit(reinterpret_cast<const int##bit##_t*>(tmp_values)); \
-    }                                                                         \
-  }                                                                           \
-  inline void Vectorized<int##bit##_t>::store(void* ptr, int64_t count)       \
-      const {                                                                 \
-    if (count == size()) {                                                    \
-      vst1q_s##bit(reinterpret_cast<int##bit##_t*>(ptr), values);             \
-    } else {                                                                  \
-      int##bit##_t tmp_values[size()];                                        \
-      vst1q_s##bit(reinterpret_cast<int##bit##_t*>(tmp_values), values);      \
-      std::memcpy(ptr, tmp_values, count * sizeof(int##bit##_t));             \
-    }                                                                         \
-  }
-
-VEC_INT_NEON_OPS(2, 64)
-VEC_INT_NEON_OPS(4, 32)
-VEC_INT_NEON_OPS(8, 16)
-VEC_INT_NEON_OPS(16, 8)
-
-template <>
-Vectorized<int64_t> inline operator*(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& b) {
-  int64x2_t x = a;
-  int64x2_t y = b;
-  return x * y;
-}
-
-template <>
-Vectorized<int64_t> inline operator/(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& b) {
-  int64x2_t x = a;
-  int64x2_t y = b;
-  return x / y;
-}
-
-template <>
-Vectorized<int32_t> inline operator/(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& b) {
-  int32x4_t x = a;
-  int32x4_t y = b;
-  return x / y;
-}
-
-inline int64_t Vectorized<int64_t>::reduce_max() const {
-  return std::max(values[0], values[1]);
-}
-
-template <>
-Vectorized<int64_t> inline minimum(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& b) {
-  int64x2_t x = a;
-  int64x2_t y = b;
-  return {std::min(x[0], y[0]), std::min(x[1], y[1])};
-}
-
-template <>
-Vectorized<int64_t> inline maximum(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& b) {
-  int64x2_t x = a;
-  int64x2_t y = b;
-  return {std::max(x[0], y[0]), std::max(x[1], y[1])};
-}
-
-template <typename step_t>
-inline Vectorized<int64_t> Vectorized<int64_t>::arange(
-    int64_t base,
-    step_t step) {
-  const Vectorized<int64_t> base_vec(base);
-  const Vectorized<int64_t> step_vec(step);
-  const int64x2_t step_sizes = {0, 1};
-  return base_vec.values + step_sizes * step_vec.values;
-}
-
-template <typename step_t>
-inline Vectorized<int32_t> Vectorized<int32_t>::arange(
-    int32_t base,
-    step_t step) {
-  const Vectorized<int32_t> base_vec(base);
-  const Vectorized<int32_t> step_vec(step);
-  const int32x4_t step_sizes = {0, 1, 2, 3};
-  return vmlaq_s32(base_vec, step_sizes, step_vec);
-}
-
-template <typename step_t>
-inline Vectorized<int16_t> Vectorized<int16_t>::arange(
-    int16_t base,
-    step_t step) {
-  const Vectorized<int16_t> base_vec(base);
-  const Vectorized<int16_t> step_vec(step);
-  const int16x8_t step_sizes = {0, 1, 2, 3, 4, 5, 6, 7};
-  return vmlaq_s16(base_vec, step_sizes, step_vec);
-}
-
-template <typename step_t>
-inline Vectorized<int8_t> Vectorized<int8_t>::arange(int8_t base, step_t step) {
-  const Vectorized<int8_t> base_vec(base);
-  const Vectorized<int8_t> step_vec(step);
-  const int8x16_t step_sizes = {
-      0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
-  return vmlaq_s8(base_vec, step_sizes, step_vec);
-}
-
-template <>
-Vectorized<int64_t> inline operator>>(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& b) {
-  int64x2_t x = a;
-  int64x2_t y = b;
-  uint64x2_t u = vreinterpretq_u64_s64(y);
-  uint64x2_t z = {std::min(u[0], (uint64_t)63), std::min(u[1], (uint64_t)63)};
-  return x >> vreinterpretq_s64_u64(z);
-}
-
-template <>
-Vectorized<int32_t> inline operator>>(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& b) {
-  int32x4_t x = a;
-  int32x4_t y = b;
-  uint32x4_t bound = vdupq_n_u32(31);
-  uint32x4_t z = vminq_u32(vreinterpretq_u32_s32(y), bound);
-  return x >> vreinterpretq_s32_u32(z);
-}
-
-template <>
-Vectorized<int16_t> inline operator>>(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& b) {
-  int16x8_t x = a;
-  int16x8_t y = b;
-  uint16x8_t bound = vdupq_n_u16(15);
-  uint16x8_t z = vminq_u16(vreinterpretq_u16_s16(y), bound);
-  return x >> vreinterpretq_s16_u16(z);
-}
-
-template <>
-Vectorized<int8_t> inline operator>>(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& b) {
-  int8x16_t x = a;
-  int8x16_t y = b;
-  uint8x16_t bound = vdupq_n_u8(7);
-  int8x16_t z = vreinterpretq_s8_u8(vminq_u8(vreinterpretq_u8_s8(y), bound));
-  return x >> z;
-}
-
-template <>
-Vectorized<int64_t> inline operator<<(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& b) {
-  int64x2_t y = b;
-  uint64x2_t u = vreinterpretq_u64_s64(y);
-  uint64x2_t z = {std::min(u[0], (uint64_t)64), std::min(u[1], (uint64_t)64)};
-  return vshlq_s64(a, vreinterpretq_s64_u64(z));
-}
-
-template <>
-Vectorized<int32_t> inline operator<<(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& b) {
-  int32x4_t y = b;
-  uint32x4_t bound = vdupq_n_u32(32);
-  uint32x4_t z = vminq_u32(vreinterpretq_u32_s32(y), bound);
-  return vshlq_s32(a, vreinterpretq_s32_u32(z));
-}
-
-template <>
-Vectorized<int16_t> inline operator<<(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& b) {
-  int16x8_t y = b;
-  uint16x8_t bound = vdupq_n_u16(16);
-  uint16x8_t z = vminq_u16(vreinterpretq_u16_s16(y), bound);
-  return vshlq_s16(a, vreinterpretq_s16_u16(z));
-}
-
-template <>
-Vectorized<int8_t> inline operator<<(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& b) {
-  int8x16_t y = b;
-  uint8x16_t bound = vdupq_n_u8(8);
-  int8x16_t z = vreinterpretq_s8_u8(vminq_u8(vreinterpretq_u8_s8(y), bound));
-  return vshlq_s8(a, z);
-}
-
-inline Vectorized<int64_t> Vectorized<int64_t>::set(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& b,
-    int64_t count) {
-  if (count == 0) {
-    return a;
-  } else if (count >= 2) {
-    return b;
-  } else {
-    int64x2_t c = {b.values[0], a.values[1]};
-    return c;
-  }
-}
-
-inline Vectorized<int32_t> Vectorized<int32_t>::set(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& b,
-    int64_t count) {
-  if (count == 0) {
-    return a;
-  } else if (count >= 4) {
-    return b;
-  } else {
-    // Build an array of flags: each bit of element is 1 if the corresponding
-    // bit in 'mask' is set, 0 otherwise.
-    uint32x4_t maskArray = {
-        (count >= 1LL) ? 0xFFFFFFFF : 0,
-        (count >= 2LL) ? 0xFFFFFFFF : 0,
-        (count >= 3LL) ? 0xFFFFFFFF : 0,
-        0};
-    // Use BSL to select elements from b where the mask is 1, else from a
-    return vbslq_s32(maskArray, b.values, a.values);
-  }
-}
-
-inline Vectorized<int16_t> Vectorized<int16_t>::set(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& b,
-    int64_t count) {
-  if (count == 0) {
-    return a;
-  } else if (count >= 8) {
-    return b;
-  } else {
-    // Build an array of flags: each bit of element is 1 if the corresponding
-    // bit in 'mask' is set, 0 otherwise.
-    uint16x8_t maskArray = {
-        static_cast<uint16_t>((count >= 1LL) ? 0xFFFF : 0),
-        static_cast<uint16_t>((count >= 2LL) ? 0xFFFF : 0),
-        static_cast<uint16_t>((count >= 3LL) ? 0xFFFF : 0),
-        static_cast<uint16_t>((count >= 4LL) ? 0xFFFF : 0),
-        static_cast<uint16_t>((count >= 5LL) ? 0xFFFF : 0),
-        static_cast<uint16_t>((count >= 6LL) ? 0xFFFF : 0),
-        static_cast<uint16_t>((count >= 7LL) ? 0xFFFF : 0),
-        0};
-    // Use BSL to select elements from b where the mask is 1, else from a
-    return vbslq_s16(maskArray, b.values, a.values);
-  }
-}
-
-inline Vectorized<int8_t> Vectorized<int8_t>::set(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& b,
-    int64_t count) {
-  if (count == 0) {
-    return a;
-  } else if (count >= 16) {
-    return b;
-  } else {
-    // Build an array of flags: each bit of element is 1 if the corresponding
-    // bit in 'mask' is set, 0 otherwise.
-    uint8x16_t maskArray = {
-        static_cast<uint8_t>((count >= 1LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 2LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 3LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 4LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 5LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 6LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 7LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 8LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 9LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 10LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 11LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 12LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 13LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 14LL) ? 0xFF : 0),
-        static_cast<uint8_t>((count >= 15LL) ? 0xFF : 0),
-        0};
-
-    // Use BSL to select elements from b where the mask is 1, else from a
-    return vbslq_s8(maskArray, b.values, a.values);
-  }
-}
-
-template <>
-Vectorized<int16_t> inline operator/(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& b) {
-  Vectorized<int32_t> highBitsA = vmovl_high_s16(a);
-  Vectorized<int32_t> highBitsB = vmovl_high_s16(b);
-  Vectorized<int32_t> lowBitsA = vmovl_s16(vget_low_s16(a));
-  Vectorized<int32_t> lowBitsB = vmovl_s16(vget_low_s16(b));
-  int32x4_t highBitsResult = highBitsA / highBitsB;
-  int32x4_t lowBitsResult = lowBitsA / lowBitsB;
-  return vuzp1q_s16(
-      vreinterpretq_s16_s32(lowBitsResult),
-      vreinterpretq_s16_s32(highBitsResult));
-}
-
-template <>
-Vectorized<int8_t> inline operator/(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& b) {
-  Vectorized<int16_t> highBitsA = vmovl_high_s8(a);
-  Vectorized<int16_t> highBitsB = vmovl_high_s8(b);
-  Vectorized<int16_t> lowBitsA = vmovl_s8(vget_low_s8(a));
-  Vectorized<int16_t> lowBitsB = vmovl_s8(vget_low_s8(b));
-  int16x8_t highBitsResult = highBitsA / highBitsB;
-  int16x8_t lowBitsResult = lowBitsA / lowBitsB;
-  return vuzp1q_s8(
-      vreinterpretq_s8_s16(lowBitsResult),
-      vreinterpretq_s8_s16(highBitsResult));
-}
-
-template <>
-Vectorized<int64_t> inline clamp(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& min,
-    const Vectorized<int64_t>& max) {
-  return minimum(max, maximum(min, a));
-}
-
-template <>
-Vectorized<int32_t> inline clamp(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& min,
-    const Vectorized<int32_t>& max) {
-  return minimum(max, maximum(min, a));
-}
-
-template <>
-Vectorized<int16_t> inline clamp(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& min,
-    const Vectorized<int16_t>& max) {
-  return minimum(max, maximum(min, a));
-}
-
-template <>
-Vectorized<int8_t> inline clamp(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& min,
-    const Vectorized<int8_t>& max) {
-  return minimum(max, maximum(min, a));
-}
-
-template <>
-Vectorized<int64_t> inline clamp_max(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& max) {
-  return minimum(max, a);
-}
-
-template <>
-Vectorized<int32_t> inline clamp_max(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& max) {
-  return minimum(max, a);
-}
-
-template <>
-Vectorized<int16_t> inline clamp_max(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& max) {
-  return minimum(max, a);
-}
-
-template <>
-Vectorized<int8_t> inline clamp_max(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& max) {
-  return minimum(max, a);
-}
-
-template <>
-Vectorized<int64_t> inline clamp_min(
-    const Vectorized<int64_t>& a,
-    const Vectorized<int64_t>& min) {
-  return maximum(min, a);
-}
-
-template <>
-Vectorized<int32_t> inline clamp_min(
-    const Vectorized<int32_t>& a,
-    const Vectorized<int32_t>& min) {
-  return maximum(min, a);
-}
-
-template <>
-Vectorized<int16_t> inline clamp_min(
-    const Vectorized<int16_t>& a,
-    const Vectorized<int16_t>& min) {
-  return maximum(min, a);
-}
-
-template <>
-Vectorized<int8_t> inline clamp_min(
-    const Vectorized<int8_t>& a,
-    const Vectorized<int8_t>& min) {
-  return maximum(min, a);
-}
-
-} // namespace CPU_CAPABILITY
-} // namespace at::vec

aten/src/ATen/cpu/vec/vec256/vec256_qint.h

@@ -1377,7 +1377,7 @@ Vectorized<c10::quint8> inline maximum(
 #if (defined(__aarch64__) && !defined(CPU_CAPABILITY_SVE256))
 std::pair<Vectorized<float>, Vectorized<float>> inline convert_int8_to_float(
     at::vec::Vectorized<int8_t> src) {
-  auto s8x8 = vget_low_s8(src);
+  auto s8x8 = vld1_s8(src.operator const int8_t*());
   auto s16x8 = vmovl_s8(s8x8);
 
   auto s32x4_hi = vmovl_s16(vget_high_s16(s16x8));
@@ -1402,7 +1402,7 @@ std::pair<Vectorized<float>, Vectorized<float>> inline convert_int8_to_float(
 
 Vectorized<float> inline convert_int8_half_register_to_float(
     at::vec::Vectorized<int8_t> src) {
-  auto s8x8 = vget_low_s8(src);
+  auto s8x8 = vld1_s8(src.operator const int8_t*());
   auto s16x8 = vmovl_s8(s8x8);
 
   auto s32x4_lo = vmovl_s16(vget_low_s16(s16x8));

aten/src/ATen/cuda/CUDABlas.cpp

@@ -16,8 +16,6 @@
 #include <c10/util/irange.h>
 #include <c10/core/ScalarType.h>
 
-#include <ATen/cuda/detail/BLASConstants.h>
-
 #ifdef USE_ROCM
 #include <c10/cuda/CUDAStream.h>
 #include <hipblaslt/hipblaslt-ext.hpp>
@@ -1956,15 +1954,13 @@ void scaled_gemm(
     const void *result_scale_ptr,
     int64_t result_ld,
     ScalarType result_dtype,
-    bool use_fast_accum,
-    const std::optional<Tensor>& alpha) {
+    bool use_fast_accum) {
   // Note: see `cublasCommonArgs` for various non-intuitive manupulations
   // of input arguments to this function.
   const auto computeType = CUBLAS_COMPUTE_32F;
   const auto scaleType = CUDA_R_32F;
-  // Note: alpha_val may change later depending on user-passed argument
-  float alpha_val = 1.0;
-  float beta_val = 0.0;
+  const float alpha_val = 1.0;
+  const float beta_val = 0.0;
   CuBlasLtMatmulDescriptor computeDesc(computeType, scaleType);
   computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_TRANSA, _cublasOpFromChar(transa));
   computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_TRANSB, _cublasOpFromChar(transb));
@@ -2035,33 +2031,6 @@ void scaled_gemm(
     computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_EPILOGUE, CUBLASLT_EPILOGUE_BIAS);
     computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE, ScalarTypeToCudaDataType(bias_dtype));
   }
-
-  // Handle user-passed alpha
-  float *alpha_ptr = &alpha_val;
-  float *beta_ptr = &beta_val;
-
-  if (alpha.has_value()) {
-    auto& a = alpha.value();
-
-    // if device-tensor
-    if (a.is_cuda()) {
-      // NOTE: there are lifetime requirements on device-side pointers for alpha/beta -- the value must be
-      //       valid & correct until the cublas call finishes (not is scheduled like host-side values). Thus
-      //       we need to use allocations for alpha/beta that have some guarantees on lifetime - a statically
-      //       managed 4B buffer for alpha that we'll copy the passed alpha value into, and constant memory
-      //       for beta respectively.
-      float *user_alpha_ptr = at::cuda::detail::get_user_alpha_ptr();
-      at::Tensor user_alpha = at::from_blob(user_alpha_ptr, {1}, TensorOptions().device(kCUDA).dtype(kFloat));
-      user_alpha.copy_(a);
-      // Tell cublasLt we're using device-side pointers for alpha/beta
-      auto pointer_mode = CUBLASLT_POINTER_MODE_DEVICE;
-      computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_POINTER_MODE, pointer_mode);
-      alpha_ptr = user_alpha.data_ptr<float>();
-      beta_ptr = at::cuda::detail::get_cublas_device_zero();
-    } else {
-      alpha_val = a.item<float>();
-    }
-  }
     // For other data types, use the get_scale_mode function based on scaling type
     // The SCALE_MODE attrs only exist in cuBLAS 12.8+/ROCm 7.0 or in recent hipblaslt,
     // but we must invoke get_scale_mode anyways to trigger the version checks.
@@ -2079,7 +2048,6 @@ void scaled_gemm(
   cublasLtMatmulHeuristicResult_t heuristicResult = {};
   int returnedResult = 0;
   cublasLtHandle_t ltHandle = at::cuda::getCurrentCUDABlasLtHandle();
-
   TORCH_CUDABLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(
       ltHandle,
       computeDesc.descriptor(),
@@ -2120,10 +2088,10 @@ void scaled_gemm(
         auto is_valid_status = hipblaslt_ext::matmulIsAlgoSupported(
                 ltHandle,
                 computeDesc.descriptor(),
-                alpha_ptr,
+                &alpha_val,
                 Adesc.descriptor(),
                 Bdesc.descriptor(),
-                beta_ptr,
+                &beta_val,
                 Cdesc.descriptor(),
                 Ddesc.descriptor(),
                 all_algos[i].algo,
@@ -2142,14 +2110,17 @@ void scaled_gemm(
   cublasStatus_t cublasStatus = cublasLtMatmul(
       ltHandle,
       computeDesc.descriptor(),
-      alpha_ptr,
+      &alpha_val,
       mat1_ptr,
       Adesc.descriptor(),
       mat2_ptr,
       Bdesc.descriptor(),
-      beta_ptr,
-      // NOTE: always use result_ptr here, because cuBLASLt w/device beta=0 can't handle nullptr either
+      &beta_val,
+#ifdef USE_ROCM
       result_ptr, // unused, since beta_val is 0, but hipblaslt can't handle nullptr
+#else
+      nullptr,
+#endif // ifdef USE_ROCM
       Cdesc.descriptor(),
       result_ptr,
       Ddesc.descriptor(),

aten/src/ATen/cuda/CUDABlas.h

@@ -161,8 +161,7 @@ void scaled_gemm(
     const void* result_scale_ptr,
     int64_t result_ld,
     ScalarType result_dtype,
-    bool use_fast_accum,
-    const std::optional<Tensor>& alpha);
+    bool use_fast_accum);
 
 #define CUDABLAS_BGEMM_ARGTYPES(Dtype)  CUDABLAS_BGEMM_ARGTYPES_AND_C_DTYPE(Dtype, Dtype)
 

aten/src/ATen/cuda/CUDAGeneratorImpl.cpp

@@ -325,9 +325,9 @@ uint64_t CUDAGeneratorImpl::seed() {
  */
 c10::intrusive_ptr<c10::TensorImpl> CUDAGeneratorImpl::get_state() const {
   // The RNG state comprises the seed, and an offset used for Philox.
-  constexpr size_t seed_size = sizeof(uint64_t);
-  constexpr size_t offset_size = sizeof(int64_t);
-  constexpr size_t total_size = seed_size + offset_size;
+  static const size_t seed_size = sizeof(uint64_t);
+  static const size_t offset_size = sizeof(int64_t);
+  static const size_t total_size = seed_size + offset_size;
 
   auto state_tensor = at::detail::empty_cpu({(int64_t)total_size}, ScalarType::Byte, std::nullopt, std::nullopt, std::nullopt, std::nullopt);
   auto rng_state = state_tensor.data_ptr<uint8_t>();
@@ -346,9 +346,9 @@ c10::intrusive_ptr<c10::TensorImpl> CUDAGeneratorImpl::get_state() const {
  * and size of the internal state.
  */
 void CUDAGeneratorImpl::set_state(const c10::TensorImpl& new_state) {
-  constexpr size_t seed_size = sizeof(uint64_t);
-  constexpr size_t offset_size = sizeof(int64_t);
-  constexpr size_t total_size = seed_size + offset_size;
+  static const size_t seed_size = sizeof(uint64_t);
+  static const size_t offset_size = sizeof(int64_t);
+  static const size_t total_size = seed_size + offset_size;
 
   detail::check_rng_state(new_state);
 

aten/src/ATen/cuda/detail/BLASConstants.cu

@@ -1,54 +0,0 @@
-#include <ATen/Functions.h>
-#include <ATen/Tensor.h>
-#include <ATen/cuda/Exceptions.h>
-
-#include <mutex>
-
-namespace at {
-namespace cuda {
-namespace detail {
-
-__device__ __constant__ float cublas_one_device;
-__device__ __constant__ float cublas_zero_device;
-
-float *get_cublas_device_one() {
-  static c10::once_flag init_flag;
-
-  c10::call_once(init_flag, []() {
-    const float one = 1.f;
-    AT_CUDA_CHECK(cudaMemcpyToSymbol(cublas_one_device, &one, sizeof(float)));
-  });
-
-  float *ptr;
-  AT_CUDA_CHECK(cudaGetSymbolAddress(reinterpret_cast<void**>(&ptr), cublas_one_device));
-  return ptr;
-}
-
-float *get_cublas_device_zero() {
-  static c10::once_flag init_flag;
-
-  c10::call_once(init_flag, []() {
-    const float zero = 0.f;
-    AT_CUDA_CHECK(cudaMemcpyToSymbol(cublas_zero_device, &zero, sizeof(float)));
-  });
-
-  float *ptr;
-  AT_CUDA_CHECK(cudaGetSymbolAddress(reinterpret_cast<void**>(&ptr), cublas_zero_device));
-  return ptr;
-}
-
-float *get_user_alpha_ptr() {
-  static float *alpha_ptr;
-
-  static c10::once_flag init_flag;
-
-  c10::call_once(init_flag, []() {
-    AT_CUDA_CHECK(cudaMalloc(&alpha_ptr, sizeof(float)));
-  });
-
-  return alpha_ptr;
-}
-
-} // namespace detail
-} // namespace cuda
-} // namespace at

aten/src/ATen/cuda/detail/BLASConstants.h

@@ -1,11 +0,0 @@
-#pragma once
-
-#include <ATen/core/TensorBase.h>
-
-namespace at::cuda::detail {
-
-float *get_cublas_device_one();
-float *get_cublas_device_zero();
-float *get_user_alpha_ptr();
-
-} // namespace at::cuda::detail

aten/src/ATen/cuda/tunable/TunableGemm.h

@@ -109,8 +109,7 @@ class DefaultScaledGemmOp : public Callable<ScaledGemmParams<T>> {
           params->c_scale_ptr,
           params->ldc,
           params->c_dtype,
-          params->use_fast_accum,
-          std::nullopt /* alpha */);
+          params->use_fast_accum);
       return OK;
     }
 };

aten/src/ATen/native/Activation.cpp

@@ -240,8 +240,8 @@ TORCH_META_FUNC(gelu_backward) (
 
 namespace at::native {
 
-static constexpr double SELU_ALPHA = 1.6732632423543772848170429916717;
-static constexpr double SELU_SCALE = 1.0507009873554804934193349852946;
+static const double SELU_ALPHA = 1.6732632423543772848170429916717;
+static const double SELU_SCALE = 1.0507009873554804934193349852946;
 
 DEFINE_DISPATCH(elu_stub);
 DEFINE_DISPATCH(elu_backward_stub);

aten/src/ATen/native/BlasKernel.cpp

@@ -286,7 +286,7 @@ template void scal_fast_path<scalar_t>(int *n, scalar_t *a, scalar_t *x, int *in
 #if AT_BUILD_WITH_BLAS()
 template <>
 bool scal_use_fast_path<double>(int64_t n, int64_t incx) {
-  auto constexpr intmax = std::numeric_limits<int>::max();
+  auto intmax = std::numeric_limits<int>::max();
   return n <= intmax && incx <= intmax;
 }
 
@@ -315,7 +315,7 @@ bool gemv_use_fast_path<float>(
     int64_t incx,
     [[maybe_unused]] float beta,
     int64_t incy) {
-  auto constexpr intmax = std::numeric_limits<int>::max();
+  auto intmax = std::numeric_limits<int>::max();
   return (m <= intmax) && (n <= intmax) && (lda <= intmax) &&
          (incx > 0) && (incx <= intmax) && (incy > 0) && (incy <= intmax);
 }

aten/src/ATen/native/Convolution.cpp

@@ -658,7 +658,6 @@ static void check_shape_forward(const at::Tensor& input,
   TORCH_CHECK(!params.is_output_padding_neg(), "negative output_padding is not supported");
   TORCH_CHECK(!params.is_stride_nonpos(), "non-positive stride is not supported");
   TORCH_CHECK(!params.is_dilation_neg(), "dilation should be greater than zero");
-  TORCH_CHECK(groups > 0, "expected groups to be greater than 0, but got groups=", groups);
 
   TORCH_CHECK(weight_dim == k,
            "Expected ", weight_dim, "-dimensional input for ", weight_dim,

aten/src/ATen/native/Distributions.h

@@ -1,6 +1,5 @@
 #pragma once
 
-#include <array>
 #include <ATen/native/Math.h>
 #include <c10/macros/Macros.h>
 #include <c10/util/MathConstants.h>
@@ -128,7 +127,7 @@ C10_DEVICE scalar_t sample_gamma(scalar_t alpha, BaseSampler<accscalar_t, unifor
 
 template<typename scalar_t>
 C10_DEVICE scalar_t stirling_approx_tail(scalar_t k) {
-  constexpr static scalar_t kTailValues[] = {
+  const static scalar_t kTailValues[] = {
     0.0810614667953272,
     0.0413406959554092,
     0.0276779256849983,
@@ -140,7 +139,7 @@ C10_DEVICE scalar_t stirling_approx_tail(scalar_t k) {
     0.00925546218271273,
     0.00833056343336287
   };
-  if (k < std::size(kTailValues)) {
+  if (k <= 9) {
     return kTailValues[static_cast<size_t>(k)];
   }
   scalar_t kp1sq = (k + 1) * (k + 1);

aten/src/ATen/native/Math.h

@@ -581,7 +581,7 @@ scalar_t ratevl(scalar_t x, const scalar_t num[], int64_t M,
 template <typename scalar_t>
 static scalar_t lanczos_sum_expg_scaled(scalar_t x) {
   // lanczos approximation
-  static constexpr scalar_t lanczos_sum_expg_scaled_num[13] = {
+  static const scalar_t lanczos_sum_expg_scaled_num[13] = {
     0.006061842346248906525783753964555936883222,
     0.5098416655656676188125178644804694509993,
     19.51992788247617482847860966235652136208,
@@ -596,7 +596,7 @@ static scalar_t lanczos_sum_expg_scaled(scalar_t x) {
     103794043.1163445451906271053616070238554,
     56906521.91347156388090791033559122686859
   };
-  static constexpr scalar_t lanczos_sum_expg_scaled_denom[13] = {
+  static const scalar_t lanczos_sum_expg_scaled_denom[13] = {
     1.,
     66.,
     1925.,
@@ -712,7 +712,7 @@ static scalar_t _igamc_helper_series(scalar_t a, scalar_t x) {
 template <typename scalar_t>
 static scalar_t _igam_helper_asymptotic_series(scalar_t a, scalar_t x, bool igam) {
   // Compute igam/igamc using DLMF 8.12.3/8.12.4 [igam1]
-  static constexpr scalar_t d[25][25] =
+  static const scalar_t d[25][25] =
     {{-3.3333333333333333e-1, 8.3333333333333333e-2, -1.4814814814814815e-2,
       1.1574074074074074e-3, 3.527336860670194e-4, -1.7875514403292181e-4,
       3.9192631785224378e-5, -2.1854485106799922e-6, -1.85406221071516e-6,

aten/src/ATen/native/Normalization.cpp

@@ -62,7 +62,7 @@
 #include <utility>
 #include <vector>
 
-static constexpr int MIOPEN_DIM_MAX = 5;
+static const int MIOPEN_DIM_MAX = 5;
 
 namespace at::meta {
 

aten/src/ATen/native/TensorAdvancedIndexingUtils.h

@@ -77,7 +77,7 @@ inline AdvancedIndex make_info(Tensor self, IOptTensorListRef orig) {
   // next broadcast all index tensors together
   try {
     indices = expand_outplace(indices);
-  } catch (std::exception&) {
+  } catch (std::exception& e) {
     TORCH_CHECK_INDEX(
         false,
         "shape mismatch: indexing tensors could not be broadcast together"

aten/src/ATen/native/cpu/UpSampleKernel.cpp

@@ -1038,7 +1038,7 @@ struct HelperInterpNearest : public HelperInterpBase {
   // We keep this structure for BC and consider as deprecated.
   // See HelperInterpNearestExact as replacement
 
-  static constexpr int interp_size = 1;
+  static const int interp_size = 1;
 
   static inline void init_indices_weights(
     at::ScalarType output_type,
@@ -1155,7 +1155,7 @@ struct HelperInterpNearestExact : public HelperInterpNearest {
 
 struct HelperInterpLinear : public HelperInterpBase {
 
-  static constexpr int interp_size = 2;
+  static const int interp_size = 2;
 
   // Compute indices and weights for each interpolated dimension
   // indices_weights = {
@@ -1275,7 +1275,7 @@ struct HelperInterpLinear : public HelperInterpBase {
 
 struct HelperInterpCubic : public HelperInterpBase {
 
-  static constexpr int interp_size = 4;
+  static const int interp_size = 4;
 
   // Compute indices and weights for each interpolated dimension
   // indices_weights = {

aten/src/ATen/native/cuda/Blas.cpp

@@ -1359,8 +1359,7 @@ _scaled_gemm(
           const ScalingType scaling_choice_a, const ScalingType scaling_choice_b,
           const std::optional<Tensor>& bias,
           const bool use_fast_accum,
-          Tensor& out,
-          const std::optional<Tensor>& alpha = std::nullopt) {
+          Tensor& out) {
   cublasCommonArgs args(mat1, mat2, out, scale_a, scale_b, std::nullopt, scaling_choice_a, scaling_choice_b);
   const auto out_dtype_ = args.result->scalar_type();
   TORCH_CHECK(args.transa == 't' && args.transb == 'n', "Only multiplication of row-major and column-major matrices is supported by cuBLASLt");
@@ -1411,8 +1410,7 @@ _scaled_gemm(
           args.scale_result_ptr,
           args.result_ld,
           out_dtype_,
-          use_fast_accum,
-          alpha);
+          use_fast_accum);
       return out;
   }
 }
@@ -1761,7 +1759,6 @@ enum class ScaledGemmImplementation {
   MXFP8_MXFP8 = 6,
   NVFP4_NVFP4 = 7,
   NVFP4_NVFP4_SINGLE_SCALE = 8,
-  MXFP4_MXFP4 = 9,
 };
 
 /**
@@ -1958,39 +1955,10 @@ bool check_mxfp8_recipe(c10::ScalarType type_a,
   return true;
 }
 
-/**
- * Both inputs must be fp4
- * A, B must have 1 scale each, {Blockwise_1x32, e8m0}
- */
-bool check_mxfp4_recipe(c10::ScalarType type_a,
-                        std::vector<ScalingType>& recipe_a,
-                        ArrayRef<Tensor>& scales_a,
-                        c10::ScalarType type_b,
-                        std::vector<ScalingType>& recipe_b,
-                        ArrayRef<Tensor>& scales_b) {
-  // both types must be fp4
-  if (type_a != ScalarType::Float4_e2m1fn_x2 || type_b != ScalarType::Float4_e2m1fn_x2) {
-    return false;
-  }
-
-  // 1 scales, 1 recipes for each input
-  if (scales_a.size() != 1 || recipe_a.size() != 1 || scales_b.size() != 1 || recipe_b.size() != 1) {
-    return false;
-  }
-
-  // Need {Blockwise_1x32, e8m0} for A & B
-  if (recipe_a[0] != ScalingType::BlockWise1x32) return false;
-  if (scales_a[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
-  if (recipe_b[0] != ScalingType::BlockWise1x32) return false;
-  if (scales_b[0].scalar_type() != ScalarType::Float8_e8m0fnu) return false;
-
-  return true;
-}
-
 using acceptance_fn = std::function<bool(c10::ScalarType, std::vector<ScalingType>&, ArrayRef<Tensor>&, c10::ScalarType, std::vector<ScalingType>&, ArrayRef<Tensor>&)>;
 using namespace std::placeholders;
 
-std::array<std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>, 9> scale_kernel_dispatch = {{
+std::array<std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>, 8> scale_kernel_dispatch = {{
   { "tensorwise_tensorwise", check_tensorwise_recipe, ScaledGemmImplementation::TENSORWISE_TENSORWISE },
   { "rowwise_rowwise", check_rowwise_recipe, ScaledGemmImplementation::ROWWISE_ROWWISE},
   { "block_1x128_128x128", std::bind(check_deepseek_recipe, ScalingType::BlockWise1x128, ScalingType::BlockWise128x128, _1, _2, _3, _4, _5, _6),
@@ -2001,8 +1969,7 @@ std::array<std::tuple<std::string, acceptance_fn, ScaledGemmImplementation>, 9>
     ScaledGemmImplementation::BLOCK_1x128_1x128},
   { "nvfp4_nvfp4", check_nvfp4_recipe, ScaledGemmImplementation::NVFP4_NVFP4},
   { "nvfp4_nvfp4_single_scale", check_nvfp4_recipe_single_scale, ScaledGemmImplementation::NVFP4_NVFP4_SINGLE_SCALE },
-  { "mxfp8_mxfp8", check_mxfp8_recipe, ScaledGemmImplementation::MXFP8_MXFP8},
-  { "mxfp4_mxfp4", check_mxfp4_recipe, ScaledGemmImplementation::MXFP4_MXFP4}}};
+  { "mxfp8_mxfp8", check_mxfp8_recipe, ScaledGemmImplementation::MXFP8_MXFP8}}};
 
 Tensor&
 _scaled_tensorwise_tensorwise(
@@ -2220,22 +2187,15 @@ _scaled_mxfp8_mxfp8(
   TORCH_CHECK_VALUE(isFloat8Type(mat_a.scalar_type()) && isFloat8Type(mat_b.scalar_type()), "mat_a and mat_b must be fp8 types, got: ",
       mat_a.scalar_type(), mat_b.scalar_type());
 
-#ifdef USE_ROCM
-  auto scale_a_elems = ceil_div<int64_t>(mat_a.size(0), 32) * mat_a.size(1);
-  auto scale_b_elems = ceil_div<int64_t>(mat_b.size(1), 32) * mat_b.size(0);
-#else
   auto scale_a_elems = round_up<int64_t>(mat_a.size(0), 128) * round_up<int64_t>(ceil_div<int64_t>(mat_a.size(1), 32), 4);
   auto scale_b_elems = round_up<int64_t>(mat_b.size(1), 128) * round_up<int64_t>(ceil_div<int64_t>(mat_b.size(0), 32), 4);
-#endif
   TORCH_CHECK_VALUE(scale_a_elems == scale_a.numel(),
          "For Blockwise scaling scale_a should have ", scale_a_elems, " elements, got: ", scale_a.numel());
   TORCH_CHECK_VALUE(scale_b_elems == scale_b.numel(),
          "For Blockwise scaling scale_b should have ", scale_b_elems, " elements, got: ", scale_b.numel());
 
-#ifndef USE_ROCM
   TORCH_CHECK_VALUE(swizzle_a == SwizzleType::SWIZZLE_32_4_4, "scale_a must be swizzled to SWIZZLE_32_4_4 format");
   TORCH_CHECK_VALUE(swizzle_b == SwizzleType::SWIZZLE_32_4_4, "scale_b must be swizzled to SWIZZLE_32_4_4 format");
-#endif
 
   TORCH_CHECK_VALUE(scale_a.is_contiguous() && scale_b.is_contiguous(),
         "For Blockwise scaling both scales should be contiguous");
@@ -2265,56 +2225,6 @@ _scaled_mxfp8_mxfp8(
   return _scaled_gemm(mat_a, mat_b, scale_a, scale_b, scaling_choice_a, scaling_choice_b, bias, false /* use_fast_accum */, out);
 }
 
-
-Tensor&
-_scaled_mxfp4_mxfp4(
-          const Tensor& mat_a, const Tensor& mat_b,
-          const Tensor& scale_a, const SwizzleType swizzle_a,
-          const Tensor& scale_b, const SwizzleType swizzle_b,
-          const std::optional<Tensor>& bias,
-          const c10::ScalarType out_dtype,
-          Tensor& out) {
-#ifndef USE_ROCM
-  TORCH_CHECK_NOT_IMPLEMENTED(false, "MXFP4 scaling supported on ROCM only");
-#endif
-  // Restrictions:
-  // A, B are FP4, scales are e8m0, A: shape K//32, B: K, N//32
-  TORCH_CHECK_VALUE(mat_a.scalar_type() == at::kFloat4_e2m1fn_x2 && mat_b.scalar_type() == at::kFloat4_e2m1fn_x2, "mat_a and mat_b must be fp4 types, got: ",
-      mat_a.scalar_type(), mat_b.scalar_type());
-
-  auto scale_a_elems = ceil_div<int64_t>(2 * mat_a.size(0), 32) * mat_a.size(1);
-  auto scale_b_elems = ceil_div<int64_t>(2 * mat_b.size(1), 32) * mat_b.size(0);
-  TORCH_CHECK_VALUE(scale_a_elems == scale_a.numel(),
-         "For Blockwise scaling scale_a should have ", scale_a_elems, " elements, got: ", scale_a.numel());
-  TORCH_CHECK_VALUE(scale_b_elems == scale_b.numel(),
-         "For Blockwise scaling scale_b should have ", scale_b_elems, " elements, got: ", scale_b.numel());
-
-  TORCH_CHECK_VALUE(scale_a.is_contiguous() && scale_b.is_contiguous(),
-        "For Blockwise scaling both scales should be contiguous");
-
-  TORCH_CHECK_VALUE(out.scalar_type() == out_dtype, "expected out.scalar_type() to be ", out_dtype, ", but got ", out_dtype);
-
-  auto scaling_choice_a = ScalingType::BlockWise1x32;
-  auto scaling_choice_b = ScalingType::BlockWise1x32;
-
-#if ROCM_VERSION >= 70000
-  TORCH_CHECK_NOT_IMPLEMENTED(at::detail::getCUDAHooks().isGPUArch({"gfx950"}),
-              "Block-wise scaling for Float8_e8m0fnu is only supported on gfx950");
-
-  TORCH_CHECK_VALUE(mat_a.size(0) % 32 == 0 && mat_a.size(1) % 32 == 0 &&
-              mat_b.size(0) % 32 == 0 && mat_b.size(1) % 32 == 0,
-              "Matrix dimensions must be multiples of 32 for block-wise scaling");
-
-  TORCH_CHECK_VALUE(out.scalar_type() == ScalarType::BFloat16 ||
-              out.scalar_type() == ScalarType::Half,
-              "Block-wise scaling only supports BFloat16 or Half output types");
-#else
-    TORCH_CHECK_NOT_IMPLEMENTED(false, "Block-wise scaling for Float8_e8m0fnu requires ROCm 7.0 or later");
-#endif
-
-  return _scaled_gemm(mat_a, mat_b, scale_a, scale_b, scaling_choice_a, scaling_choice_b, bias, false /* use_fast_accum */, out);
-}
-
 Tensor&
 _scaled_nvfp4_nvfp4(
           const Tensor& mat_a, const Tensor& mat_b,
@@ -2558,8 +2468,6 @@ _scaled_mm_cuda_v2_out(
     TORCH_CHECK_NOT_IMPLEMENTED(false, "Only single-scale NVFP4 currently supported");
   } else if (gemm_impl == ScaledGemmImplementation::NVFP4_NVFP4_SINGLE_SCALE) {
     return _scaled_nvfp4_nvfp4(mat_a, mat_b, scale_a[0], swizzle_a_enum[0], scale_b[0], swizzle_b_enum[0], bias, out_dtype_, true /* single_scale */, out);
-  } else if (gemm_impl == ScaledGemmImplementation::MXFP4_MXFP4) {
-    return _scaled_mxfp4_mxfp4(mat_a, mat_b, scale_a[0], swizzle_a_enum[0], scale_b[0], swizzle_b_enum[0], bias, out_dtype_, out);
   } else {
     TORCH_CHECK_VALUE(false, "Invalid state - found an implementation, but not really");
   }

aten/src/ATen/native/cuda/DilatedMaxPool2d.cu

@@ -249,7 +249,7 @@ __global__ void max_pool_forward_nhwc(
 }
 
 
-static constexpr int BLOCK_THREADS = 256;
+static const int BLOCK_THREADS = 256;
 
 template <typename scalar_t, typename accscalar_t>
 #if defined (USE_ROCM)

aten/src/ATen/native/cuda/Embedding.cu

@@ -36,9 +36,9 @@ namespace at::native {
 namespace {
 
 #if defined(USE_ROCM)
-static constexpr int BLOCKDIMY = 16;
+static const int BLOCKDIMY = 16;
 #else
-static constexpr int BLOCKDIMY = 32;
+static const int BLOCKDIMY = 32;
 #endif
 
 template

aten/src/ATen/native/cuda/IGammaKernel.cu

@@ -82,7 +82,7 @@ __host__ __device__ scalar_t lanczos_sum_expg_scaled(scalar_t x) {
   // lanczos approximation
   using accscalar_t = at::acc_type<scalar_t, /*is_cuda=*/true>;
 
-  constexpr accscalar_t lanczos_sum_expg_scaled_num[13] = {
+  static const accscalar_t lanczos_sum_expg_scaled_num[13] = {
     0.006061842346248906525783753964555936883222,
     0.5098416655656676188125178644804694509993,
     19.51992788247617482847860966235652136208,
@@ -97,7 +97,7 @@ __host__ __device__ scalar_t lanczos_sum_expg_scaled(scalar_t x) {
     103794043.1163445451906271053616070238554,
     56906521.91347156388090791033559122686859
   };
-  constexpr accscalar_t lanczos_sum_expg_scaled_denom[13] = {
+  static const accscalar_t lanczos_sum_expg_scaled_denom[13] = {
     1.,
     66.,
     1925.,
@@ -126,10 +126,10 @@ __host__ __device__ scalar_t _igam_helper_fac(scalar_t a, scalar_t x) {
 
   using accscalar_t = at::acc_type<scalar_t, /*is_cuda=*/true>;
   accscalar_t ax, fac, res, num, numfac;
-  constexpr accscalar_t MAXLOG = std::is_same_v<accscalar_t,double> ?
+  static const accscalar_t MAXLOG = std::is_same_v<accscalar_t,double> ?
     7.09782712893383996843E2 : 88.72283905206835;
-  constexpr accscalar_t EXP1 = 2.718281828459045;
-  constexpr accscalar_t lanczos_g = 6.024680040776729583740234375;
+  static const accscalar_t EXP1 = 2.718281828459045;
+  static const accscalar_t lanczos_g = 6.024680040776729583740234375;
 
   if (::fabs(a - x) > 0.4 * ::fabs(a)) {
     ax = a * ::log(x) - x - ::lgamma(a);
@@ -158,9 +158,9 @@ __host__ __device__ scalar_t _igam_helper_series(scalar_t a, scalar_t x) {
   // Compute igam using DLMF 8.11.4. [igam1]
 
   using accscalar_t = at::acc_type<scalar_t, /*is_cuda=*/true>;
-  constexpr accscalar_t MACHEP = std::is_same_v<accscalar_t, double> ?
+  static const accscalar_t MACHEP = std::is_same_v<accscalar_t, double> ?
     1.11022302462515654042E-16 : 5.9604644775390625E-8;
-  constexpr int MAXITER = 2000;
+  static const int MAXITER = 2000;
 
   int i;
   accscalar_t ans, ax, c, r;
@@ -196,8 +196,8 @@ __host__ __device__ scalar_t _igamc_helper_series(scalar_t a, scalar_t x) {
   accscalar_t fac = 1;
   accscalar_t sum = 0;
   accscalar_t term, logx;
-  constexpr int MAXITER = 2000;
-  constexpr accscalar_t MACHEP = std::is_same_v<accscalar_t, double> ?
+  static const int MAXITER = 2000;
+  static const accscalar_t MACHEP = std::is_same_v<accscalar_t, double> ?
     1.11022302462515654042E-16 : 5.9604644775390625E-8;
 
   for (n = 1; n < MAXITER; n++) {
@@ -219,7 +219,7 @@ __host__ __device__ scalar_t _igam_helper_asymptotic_series(scalar_t a, scalar_t
   // Compute igam/igamc using DLMF 8.12.3/8.12.4 [igam1]
 
   using accscalar_t = at::acc_type<scalar_t, /*is_cuda=*/true>;
-  constexpr accscalar_t d[25][25] =
+  static const accscalar_t d[25][25] =
     {{-3.3333333333333333e-1, 8.3333333333333333e-2, -1.4814814814814815e-2, 1.1574074074074074e-3, 3.527336860670194e-4, -1.7875514403292181e-4, 3.9192631785224378e-5, -2.1854485106799922e-6, -1.85406221071516e-6, 8.296711340953086e-7, -1.7665952736826079e-7, 6.7078535434014986e-9, 1.0261809784240308e-8, -4.3820360184533532e-9, 9.1476995822367902e-10, -2.551419399494625e-11, -5.8307721325504251e-11, 2.4361948020667416e-11, -5.0276692801141756e-12, 1.1004392031956135e-13, 3.3717632624009854e-13, -1.3923887224181621e-13, 2.8534893807047443e-14, -5.1391118342425726e-16, -1.9752288294349443e-15},
     {-1.8518518518518519e-3, -3.4722222222222222e-3, 2.6455026455026455e-3, -9.9022633744855967e-4, 2.0576131687242798e-4, -4.0187757201646091e-7, -1.8098550334489978e-5, 7.6491609160811101e-6, -1.6120900894563446e-6, 4.6471278028074343e-9, 1.378633446915721e-7, -5.752545603517705e-8, 1.1951628599778147e-8, -1.7543241719747648e-11, -1.0091543710600413e-9, 4.1627929918425826e-10, -8.5639070264929806e-11, 6.0672151016047586e-14, 7.1624989648114854e-12, -2.9331866437714371e-12, 5.9966963656836887e-13, -2.1671786527323314e-16, -4.9783399723692616e-14, 2.0291628823713425e-14, -4.13125571381061e-15},
     {4.1335978835978836e-3, -2.6813271604938272e-3, 7.7160493827160494e-4, 2.0093878600823045e-6, -1.0736653226365161e-4, 5.2923448829120125e-5, -1.2760635188618728e-5, 3.4235787340961381e-8, 1.3721957309062933e-6, -6.298992138380055e-7, 1.4280614206064242e-7, -2.0477098421990866e-10, -1.4092529910867521e-8, 6.228974084922022e-9, -1.3670488396617113e-9, 9.4283561590146782e-13, 1.2872252400089318e-10, -5.5645956134363321e-11, 1.1975935546366981e-11, -4.1689782251838635e-15, -1.0940640427884594e-12, 4.6622399463901357e-13, -9.905105763906906e-14, 1.8931876768373515e-17, 8.8592218725911273e-15},
@@ -248,7 +248,7 @@ __host__ __device__ scalar_t _igam_helper_asymptotic_series(scalar_t a, scalar_t
 
   int k, n, sgn;
   int maxpow = 0;
-  constexpr accscalar_t MACHEP = std::is_same_v<accscalar_t, double> ?
+  static const accscalar_t MACHEP = std::is_same_v<accscalar_t, double> ?
     1.11022302462515654042E-16 : 5.9604644775390625E-8;
   accscalar_t lambda = x / a;
   accscalar_t sigma = (x - a) / a;
@@ -314,12 +314,12 @@ __host__ __device__ scalar_t _igamc_helper_continued_fraction(scalar_t a, scalar
   int i;
   accscalar_t ans, ax, c, yc, r, t, y, z;
   accscalar_t pk, pkm1, pkm2, qk, qkm1, qkm2;
-  constexpr int MAXITER = 2000;
-  constexpr accscalar_t MACHEP = std::is_same_v<accscalar_t, double> ?
+  static const int MAXITER = 2000;
+  static const accscalar_t MACHEP = std::is_same_v<accscalar_t, double> ?
     1.11022302462515654042E-16 : 5.9604644775390625E-8;
-  constexpr accscalar_t BIG = std::is_same_v<accscalar_t,double> ?
+  static const accscalar_t BIG = std::is_same_v<accscalar_t,double> ?
     4.503599627370496e15 : 16777216.;
-  constexpr accscalar_t BIGINV = std::is_same_v<accscalar_t,double> ?
+  static const accscalar_t BIGINV = std::is_same_v<accscalar_t,double> ?
     2.22044604925031308085e-16 : 5.9604644775390625E-8;
 
   ax = _igam_helper_fac(a, x);
@@ -385,10 +385,10 @@ __noinline__ __host__ __device__ scalar_t calc_igammac(scalar_t a, scalar_t x) {
   using accscalar_t = at::acc_type<scalar_t, /*is_cuda=*/true>;
   accscalar_t absxma_a;
 
-  constexpr accscalar_t SMALL = 20.0;
-  constexpr accscalar_t LARGE = 200.0;
-  constexpr accscalar_t SMALLRATIO = 0.3;
-  constexpr accscalar_t LARGERATIO = 4.5;
+  static const accscalar_t SMALL = 20.0;
+  static const accscalar_t LARGE = 200.0;
+  static const accscalar_t SMALLRATIO = 0.3;
+  static const accscalar_t LARGERATIO = 4.5;
 
   if ((x < 0) || (a < 0)) {
     // out of defined-region of the function
@@ -467,10 +467,10 @@ __noinline__ __host__ __device__ scalar_t calc_igamma(scalar_t a, scalar_t x) {
 
   using accscalar_t = at::acc_type<scalar_t, /*is_cuda=*/true>;
   accscalar_t absxma_a;
-  constexpr accscalar_t SMALL = 20.0;
-  constexpr accscalar_t LARGE = 200.0;
-  constexpr accscalar_t SMALLRATIO = 0.3;
-  constexpr accscalar_t LARGERATIO = 4.5;
+  static const accscalar_t SMALL = 20.0;
+  static const accscalar_t LARGE = 200.0;
+  static const accscalar_t SMALLRATIO = 0.3;
+  static const accscalar_t LARGERATIO = 4.5;
 
   // boundary values following SciPy
   if ((x < 0) || (a < 0)) {

aten/src/ATen/native/cuda/Math.cuh

@@ -231,7 +231,7 @@ const auto lcm_string = jiterator_stringify(
 const auto digamma_string = jiterator_stringify(
   template <typename T>
   T digamma(T x) {
-    static constexpr double PI_f64 = 3.14159265358979323846;
+    static const double PI_f64 = 3.14159265358979323846;
 
     // Short-circuits if x is +/- 0 and returns -/+ ∞ per the C++ standard
     if (x == 0) {
@@ -3072,9 +3072,9 @@ template <typename scalar_t>
 static inline C10_HOST_DEVICE scalar_t calc_digamma(scalar_t in) {
   // [C++ Standard Reference: Gamma Function] https://en.cppreference.com/w/cpp/numeric/math/tgamma
   using accscalar_t = at::acc_type<scalar_t, /*is_cuda=*/true>;
-  static constexpr double PI_f64 = 3.14159265358979323846;
-  constexpr accscalar_t PSI_10 = 2.25175258906672110764;
-  constexpr accscalar_t A[] = {
+  static const double PI_f64 = 3.14159265358979323846;
+  const accscalar_t PSI_10 = 2.25175258906672110764;
+  const accscalar_t A[] = {
       8.33333333333333333333E-2,
       -2.10927960927960927961E-2,
       7.57575757575757575758E-3,

aten/src/ATen/native/cuda/Reduce.cuh

@@ -1097,7 +1097,11 @@ ReduceConfig setReduceConfig(const TensorIterator& iter){
   // threads with different threadIdx.x are independent and will produce results for different outputs.
   // In such case, values in each loaded vector always correspond to different outputs.
   if (fastest_moving_stride == sizeof(scalar_t)) {
+#ifdef USE_ROCM
     if (reduction_on_fastest_striding_dimension && dim0 >= 128 && iter.num_reduce_dims() == 1) {
+#else
+    if (reduction_on_fastest_striding_dimension && dim0 > 128 && iter.num_reduce_dims() == 1 && vt0 >= input_vec_size) {
+#endif
       // Case 1: "vectorize along input"
       // Note that if vt0 < ReduceConfig::vec_size, then this means the register pressure could be high, in such case,
       // we should avoid vectorization.

aten/src/ATen/native/cuda/ReduceMomentKernel.cu

@@ -39,14 +39,9 @@ static void std_var_kernel_cuda(TensorIterator& iter, double correction, bool ta
 template <typename scalar_t, typename acc_t=scalar_t, typename out_t=scalar_t>
 void mean_kernel_impl(TensorIterator& iter) {
   //  returns acc_t for all non-complex dtypes and returns T for c10::complex<T>
-  constexpr bool is_16_bits = sizeof(scalar_t) == 2;
   using factor_t = typename c10::scalar_value_type<acc_t>::type;
   factor_t factor = static_cast<factor_t>(iter.num_output_elements()) / iter.numel();
-  if constexpr (is_16_bits) {
-    gpu_reduce_kernel<scalar_t, out_t, /*vt0=*/4, /*input_vec_size=*/8>(iter, MeanOps<scalar_t, acc_t, factor_t, out_t> {factor});
-  } else {
-    gpu_reduce_kernel<scalar_t, out_t>(iter, MeanOps<scalar_t, acc_t, factor_t, out_t> {factor});
-  }
+  gpu_reduce_kernel<scalar_t, out_t>(iter, MeanOps<scalar_t, acc_t, factor_t, out_t> {factor});
 }
 
 static void mean_kernel_cuda(TensorIterator& iter) {

aten/src/ATen/native/cuda/ReduceSumProdKernel.cu

@@ -13,19 +13,24 @@ namespace at::native {
 template <typename scalar_t, typename acc_t = scalar_t, typename out_t = scalar_t>
 struct sum_functor {
   void operator()(TensorIterator& iter) {
-    const auto sum_combine = [] GPU_LAMBDA(acc_t a, acc_t b) -> acc_t {
-      return a + b;
-    };
-    constexpr bool is_16_bits = sizeof(scalar_t) == 2;
-    if constexpr (is_16_bits) {
+#ifdef USE_ROCM
+    // Half and BFloat16 can be packed in groups of up to 8 elements and
+    // can use *_DWORDX4 instructions to achieve that.
+    const bool is_16_bits =
+      ( (std::is_same<at::Half, scalar_t>::value) ||
+        (std::is_same<at::BFloat16, scalar_t>::value) );
+    if (is_16_bits) {
       gpu_reduce_kernel<scalar_t, out_t, /*vt0=*/4, /*input_vec_size=*/8>(
-        iter, func_wrapper<out_t>(sum_combine)
-      );
-    } else {
-      gpu_reduce_kernel<scalar_t, out_t>(
-        iter, func_wrapper<out_t>(sum_combine)
-      );
+        iter, func_wrapper<out_t>([] GPU_LAMBDA(acc_t a, acc_t b) -> acc_t {
+          return a + b;
+        }));
+      return;
     }
+#endif
+    gpu_reduce_kernel<scalar_t, out_t>(
+        iter, func_wrapper<out_t>([] GPU_LAMBDA(acc_t a, acc_t b) -> acc_t {
+          return a + b;
+        }));
   }
 };
 

aten/src/ATen/native/cuda/UpSample.cuh

@@ -277,7 +277,7 @@ struct BilinearFilterFunctor {
     return 0;
   }
 
-  static constexpr int size = 2;
+  static const int size = 2;
 };
 
 // taken from
@@ -301,7 +301,7 @@ struct BicubicFilterFunctor {
     return 0;
   }
 
-  static constexpr int size = 4;
+  static const int size = 4;
 };
 
 template <typename accscalar_t>

aten/src/ATen/native/cuda/UpSampleBilinear2d.cu

@@ -127,6 +127,29 @@ __global__ void upsample_bilinear2d_nhwc_out_frame(
   }
 }
 
+#ifdef USE_ROCM
+// Helper function to compute output pixel range that can contribute to input pixel
+template <typename accscalar_t>
+__device__ __forceinline__ void compute_output_range(
+    int input_pos,
+    accscalar_t scale,
+    int output_size,
+    bool align_corners,
+    int& min_output,
+    int& max_output) {
+  accscalar_t lo, hi;
+  if (align_corners) {
+      lo = static_cast<accscalar_t>(input_pos - 1) / scale;
+      hi = static_cast<accscalar_t>(input_pos + 1) / scale;
+  } else {
+      lo = (input_pos - static_cast<accscalar_t>(0.5)) / scale - static_cast<accscalar_t>(0.5);
+      hi = (input_pos + static_cast<accscalar_t>(1.5)) / scale - static_cast<accscalar_t>(0.5);
+  }
+  min_output = max(0, static_cast<int>(ceil(lo)));
+  max_output = min(output_size - 1, static_cast<int>(floor(hi)));
+}
+#endif
+
 // Backward (adjoint) operation 1 <- 2 (accumulates)
 template <typename scalar_t, typename accscalar_t>
 C10_LAUNCH_BOUNDS_1(1024)
@@ -141,8 +164,74 @@ __global__ void upsample_bilinear2d_backward_out_frame(
     const bool align_corners,
     scalar_t* __restrict__ idata,
     const scalar_t* __restrict__ odata) {
-  const size_t o_numel = nc * width2 * height2;
+  // In C++, integer multiplication, like in standard arithmetic, is generally commutative.
   const size_t i_numel = nc * width1 * height1;
+#ifdef USE_ROCM
+  for (size_t index = blockDim.x * blockIdx.x + threadIdx.x; index < i_numel;
+       index += blockDim.x * gridDim.x) {
+    // Decode input pixel coordinates
+    size_t index_temp = index;
+    const int w1 = index_temp % width1;
+    index_temp /= width1;
+    const int h1 = index_temp % height1;
+    const size_t nc_idx = index_temp / height1;
+
+    accscalar_t grad_sum = 0;
+
+    // Find range of output pixels that could interpolate from this input pixel
+    int h2_min, h2_max, w2_min, w2_max;
+    compute_output_range<accscalar_t>(h1, rheight, height2, align_corners, h2_min, h2_max);
+    compute_output_range<accscalar_t>(w1, rwidth, width2, align_corners, w2_min, w2_max);
+
+    // Iterate over potential output pixels
+    for (int h2 = h2_min; h2 <= h2_max; h2++) {
+      for (int w2 = w2_min; w2 <= w2_max; w2++) {
+        // Compute source coordinates for this output pixel
+        const accscalar_t h1r = area_pixel_compute_source_index<accscalar_t>(
+            rheight, h2, align_corners, /*cubic=*/false);
+        const int h1_base = (int)h1r;
+        const int h1p = (h1_base < height1 - 1) ? 1 : 0;
+        const accscalar_t h1lambda = h1r - h1_base;
+        const accscalar_t h0lambda = static_cast<accscalar_t>(1) - h1lambda;
+
+        const accscalar_t w1r = area_pixel_compute_source_index<accscalar_t>(
+            rwidth, w2, align_corners, /*cubic=*/false);
+        const int w1_base = (int)w1r;
+        const int w1p = (w1_base < width1 - 1) ? 1 : 0;
+        const accscalar_t w1lambda = w1r - w1_base;
+        const accscalar_t w0lambda = static_cast<accscalar_t>(1) - w1lambda;
+
+        // Check if our input pixel participates in this interpolation and accumulate all weights
+        // At boundaries, h1p=0 or w1p=0 causes some sampling positions to collapse
+        // to the same pixel, so we need to accumulate weights from all matching positions
+        accscalar_t weight = 0;
+
+        // Check all four interpolation positions and accumulate weights
+        if (h1 == h1_base && w1 == w1_base) {
+          weight += h0lambda * w0lambda;  // top-left
+        }
+        if (h1 == h1_base && w1 == w1_base + w1p) {
+          weight += h0lambda * w1lambda;  // top-right (may be same as top-left if w1p=0)
+        }
+        if (h1 == h1_base + h1p && w1 == w1_base) {
+          weight += h1lambda * w0lambda;  // bottom-left (may be same as top-left if h1p=0)
+        }
+        if (h1 == h1_base + h1p && w1 == w1_base + w1p) {
+          weight += h1lambda * w1lambda;  // bottom-right (may collapse to other positions)
+        }
+
+        if (weight > 0) {
+          const size_t output_idx = nc_idx * height2 * width2 + h2 * width2 + w2;
+          grad_sum += weight * static_cast<accscalar_t>(odata[output_idx]);
+        }
+      }
+    }
+
+    // Write accumulated gradient (no atomics needed)
+    idata[index] = static_cast<scalar_t>(grad_sum);
+  }
+#else
+  const size_t o_numel = nc * width2 * height2;
   for (size_t index = blockDim.x * blockIdx.x + threadIdx.x; index < o_numel;
        index += blockDim.x * gridDim.x) {
     size_t index_temp = index;
@@ -191,6 +280,7 @@ __global__ void upsample_bilinear2d_backward_out_frame(
         static_cast<scalar_t>(h1lambda * w1lambda * d2val),
         true);
   }
+#endif
 }
 
 template <typename scalar_t, typename accscalar_t>
@@ -387,7 +477,6 @@ static void upsample_bilinear2d_backward_out_cuda_template(
   // threads are not covering the whole input tensor.
   grad_input.zero_();
 
-  const size_t num_kernels = nbatch * channels * output_height * output_width;
   const int num_threads = std::min(
       at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock, 1024);
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
@@ -397,6 +486,12 @@ static void upsample_bilinear2d_backward_out_cuda_template(
     return;
   }
 
+#ifdef USE_ROCM
+  constexpr bool use_input = true;
+#else
+  constexpr bool use_input = false;
+#endif
+
   AT_DISPATCH_FLOATING_TYPES_AND2(
       at::ScalarType::Half, at::ScalarType::BFloat16,
       grad_output_.scalar_type(), "upsample_bilinear2d_backward_out_frame", [&] {
@@ -414,6 +509,8 @@ static void upsample_bilinear2d_backward_out_cuda_template(
       const accscalar_t rwidth = area_pixel_compute_scale<accscalar_t>(
           input_width, output_width, align_corners, scales_w);
 
+      const size_t num_kernels = nbatch * channels * output_height * output_width;
+
       upsample_bilinear2d_backward_nhwc_out_frame<scalar_t, accscalar_t>
           <<<ceil_div(num_kernels, static_cast<size_t>(num_threads)), num_threads, 0, stream>>>(
               input_height,
@@ -444,6 +541,8 @@ static void upsample_bilinear2d_backward_out_cuda_template(
       const accscalar_t rwidth = area_pixel_compute_scale<accscalar_t>(
           input_width, output_width, align_corners, scales_w);
 
+      const size_t num_kernels = nbatch * channels * (use_input ? input_height * input_width : output_height * output_width);
+
       upsample_bilinear2d_backward_out_frame<scalar_t, accscalar_t>
           <<<ceil_div(num_kernels, static_cast<size_t>(num_threads)),
              num_threads,

aten/src/ATen/native/cuda/layer_norm_kernel.cu

@@ -141,11 +141,7 @@ WelfordDataLN cuWelfordOnlineSum(
   if constexpr (!rms_norm){
     U delta = val - curr_sum.mean;
     U new_count = curr_sum.count + 1.f;
-#if defined(USE_ROCM) && defined(USE_LAYERNORM_FAST_RECIPROCAL)
-    U new_mean = curr_sum.mean + delta * __builtin_amdgcn_rcpf(new_count);
-#else
     U new_mean = curr_sum.mean + delta * (1.f/new_count); //proper division is slow, this is less accurate but noticeably faster
-#endif
     return {new_mean, curr_sum.sigma2 + delta * (val - new_mean), new_count};
   } else{
     return {0.f, curr_sum.sigma2 + val * val, 0};
@@ -163,11 +159,7 @@ WelfordDataLN cuWelfordCombine(
     U count = dataA.count + dataB.count;
     U mean, sigma2;
     if (count > decltype(dataB.count){0}) {
-#if defined(USE_ROCM) && defined(USE_LAYERNORM_FAST_RECIPROCAL)
-      auto coef = __builtin_amdgcn_rcpf(count);
-#else
       auto coef = 1.f/count; //NB we don't use --use_fast_math, but this is emulation, 1./count goes to intrinsic, `* coef` is multiplication, instead of slow fp division
-#endif
       auto nA = dataA.count * coef;
       auto nB = dataB.count * coef;
       mean = nA*dataA.mean + nB*dataB.mean;

aten/src/ATen/native/mkldnn/Matmul.cpp

@@ -416,7 +416,7 @@ static inline bool checksize(const Tensor& mat1, const Tensor& mat2){
   // else if dim = 3, mat1's size = (b * m * n), mat2's size = (b * n * k)
   // else called from aten::mv, mat1.size = (m * n), mat2.size = (n)
   // only m * n * b * k(if exist) are large enough we can get benefit from mkldnn optimized gemm kernel
-  constexpr int64_t mkldnn_gemm_min_size = 16 * 16 * 16;
+  static const int64_t mkldnn_gemm_min_size = 16 * 16 * 16;
   if (mat1.dim() == 1 && mat2.dim() == 1) {
     // aten::dot
     return mat1.size(0) > mkldnn_gemm_min_size;

aten/src/ATen/native/mps/kernels/Shape.h

@@ -1,16 +1,16 @@
 #pragma once
 #include <c10/metal/common.h>
 
-template <typename idx_type_t = int64_t, unsigned N = c10::metal::max_ndim>
-struct CatSharedParams {
+template <unsigned N = c10::metal::max_ndim, typename idx_type_t = int64_t>
+struct CatLargeSharedParams {
   int32_t ndim;
   int32_t cat_dim;
   ::c10::metal::array<idx_type_t, N> output_strides;
   ::c10::metal::array<idx_type_t, N> output_sizes;
 };
 
-template <typename idx_type_t = int64_t, unsigned N = c10::metal::max_ndim>
-struct CatInputParams {
+template <unsigned N = c10::metal::max_ndim, typename idx_type_t = int64_t>
+struct CatLargeInputParams {
   idx_type_t cat_dim_offset;
   idx_type_t input_element_offset;
   ::c10::metal::array<idx_type_t, N> input_strides;

aten/src/ATen/native/mps/kernels/Shape.metal

@@ -6,25 +6,26 @@
 using namespace metal;
 using namespace c10::metal;
 
-template <typename I, typename T_in, typename T_out>
-kernel void cat(
+template <typename T_in, typename T_out>
+kernel void cat_large(
     constant T_in* input [[buffer(0)]],
     device T_out* output [[buffer(1)]],
-    constant CatSharedParams<I>& shared_params [[buffer(2)]],
-    constant CatInputParams<I>& input_params [[buffer(3)]],
+    constant CatLargeSharedParams<>& shared_params [[buffer(2)]],
+    constant CatLargeInputParams<>& input_params [[buffer(3)]],
     uint tid [[thread_position_in_grid]]) {
   auto ndim = shared_params.ndim;
   auto cat_dim = shared_params.cat_dim;
   constant auto& output_strides = shared_params.output_strides;
+  constant auto& output_sizes = shared_params.output_sizes;
 
   auto cat_dim_offset = input_params.cat_dim_offset;
   auto input_element_offset = input_params.input_element_offset;
   constant auto& input_strides = input_params.input_strides;
   constant auto& input_sizes = input_params.input_sizes;
 
-  auto input_element_idx = static_cast<I>(tid) + input_element_offset;
-  I input_offset = 0;
-  I output_offset = 0;
+  auto input_element_idx = static_cast<int64_t>(tid) + input_element_offset;
+  int64_t input_offset = 0;
+  int64_t output_offset = 0;
 
   for (auto dim = ndim - 1; dim >= 0; dim--) {
     auto dim_size = input_sizes[dim];
@@ -41,45 +42,41 @@ kernel void cat(
   output[output_offset] = static_cast<T_out>(input[input_offset]);
 }
 
-#define REGISTER_CAT_OP(I, T_in, T_out)                          \
-  template [[host_name("cat_" #I "_" #T_in "_" #T_out)]]         \
-  kernel void cat<I, T_in, T_out>(                               \
-      constant T_in * input [[buffer(0)]],                       \
-      device T_out * output [[buffer(1)]],                       \
-      constant CatSharedParams<I> & shared_params [[buffer(2)]], \
-      constant CatInputParams<I> & input_params [[buffer(3)]],   \
+#define REGISTER_CAT_LARGE_OP(T_in, T_out)                           \
+  template [[host_name("cat_large_" #T_in "_" #T_out)]]              \
+  kernel void cat_large<T_in, T_out>(                                \
+      constant T_in * input [[buffer(0)]],                           \
+      device T_out * output [[buffer(1)]],                           \
+      constant CatLargeSharedParams<> & shared_params [[buffer(2)]], \
+      constant CatLargeInputParams<> & input_params [[buffer(3)]],   \
       uint tid [[thread_position_in_grid]]);
 
-#define REGISTER_CAT_OP_ALL_INPUT_TYPES(I, T_out) \
-  REGISTER_CAT_OP(I, float, T_out);               \
-  REGISTER_CAT_OP(I, half, T_out);                \
-  REGISTER_CAT_OP(I, bfloat, T_out);              \
-  REGISTER_CAT_OP(I, int, T_out);                 \
-  REGISTER_CAT_OP(I, uint, T_out);                \
-  REGISTER_CAT_OP(I, long, T_out);                \
-  REGISTER_CAT_OP(I, ulong, T_out);               \
-  REGISTER_CAT_OP(I, short, T_out);               \
-  REGISTER_CAT_OP(I, ushort, T_out);              \
-  REGISTER_CAT_OP(I, char, T_out);                \
-  REGISTER_CAT_OP(I, uchar, T_out);               \
-  REGISTER_CAT_OP(I, bool, T_out);
+#define REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(T_out) \
+  REGISTER_CAT_LARGE_OP(float, T_out);               \
+  REGISTER_CAT_LARGE_OP(half, T_out);                \
+  REGISTER_CAT_LARGE_OP(bfloat, T_out);              \
+  REGISTER_CAT_LARGE_OP(int, T_out);                 \
+  REGISTER_CAT_LARGE_OP(uint, T_out);                \
+  REGISTER_CAT_LARGE_OP(long, T_out);                \
+  REGISTER_CAT_LARGE_OP(ulong, T_out);               \
+  REGISTER_CAT_LARGE_OP(short, T_out);               \
+  REGISTER_CAT_LARGE_OP(ushort, T_out);              \
+  REGISTER_CAT_LARGE_OP(char, T_out);                \
+  REGISTER_CAT_LARGE_OP(uchar, T_out);               \
+  REGISTER_CAT_LARGE_OP(bool, T_out);
 
-#define REGISTER_CAT_FOR_INDEX_TYPE(I)        \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, float);  \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, half);   \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, bfloat); \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, int);    \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, uint);   \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, long);   \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, ulong);  \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, short);  \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, ushort); \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, char);   \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, uchar);  \
-  REGISTER_CAT_OP_ALL_INPUT_TYPES(I, bool);   \
-                                              \
-  REGISTER_CAT_OP(I, float2, float2);         \
-  REGISTER_CAT_OP(I, half2, half2);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(float);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(half);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(bfloat);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(int);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(uint);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(long);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(ulong);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(short);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(ushort);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(char);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(uchar);
+REGISTER_CAT_LARGE_OP_ALL_INPUT_TYPES(bool);
 
-REGISTER_CAT_FOR_INDEX_TYPE(int64_t);
-REGISTER_CAT_FOR_INDEX_TYPE(int32_t);
+REGISTER_CAT_LARGE_OP(float2, float2);
+REGISTER_CAT_LARGE_OP(half2, half2);

aten/src/ATen/native/mps/operations/LinearAlgebra.mm

@@ -196,28 +196,6 @@ bool use_metal_mm(const Tensor& self, const Tensor& other, const Tensor& output)
        other.size(0) > max_stride_size || other.size(1) > max_stride_size);
 }
 
-void map_mps_decomposition_error_code_to_blas(const Tensor& status) {
-  const auto& status_flat = status.view(-1);
-
-  for (const auto i : c10::irange(status_flat.size(0))) {
-    int code = status_flat[i].item<int>();
-    switch (code) {
-      case MPSMatrixDecompositionStatusSuccess:
-        status_flat[i] = 0;
-        break;
-      case MPSMatrixDecompositionStatusNonPositiveDefinite:
-      case MPSMatrixDecompositionStatusSingular:
-        status_flat[i] = 2;
-        break;
-      case MPSMatrixDecompositionStatusFailure:
-        status_flat[i] = -1;
-        break;
-      default:
-        TORCH_INTERNAL_ASSERT(false, "Unknown MPSMatrixDecompositionStatus enum value: ", code);
-    }
-  }
-}
-
 } // anonymous namespace
 
 static void linalg_lu_factor_ex_out_mps_impl(const Tensor& A,
@@ -509,9 +487,6 @@ static void linalg_solve_out_mps_impl(const Tensor& A,
                   "mpsmatrixdecompositionstatus for details.");
     }
   }
-
-  map_mps_decomposition_error_code_to_blas(info);
-
   if (!left) {
     // If this was a right solve, transpose the result back
     result.copy_(result_t.transpose(-2, -1).contiguous());

aten/src/ATen/native/mps/operations/Shape.mm

@@ -3,7 +3,6 @@
 #include <ATen/MemoryOverlap.h>
 #include <ATen/WrapDimUtils.h>
 #include <ATen/mps/MPSProfiler.h>
-#include <ATen/native/Pool.h>
 #include <ATen/native/TensorShape.h>
 #include <ATen/native/TypeProperties.h>
 #include <ATen/native/mps/OperationUtils.h>
@@ -70,40 +69,29 @@ static void check_shape_except_dim(const Tensor& first, const Tensor& second, in
   }
 }
 
-template <typename T>
-std::string get_type_str();
-
-template <>
-std::string get_type_str<int64_t>() {
-  return "int64_t";
-}
-
-template <>
-std::string get_type_str<int32_t>() {
-  return "int32_t";
-}
-
+// This implementation of cat is used only if one of the inputs or the output is
+// too large to use MPSGraph.
 // NOTE: `output` is expected to already have the correct size.
-template <typename idx_type_t>
-static void cat_out_mps_impl(const ITensorListRef& inputs, int64_t dimension, const Tensor& output) {
-  CatSharedParams<idx_type_t> shared_params;
+static void cat_out_large_tensor_mps(const ITensorListRef& inputs, int64_t dimension, const Tensor& output) {
+  CatLargeSharedParams shared_params;
 
   shared_params.ndim = output.dim();
   shared_params.cat_dim = dimension;
 
   for (const auto dim : c10::irange(output.dim())) {
-    shared_params.output_strides[dim] = safe_downcast<idx_type_t, int64_t>(output.stride(dim));
-    shared_params.output_sizes[dim] = safe_downcast<idx_type_t, int64_t>(output.size(dim));
+    shared_params.output_strides[dim] = output.stride(dim);
+    shared_params.output_sizes[dim] = output.size(dim);
   }
 
-  idx_type_t cat_dim_offset = 0;
+  int64_t cat_dim_offset = 0;
   size_t input_idx = 0;
   MPSStream* stream = getCurrentMPSStream();
 
-  // Launch a separate kernels for each input. This will produce some overhead.
-  // In order to launch only one kernel to process all inputs, we would have to
-  // copy all the input tensor data into a packed buffer, which would not be
-  // ideal.
+  // Launch a separate kernels for each input. This will produce some overhead,
+  // but that should be relatively minimal since at least one of the inputs is
+  // very large. In order to launch only one kernel to process all inputs, we
+  // would have to copy all the input tensor data into a packed buffer, which
+  // would not be ideal.
   for (const Tensor& input : inputs) {
     if (input.numel() == 0) {
       continue;
@@ -116,23 +104,21 @@ static void cat_out_mps_impl(const ITensorListRef& inputs, int64_t dimension, co
 
     for (int64_t numel_remaining = input.numel(); numel_remaining > 0; numel_remaining -= max_num_threads) {
       auto num_threads = std::min(max_num_threads, numel_remaining);
-      CatInputParams<idx_type_t> input_params;
+      CatLargeInputParams input_params;
 
-      input_params.cat_dim_offset = safe_downcast<idx_type_t, int64_t>(cat_dim_offset);
-      input_params.input_element_offset = safe_downcast<idx_type_t, int64_t>(input.numel() - numel_remaining);
+      input_params.cat_dim_offset = cat_dim_offset;
+      input_params.input_element_offset = input.numel() - numel_remaining;
 
       for (const auto dim : c10::irange(input.dim())) {
-        input_params.input_strides[dim] = safe_downcast<idx_type_t, int64_t>(input.stride(dim));
-        input_params.input_sizes[dim] = safe_downcast<idx_type_t, int64_t>(input.size(dim));
+        input_params.input_strides[dim] = input.stride(dim);
+        input_params.input_sizes[dim] = input.size(dim);
       }
 
       dispatch_sync_with_rethrow(stream->queue(), ^() {
         @autoreleasepool {
           id<MTLComputeCommandEncoder> computeEncoder = stream->commandEncoder();
-          auto pipeline_state = lib.getPipelineStateForFunc(fmt::format("cat_{}_{}_{}",
-                                                                        get_type_str<idx_type_t>(),
-                                                                        scalarToMetalTypeString(input),
-                                                                        scalarToMetalTypeString(output)));
+          auto pipeline_state = lib.getPipelineStateForFunc(
+              fmt::format("cat_large_{}_{}", scalarToMetalTypeString(input), scalarToMetalTypeString(output)));
           getMPSProfiler().beginProfileKernel(pipeline_state, "cat", {input});
           [computeEncoder setComputePipelineState:pipeline_state];
           mtl_setArgs(computeEncoder, input, output, shared_params, input_params);
@@ -308,6 +294,13 @@ TORCH_IMPL_FUNC(cat_out_mps)
               " and out is on ",
               out.device());
 
+  // TODO: For better performance by eliminating input tensor gathering and post transpose,
+  // TODO: it is better to keep the out tensor's memory format.
+  // TODO: dimension needs to be recomputed as:
+  // TODO: dim = 0 --> dim = 0; dim = 1 or 2 --> dim = out.dim()- dim; otherwise dim = dim-1
+  if (needsGather(out)) {
+    out.unsafeGetTensorImpl()->empty_tensor_restride(MemoryFormat::Contiguous);
+  }
   std::vector<int64_t> size(notSkippedTensor.sizes().vec());
 
   // Compute size of the result in the cat dimension
@@ -338,9 +331,82 @@ TORCH_IMPL_FUNC(cat_out_mps)
   has_large_tensor |= isTooLargeForMPSGraph(out);
 
   if (has_large_tensor) {
-    return mps::cat_out_mps_impl<int64_t>(materialized_inputs, dimension, out);
-  } else {
-    return mps::cat_out_mps_impl<int32_t>(materialized_inputs, dimension, out);
+    return mps::cat_out_large_tensor_mps(materialized_inputs, dimension, out);
+  }
+
+  struct CachedGraph : public MPSCachedGraph {
+    CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {}
+    std::vector<MPSGraphTensor*> inputTensors_;
+    MPSGraphTensor* outputTensor_ = nil;
+  };
+
+  @autoreleasepool {
+    std::string key = "cat_out_mps:" + std::to_string(dimension) + ":" +
+        (memory_format == MemoryFormat::ChannelsLast ? "NHWC" : "NCHW");
+    if (!all_same_dtype) {
+      key += getTensorsStringKey(input_tensors, true, all_same_sizes_and_stride);
+    } else {
+      key += ":" + getMPSTypeString(input_tensors[0].scalar_type(), true) + ":" + std::to_string(inputs.size());
+    }
+    for (auto idx : skipped_tensor_indices) {
+      key += "," + std::to_string(idx);
+    }
+
+    auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
+      auto len_tensor_array = inputs.size() - skipped_tensor_indices.size();
+      std::vector<MPSGraphTensor*> castInputTensors(len_tensor_array);
+      newCachedGraph->inputTensors_.reserve(len_tensor_array);
+
+      for (const auto idx : c10::irange(len_tensor_array)) {
+        const Tensor& tensor = input_tensors[idx];
+        auto scalar_type = getMPSScalarType(tensor.scalar_type());
+        if (tensor.scalar_type() == kBool) {
+          scalar_type = MPSDataTypeInt8;
+        }
+        newCachedGraph->inputTensors_[idx] = mpsGraphUnrankedPlaceHolder(mpsGraph, scalar_type);
+        if (tensor.scalar_type() != out_dtype) {
+          castInputTensors[idx] = [mpsGraph castTensor:newCachedGraph->inputTensors_[idx]
+                                                toType:getMPSDataType(out_dtype)
+                                                  name:@"castInput"];
+        } else {
+          castInputTensors[idx] = newCachedGraph->inputTensors_[idx];
+        }
+      }
+
+      auto inputTensorsArray = [NSArray arrayWithObjects:castInputTensors.data() count:len_tensor_array];
+      MPSGraphTensor* outputTensor = [mpsGraph concatTensors:inputTensorsArray
+                                                   dimension:dimension // Maybe convert this from int64_t -> int32
+                                                        name:nil];
+      if (getMPSDataType(out_dtype) == MPSDataTypeBool) {
+        outputTensor = [mpsGraph castTensor:outputTensor toType:MPSDataTypeBool name:@"outputTensor"];
+      }
+      newCachedGraph->outputTensor_ = outputTensor;
+    });
+
+    std::vector<Placeholder> inputPlaceholders;
+    int i = 0;
+    int t_idx = 0;
+    for (const Tensor& tensor : materialized_inputs) {
+      if (std::find(skipped_tensor_indices.begin(), skipped_tensor_indices.end(), i) == skipped_tensor_indices.end()) {
+        auto scalar_type = getMPSScalarType(tensor.scalar_type());
+        if (tensor.scalar_type() == kBool) {
+          scalar_type = MPSDataTypeInt8;
+        }
+        inputPlaceholders.emplace_back(cachedGraph->inputTensors_[t_idx], tensor, nullptr, true, scalar_type);
+        t_idx++;
+      }
+      i++;
+    }
+
+    auto outputDataType = getMPSScalarType(out.scalar_type());
+    Placeholder outputPlaceholder =
+        Placeholder(cachedGraph->outputTensor_, out, /*mpsShape=*/nil, /*gatherTensorData=*/false, outputDataType);
+
+    NSMutableDictionary* feeds = [[NSMutableDictionary new] autorelease];
+    for (auto& inputPlaceholder : inputPlaceholders) {
+      feeds[inputPlaceholder.getMPSGraphTensor()] = inputPlaceholder.getMPSGraphTensorData();
+    }
+    runMPSGraph(getCurrentMPSStream(), cachedGraph->graph(), feeds, outputPlaceholder);
   }
 }
 

aten/src/ATen/native/native_functions.yaml

@@ -6531,7 +6531,6 @@
   dispatch:
     CPU, CUDA: var
     MPS: var_mps
-    MTIA: var_mtia
   tags: core
 
 - func: var.out(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!)

aten/src/ATen/native/quantized/FakeQuantPerTensorAffine.cpp

@@ -184,23 +184,15 @@ std::tuple<Tensor, Tensor, Tensor> _fake_quantize_learnable_per_tensor_affine_ba
           0 & \text{ else }
         \end{cases}
   */
-
-  bool is_bfloat16 = (X.scalar_type() == at::kBFloat16);
-
-  at::Tensor X_ = is_bfloat16 ? X.to(ScalarType::Float) : X;
-  at::Tensor dY_ = is_bfloat16 ? dY.to(ScalarType::Float) : dY;
-  at::Tensor scale_ = is_bfloat16 ? scale.to(ScalarType::Float) : scale;
-  at::Tensor zero_point_ = is_bfloat16 ? zero_point.to(ScalarType::Float) : zero_point;
-
-  float scale_val = scale_[0].item<float>();
+  float scale_val = scale[0].item<float>();
   float inv_scale_val = 1.0f / scale_val;
-  int64_t zero_point_val = native::_get_zero_point_from_tensor(zero_point_, quant_min, quant_max, false);
+  int64_t zero_point_val = native::_get_zero_point_from_tensor(zero_point, quant_min, quant_max, false);
 
-  TORCH_CHECK(dY_.scalar_type() == ScalarType::Float);
-  TORCH_CHECK(X_.scalar_type() == ScalarType::Float);
-  TORCH_CHECK(scale_.scalar_type() == ScalarType::Float);
-  TORCH_CHECK(zero_point_.scalar_type() == ScalarType::Float);
-  TORCH_CHECK(X_.numel() == dY_.numel(), "`X` and `dY` are not the same size");
+  TORCH_CHECK(dY.scalar_type() == ScalarType::Float);
+  TORCH_CHECK(X.scalar_type() == ScalarType::Float);
+  TORCH_CHECK(scale.scalar_type() == ScalarType::Float);
+  TORCH_CHECK(zero_point.scalar_type() == ScalarType::Float);
+  TORCH_CHECK(X.numel() == dY.numel(), "`X` and `dY` are not the same size");
   TORCH_CHECK(
       quant_min <= 0 && quant_max >= 0,
       "`quant_min` should be less than or \
@@ -208,28 +200,28 @@ std::tuple<Tensor, Tensor, Tensor> _fake_quantize_learnable_per_tensor_affine_ba
   TORCH_CHECK(
       zero_point_val >= quant_min && zero_point_val <= quant_max,
       "`zero_point` must be between `quant_min` and `quant_max`.");
-  if (X_.numel() <= 0) {
+  if (X.numel() <= 0) {
     return std::make_tuple(X, scale, zero_point);
   }
 
-  auto dX = at::empty_like(X_, X_.options(), MemoryFormat::Preserve);
-  auto dScale_vec = at::empty_like(X_, X_.options(), MemoryFormat::Preserve);
-  auto dZeroPoint_vec = at::empty_like(X_, X_.options(), MemoryFormat::Preserve);
+  auto dX = at::empty_like(X, X.options(), MemoryFormat::Preserve);
+  auto dScale_vec = at::empty_like(X, X.options(), MemoryFormat::Preserve);
+  auto dZeroPoint_vec = at::empty_like(X, X.options(), MemoryFormat::Preserve);
 
   auto iter = TensorIteratorConfig()
     .add_output(dX)
     .add_output(dScale_vec)
     .add_output(dZeroPoint_vec)
-    .add_input(X_)
-    .add_input(dY_)
+    .add_input(X)
+    .add_input(dY)
     .build();
 
   fake_quant_grad_learnable_tensor_stub(
-    X_.device().type(), iter, scale_val, inv_scale_val, zero_point_val, quant_min, quant_max, grad_factor);
+    X.device().type(), iter, scale_val, inv_scale_val, zero_point_val, quant_min, quant_max, grad_factor);
 
   // The total sums over the scale and zero point gradient vectors are what will be returned in the end.
-  auto dScale = dScale_vec.sum().unsqueeze(0).to(scale_.device());
-  auto dZeroPoint = dZeroPoint_vec.sum().unsqueeze(0).to(zero_point_.device());
+  auto dScale = dScale_vec.sum().unsqueeze(0).to(scale.device());
+  auto dZeroPoint = dZeroPoint_vec.sum().unsqueeze(0).to(zero_point.device());
 
   return std::make_tuple(dX, dScale, dZeroPoint);
 }

aten/src/ATen/native/quantized/cpu/kernels/QuantizedOpKernels.cpp

@@ -3551,7 +3551,7 @@ void dequantize_tensor_per_tensor_affine_cpu(
 
 #if defined(__ARM_NEON__) || defined(__aarch64__)
 
-constexpr static int PARALLEL_THRESHOLD = 1 << 20;
+const static int PARALLEL_THRESHOLD = 1 << 20;
 
 // Generic template defaults to naive quantize implementation
 template <typename T>

aten/src/ATen/native/quantized/cpu/qlinear.cpp

@@ -1388,7 +1388,7 @@ namespace at::native {
     TORCH_CHECK(act_scale.numel() == 1 && act_zero_point.numel() <= 1,
         "onednn int8 linear: act scale/zp size should be 1/<=1");
     static std::optional<at::Tensor> other = std::nullopt;
-    constexpr std::string_view binary_post_op = "none";
+    static const std::string_view binary_post_op = "none";
     int64_t act_zp = act_zero_point.numel() == 1 ? act_zero_point.item().toLong() : 0;
     return linear_int8_with_onednn_weight(
         act, act_scale.item().toDouble(), act_zp,

aten/src/ATen/native/quantized/cpu/qsoftmax.cpp

@@ -16,8 +16,8 @@ namespace {
 
 #ifdef USE_PYTORCH_QNNPACK
 
-constexpr static float qnnpack_softmax_output_scale = 0x1.0p-8f;
-constexpr static int qnnpack_softmax_output_zero_point = 0;
+const static float qnnpack_softmax_output_scale = 0x1.0p-8f;
+const static int qnnpack_softmax_output_zero_point = 0;
 
 bool is_qnnpack_compatible(
     const Tensor& qx,

aten/src/ATen/native/transformers/cuda/mem_eff_attention/epilogue/epilogue_thread_apply_logsumexp.h

@@ -110,9 +110,9 @@ class ApplyLogSumExp {
   using ElementCompute = ElementCompute_;
   using ElementLSE = ElementLSE_;
 
-  static int constexpr kElementsPerAccess = ElementsPerAccess;
-  static int constexpr kCount = kElementsPerAccess;
-  static constexpr ScaleType::Kind kScale =
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kCount = kElementsPerAccess;
+  static const ScaleType::Kind kScale =
       cutlass::epilogue::thread::ScaleType::NoBetaScaling;
 
   using FragmentOutput = Array<ElementOutput, kCount>;

aten/src/ATen/test/pow_test.cpp

@@ -14,16 +14,16 @@ using namespace at;
 
 namespace {
 
-constexpr auto int_min = std::numeric_limits<int>::min();
-constexpr auto int_max = std::numeric_limits<int>::max();
-constexpr auto long_min = std::numeric_limits<int64_t>::min();
-constexpr auto long_max = std::numeric_limits<int64_t>::max();
-constexpr auto float_lowest = std::numeric_limits<float>::lowest();
-constexpr auto float_min = std::numeric_limits<float>::min();
-constexpr auto float_max = std::numeric_limits<float>::max();
-constexpr auto double_lowest = std::numeric_limits<double>::lowest();
-constexpr auto double_min = std::numeric_limits<double>::min();
-constexpr auto double_max = std::numeric_limits<double>::max();
+const auto int_min = std::numeric_limits<int>::min();
+const auto int_max = std::numeric_limits<int>::max();
+const auto long_min = std::numeric_limits<int64_t>::min();
+const auto long_max = std::numeric_limits<int64_t>::max();
+const auto float_lowest = std::numeric_limits<float>::lowest();
+const auto float_min = std::numeric_limits<float>::min();
+const auto float_max = std::numeric_limits<float>::max();
+const auto double_lowest = std::numeric_limits<double>::lowest();
+const auto double_min = std::numeric_limits<double>::min();
+const auto double_max = std::numeric_limits<double>::max();
 
 const std::vector<int> ints {
   int_min,

aten/src/ATen/xpu/XPUGeneratorImpl.cpp

@@ -146,9 +146,9 @@ uint64_t XPUGeneratorImpl::seed() {
 
 c10::intrusive_ptr<c10::TensorImpl> XPUGeneratorImpl::get_state() const {
   // The RNG state comprises the seed, and an offset used for Philox.
-  constexpr size_t seed_size = sizeof(uint64_t);
-  constexpr size_t offset_size = sizeof(uint64_t);
-  constexpr size_t total_size = seed_size + offset_size;
+  static const size_t seed_size = sizeof(uint64_t);
+  static const size_t offset_size = sizeof(uint64_t);
+  static const size_t total_size = seed_size + offset_size;
 
   // The internal state is returned as a CPU byte tensor.
   auto state_tensor = at::detail::empty_cpu(
@@ -170,9 +170,9 @@ c10::intrusive_ptr<c10::TensorImpl> XPUGeneratorImpl::get_state() const {
 void XPUGeneratorImpl::set_state(const c10::TensorImpl& new_state) {
   at::xpu::assertNotCapturing(
       "Please ensure to utilize the XPUGeneratorImpl::set_state_index method during capturing.");
-  constexpr size_t seed_size = sizeof(uint64_t);
-  constexpr size_t offset_size = sizeof(uint64_t);
-  constexpr size_t total_size = seed_size + offset_size;
+  static const size_t seed_size = sizeof(uint64_t);
+  static const size_t offset_size = sizeof(uint64_t);
+  static const size_t total_size = seed_size + offset_size;
 
   at::detail::check_rng_state(new_state);
 

benchmarks/dynamo/cachebench.py

@@ -6,7 +6,7 @@ import os
 import subprocess
 import sys
 import tempfile
-from collections.abc import Callable
+from typing import Callable
 
 from torch._inductor.utils import fresh_cache
 

benchmarks/dynamo/common.py

@@ -2284,11 +2284,9 @@ class BenchmarkRunner:
                     )
                 ):
                     is_same = False
-            except Exception as e:
+            except Exception:
                 # Sometimes torch.allclose may throw RuntimeError
-                exception_string = str(e)
-                accuracy_status = f"fail_exception: {exception_string}"
-                return record_status(accuracy_status, dynamo_start_stats=start_stats)
+                is_same = False
 
             if not is_same:
                 accuracy_status = "eager_two_runs_differ"
@@ -2405,11 +2403,9 @@ class BenchmarkRunner:
                     force_max_multiplier=force_max_multiplier,
                 ):
                     is_same = False
-            except Exception as e:
+            except Exception:
                 # Sometimes torch.allclose may throw RuntimeError
-                exception_string = str(e)
-                accuracy_status = f"fail_exception: {exception_string}"
-                return record_status(accuracy_status, dynamo_start_stats=start_stats)
+                is_same = False
 
             if not is_same:
                 if self.args.skip_accuracy_check:
@@ -4064,7 +4060,7 @@ def run(runner, args, original_dir=None):
         else:
             optimize_ctx = torch._dynamo.optimize(args.backend, nopython=args.nopython)
         experiment = (
-            speedup_experiment if args.backend != "torchao" else latency_experiment
+            speedup_experiment if not args.backend == "torchao" else latency_experiment
         )
         if args.accuracy:
             output_filename = f"accuracy_{args.backend}.csv"

benchmarks/dynamo/genai_layers/utils.py

@@ -1,8 +1,7 @@
 import os
 from collections import defaultdict
-from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import matplotlib.pyplot as plt
 

benchmarks/dynamo/torchao_backend.py

@@ -1,5 +1,4 @@
-from collections.abc import Callable
-from typing import Any
+from typing import Any, Callable
 
 import torch
 

benchmarks/functional_autograd_benchmark/functional_autograd_benchmark.py

@@ -1,8 +1,7 @@
 import time
 from argparse import ArgumentParser
 from collections import defaultdict
-from collections.abc import Callable
-from typing import Any, NamedTuple
+from typing import Any, Callable, NamedTuple
 
 import torch
 from torch.autograd import functional

benchmarks/functional_autograd_benchmark/utils.py

@@ -1,6 +1,5 @@
 from collections import defaultdict
-from collections.abc import Callable
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 import torch
 from torch import nn, Tensor

benchmarks/gpt_fast/common.py

@@ -1,6 +1,5 @@
 import dataclasses
-from collections.abc import Callable
-from typing import Optional
+from typing import Callable, Optional
 
 
 all_experiments: dict[str, Callable] = {}

benchmarks/inductor_backends/cutlass.py

@@ -9,9 +9,8 @@ import logging
 import time
 from abc import abstractmethod
 from collections import defaultdict
-from collections.abc import Callable
 from dataclasses import asdict, dataclass, field
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 from tabulate import tabulate
 from tqdm import tqdm

benchmarks/operator_benchmark/benchmark_all_other_test.py

@@ -7,7 +7,6 @@ from pt import (  # noqa: F401
     binary_inplace_test,
     binary_test,
     bmm_test,
-    boolean_test,
     cat_test,
     channel_shuffle_test,
     chunk_test,

benchmarks/operator_benchmark/pt/binary_test.py

@@ -56,9 +56,6 @@ binary_ops_list = op_bench.op_list(
         ["sub", torch.sub],
         ["div", torch.div],
         ["mul", torch.mul],
-        ["asr", torch.bitwise_right_shift],
-        ["lsl", torch.bitwise_left_shift],
-        ["xor", torch.bitwise_xor],
     ],
 )
 

benchmarks/operator_benchmark/pt/boolean_test.py

@@ -1,73 +0,0 @@
-import operator_benchmark as op_bench
-
-import torch
-
-
-"""Microbenchmarks for boolean operators. Supports both Caffe2/PyTorch."""
-
-# Configs for PT all operator
-all_long_configs = op_bench.cross_product_configs(
-    M=[8, 128], N=[32, 64], K=[256, 512], device=["cpu", "cuda"], tags=["long"]
-)
-
-
-all_short_configs = op_bench.config_list(
-    attr_names=["M", "N", "K"],
-    attrs=[
-        [1, 1, 1],
-        [64, 64, 64],
-        [64, 64, 128],
-    ],
-    cross_product_configs={
-        "device": ["cpu", "cuda"],
-    },
-    tags=["short"],
-)
-
-
-class AllBenchmark(op_bench.TorchBenchmarkBase):
-    def init(self, M, N, K, device):
-        self.inputs = {
-            "input_one": torch.randint(0, 2, (M, N, K), device=device, dtype=torch.bool)
-        }
-        self.set_module_name("all")
-
-    def forward(self, input_one):
-        return torch.all(input_one)
-
-
-# The generated test names based on all_short_configs will be in the following pattern:
-# all_M8_N16_K32_devicecpu
-# all_M8_N16_K32_devicecpu_bwdall
-# all_M8_N16_K32_devicecpu_bwd1
-# all_M8_N16_K32_devicecpu_bwd2
-# ...
-# Those names can be used to filter tests.
-
-op_bench.generate_pt_test(all_long_configs + all_short_configs, AllBenchmark)
-
-"""Mircobenchmark for any operator."""
-
-
-class AnyBenchmark(op_bench.TorchBenchmarkBase):
-    def init(self, M, N, device):
-        self.inputs = {
-            "input_one": torch.randint(0, 2, (M, N), device=device, dtype=torch.bool)
-        }
-        self.set_module_name("any")
-
-    def forward(self, input_one):
-        return torch.any(input_one)
-
-
-any_configs = op_bench.cross_product_configs(
-    M=[8, 256],
-    N=[256, 16],
-    device=["cpu", "cuda"],
-    tags=["any"],
-)
-
-op_bench.generate_pt_test(any_configs, AnyBenchmark)
-
-if __name__ == "__main__":
-    op_bench.benchmark_runner.main()

benchmarks/operator_benchmark/pt/conv_test.py

@@ -38,16 +38,12 @@ class ConvTranspose1dBenchmark(op_bench.TorchBenchmarkBase):
 op_bench.generate_pt_test(
     configs.conv_1d_configs_short + configs.conv_1d_configs_long, Conv1dBenchmark
 )
-
-
-if not torch.backends.mkldnn.is_acl_available():
-    # convtranpose1d crashes with ACL, see https://github.com/pytorch/pytorch/issues/165654
-    op_bench.generate_pt_test(
-        configs.convtranspose_1d_configs_short
-        + configs.conv_1d_configs_short
-        + configs.conv_1d_configs_long,
-        ConvTranspose1dBenchmark,
-    )
+op_bench.generate_pt_test(
+    configs.convtranspose_1d_configs_short
+    + configs.conv_1d_configs_short
+    + configs.conv_1d_configs_long,
+    ConvTranspose1dBenchmark,
+)
 
 
 """

benchmarks/transformer/attention_bias_benchmarks.py

@@ -1,8 +1,7 @@
 import itertools
-from collections.abc import Callable
 from dataclasses import asdict, dataclass
 from functools import partial
-from typing import Union
+from typing import Callable, Union
 
 import numpy as np
 from tabulate import tabulate

benchmarks/transformer/score_mod.py

@@ -3,11 +3,10 @@ import csv
 import itertools
 import random
 from collections import defaultdict
-from collections.abc import Callable
 from contextlib import nullcontext
 from dataclasses import asdict, dataclass
 from functools import partial
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 import numpy as np
 from tabulate import tabulate
@@ -271,7 +270,7 @@ def run_single_backend_sdpa(
 
         if config.calculate_bwd_time:
             # TODO: debug backward pass for njt
-            if eager_sdpa and config.attn_type != "document_mask":
+            if eager_sdpa and not config.attn_type == "document_mask":
                 d_out = torch.randn_like(out_eager.transpose(1, 2)).transpose(1, 2)
                 backward_eager_time = benchmark_torch_function_in_microseconds(
                     out_eager.backward, d_out, retain_graph=True

benchmarks/transformer/sdpa.py

@@ -1,8 +1,8 @@
 import itertools
 from collections import defaultdict
-from collections.abc import Callable
 from contextlib import nullcontext
 from dataclasses import asdict, dataclass
+from typing import Callable
 
 from tabulate import tabulate
 from tqdm import tqdm

c10/core/AllocatorConfig.cpp

@@ -13,22 +13,20 @@ constexpr size_t kRoundUpPowerOfTwoEnd = 64 * 1024ul * kMB; // 64GB
 
 AcceleratorAllocatorConfig& AcceleratorAllocatorConfig::instance() {
   static AcceleratorAllocatorConfig instance;
-#define C10_ALLOCATOR_CONFIG_PARSE_ENV(env)    \
-  auto env##_name = c10::utils::get_env(#env); \
-  if (env##_name.has_value()) {                \
-    instance.parseArgs(env##_name.value());    \
-    return true;                               \
+#define C10_ALLOCATOR_CONFIG_PARSE_ENV(env, deprecated)                       \
+  auto env##_name = c10::utils::get_env(#env);                                \
+  if (env##_name.has_value()) {                                               \
+    if (deprecated) {                                                         \
+      TORCH_WARN_ONCE(#env " is deprecated, use PYTORCH_ALLOC_CONF instead"); \
+    }                                                                         \
+    instance.parseArgs(env##_name.value());                                   \
+    return true;                                                              \
   }
   static bool env_flag [[maybe_unused]] = []() {
-    // Parse allocator configuration from environment variables.
-    // The first two entries are kept for backward compatibility with legacy
-    // CUDA and HIP environment variable names. The new unified variable
-    // (PYTORCH_ALLOC_CONF) should be used going forward.
-    // Note: keep the parsing order and logic stable to avoid potential
-    // performance regressions in internal tests.
-    C10_ALLOCATOR_CONFIG_PARSE_ENV(PYTORCH_CUDA_ALLOC_CONF)
-    C10_ALLOCATOR_CONFIG_PARSE_ENV(PYTORCH_HIP_ALLOC_CONF)
-    C10_ALLOCATOR_CONFIG_PARSE_ENV(PYTORCH_ALLOC_CONF)
+    C10_ALLOCATOR_CONFIG_PARSE_ENV(PYTORCH_ALLOC_CONF, false)
+    // Keep this for backwards compatibility
+    C10_ALLOCATOR_CONFIG_PARSE_ENV(PYTORCH_CUDA_ALLOC_CONF, /*deprecated=*/true)
+    C10_ALLOCATOR_CONFIG_PARSE_ENV(PYTORCH_HIP_ALLOC_CONF, /*deprecated=*/true)
     return false;
   }();
 #undef C10_ALLOCATOR_CONFIG_PARSE_ENV
@@ -129,7 +127,8 @@ size_t AcceleratorAllocatorConfig::parseRoundUpPower2Divisions(
         std::fill(
             std::next(
                 roundup_power2_divisions_.begin(),
-                static_cast<std::vector<size_t>::difference_type>(last_index)),
+                static_cast<std::vector<size_t>::difference_type>(
+                    last_index + 1)),
             roundup_power2_divisions_.end(),
             value);
       } else {

c10/cuda/CUDAAllocatorConfig.cpp

@@ -1,5 +1,6 @@
 #include <c10/cuda/CUDAAllocatorConfig.h>
 #include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/util/llvmMathExtras.h>
 
 #if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED)
 #include <c10/cuda/driver_api.h>
@@ -7,119 +8,386 @@
 
 namespace c10::cuda::CUDACachingAllocator {
 
+constexpr size_t kRoundUpPowerOfTwoIntervals = 16;
+
+CUDAAllocatorConfig::CUDAAllocatorConfig()
+    : m_max_split_size(std::numeric_limits<size_t>::max()),
+      m_max_non_split_rounding_size(kLargeBuffer),
+      m_garbage_collection_threshold(0),
+      m_pinned_num_register_threads(1),
+      m_pinned_reserve_segment_size_mb(0),
+      m_expandable_segments(false),
+#if CUDA_VERSION >= 12030
+      m_expandable_segments_handle_type(
+          Expandable_Segments_Handle_Type::UNSPECIFIED),
+#else
+      m_expandable_segments_handle_type(
+          Expandable_Segments_Handle_Type::POSIX_FD),
+#endif
+      m_release_lock_on_cudamalloc(false),
+      m_pinned_use_cuda_host_register(false),
+      m_graph_capture_record_stream_reuse(false),
+      m_pinned_use_background_threads(false) {
+  m_roundup_power2_divisions.assign(kRoundUpPowerOfTwoIntervals, 0);
+}
+
+size_t CUDAAllocatorConfig::roundup_power2_divisions(size_t size) {
+  size_t log_size = (63 - llvm::countLeadingZeros(size));
+
+  // Our intervals start at 1MB and end at 64GB
+  const size_t interval_start =
+      63 - llvm::countLeadingZeros(static_cast<size_t>(1048576));
+  const size_t interval_end =
+      63 - llvm::countLeadingZeros(static_cast<size_t>(68719476736));
+  TORCH_CHECK(
+      (interval_end - interval_start == kRoundUpPowerOfTwoIntervals),
+      "kRoundUpPowerOfTwoIntervals mismatch");
+
+  int index = static_cast<int>(log_size) - static_cast<int>(interval_start);
+
+  index = std::max(0, index);
+  index = std::min(index, static_cast<int>(kRoundUpPowerOfTwoIntervals) - 1);
+  return instance().m_roundup_power2_divisions[index];
+}
+
+void CUDAAllocatorConfig::lexArgs(
+    const std::string& env,
+    std::vector<std::string>& config) {
+  std::vector<char> buf;
+
+  for (char ch : env) {
+    if (ch == ',' || ch == ':' || ch == '[' || ch == ']') {
+      if (!buf.empty()) {
+        config.emplace_back(buf.begin(), buf.end());
+        buf.clear();
+      }
+      config.emplace_back(1, ch);
+    } else if (ch != ' ') {
+      buf.emplace_back(ch);
+    }
+  }
+  if (!buf.empty()) {
+    config.emplace_back(buf.begin(), buf.end());
+  }
+}
+
+void CUDAAllocatorConfig::consumeToken(
+    const std::vector<std::string>& config,
+    size_t i,
+    const char c) {
+  TORCH_CHECK(
+      i < config.size() && config[i] == std::string(1, c),
+      "Error parsing CachingAllocator settings, expected ",
+      c,
+      "");
+}
+
+size_t CUDAAllocatorConfig::parseMaxSplitSize(
+    const std::vector<std::string>& config,
+    size_t i) {
+  consumeToken(config, ++i, ':');
+  constexpr int mb = 1024 * 1024;
+  if (++i < config.size()) {
+    size_t val1 = stoi(config[i]);
+    TORCH_CHECK(
+        val1 > kLargeBuffer / mb,
+        "CachingAllocator option max_split_size_mb too small, must be > ",
+        kLargeBuffer / mb,
+        "");
+    val1 = std::max(val1, kLargeBuffer / mb);
+    val1 = std::min(val1, (std::numeric_limits<size_t>::max() / mb));
+    m_max_split_size = val1 * 1024 * 1024;
+  } else {
+    TORCH_CHECK(false, "Error, expecting max_split_size_mb value", "");
+  }
+  return i;
+}
+
+size_t CUDAAllocatorConfig::parseMaxNonSplitRoundingSize(
+    const std::vector<std::string>& config,
+    size_t i) {
+  consumeToken(config, ++i, ':');
+  constexpr int mb = 1024 * 1024;
+  if (++i < config.size()) {
+    size_t val1 = stoi(config[i]);
+    TORCH_CHECK(
+        val1 > kLargeBuffer / mb,
+        "CachingAllocator option max_non_split_rounding_mb too small, must be > ",
+        kLargeBuffer / mb,
+        "");
+    val1 = std::max(val1, kLargeBuffer / mb);
+    val1 = std::min(val1, (std::numeric_limits<size_t>::max() / mb));
+    m_max_non_split_rounding_size = val1 * 1024 * 1024;
+  } else {
+    TORCH_CHECK(false, "Error, expecting max_non_split_rounding_mb value", "");
+  }
+  return i;
+}
+
+size_t CUDAAllocatorConfig::parseGarbageCollectionThreshold(
+    const std::vector<std::string>& config,
+    size_t i) {
+  consumeToken(config, ++i, ':');
+  if (++i < config.size()) {
+    double val1 = stod(config[i]);
+    TORCH_CHECK(
+        val1 > 0, "garbage_collect_threshold too small, set it 0.0~1.0", "");
+    TORCH_CHECK(
+        val1 < 1.0, "garbage_collect_threshold too big, set it 0.0~1.0", "");
+    m_garbage_collection_threshold = val1;
+  } else {
+    TORCH_CHECK(
+        false, "Error, expecting garbage_collection_threshold value", "");
+  }
+  return i;
+}
+
+size_t CUDAAllocatorConfig::parseRoundUpPower2Divisions(
+    const std::vector<std::string>& config,
+    size_t i) {
+  consumeToken(config, ++i, ':');
+  bool first_value = true;
+
+  if (++i < config.size()) {
+    if (std::string_view(config[i]) == "[") {
+      size_t last_index = 0;
+      // NOLINTNEXTLINE(bugprone-inc-dec-in-conditions)
+      while (++i < config.size() && std::string_view(config[i]) != "]") {
+        const std::string& val1 = config[i];
+        size_t val2 = 0;
+
+        consumeToken(config, ++i, ':');
+        if (++i < config.size()) {
+          val2 = stoi(config[i]);
+        } else {
+          TORCH_CHECK(
+              false, "Error parsing roundup_power2_divisions value", "");
+        }
+        TORCH_CHECK(
+            val2 == 0 || llvm::isPowerOf2_64(val2),
+            "For roundups, the divisions has to be power of 2 or 0 to disable roundup ",
+            "");
+
+        if (std::string_view(val1) == ">") {
+          std::fill(
+              std::next(
+                  m_roundup_power2_divisions.begin(),
+                  static_cast<std::vector<unsigned long>::difference_type>(
+                      last_index)),
+              m_roundup_power2_divisions.end(),
+              val2);
+        } else {
+          size_t val1_long = stoul(val1);
+          TORCH_CHECK(
+              llvm::isPowerOf2_64(val1_long),
+              "For roundups, the intervals have to be power of 2 ",
+              "");
+
+          size_t index = 63 - llvm::countLeadingZeros(val1_long);
+          index = std::max((size_t)0, index);
+          index = std::min(index, m_roundup_power2_divisions.size() - 1);
+
+          if (first_value) {
+            std::fill(
+                m_roundup_power2_divisions.begin(),
+                std::next(
+                    m_roundup_power2_divisions.begin(),
+                    static_cast<std::vector<unsigned long>::difference_type>(
+                        index)),
+                val2);
+            first_value = false;
+          }
+          if (index < m_roundup_power2_divisions.size()) {
+            m_roundup_power2_divisions[index] = val2;
+          }
+          last_index = index;
+        }
+
+        if (std::string_view(config[i + 1]) != "]") {
+          consumeToken(config, ++i, ',');
+        }
+      }
+    } else { // Keep this for backwards compatibility
+      size_t val1 = stoi(config[i]);
+      TORCH_CHECK(
+          llvm::isPowerOf2_64(val1),
+          "For roundups, the divisions has to be power of 2 ",
+          "");
+      std::fill(
+          m_roundup_power2_divisions.begin(),
+          m_roundup_power2_divisions.end(),
+          val1);
+    }
+  } else {
+    TORCH_CHECK(false, "Error, expecting roundup_power2_divisions value", "");
+  }
+  return i;
+}
+
 size_t CUDAAllocatorConfig::parseAllocatorConfig(
-    const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+    const std::vector<std::string>& config,
     size_t i,
     bool& used_cudaMallocAsync) {
+  // For ease of maintenance and understanding, the CUDA and ROCm
+  // implementations of this function are separated. This avoids having many
+  // #ifdef's throughout.
+#ifdef USE_ROCM
   // Ease burden on ROCm users by allowing either cuda or hip tokens.
   // cuda token is broken up to prevent hipify matching it.
 #define PYTORCH_TOKEN1 \
   "cud"                \
   "aMallocAsync"
 #define PYTORCH_TOKEN2 "hipMallocAsync"
-  tokenizer.checkToken(++i, ":");
-  i++; // Move to the value after the colon
-#ifdef USE_ROCM
-  TORCH_CHECK(
-      ((tokenizer[i] == "native") || (tokenizer[i] == PYTORCH_TOKEN1) ||
-       (tokenizer[i] == PYTORCH_TOKEN2)),
-      "Unknown allocator backend, "
-      "options are native, " PYTORCH_TOKEN1 ", and " PYTORCH_TOKEN2);
-  used_cudaMallocAsync =
-      (tokenizer[i] == PYTORCH_TOKEN1 || tokenizer[i] == PYTORCH_TOKEN2);
-  TORCH_INTERNAL_ASSERT(
-      tokenizer[i] == get()->name() ||
-          (tokenizer[i] == PYTORCH_TOKEN1 && get()->name() == PYTORCH_TOKEN2),
-      "Allocator backend parsed at runtime != "
-      "allocator backend parsed at load time, ",
-      tokenizer[i],
-      " != ",
-      get()->name());
-#else // USE_ROCM
-  TORCH_CHECK(
-      ((tokenizer[i] == "native") || (tokenizer[i] == PYTORCH_TOKEN1)),
-      "Unknown allocator backend, "
-      "options are native and " PYTORCH_TOKEN1);
-  used_cudaMallocAsync = (tokenizer[i] == PYTORCH_TOKEN1);
-  TORCH_INTERNAL_ASSERT(
-      tokenizer[i] == get()->name(),
-      "Allocator backend parsed at runtime != "
-      "allocator backend parsed at load time, ",
-      tokenizer[i],
-      " != ",
-      get()->name());
-  if (used_cudaMallocAsync) {
-#if CUDA_VERSION >= 11040
-    int version = 0;
-    C10_CUDA_CHECK(cudaDriverGetVersion(&version));
+  consumeToken(config, ++i, ':');
+  if (++i < config.size()) {
     TORCH_CHECK(
-        version >= 11040,
-        "backend:cudaMallocAsync requires CUDA runtime "
-        "11.4 or newer, but cudaDriverGetVersion returned ",
-        version);
-#else // CUDA_VERSION >= 11040
-    TORCH_CHECK(
-        false,
-        "backend:cudaMallocAsync requires PyTorch to be built with "
-        "CUDA 11.4 or newer, but CUDA_VERSION is ",
-        CUDA_VERSION);
-#endif // CUDA_VERSION >= 11040
+        ((config[i] == "native") || (config[i] == PYTORCH_TOKEN1) ||
+         (config[i] == PYTORCH_TOKEN2)),
+        "Unknown allocator backend, "
+        "options are native, " PYTORCH_TOKEN1 ", and " PYTORCH_TOKEN2);
+    used_cudaMallocAsync =
+        (config[i] == PYTORCH_TOKEN1 || config[i] == PYTORCH_TOKEN2);
+    TORCH_INTERNAL_ASSERT(
+        config[i] == get()->name() ||
+            (config[i] == PYTORCH_TOKEN1 && get()->name() == PYTORCH_TOKEN2),
+        "Allocator backend parsed at runtime != "
+        "allocator backend parsed at load time, ",
+        config[i],
+        " != ",
+        get()->name());
+  } else {
+    TORCH_CHECK(false, "Error parsing backend value", "");
   }
-#endif // USE_ROCM
   return i;
+#undef PYTORCH_TOKEN1
+#undef PYTORCH_TOKEN2
+#else // USE_ROCM
+  consumeToken(config, ++i, ':');
+  if (++i < config.size()) {
+    TORCH_CHECK(
+        ((config[i] == "native") || (config[i] == "cudaMallocAsync")),
+        "Unknown allocator backend, "
+        "options are native and cudaMallocAsync");
+    used_cudaMallocAsync = (config[i] == "cudaMallocAsync");
+    if (used_cudaMallocAsync) {
+#if CUDA_VERSION >= 11040
+      int version = 0;
+      C10_CUDA_CHECK(cudaDriverGetVersion(&version));
+      TORCH_CHECK(
+          version >= 11040,
+          "backend:cudaMallocAsync requires CUDA runtime "
+          "11.4 or newer, but cudaDriverGetVersion returned ",
+          version);
+#else
+      TORCH_CHECK(
+          false,
+          "backend:cudaMallocAsync requires PyTorch to be built with "
+          "CUDA 11.4 or newer, but CUDA_VERSION is ",
+          CUDA_VERSION);
+#endif
+    }
+    TORCH_INTERNAL_ASSERT(
+        config[i] == get()->name(),
+        "Allocator backend parsed at runtime != "
+        "allocator backend parsed at load time");
+  } else {
+    TORCH_CHECK(false, "Error parsing backend value", "");
+  }
+  return i;
+#endif // USE_ROCM
 }
 
-void CUDAAllocatorConfig::parseArgs(const std::string& env) {
+void CUDAAllocatorConfig::parseArgs(const std::optional<std::string>& env) {
+  // If empty, set the default values
+  m_max_split_size = std::numeric_limits<size_t>::max();
+  m_roundup_power2_divisions.assign(kRoundUpPowerOfTwoIntervals, 0);
+  m_garbage_collection_threshold = 0;
   bool used_cudaMallocAsync = false;
   bool used_native_specific_option = false;
 
-  c10::CachingAllocator::ConfigTokenizer tokenizer(env);
-  for (size_t i = 0; i < tokenizer.size(); i++) {
-    const auto& key = tokenizer[i];
-    if (key == "backend") {
-      i = parseAllocatorConfig(tokenizer, i, used_cudaMallocAsync);
+  if (!env.has_value()) {
+    return;
+  }
+  {
+    std::lock_guard<std::mutex> lock(m_last_allocator_settings_mutex);
+    m_last_allocator_settings = env.value();
+  }
+
+  std::vector<std::string> config;
+  lexArgs(env.value(), config);
+
+  for (size_t i = 0; i < config.size(); i++) {
+    std::string_view config_item_view(config[i]);
+    if (config_item_view == "max_split_size_mb") {
+      i = parseMaxSplitSize(config, i);
+      used_native_specific_option = true;
+    } else if (config_item_view == "max_non_split_rounding_mb") {
+      i = parseMaxNonSplitRoundingSize(config, i);
+      used_native_specific_option = true;
+    } else if (config_item_view == "garbage_collection_threshold") {
+      i = parseGarbageCollectionThreshold(config, i);
+      used_native_specific_option = true;
+    } else if (config_item_view == "roundup_power2_divisions") {
+      i = parseRoundUpPower2Divisions(config, i);
+      used_native_specific_option = true;
+    } else if (config_item_view == "backend") {
+      i = parseAllocatorConfig(config, i, used_cudaMallocAsync);
+    } else if (config_item_view == "expandable_segments") {
+      used_native_specific_option = true;
+      consumeToken(config, ++i, ':');
+      ++i;
+      TORCH_CHECK(
+          i < config.size() &&
+              (std::string_view(config[i]) == "True" ||
+               std::string_view(config[i]) == "False"),
+          "Expected a single True/False argument for expandable_segments");
+      config_item_view = config[i];
+      m_expandable_segments = (config_item_view == "True");
     } else if (
         // ROCm build's hipify step will change "cuda" to "hip", but for ease of
         // use, accept both. We must break up the string to prevent hipify here.
-        key == "release_lock_on_hipmalloc" ||
-        key ==
+        config_item_view == "release_lock_on_hipmalloc" ||
+        config_item_view ==
             "release_lock_on_c"
             "udamalloc") {
       used_native_specific_option = true;
-      tokenizer.checkToken(++i, ":");
-      m_release_lock_on_cudamalloc = tokenizer.toBool(++i);
+      consumeToken(config, ++i, ':');
+      ++i;
+      TORCH_CHECK(
+          i < config.size() &&
+              (std::string_view(config[i]) == "True" ||
+               std::string_view(config[i]) == "False"),
+          "Expected a single True/False argument for release_lock_on_cudamalloc");
+      config_item_view = config[i];
+      m_release_lock_on_cudamalloc = (config_item_view == "True");
     } else if (
         // ROCm build's hipify step will change "cuda" to "hip", but for ease of
         // use, accept both. We must break up the string to prevent hipify here.
-        key == "pinned_use_hip_host_register" ||
-        key ==
+        config_item_view == "pinned_use_hip_host_register" ||
+        config_item_view ==
             "pinned_use_c"
             "uda_host_register") {
-      i = parsePinnedUseCudaHostRegister(tokenizer, i);
+      i = parsePinnedUseCudaHostRegister(config, i);
       used_native_specific_option = true;
-    } else if (key == "pinned_num_register_threads") {
-      i = parsePinnedNumRegisterThreads(tokenizer, i);
+    } else if (config_item_view == "pinned_num_register_threads") {
+      i = parsePinnedNumRegisterThreads(config, i);
       used_native_specific_option = true;
-    } else if (key == "pinned_reserve_segment_size_mb") {
-      i = parsePinnedReserveSegmentSize(tokenizer, i);
+    } else if (config_item_view == "pinned_reserve_segment_size_mb") {
+      i = parsePinnedReserveSegmentSize(config, i);
       used_native_specific_option = true;
-    } else if (key == "graph_capture_record_stream_reuse") {
-      i = parseGraphCaptureRecordStreamReuse(tokenizer, i);
+    } else if (config_item_view == "pinned_use_background_threads") {
+      i = parsePinnedUseBackgroundThreads(config, i);
+      used_native_specific_option = true;
+    } else if (config_item_view == "graph_capture_record_stream_reuse") {
+      i = parseGraphCaptureRecordStreamReuse(config, i);
       used_native_specific_option = true;
     } else {
-      const auto& keys =
-          c10::CachingAllocator::AcceleratorAllocatorConfig::getKeys();
       TORCH_CHECK(
-          keys.find(key) != keys.end(),
-          "Unrecognized key '",
-          key,
-          "' in CUDA allocator config.");
-      // Skip the key and its value
-      i = tokenizer.skipKey(i);
+          false, "Unrecognized CachingAllocator option: ", config_item_view);
     }
 
-    if (i + 1 < tokenizer.size()) {
-      tokenizer.checkToken(++i, ",");
+    if (i + 1 < config.size()) {
+      consumeToken(config, ++i, ',');
     }
   }
 
@@ -131,51 +399,97 @@ void CUDAAllocatorConfig::parseArgs(const std::string& env) {
 }
 
 size_t CUDAAllocatorConfig::parsePinnedUseCudaHostRegister(
-    const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+    const std::vector<std::string>& config,
     size_t i) {
-  tokenizer.checkToken(++i, ":");
-  m_pinned_use_cuda_host_register = tokenizer.toBool(++i);
+  consumeToken(config, ++i, ':');
+  if (++i < config.size()) {
+    TORCH_CHECK(
+        (config[i] == "True" || config[i] == "False"),
+        "Expected a single True/False argument for pinned_use_cuda_host_register");
+    m_pinned_use_cuda_host_register = (config[i] == "True");
+  } else {
+    TORCH_CHECK(
+        false, "Error, expecting pinned_use_cuda_host_register value", "");
+  }
   return i;
 }
 
 size_t CUDAAllocatorConfig::parseGraphCaptureRecordStreamReuse(
-    const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+    const std::vector<std::string>& config,
     size_t i) {
-  tokenizer.checkToken(++i, ":");
-  m_graph_capture_record_stream_reuse = tokenizer.toBool(++i);
+  consumeToken(config, ++i, ':');
+  if (++i < config.size()) {
+    TORCH_CHECK(
+        (config[i] == "True" || config[i] == "False"),
+        "Expected a single True/False argument for graph_capture_record_stream_reuse");
+    m_graph_capture_record_stream_reuse = (config[i] == "True");
+  } else {
+    TORCH_CHECK(
+        false, "Error, expecting graph_capture_record_stream_reuse value", "");
+  }
+
   return i;
 }
 
 size_t CUDAAllocatorConfig::parsePinnedNumRegisterThreads(
-    const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+    const std::vector<std::string>& config,
     size_t i) {
-  tokenizer.checkToken(++i, ":");
-  size_t val2 = tokenizer.toSizeT(++i);
-  TORCH_CHECK(
-      llvm::isPowerOf2_64(val2),
-      "Number of register threads has to be power of 2, got ",
-      val2);
-  auto maxThreads = CUDAAllocatorConfig::pinned_max_register_threads();
-  TORCH_CHECK(
-      val2 <= maxThreads,
-      "Number of register threads should be less than or equal to ",
-      maxThreads,
-      ", got ",
-      val2);
-  m_pinned_num_register_threads = val2;
+  consumeToken(config, ++i, ':');
+  if (++i < config.size()) {
+    size_t val2 = stoi(config[i]);
+    TORCH_CHECK(
+        llvm::isPowerOf2_64(val2),
+        "Number of register threads has to be power of 2 ",
+        "");
+    auto maxThreads = CUDAAllocatorConfig::pinned_max_register_threads();
+    TORCH_CHECK(
+        val2 <= maxThreads,
+        "Number of register threads should be less than or equal to " +
+            std::to_string(maxThreads),
+        "");
+    m_pinned_num_register_threads = val2;
+  } else {
+    TORCH_CHECK(
+        false, "Error, expecting pinned_num_register_threads value", "");
+  }
   return i;
 }
 
 size_t CUDAAllocatorConfig::parsePinnedReserveSegmentSize(
-    const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+    const std::vector<std::string>& config,
     size_t i) {
-  tokenizer.checkToken(++i, ":");
-  size_t val2 = tokenizer.toSizeT(++i);
-  TORCH_CHECK(val2 > 0, "Pinned reserve segment size has to be greater than 0");
-  m_pinned_reserve_segment_size_mb = val2;
+  consumeToken(config, ++i, ':');
+  if (++i < config.size()) {
+    size_t val2 = stoi(config[i]);
+    TORCH_CHECK(
+        val2 > 0, "Pinned reserve segment size has to be greater than 0 ", "");
+    m_pinned_reserve_segment_size_mb = val2;
+  } else {
+    TORCH_CHECK(
+        false, "Error, expecting pinned_reserve_segment_size_mb value", "");
+  }
   return i;
 }
 
-REGISTER_ALLOCATOR_CONFIG_PARSE_HOOK(CUDAAllocatorConfig)
+size_t CUDAAllocatorConfig::parsePinnedUseBackgroundThreads(
+    const std::vector<std::string>& config,
+    size_t i) {
+  consumeToken(config, ++i, ':');
+  if (++i < config.size()) {
+    TORCH_CHECK(
+        (config[i] == "True" || config[i] == "False"),
+        "Expected a single True/False argument for pinned_use_background_threads");
+    m_pinned_use_background_threads = (config[i] == "True");
+  } else {
+    TORCH_CHECK(
+        false, "Error, expecting pinned_use_background_threads value", "");
+  }
+  return i;
+}
+
+// General caching allocator utilities
+void setAllocatorSettings(const std::string& env) {
+  CUDACachingAllocator::CUDAAllocatorConfig::instance().parseArgs(env.c_str());
+}
 
 } // namespace c10::cuda::CUDACachingAllocator

c10/cuda/CUDAAllocatorConfig.h

@@ -1,11 +1,16 @@
 #pragma once
 
-#include <c10/core/AllocatorConfig.h>
-#include <c10/cuda/CUDAException.h>
 #include <c10/cuda/CUDAMacros.h>
 #include <c10/util/Exception.h>
 #include <c10/util/env.h>
 
+#include <atomic>
+#include <cstddef>
+#include <cstdlib>
+#include <mutex>
+#include <string>
+#include <vector>
+
 namespace c10::cuda::CUDACachingAllocator {
 
 enum class Expandable_Segments_Handle_Type : int {
@@ -18,23 +23,20 @@ enum class Expandable_Segments_Handle_Type : int {
 class C10_CUDA_API CUDAAllocatorConfig {
  public:
   static size_t max_split_size() {
-    return c10::CachingAllocator::AcceleratorAllocatorConfig::max_split_size();
+    return instance().m_max_split_size;
   }
   static double garbage_collection_threshold() {
-    return c10::CachingAllocator::AcceleratorAllocatorConfig::
-        garbage_collection_threshold();
+    return instance().m_garbage_collection_threshold;
   }
 
   static bool expandable_segments() {
-    bool enabled = c10::CachingAllocator::AcceleratorAllocatorConfig::
-        use_expandable_segments();
 #ifndef PYTORCH_C10_DRIVER_API_SUPPORTED
-    if (enabled) {
+    if (instance().m_expandable_segments) {
       TORCH_WARN_ONCE("expandable_segments not supported on this platform")
     }
     return false;
 #else
-    return enabled;
+    return instance().m_expandable_segments;
 #endif
   }
 
@@ -65,8 +67,7 @@ class C10_CUDA_API CUDAAllocatorConfig {
   }
 
   static bool pinned_use_background_threads() {
-    return c10::CachingAllocator::AcceleratorAllocatorConfig::
-        pinned_use_background_threads();
+    return instance().m_pinned_use_background_threads;
   }
 
   static size_t pinned_reserve_segment_size_mb() {
@@ -80,23 +81,24 @@ class C10_CUDA_API CUDAAllocatorConfig {
     return 128;
   }
 
-  static size_t roundup_power2_divisions(size_t size) {
-    return c10::CachingAllocator::AcceleratorAllocatorConfig::
-        roundup_power2_divisions(size);
-  }
+  // This is used to round-up allocation size to nearest power of 2 divisions.
+  // More description below in function roundup_power2_next_division
+  // As an example, if we want 4 divisions between 2's power, this can be done
+  // using env variable: PYTORCH_CUDA_ALLOC_CONF=roundup_power2_divisions:4
+  static size_t roundup_power2_divisions(size_t size);
 
   static std::vector<size_t> roundup_power2_divisions() {
-    return c10::CachingAllocator::AcceleratorAllocatorConfig::
-        roundup_power2_divisions();
+    return instance().m_roundup_power2_divisions;
   }
 
   static size_t max_non_split_rounding_size() {
-    return c10::CachingAllocator::AcceleratorAllocatorConfig::
-        max_non_split_rounding_size();
+    return instance().m_max_non_split_rounding_size;
   }
 
   static std::string last_allocator_settings() {
-    return c10::CachingAllocator::getAllocatorSettings();
+    std::lock_guard<std::mutex> lock(
+        instance().m_last_allocator_settings_mutex);
+    return instance().m_last_allocator_settings;
   }
 
   static CUDAAllocatorConfig& instance() {
@@ -109,75 +111,70 @@ class C10_CUDA_API CUDAAllocatorConfig {
         env = c10::utils::get_env("PYTORCH_HIP_ALLOC_CONF");
       }
 #endif
-      // Note: keep the parsing order and logic stable to avoid potential
-      // performance regressions in internal tests.
-      if (!env.has_value()) {
-        env = c10::utils::get_env("PYTORCH_ALLOC_CONF");
-      }
-      if (env.has_value()) {
-        inst->parseArgs(env.value());
-      }
+      inst->parseArgs(env);
       return inst;
     })();
     return *s_instance;
   }
 
-  // Use `Construct On First Use Idiom` to avoid `Static Initialization Order`
-  // issue.
-  static const std::unordered_set<std::string>& getKeys() {
-    static std::unordered_set<std::string> keys{
-        "backend",
-        // keep BC for Rocm: `cuda` -> `cud` `a`, to avoid hipify issues
-        // NOLINTBEGIN(bugprone-suspicious-missing-comma,-warnings-as-errors)
-        "release_lock_on_cud"
-        "amalloc",
-        "pinned_use_cud"
-        "a_host_register",
-        // NOLINTEND(bugprone-suspicious-missing-comma,-warnings-as-errors)
-        "release_lock_on_hipmalloc",
-        "pinned_use_hip_host_register",
-        "graph_capture_record_stream_reuse",
-        "pinned_reserve_segment_size_mb",
-        "pinned_num_register_threads"};
-    return keys;
-  }
-
-  void parseArgs(const std::string& env);
+  void parseArgs(const std::optional<std::string>& env);
 
  private:
-  CUDAAllocatorConfig() = default;
+  CUDAAllocatorConfig();
 
+  static void lexArgs(const std::string& env, std::vector<std::string>& config);
+  static void consumeToken(
+      const std::vector<std::string>& config,
+      size_t i,
+      const char c);
+  size_t parseMaxSplitSize(const std::vector<std::string>& config, size_t i);
+  size_t parseMaxNonSplitRoundingSize(
+      const std::vector<std::string>& config,
+      size_t i);
+  size_t parseGarbageCollectionThreshold(
+      const std::vector<std::string>& config,
+      size_t i);
+  size_t parseRoundUpPower2Divisions(
+      const std::vector<std::string>& config,
+      size_t i);
   size_t parseAllocatorConfig(
-      const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+      const std::vector<std::string>& config,
       size_t i,
       bool& used_cudaMallocAsync);
   size_t parsePinnedUseCudaHostRegister(
-      const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+      const std::vector<std::string>& config,
       size_t i);
   size_t parsePinnedNumRegisterThreads(
-      const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+      const std::vector<std::string>& config,
       size_t i);
   size_t parsePinnedReserveSegmentSize(
-      const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+      const std::vector<std::string>& config,
+      size_t i);
+  size_t parsePinnedUseBackgroundThreads(
+      const std::vector<std::string>& config,
       size_t i);
   size_t parseGraphCaptureRecordStreamReuse(
-      const c10::CachingAllocator::ConfigTokenizer& tokenizer,
+      const std::vector<std::string>& config,
       size_t i);
 
-  std::atomic<size_t> m_pinned_num_register_threads{1};
-  std::atomic<size_t> m_pinned_reserve_segment_size_mb{0};
-  std::atomic<Expandable_Segments_Handle_Type> m_expandable_segments_handle_type
-#if CUDA_VERSION >= 12030
-      {Expandable_Segments_Handle_Type::UNSPECIFIED};
-#else
-      {Expandable_Segments_Handle_Type::POSIX_FD};
-#endif
-  std::atomic<bool> m_release_lock_on_cudamalloc{false};
-  std::atomic<bool> m_pinned_use_cuda_host_register{false};
-  std::atomic<bool> m_graph_capture_record_stream_reuse{false};
+  std::atomic<size_t> m_max_split_size;
+  std::atomic<size_t> m_max_non_split_rounding_size;
+  std::vector<size_t> m_roundup_power2_divisions;
+  std::atomic<double> m_garbage_collection_threshold;
+  std::atomic<size_t> m_pinned_num_register_threads;
+  std::atomic<size_t> m_pinned_reserve_segment_size_mb;
+  std::atomic<bool> m_expandable_segments;
+  std::atomic<Expandable_Segments_Handle_Type>
+      m_expandable_segments_handle_type;
+  std::atomic<bool> m_release_lock_on_cudamalloc;
+  std::atomic<bool> m_pinned_use_cuda_host_register;
+  std::atomic<bool> m_graph_capture_record_stream_reuse;
+  std::atomic<bool> m_pinned_use_background_threads;
+  std::string m_last_allocator_settings;
+  std::mutex m_last_allocator_settings_mutex;
 };
 
-// Keep this for backwards compatibility
-using c10::CachingAllocator::setAllocatorSettings;
+// General caching allocator utilities
+C10_CUDA_API void setAllocatorSettings(const std::string& env);
 
 } // namespace c10::cuda::CUDACachingAllocator

c10/cuda/CUDACachingAllocator.cpp

@@ -64,6 +64,10 @@ namespace cuda::CUDACachingAllocator {
 using namespace c10::CachingAllocator;
 using namespace c10::CachingDeviceAllocator;
 
+// Included here as this is externally used in CUDAAllocatorConfig
+const size_t kLargeBuffer =
+    20971520; // "large" allocations may be packed in 20 MiB blocks
+
 namespace Native {
 
 //

c10/cuda/CUDACachingAllocator.h

@@ -1,6 +1,5 @@
 #pragma once
 
-#include <c10/core/AllocatorConfig.h>
 #include <c10/core/CachingDeviceAllocator.h>
 #include <c10/cuda/CUDAGraphsC10Utils.h>
 #include <c10/cuda/CUDAMacros.h>
@@ -50,9 +49,10 @@ namespace c10::cuda::CUDACachingAllocator {
 
 // Preserved only for BC reasons
 // NOLINTNEXTLINE(misc-unused-using-decls)
-using c10::CachingAllocator::kLargeBuffer;
 using c10::CachingDeviceAllocator::DeviceStats;
 
+extern const size_t kLargeBuffer;
+
 typedef std::shared_ptr<GatheredContext> (*CreateContextFn)();
 
 // Struct containing info of an allocation block (i.e. a fractional part of a

c10/test/core/AllocatorConfig_test.cpp

@@ -67,8 +67,8 @@ TEST(AllocatorConfigTest, allocator_config_test) {
   EXPECT_EQ(AcceleratorAllocatorConfig::roundup_power2_divisions(128 * kMB), 2);
   EXPECT_EQ(AcceleratorAllocatorConfig::roundup_power2_divisions(256 * kMB), 4);
   EXPECT_EQ(AcceleratorAllocatorConfig::roundup_power2_divisions(512 * kMB), 2);
-  // EXPECT_EQ(
-  //     AcceleratorAllocatorConfig::roundup_power2_divisions(1024 * kMB), 4);
+  EXPECT_EQ(
+      AcceleratorAllocatorConfig::roundup_power2_divisions(1024 * kMB), 4);
   EXPECT_EQ(
       AcceleratorAllocatorConfig::roundup_power2_divisions(2048 * kMB), 1);
   EXPECT_EQ(
@@ -101,8 +101,8 @@ TEST(AllocatorConfigTest, allocator_config_test) {
   EXPECT_EQ(AcceleratorAllocatorConfig::roundup_power2_divisions(512 * kMB), 1);
   EXPECT_EQ(
       AcceleratorAllocatorConfig::roundup_power2_divisions(1024 * kMB), 0);
-  // EXPECT_EQ(
-  //     AcceleratorAllocatorConfig::roundup_power2_divisions(2048 * kMB), 8);
+  EXPECT_EQ(
+      AcceleratorAllocatorConfig::roundup_power2_divisions(2048 * kMB), 8);
   EXPECT_EQ(
       AcceleratorAllocatorConfig::roundup_power2_divisions(4096 * kMB), 2);
 

c10/util/ArrayRef.h

@@ -18,7 +18,6 @@
 #include <c10/macros/Macros.h>
 #include <c10/util/Exception.h>
 #include <c10/util/SmallVector.h>
-#include <torch/headeronly/util/HeaderOnlyArrayRef.h>
 
 #include <array>
 #include <cstddef>
@@ -41,106 +40,200 @@ namespace c10 {
 ///
 /// This is intended to be trivially copyable, so it should be passed by
 /// value.
-///
-/// NOTE: We have refactored out the headeronly parts of the ArrayRef struct
-/// into HeaderOnlyArrayRef. As adding `virtual` would change the performance of
-/// the underlying constexpr calls, we rely on apparent-type dispatch for
-/// inheritance. This should be fine because their memory format is the same,
-/// and it is never incorrect for ArrayRef to call HeaderOnlyArrayRef methods.
-/// However, you should prefer to use ArrayRef when possible, because its use
-/// of TORCH_CHECK will lead to better user-facing error messages.
 template <typename T>
-class ArrayRef final : public HeaderOnlyArrayRef<T> {
+class ArrayRef final {
  public:
-  /// @name Constructors, all inherited from HeaderOnlyArrayRef except for
-  /// SmallVector.
+  using iterator = const T*;
+  using const_iterator = const T*;
+  using size_type = size_t;
+  using value_type = T;
+
+  using reverse_iterator = std::reverse_iterator<iterator>;
+
+ private:
+  /// The start of the array, in an external buffer.
+  const T* Data;
+
+  /// The number of elements.
+  size_type Length;
+
+  void debugCheckNullptrInvariant() {
+    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
+        Data != nullptr || Length == 0,
+        "created ArrayRef with nullptr and non-zero length! std::optional relies on this being illegal");
+  }
+
+ public:
+  /// @name Constructors
   /// @{
 
-  using HeaderOnlyArrayRef<T>::HeaderOnlyArrayRef;
+  /// Construct an empty ArrayRef.
+  /* implicit */ constexpr ArrayRef() : Data(nullptr), Length(0) {}
 
-  /// Construct an ArrayRef from a std::vector.
-  /// This constructor is identical to the one in HeaderOnlyArrayRef, but we
-  /// include it to help with Class Template Argument Deduction (CTAD).
-  /// Without it, CTAD can fail sometimes due to the indirect constructor
-  /// inheritance. So we explicitly include this constructor.
-  template <typename A>
-  /* implicit */ ArrayRef(const std::vector<T, A>& Vec)
-      : HeaderOnlyArrayRef<T>(Vec.data(), Vec.size()) {}
+  /// Construct an ArrayRef from a single element.
+  // TODO Make this explicit
+  constexpr ArrayRef(const T& OneElt) : Data(&OneElt), Length(1) {}
+
+  /// Construct an ArrayRef from a pointer and length.
+  constexpr ArrayRef(const T* data, size_t length)
+      : Data(data), Length(length) {
+    debugCheckNullptrInvariant();
+  }
+
+  /// Construct an ArrayRef from a range.
+  constexpr ArrayRef(const T* begin, const T* end)
+      : Data(begin), Length(end - begin) {
+    debugCheckNullptrInvariant();
+  }
 
   /// Construct an ArrayRef from a SmallVector. This is templated in order to
   /// avoid instantiating SmallVectorTemplateCommon<T> whenever we
   /// copy-construct an ArrayRef.
-  /// NOTE: this is the only constructor that is not inherited from
-  /// HeaderOnlyArrayRef.
   template <typename U>
   /* implicit */ ArrayRef(const SmallVectorTemplateCommon<T, U>& Vec)
-      : HeaderOnlyArrayRef<T>(Vec.data(), Vec.size()) {}
+      : Data(Vec.data()), Length(Vec.size()) {
+    debugCheckNullptrInvariant();
+  }
+
+  template <
+      typename Container,
+      typename U = decltype(std::declval<Container>().data()),
+      typename = std::enable_if_t<
+          (std::is_same_v<U, T*> || std::is_same_v<U, T const*>)>>
+  /* implicit */ ArrayRef(const Container& container)
+      : Data(container.data()), Length(container.size()) {
+    debugCheckNullptrInvariant();
+  }
+
+  /// Construct an ArrayRef from a std::vector.
+  // The enable_if stuff here makes sure that this isn't used for
+  // std::vector<bool>, because ArrayRef can't work on a std::vector<bool>
+  // bitfield.
+  template <typename A>
+  /* implicit */ ArrayRef(const std::vector<T, A>& Vec)
+      : Data(Vec.data()), Length(Vec.size()) {
+    static_assert(
+        !std::is_same_v<T, bool>,
+        "ArrayRef<bool> cannot be constructed from a std::vector<bool> bitfield.");
+  }
+
+  /// Construct an ArrayRef from a std::array
+  template <size_t N>
+  /* implicit */ constexpr ArrayRef(const std::array<T, N>& Arr)
+      : Data(Arr.data()), Length(N) {}
+
+  /// Construct an ArrayRef from a C array.
+  template <size_t N>
+  // NOLINTNEXTLINE(*c-arrays*)
+  /* implicit */ constexpr ArrayRef(const T (&Arr)[N]) : Data(Arr), Length(N) {}
+
+  /// Construct an ArrayRef from a std::initializer_list.
+  /* implicit */ constexpr ArrayRef(const std::initializer_list<T>& Vec)
+      : Data(
+            std::begin(Vec) == std::end(Vec) ? static_cast<T*>(nullptr)
+                                             : std::begin(Vec)),
+        Length(Vec.size()) {}
 
   /// @}
-  /// @name Simple Operations, mostly inherited from HeaderOnlyArrayRef
+  /// @name Simple Operations
   /// @{
 
+  constexpr iterator begin() const {
+    return Data;
+  }
+  constexpr iterator end() const {
+    return Data + Length;
+  }
+
+  // These are actually the same as iterator, since ArrayRef only
+  // gives you const iterators.
+  constexpr const_iterator cbegin() const {
+    return Data;
+  }
+  constexpr const_iterator cend() const {
+    return Data + Length;
+  }
+
+  constexpr reverse_iterator rbegin() const {
+    return reverse_iterator(end());
+  }
+  constexpr reverse_iterator rend() const {
+    return reverse_iterator(begin());
+  }
+
+  /// Check if all elements in the array satisfy the given expression
+  constexpr bool allMatch(const std::function<bool(const T&)>& pred) const {
+    return std::all_of(cbegin(), cend(), pred);
+  }
+
+  /// empty - Check if the array is empty.
+  constexpr bool empty() const {
+    return Length == 0;
+  }
+
+  constexpr const T* data() const {
+    return Data;
+  }
+
+  /// size - Get the array size.
+  constexpr size_t size() const {
+    return Length;
+  }
+
   /// front - Get the first element.
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr const T& front() const {
     TORCH_CHECK(
-        !this->empty(), "ArrayRef: attempted to access front() of empty list");
-    return this->Data[0];
+        !empty(), "ArrayRef: attempted to access front() of empty list");
+    return Data[0];
   }
 
   /// back - Get the last element.
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr const T& back() const {
-    TORCH_CHECK(
-        !this->empty(), "ArrayRef: attempted to access back() of empty list");
-    return this->Data[this->Length - 1];
+    TORCH_CHECK(!empty(), "ArrayRef: attempted to access back() of empty list");
+    return Data[Length - 1];
+  }
+
+  /// equals - Check for element-wise equality.
+  constexpr bool equals(ArrayRef RHS) const {
+    return Length == RHS.Length && std::equal(begin(), end(), RHS.begin());
   }
 
   /// slice(n, m) - Take M elements of the array starting at element N
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr ArrayRef<T> slice(size_t N, size_t M) const {
     TORCH_CHECK(
-        N + M <= this->size(),
+        N + M <= size(),
         "ArrayRef: invalid slice, N = ",
         N,
         "; M = ",
         M,
         "; size = ",
-        this->size());
-    return ArrayRef<T>(this->data() + N, M);
+        size());
+    return ArrayRef<T>(data() + N, M);
   }
 
   /// slice(n) - Chop off the first N elements of the array.
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr ArrayRef<T> slice(size_t N) const {
     TORCH_CHECK(
-        N <= this->size(),
-        "ArrayRef: invalid slice, N = ",
-        N,
-        "; size = ",
-        this->size());
-    return slice(N, this->size() - N); // should this slice be this->slice?
+        N <= size(), "ArrayRef: invalid slice, N = ", N, "; size = ", size());
+    return slice(N, size() - N);
   }
 
   /// @}
   /// @name Operator Overloads
   /// @{
+  constexpr const T& operator[](size_t Index) const {
+    return Data[Index];
+  }
 
   /// Vector compatibility
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr const T& at(size_t Index) const {
     TORCH_CHECK(
-        Index < this->Length,
+        Index < Length,
         "ArrayRef: invalid index Index = ",
         Index,
         "; Length = ",
-        this->Length);
-    return this->Data[Index];
+        Length);
+    return Data[Index];
   }
 
   /// Disallow accidental assignment from a temporary.
@@ -160,6 +253,13 @@ class ArrayRef final : public HeaderOnlyArrayRef<T> {
   std::enable_if_t<std::is_same_v<U, T>, ArrayRef<T>>& operator=(
       std::initializer_list<U>) = delete;
 
+  /// @}
+  /// @name Expensive Operations
+  /// @{
+  std::vector<T> vec() const {
+    return std::vector<T>(Data, Data + Length);
+  }
+
   /// @}
 };
 

c10/xpu/XPUFunctions.cpp

@@ -120,23 +120,17 @@ inline void initGlobalDevicePoolState() {
   TORCH_CHECK(
       gDevicePool.devices.size() <= std::numeric_limits<DeviceIndex>::max(),
       "Too many XPU devices, DeviceIndex overflowed!");
-  // Check each device's architecture and issue a warning if it is older than
-  // the officially supported range (Intel GPUs starting from Arc (Alchemist)
-  // series).
-  namespace syclex = sycl::ext::oneapi::experimental;
-  for (const auto& device : gDevicePool.devices) {
-    auto architecture = device->get_info<syclex::info::device::architecture>();
-    if (architecture < syclex::architecture::intel_gpu_acm_g10) {
-      TORCH_WARN(
-          "The detected GPU (",
-          device->get_info<sycl::info::device::name>(),
-          ") is not officially supported by PyTorch XPU. Running workloads on this device may result in unexpected behavior.\n",
-          "For stable and fully supported execution, please use GPUs based on Intel Arc (Alchemist) series or newer.\n",
-          "Refer to the hardware prerequisites for more information: ",
-          "https://github.com/pytorch/pytorch/blob/main/docs/source/notes/get_start_xpu.rst#hardware-prerequisite");
-    }
-  }
 
+#if defined(_WIN32) && SYCL_COMPILER_VERSION < 20250000
+  // The default context feature is disabled by default on Windows for SYCL
+  // compiler versions earlier than 2025.0.0.
+  std::vector<sycl::device> deviceList;
+  for (auto it = gDevicePool.devices.begin(); it != gDevicePool.devices.end();
+       ++it) {
+    deviceList.push_back(*(*it));
+  }
+  gDevicePool.context = std::make_unique<sycl::context>(deviceList);
+#else
   // The default context is utilized for each Intel GPU device, allowing the
   // retrieval of the context from any GPU device.
   const auto& platform = gDevicePool.devices[0]->get_platform();
@@ -146,6 +140,7 @@ inline void initGlobalDevicePoolState() {
 #else
       platform.ext_oneapi_get_default_context());
 #endif
+#endif
 }
 
 inline void initDevicePoolCallOnce() {
@@ -170,9 +165,9 @@ void initDeviceProperties(DeviceProp* device_prop, DeviceIndex device) {
 #define ASSIGN_DEVICE_ASPECT(member) \
   device_prop->has_##member = raw_device.has(sycl::aspect::member);
 
-#define ASSIGN_EXP_CL_ASPECT(member) \
-  device_prop->has_##member =        \
-      raw_device.ext_oneapi_supports_cl_extension("cl_intel_" #member);
+#define ASSIGN_EXP_CL_ASPECT(member)                                       \
+  device_prop->has_##member = raw_device.ext_oneapi_supports_cl_extension( \
+      "cl_intel_" #member, &cl_version);
 
 #define ASSIGN_EXP_DEVICE_PROP(property) \
   device_prop->property =                \
@@ -187,6 +182,8 @@ void initDeviceProperties(DeviceProp* device_prop, DeviceIndex device) {
 
   AT_FORALL_XPU_DEVICE_ASPECT(ASSIGN_DEVICE_ASPECT);
 
+  // TODO: Remove cl_version since it is unnecessary.
+  sycl::ext::oneapi::experimental::cl_version cl_version;
   AT_FORALL_XPU_EXP_CL_ASPECT(ASSIGN_EXP_CL_ASPECT);
 
 #if SYCL_COMPILER_VERSION >= 20250000

cmake/Dependencies.cmake

@@ -1044,17 +1044,6 @@ if(USE_ROCM)
        list(APPEND HIP_HIPCC_FLAGS -fdebug-info-for-profiling)
     endif(CMAKE_BUILD_TYPE MATCHES Debug)
 
-    # Get EnVar 'USE_LAYERNORM_FAST_RECIPROCAL' (or default to on).
-    if(DEFINED ENV{USE_LAYERNORM_FAST_RECIPROCAL})
-      set(USE_LAYERNORM_FAST_RECIPROCAL $ENV{USE_LAYERNORM_FAST_RECIPROCAL})
-    else()
-      set(USE_LAYERNORM_FAST_RECIPROCAL ON)
-    endif()
-
-    if(USE_LAYERNORM_FAST_RECIPROCAL)
-      add_definitions(-DUSE_LAYERNORM_FAST_RECIPROCAL)
-    endif()
-
     # needed for compat with newer versions of hip-clang that introduced C++20 mangling rules
     list(APPEND HIP_HIPCC_FLAGS -fclang-abi-compat=17)
 

cmake/External/aotriton.cmake

@@ -244,8 +244,7 @@ if(NOT __AOTRITON_INCLUDED)
   else()
     set(__AOTRITON_SYSTEM_ROCM "${HIP_VERSION_MAJOR}.${HIP_VERSION_MINOR}")
     list(FIND __AOTRITON_ROCM_LIST "rocm${__AOTRITON_SYSTEM_ROCM}" __AOTRITON_RUNTIME_INDEX)
-    # Always build aotriton runtime from source on Windows due to lack of pre-built binaries
-    if(${__AOTRITON_RUNTIME_INDEX} LESS 0 OR WIN32)
+    if(${__AOTRITON_RUNTIME_INDEX} LESS 0)
       message(STATUS "Cannot find AOTriton runtime for ROCM ${__AOTRITON_SYSTEM_ROCM}. \
       Build runtime from source")
       aotriton_build_from_source(ON aotriton_runtime)

cmake/Summary.cmake

@@ -128,12 +128,11 @@ function(caffe2_print_configuration_summary)
   endif()
   message(STATUS "  USE_ROCM              : ${USE_ROCM}")
   if(${USE_ROCM})
-    message(STATUS "    ROCM_VERSION                  : ${ROCM_VERSION}")
-    message(STATUS "    USE_FLASH_ATTENTION           : ${USE_FLASH_ATTENTION}")
-    message(STATUS "    USE_MEM_EFF_ATTENTION         : ${USE_MEM_EFF_ATTENTION}")
-    message(STATUS "    USE_ROCM_CK_SDPA              : ${USE_ROCM_CK_SDPA}")
-    message(STATUS "    USE_ROCM_CK_GEMM              : ${USE_ROCM_CK_GEMM}")
-    message(STATUS "    USE_LAYERNORM_FAST_RECIPROCAL : ${USE_LAYERNORM_FAST_RECIPROCAL}")
+    message(STATUS "    ROCM_VERSION          : ${ROCM_VERSION}")
+    message(STATUS "    USE_FLASH_ATTENTION   : ${USE_FLASH_ATTENTION}")
+    message(STATUS "    USE_MEM_EFF_ATTENTION : ${USE_MEM_EFF_ATTENTION}")
+    message(STATUS "    USE_ROCM_CK_SDPA      : ${USE_ROCM_CK_SDPA}")
+    message(STATUS "    USE_ROCM_CK_GEMM      : ${USE_ROCM_CK_GEMM}")
   endif()
   message(STATUS "  BUILD_NVFUSER         : ${BUILD_NVFUSER}")
   message(STATUS "  USE_EIGEN_FOR_BLAS    : ${CAFFE2_USE_EIGEN_FOR_BLAS}")

functorch/dim/__init__.py

@@ -3,11 +3,11 @@ from __future__ import annotations
 import dis
 import inspect
 import sys
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import Any, Callable, Optional, TYPE_CHECKING, Union
 
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Sequence
+    from collections.abc import Sequence
 
 import torch
 from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten

functorch/dim/_wrap.py

@@ -5,7 +5,7 @@ Python implementation of function wrapping functionality for functorch.dim.
 from __future__ import annotations
 
 import functools
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, Callable, Optional
 
 import torch
 from torch.utils._pytree import tree_map
@@ -15,10 +15,6 @@ from ._enable_all_layers import EnableAllLayers
 from ._tensor_info import TensorInfo
 
 
-if TYPE_CHECKING:
-    from collections.abc import Callable
-
-
 def handle_from_tensor(tensor: torch.Tensor) -> torch.Tensor:
     """Handle tensor conversion for torch function integration."""
     return tensor

functorch/dim/wrap_type.py

@@ -5,7 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import functools
-from collections.abc import Callable
 from types import (
     BuiltinMethodType,
     FunctionType,
@@ -13,7 +12,7 @@ from types import (
     MethodDescriptorType,
     WrapperDescriptorType,
 )
-from typing import Any
+from typing import Any, Callable
 
 
 FUNC_TYPES = (

functorch/einops/rearrange.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import functools
-from typing import TYPE_CHECKING, Union
+from typing import Callable, TYPE_CHECKING, Union
 
 import torch
 from functorch.dim import dims  # noqa: F401
@@ -16,7 +16,7 @@ from ._parsing import (
 
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Sequence
+    from collections.abc import Sequence
 
 __all__ = ["rearrange"]
 

pyproject.toml

@@ -180,7 +180,6 @@ ignore = [
     "SIM116", # Disable Use a dictionary instead of consecutive `if` statements
     "SIM117",
     "SIM118",
-    "SIM300", # Yoda condition detected
     "UP007", # keep-runtime-typing
     "UP045", # keep-runtime-typing
     "TC006",
@@ -196,7 +195,8 @@ select = [
     "E",
     "EXE",
     "F",
-    "SIM",
+    "SIM1",
+    "SIM911",
     "W",
     # Not included in flake8
     "FURB",

pyrefly.toml

@@ -1,7 +1,5 @@
 # A Pyrefly configuration for PyTorch
 # Based on https://github.com/pytorch/pytorch/blob/main/mypy.ini
-python-version = "3.12"
-
 project-includes = [
     "torch",
     "caffe2",
@@ -23,7 +21,7 @@ project-excludes = [
   # ==== below will be enabled directory by directory ====
   # ==== to test Pyrefly on a specific directory, simply comment it out ====
   "torch/_inductor/runtime",
-  "torch/_inductor/codegen/triton.py",
+  "torch/_inductor/codegen",
   # formatting issues, will turn on after adjusting where suppressions can be
   # in import statements
   "torch/linalg/__init__.py",
@@ -38,7 +36,6 @@ project-excludes = [
   "torch/nn/modules/rnn.py", # only remove when parsing errors are fixed
   "torch/_inductor/codecache.py",
   "torch/distributed/elastic/metrics/__init__.py",
-  "torch/_inductor/fx_passes/bucketing.py",
   # ====
   "benchmarks/instruction_counts/main.py",
   "benchmarks/instruction_counts/definitions/setup.py",

setup.py

@@ -156,10 +156,6 @@
 #   USE_ROCM_KERNEL_ASSERT=1
 #     Enable kernel assert in ROCm platform
 #
-#   USE_LAYERNORM_FAST_RECIPROCAL
-#     If set, enables the use of builtin functions for fast reciprocals (1/x) w.r.t.
-#     layer normalization. Default: enabled.
-#
 #   USE_ROCM_CK_GEMM=1
 #     Enable building CK GEMM backend in ROCm platform
 #