Updated docs to add the error case for torch.multinomial

Summary: Updated docs to add the error condition for torch.multinomial

Test Plan: No change in code

Reviewers:

Subscribers: @drisspg

Tasks:

Tags:

.bazelignore

@@ -1,3 +1,4 @@
 # We do not use this library in our Bazel build. It contains an
 # infinitely recursing symlink that makes Bazel very unhappy.
 third_party/ittapi/
+third_party/opentelemetry-cpp

.ci/docker/build.sh

@@ -229,6 +229,7 @@ case "$image" in
     BASEKIT_VERSION=2024.0.0-49522
     NINJA_VERSION=1.9.0
     CONDA_CMAKE=yes
+    TRITON=yes
     ;;
     pytorch-linux-jammy-py3.8-gcc11-inductor-benchmarks)
     ANACONDA_PYTHON_VERSION=3.8

.ci/docker/ci_commit_pins/executorch.txt

@@ -1 +1 @@
-e2a8f9548aecb62a68e264607174a7d207ed2929
+d4b3e5cc607e97afdba79dc90f8ef968142f347c

.ci/docker/ci_commit_pins/triton-rocm.txt

@@ -1 +1 @@
-0a22a91d04c2b4a029a69a198eac390089c3e891
+bbe6246e37d8aa791c67daaf9d9d61b26c9ccfdc

.ci/docker/ci_commit_pins/triton-xpu.txt

@@ -0,0 +1 @@
+b8c64f64c18d8cac598b3adb355c21e7439c21de

.ci/docker/ci_commit_pins/triton.txt

@@ -1 +1 @@
-a9bc1a36470eefafe0e2ab2503b8698f1e89e7e3
+45fff310c891f5a92d55445adf8cc9d29df5841e

.ci/docker/common/install_conda.sh

@@ -57,8 +57,21 @@ fi
   # Uncomment the below when resolved to track the latest conda update
   # as_jenkins conda update -y -n base conda
 
+  if [[ $(uname -m) == "aarch64" ]]; then
+    export SYSROOT_DEP="sysroot_linux-aarch64=2.17"
+  else
+    export SYSROOT_DEP="sysroot_linux-64=2.17"
+  fi
+
   # Install correct Python version
-  as_jenkins conda create -n py_$ANACONDA_PYTHON_VERSION -y python="$ANACONDA_PYTHON_VERSION"
+  # Also ensure sysroot is using a modern GLIBC to match system compilers
+  as_jenkins conda create -n py_$ANACONDA_PYTHON_VERSION -y\
+             python="$ANACONDA_PYTHON_VERSION" \
+             ${SYSROOT_DEP}
+
+  # libstdcxx from conda default channels are too old, we need GLIBCXX_3.4.30
+  # which is provided in libstdcxx 12 and up.
+  conda_install libstdcxx-ng=12.3.0 -c conda-forge
 
   # Install PyTorch conda deps, as per https://github.com/pytorch/pytorch README
   if [[ $(uname -m) == "aarch64" ]]; then
@@ -110,14 +123,5 @@ fi
     pip_install -r /opt/conda/requirements-docs.txt
   fi
 
-  # HACK HACK HACK
-  # gcc-9 for ubuntu-18.04 from http://ppa.launchpad.net/ubuntu-toolchain-r/test/ubuntu
-  # Pulls llibstdc++6 13.1.0-8ubuntu1~18.04 which is too new for conda
-  # So remove libstdc++6.so.3.29 installed by https://anaconda.org/anaconda/libstdcxx-ng/files?version=11.2.0
-  # Same is true for gcc-12 from Ubuntu-22.04
-  if grep -e [12][82].04.[623] /etc/issue >/dev/null; then
-    rm /opt/conda/envs/py_$ANACONDA_PYTHON_VERSION/lib/libstdc++.so.6
-  fi
-
   popd
 fi

.ci/docker/common/install_onnx.sh

@@ -33,12 +33,12 @@ pip_install coloredlogs packaging
 pip_install onnxruntime==1.17.0
 pip_install onnx==1.15.0
 # pip_install "onnxscript@git+https://github.com/microsoft/onnxscript@3e869ef8ccf19b5ebd21c10d3e9c267c9a9fa729" --no-deps
-pip_install onnxscript==0.1.0.dev20240301 --no-deps
+pip_install onnxscript==0.1.0.dev20240315 --no-deps
 
 # 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/
 IMPORT_SCRIPT_FILENAME="/tmp/onnx_import_script.py"
-as_jenkins echo 'import transformers; transformers.AutoModel.from_pretrained("sshleifer/tiny-gpt2"); transformers.AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2");' > "${IMPORT_SCRIPT_FILENAME}"
+as_jenkins echo 'import transformers; transformers.AutoModel.from_pretrained("sshleifer/tiny-gpt2"); transformers.AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2"); transformers.AutoModelForSpeechSeq2Seq.from_pretrained("openai/whisper-large-v3");' > "${IMPORT_SCRIPT_FILENAME}"
 
 # Need a PyTorch version for transformers to work
 pip_install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu

.ci/docker/common/install_protobuf.sh

@@ -11,7 +11,8 @@ mkdir -p $pb_dir
 ln -s /usr/lib64 "$pb_dir/lib64"
 
 curl -LO "https://github.com/protocolbuffers/protobuf/releases/download/v3.17.3/protobuf-all-3.17.3.tar.gz" --retry 3
-tar -xvz -C "$pb_dir" --strip-components 1 -f protobuf-all-3.17.3.tar.gz
+
+tar -xvz --no-same-owner -C "$pb_dir" --strip-components 1 -f protobuf-all-3.17.3.tar.gz
 NPROC=$[$(nproc) - 2]
 pushd "$pb_dir" && ./configure && make -j${NPROC} && make -j${NPROC} check && sudo make -j${NRPOC} install && sudo ldconfig
 popd

.ci/docker/common/install_triton.sh

@@ -13,8 +13,11 @@ conda_reinstall() {
 }
 
 if [ -n "${ROCM_VERSION}" ]; then
-  TRITON_REPO="https://github.com/ROCmSoftwarePlatform/triton"
+  TRITON_REPO="https://github.com/openai/triton"
   TRITON_TEXT_FILE="triton-rocm"
+elif [ -n "${BASEKIT_VERSION}" ]; then
+  TRITON_REPO="https://github.com/intel/intel-xpu-backend-for-triton"
+  TRITON_TEXT_FILE="triton-xpu"
 else
   TRITON_REPO="https://github.com/openai/triton"
   TRITON_TEXT_FILE="triton"

.ci/docker/common/install_xpu.sh

@@ -3,7 +3,7 @@ set -xe
 
 
 # Intel® software for general purpose GPU capabilities.
-# Refer to https://dgpu-docs.intel.com/releases/stable_647_21_20230714.html
+# Refer to https://dgpu-docs.intel.com/releases/LTS_803.29_20240131.html
 
 # Intel® oneAPI Base Toolkit (version 2024.0.0) has been updated to include functional and security updates.
 # Refer to https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit-download.html
@@ -21,7 +21,7 @@ function install_ubuntu() {
         | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null
 
     # Add the signed entry to APT sources and configure the APT client to use the Intel repository
-    echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/gpu/ubuntu jammy/production/2328 unified" \
+    echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/gpu/ubuntu jammy/lts/2350 unified" \
         | tee /etc/apt/sources.list.d/intel-gpu-jammy.list
     echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" \
         | tee /etc/apt/sources.list.d/oneAPI.list

.ci/docker/requirements-ci.txt

@@ -85,10 +85,10 @@ librosa>=0.6.2 ; python_version < "3.11"
 #Pinned versions:
 #test that import:
 
-mypy==1.8.0
+mypy==1.9.0
 # Pin MyPy version because new errors are likely to appear with each release
 #Description: linter
-#Pinned versions: 1.8.0
+#Pinned versions: 1.9.0
 #test that import: test_typing.py, test_type_hints.py
 
 networkx==2.8.8
@@ -134,9 +134,9 @@ opt-einsum==3.3
 #Pinned versions: 3.3
 #test that import: test_linalg.py
 
-optree==0.9.1
+optree==0.11.0
 #Description: A library for tree manipulation
-#Pinned versions: 0.9.1
+#Pinned versions: 0.11.0
 #test that import: test_vmap.py, test_aotdispatch.py, test_dynamic_shapes.py,
 #test_pytree.py, test_ops.py, test_control_flow.py, test_modules.py,
 #common_utils.py, test_eager_transforms.py, test_python_dispatch.py,
@@ -147,9 +147,9 @@ optree==0.9.1
 #test_pointwise_ops.py, test_dtensor_ops.py, test_torchinductor.py, test_fx.py,
 #test_fake_tensor.py, test_mps.py
 
-pillow==10.2.0
+pillow==10.3.0
 #Description:  Python Imaging Library fork
-#Pinned versions: 10.2.0
+#Pinned versions: 10.3.0
 #test that import:
 
 protobuf==3.20.2
@@ -228,12 +228,11 @@ scikit-image==0.20.0 ; python_version >= "3.10"
 #Pinned versions: 0.20.3
 #test that import:
 
-scipy==1.6.3 ; python_version < "3.10"
-scipy==1.8.1 ; python_version == "3.10"
-scipy==1.10.1 ; python_version == "3.11"
+scipy==1.10.1 ; python_version <= "3.11"
+scipy==1.12.0 ; python_version == "3.12"
 # Pin SciPy because of failing distribution tests (see #60347)
 #Description: scientific python
-#Pinned versions: 1.6.3
+#Pinned versions: 1.10.1
 #test that import: test_unary_ufuncs.py, test_torch.py,test_tensor_creation_ops.py
 #test_spectral_ops.py, test_sparse_csr.py, test_reductions.py,test_nn.py
 #test_linalg.py, test_binary_ufuncs.py

.ci/docker/ubuntu-xpu/Dockerfile

@@ -61,15 +61,20 @@ COPY ci_commit_pins/timm.txt timm.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.txt
 
+# Install XPU Dependencies
+ARG BASEKIT_VERSION
+COPY ./common/install_xpu.sh install_xpu.sh
+RUN bash ./install_xpu.sh && rm install_xpu.sh
+
 ARG TRITON
 # Install triton, this needs to be done before sccache because the latter will
 # try to reach out to S3, which docker build runners don't have access
 COPY ./common/install_triton.sh install_triton.sh
 COPY ./common/common_utils.sh common_utils.sh
-# TODO: will add triton xpu commit
-COPY ci_commit_pins/triton.txt triton.txt
+COPY ci_commit_pins/triton-xpu.txt triton-xpu.txt
+COPY triton_version.txt triton_version.txt
 RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
-RUN rm install_triton.sh common_utils.sh triton.txt
+RUN rm install_triton.sh common_utils.sh triton-xpu.txt triton_version.txt
 
 # (optional) Install database packages like LMDB and LevelDB
 ARG DB
@@ -85,11 +90,6 @@ RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
 RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
 ENV INSTALLED_VISION ${VISION}
 
-# Install XPU Dependencies
-ARG BASEKIT_VERSION
-COPY ./common/install_xpu.sh install_xpu.sh
-RUN bash ./install_xpu.sh && rm install_xpu.sh
-
 # (optional) Install non-default CMake version
 ARG CMAKE_VERSION
 COPY ./common/install_cmake.sh install_cmake.sh

.ci/onnx/common.sh

@@ -1,5 +1,9 @@
+#!/bin/bash
+
 set -ex
 
+source "$(dirname "${BASH_SOURCE[0]}")/../pytorch/common_utils.sh"
+
 LOCAL_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
 ROOT_DIR=$(cd "$LOCAL_DIR"/../.. && pwd)
 TEST_DIR="$ROOT_DIR/test"

.ci/onnx/test.sh

@@ -3,6 +3,20 @@
 # shellcheck source=./common.sh
 source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
 
+# Workaround for dind-rootless userid mapping (https://github.com/pytorch/ci-infra/issues/96)
+WORKSPACE_ORIGINAL_OWNER_ID=$(stat -c '%u' "/var/lib/jenkins/workspace")
+cleanup_workspace() {
+  echo "sudo may print the following warning message that can be ignored. The chown command will still run."
+  echo "    sudo: setrlimit(RLIMIT_STACK): Operation not permitted"
+  echo "For more details refer to https://github.com/sudo-project/sudo/issues/42"
+  sudo chown -R "$WORKSPACE_ORIGINAL_OWNER_ID" /var/lib/jenkins/workspace
+}
+# Disable shellcheck SC2064 as we want to parse the original owner immediately.
+# shellcheck disable=SC2064
+trap_add cleanup_workspace EXIT
+sudo chown -R jenkins /var/lib/jenkins/workspace
+git config --global --add safe.directory /var/lib/jenkins/workspace
+
 if [[ "$BUILD_ENVIRONMENT" == *onnx* ]]; then
   # TODO: This can be removed later once vision is also part of the Docker image
   pip install -q --user --no-use-pep517 "git+https://github.com/pytorch/vision.git@$(cat .github/ci_commit_pins/vision.txt)"

.ci/pytorch/build.sh

@@ -223,6 +223,24 @@ if [[ "${BUILD_ENVIRONMENT}" != *android* && "${BUILD_ENVIRONMENT}" != *cuda* ]]
   export BUILD_STATIC_RUNTIME_BENCHMARK=ON
 fi
 
+# Do not change workspace permissions for ROCm CI jobs
+# as it can leave workspace with bad permissions for cancelled jobs
+if [[ "$BUILD_ENVIRONMENT" != *rocm* ]]; then
+  # Workaround for dind-rootless userid mapping (https://github.com/pytorch/ci-infra/issues/96)
+  WORKSPACE_ORIGINAL_OWNER_ID=$(stat -c '%u' "/var/lib/jenkins/workspace")
+  cleanup_workspace() {
+    echo "sudo may print the following warning message that can be ignored. The chown command will still run."
+    echo "    sudo: setrlimit(RLIMIT_STACK): Operation not permitted"
+    echo "For more details refer to https://github.com/sudo-project/sudo/issues/42"
+    sudo chown -R "$WORKSPACE_ORIGINAL_OWNER_ID" /var/lib/jenkins/workspace
+  }
+  # Disable shellcheck SC2064 as we want to parse the original owner immediately.
+  # shellcheck disable=SC2064
+  trap_add cleanup_workspace EXIT
+  sudo chown -R jenkins /var/lib/jenkins/workspace
+  git config --global --add safe.directory /var/lib/jenkins/workspace
+fi
+
 if [[ "$BUILD_ENVIRONMENT" == *-bazel-* ]]; then
   set -e
 
@@ -248,13 +266,17 @@ else
   ( ! get_exit_code python setup.py clean bad_argument )
 
   if [[ "$BUILD_ENVIRONMENT" != *libtorch* ]]; then
-
     # rocm builds fail when WERROR=1
     # XLA test build fails when WERROR=1
     # set only when building other architectures
     # or building non-XLA tests.
     if [[ "$BUILD_ENVIRONMENT" != *rocm*  &&
           "$BUILD_ENVIRONMENT" != *xla* ]]; then
+      if [[ "$BUILD_ENVIRONMENT" != *py3.8* ]]; then
+        # Install numpy-2.0 release candidate for builds
+        # Which should be backward compatible with Numpy-1.X
+        python -mpip install --pre numpy==2.0.0rc1
+      fi
       WERROR=1 python setup.py bdist_wheel
     else
       python setup.py bdist_wheel

.ci/pytorch/common_utils.sh

@@ -159,7 +159,7 @@ function install_torchvision() {
 }
 
 function install_tlparse() {
-  pip_install --user "tlparse==0.3.5"
+  pip_install --user "tlparse==0.3.7"
   PATH="$(python -m site --user-base)/bin:$PATH"
 }
 

.ci/pytorch/multigpu-test.sh

@@ -45,6 +45,7 @@ time python test/run_test.py --verbose -i distributed/test_device_mesh
 time python test/run_test.py --verbose -i distributed/tensor/parallel/test_ddp_2d_parallel
 time python test/run_test.py --verbose -i distributed/tensor/parallel/test_fsdp_2d_parallel
 time python test/run_test.py --verbose -i distributed/tensor/parallel/test_tp_examples
+time python test/run_test.py --verbose -i distributed/tensor/parallel/test_tp_random_state.py
 
 # Other tests
 time python test/run_test.py --verbose -i test_cuda_primary_ctx

.ci/pytorch/perf_test/compare_with_baseline.py

@@ -59,16 +59,16 @@ print("sample mean: ", sample_mean)
 print("sample sigma: ", sample_sigma)
 
 if math.isnan(sample_mean):
-    raise Exception("""Error: sample mean is NaN""")
+    raise Exception("""Error: sample mean is NaN""")  # noqa: TRY002
 elif math.isnan(sample_sigma):
-    raise Exception("""Error: sample sigma is NaN""")
+    raise Exception("""Error: sample sigma is NaN""")  # noqa: TRY002
 
 z_value = (sample_mean - mean) / sigma
 
 print("z-value: ", z_value)
 
 if z_value >= 3:
-    raise Exception(
+    raise Exception(  # noqa: TRY002
         f"""\n
 z-value >= 3, there is high chance of perf regression.\n
 To reproduce this regression, run

.ci/pytorch/python_doc_push_script.sh

@@ -26,8 +26,8 @@ echo "error: python_doc_push_script.sh: version (arg2) not specified"
 fi
 
 # Argument 1: Where to copy the built documentation to
-# (pytorch.github.io/$install_path)
-install_path="${1:-${DOCS_INSTALL_PATH:-docs/${DOCS_VERSION}}}"
+# (pytorch_docs/$install_path)
+install_path="${1:-${DOCS_INSTALL_PATH:-${DOCS_VERSION}}}"
 if [ -z "$install_path" ]; then
 echo "error: python_doc_push_script.sh: install_path (arg1) not specified"
   exit 1
@@ -68,8 +68,8 @@ build_docs () {
 }
 
 
-git clone https://github.com/pytorch/pytorch.github.io -b "$branch" --depth 1
-pushd pytorch.github.io
+git clone https://github.com/pytorch/docs pytorch_docs -b "$branch" --depth 1
+pushd pytorch_docs
 
 export LC_ALL=C
 export PATH=/opt/conda/bin:$PATH
@@ -105,6 +105,7 @@ if [ "$is_main_doc" = true ]; then
     echo undocumented objects found:
     cat build/coverage/python.txt
     echo "Make sure you've updated relevant .rsts in docs/source!"
+    echo "You can reproduce locally by running 'cd docs && make coverage && cat build/coverage/python.txt'"
     exit 1
   fi
 else

.ci/pytorch/test.sh

@@ -6,6 +6,27 @@
 
 set -ex
 
+# shellcheck source=./common.sh
+source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
+
+# Do not change workspace permissions for ROCm CI jobs
+# as it can leave workspace with bad permissions for cancelled jobs
+if [[ "$BUILD_ENVIRONMENT" != *rocm* ]]; then
+  # Workaround for dind-rootless userid mapping (https://github.com/pytorch/ci-infra/issues/96)
+  WORKSPACE_ORIGINAL_OWNER_ID=$(stat -c '%u' "/var/lib/jenkins/workspace")
+  cleanup_workspace() {
+    echo "sudo may print the following warning message that can be ignored. The chown command will still run."
+    echo "    sudo: setrlimit(RLIMIT_STACK): Operation not permitted"
+    echo "For more details refer to https://github.com/sudo-project/sudo/issues/42"
+    sudo chown -R "$WORKSPACE_ORIGINAL_OWNER_ID" /var/lib/jenkins/workspace
+  }
+  # Disable shellcheck SC2064 as we want to parse the original owner immediately.
+  # shellcheck disable=SC2064
+  trap_add cleanup_workspace EXIT
+  sudo chown -R jenkins /var/lib/jenkins/workspace
+  git config --global --add safe.directory /var/lib/jenkins/workspace
+fi
+
 echo "Environment variables:"
 env
 
@@ -90,9 +111,6 @@ if [[ -n $TESTS_TO_INCLUDE ]]; then
   INCLUDE_CLAUSE="--include $TESTS_TO_INCLUDE"
 fi
 
-# shellcheck source=./common.sh
-source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
-
 echo "Environment variables"
 env
 
@@ -211,8 +229,6 @@ if [[ "$BUILD_ENVIRONMENT" == *asan* ]]; then
     export LD_PRELOAD=/usr/lib/llvm-15/lib/clang/15.0.7/lib/linux/libclang_rt.asan-x86_64.so
     # Disable valgrind for asan
     export VALGRIND=OFF
-    # Increase stack size, because ASAN red zones use more stack
-    ulimit -s 81920
 
     (cd test && python -c "import torch; print(torch.__version__, torch.version.git_version)")
     echo "The next four invocations are expected to crash; if they don't that means ASAN/UBSAN is misconfigured"
@@ -299,9 +315,12 @@ test_inductor_distributed() {
   pytest test/distributed/_composable/fsdp/test_fully_shard_training.py -k test_train_parity_multi_group
   pytest test/distributed/_composable/fsdp/test_fully_shard_training.py -k test_train_parity_with_activation_checkpointing
   pytest test/distributed/_composable/fsdp/test_fully_shard_training.py -k test_train_parity_2d_mlp
+  pytest test/distributed/_composable/fsdp/test_fully_shard_training.py -k test_train_parity_hsdp
+  pytest test/distributed/_composable/fsdp/test_fully_shard_training.py -k test_train_parity_2d_transformer_checkpoint_resume
   pytest test/distributed/_composable/fsdp/test_fully_shard_frozen.py
   pytest test/distributed/_composable/fsdp/test_fully_shard_mixed_precision.py -k test_compute_dtype
   pytest test/distributed/_composable/fsdp/test_fully_shard_mixed_precision.py -k test_reduce_dtype
+  pytest test/distributed/fsdp/test_fsdp_tp_integration.py -k test_fsdp_tp_integration
 
   # this runs on both single-gpu and multi-gpu instance. It should be smart about skipping tests that aren't supported
   # with if required # gpus aren't available
@@ -313,13 +332,13 @@ test_inductor() {
   python tools/dynamo/verify_dynamo.py
   python test/run_test.py --inductor --include test_modules test_ops test_ops_gradients test_torch --verbose
   # Do not add --inductor for the following inductor unit tests, otherwise we will fail because of nested dynamo state
-  python test/run_test.py --include inductor/test_torchinductor inductor/test_torchinductor_opinfo --verbose
+  python test/run_test.py --include inductor/test_torchinductor inductor/test_torchinductor_opinfo inductor/test_aot_inductor --verbose
 
   # docker build uses bdist_wheel which does not work with test_aot_inductor
   # TODO: need a faster way to build
   if [[ "$BUILD_ENVIRONMENT" != *rocm* ]]; then
       BUILD_AOT_INDUCTOR_TEST=1 python setup.py develop
-      CPP_TESTS_DIR="${BUILD_BIN_DIR}" LD_LIBRARY_PATH="${TORCH_LIB_DIR}" python test/run_test.py --cpp --verbose -i cpp/test_aot_inductor
+      CPP_TESTS_DIR="${BUILD_BIN_DIR}" LD_LIBRARY_PATH="${TORCH_LIB_DIR}" python test/run_test.py --cpp --verbose -i cpp/test_aoti_abi_check cpp/test_aoti_inference
   fi
 }
 
@@ -432,6 +451,17 @@ test_perf_for_dashboard() {
             "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" "$@" \
             --output "$TEST_REPORTS_DIR/${backend}_max_autotune_${suite}_${dtype}_${mode}_cuda_${target}.csv"
       fi
+      if [[ "$DASHBOARD_TAG" == *cudagraphs_low_precision-true* ]] && [[ "$mode" == "inference" ]]; then
+        # TODO: This has a new dtype called quant and the benchmarks script needs to be updated to support this.
+        # The tentative command is as follows. It doesn't work now, but it's ok because we only need mock data
+        # to fill the dashboard.
+        python "benchmarks/dynamo/$suite.py" \
+          "${target_flag[@]}" --"$mode" --quant --backend "$backend" "$@" \
+          --output "$TEST_REPORTS_DIR/${backend}_cudagraphs_low_precision_${suite}_quant_${mode}_cuda_${target}.csv" || true
+        # Copy cudagraph results as mock data, easiest choice?
+        cp "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_${suite}_${dtype}_${mode}_cuda_${target}.csv" \
+          "$TEST_REPORTS_DIR/${backend}_cudagraphs_low_precision_${suite}_quant_${mode}_cuda_${target}.csv"
+      fi
     done
   done
 }
@@ -593,6 +623,12 @@ test_inductor_torchbench_cpu_smoketest_perf(){
   done
 }
 
+test_torchbench_gcp_smoketest(){
+  pushd "${TORCHBENCHPATH}"
+  python test.py -v
+  popd
+}
+
 test_python_gloo_with_tls() {
   source "$(dirname "${BASH_SOURCE[0]}")/run_glootls_test.sh"
   assert_git_not_dirty
@@ -1172,6 +1208,9 @@ elif [[ "${TEST_CONFIG}" == *torchbench* ]]; then
       llama_v2_7b_16h resnet50 timm_efficientnet mobilenet_v3_large timm_resnest \
       shufflenet_v2_x1_0 hf_GPT2
     PYTHONPATH=$(pwd)/torchbench test_inductor_torchbench_cpu_smoketest_perf
+  elif [[ "${TEST_CONFIG}" == *torchbench_gcp_smoketest* ]]; then
+    checkout_install_torchbench
+    TORCHBENCHPATH=$(pwd)/torchbench test_torchbench_gcp_smoketest
   else
     checkout_install_torchbench
     # Do this after checkout_install_torchbench to ensure we clobber any

.clang-tidy

@@ -36,6 +36,7 @@ hicpp-exception-baseclass,
 hicpp-avoid-goto,
 misc-*,
 -misc-const-correctness,
+-misc-include-cleaner,
 -misc-use-anonymous-namespace,
 -misc-unused-parameters,
 -misc-no-recursion,

.gitattributes

@@ -4,3 +4,4 @@
 .github/generated-* linguist-generated=true
 .github/scripts/gql_mocks.json linguist-generated=true
 third_party/LICENSES_BUNDLED.txt linguist-generated=true
+tools/build/bazel/requirements.txt linguist-generated=true

.github/ISSUE_TEMPLATE/pt2-bug-report.yml

@@ -8,7 +8,18 @@ body:
       value: >
         #### Before submitting a bug, please make sure the issue hasn't been already addressed by searching through [the
         existing and past issues](https://github.com/pytorch/pytorch/issues)
-        It's likely that your bug will be resolved by checking our FAQ or troubleshooting guide [documentation](https://pytorch.org/docs/master/dynamo/index.html)
+        It's likely that your bug will be resolved by checking our FAQ or troubleshooting guide [documentation](https://pytorch.org/docs/main/dynamo/index.html)
+
+        Note: if you're submitting an issue that you generated from a fuzzer. Please do the following:
+
+        - Ensure rtol/atol are at default tolerances
+
+        - Dont compare indices of max/min etc, because that avoids the above requirement
+
+        - If comparing eager and torch.compile at fp16/bf16, you should use fp32 as baseline
+
+        If the above requirements are met, add the label "topic: fuzzer" to your issue.
+
   - type: textarea
     attributes:
       label: 🐛 Describe the bug
@@ -33,7 +44,7 @@ body:
       label: Minified repro
       description: |
         Please run the minifier on your example and paste the minified code below
-        Learn more here https://pytorch.org/docs/master/compile/troubleshooting.html
+        Learn more here https://pytorch.org/docs/main/torch.compiler_troubleshooting.html
       placeholder: |
         env TORCHDYNAMO_REPRO_AFTER="aot" python your_model.py
         or

.github/actions/download-build-artifacts/action.yml

@@ -9,6 +9,10 @@ inputs:
   use-gha:
     description: If set to any value, use GHA to download the artifact. Otherwise use s3.
     required: false
+  s3-bucket:
+    description: S3 bucket to download builds
+    required: false
+    default: "gha-artifacts"
 
 runs:
   using: composite
@@ -18,9 +22,10 @@ runs:
       uses: seemethere/download-artifact-s3@v4
       with:
         name: ${{ inputs.name }}
+        s3-bucket: ${{ inputs.s3-bucket }}
 
     - name: Download PyTorch Build Artifacts from GHA
-      if: inputs.use-gha
+      if: ${{ inputs.use-gha }}
       uses: actions/download-artifact@v3
       with:
         name: ${{ inputs.name }}
@@ -29,6 +34,10 @@ runs:
       shell: bash
       run: unzip -o artifacts.zip
 
+    - name: Remove artifacts.zip
+      shell: bash
+      run: rm artifacts.zip
+
     - name: Output disk space left
       shell: bash
       run: df -H

.github/actions/filter-test-configs/action.yml

@@ -13,6 +13,13 @@ inputs:
     required: true
     type: string
     description: JSON description of what test configs to run.
+  selected-test-configs:
+    required: false
+    type: string
+    description: |
+      A comma-separated list of test configurations from the test matrix to keep,
+      The empty list means we are going to keep every configurations by defaults
+    default: ""
   job-name:
     type: string
     required: false
@@ -40,6 +47,9 @@ outputs:
   ci-no-td:
     description: True if ci-no-td label was on PR or [ci-no-td] in PR body.
     value: ${{ steps.filter.outputs.ci-no-td }}
+  ci-td-distributed:
+    description: True if ci-td-distributed label was on PR or [ci-td-distributed] in PR body.
+    value: ${{ steps.filter.outputs.ci-td-distributed }}
 
 runs:
   using: composite
@@ -123,6 +133,7 @@ runs:
           --workflow "${GITHUB_WORKFLOW}" \
           --job-name "${JOB_NAME}" \
           --test-matrix "${{ inputs.test-matrix }}" \
+          --selected-test-configs "${{ inputs.selected-test-configs }}" \
           --pr-number "${PR_NUMBER}" \
           --tag "${TAG}" \
           --event-name "${EVENT_NAME}" \

.github/actions/linux-build/action.yml

@@ -0,0 +1,207 @@
+name: linux-build
+
+inputs:
+  build-environment:
+    required: true
+    description: Top-level label for what's being built/tested.
+  docker-image-name:
+    required: true
+    description: Name of the base docker image to build with.
+  build-generates-artifacts:
+    required: false
+    default: "true"
+    description: If set, upload generated build artifacts.
+  build-with-debug:
+    required: false
+    default: "false"
+    description: If set, build in debug mode.
+  sync-tag:
+    required: false
+    default: ""
+    description: |
+      If this is set, our linter will use this to make sure that every other
+      job with the same `sync-tag` is identical.
+  cuda-arch-list:
+    required: false
+    default: "5.2"
+    description: Runner label to select worker type
+  runner:
+    required: false
+    default: "linux.2xlarge"
+    description: |
+      List of CUDA architectures CI build should target.
+  test-matrix:
+    required: false
+    type: string
+    description: |
+      An option JSON description of what test configs to run later on. This
+      is moved here from the Linux test workflow so that we can apply filter
+      logic using test-config labels earlier and skip unnecessary builds
+  s3-bucket:
+    description: S3 bucket to download artifact
+    required: false
+    default: "gha-artifacts"
+  aws-role-to-assume:
+    description: role to assume for downloading artifacts
+    required: false
+    default: ""
+  GITHUB_TOKEN:
+    description: GitHub token
+    required: true
+  HUGGING_FACE_HUB_TOKEN:
+    description: Hugging Face Hub token
+    required: false
+    default: ""
+outputs:
+  docker-image:
+    value: ${{ steps.calculate-docker-image.outputs.docker-image }}
+    description: The docker image containing the built PyTorch.
+  test-matrix:
+    value: ${{ steps.filter.outputs.test-matrix }}
+    description: An optional JSON description of what test configs to run later on.
+
+runs:
+  using: composite
+  steps:
+    - name: Setup Linux
+      uses: ./.github/actions/setup-linux
+
+    - name: configure aws credentials
+      uses: aws-actions/configure-aws-credentials@v3
+      if: ${{ inputs.aws-role-to-assume != '' }}
+      with:
+        role-to-assume: ${{ inputs.aws-role-to-assume }}
+        role-session-name: gha-linux-build
+        role-duration-seconds: 10800
+        aws-region: us-east-1
+
+    - name: Calculate docker image
+      id: calculate-docker-image
+      uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+      with:
+        docker-image-name: ${{ inputs.docker-image-name }}
+
+    - name: Use following to pull public copy of the image
+      id: print-ghcr-mirror
+      env:
+        ECR_DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
+      shell: bash
+      run: |
+        tag=${ECR_DOCKER_IMAGE##*/}
+        echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}"
+
+    - name: Pull docker image
+      uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+      with:
+        docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
+
+    - name: Parse ref
+      id: parse-ref
+      shell: bash
+      run: .github/scripts/parse_ref.py
+
+    - name: Get workflow job id
+      id: get-job-id
+      uses: ./.github/actions/get-workflow-job-id
+      if: always()
+      with:
+        github-token: ${{ inputs.GITHUB_TOKEN }}
+
+    # Apply the filter logic to the build step too if the test-config label is already there
+    - name: Select all requested test configurations (if the test matrix is available)
+      id: filter
+      uses: ./.github/actions/filter-test-configs
+      with:
+        github-token: ${{ inputs.GITHUB_TOKEN }}
+        test-matrix: ${{ inputs.test-matrix }}
+        job-name: ${{ steps.get-job-id.outputs.job-name }}
+
+    - name: Download pytest cache
+      uses: ./.github/actions/pytest-cache-download
+      continue-on-error: true
+      with:
+        cache_dir: .pytest_cache
+        job_identifier: ${{ github.workflow }}_${{ inputs.build-environment }}
+        s3_bucket: ${{ inputs.s3-bucket }}
+
+    - name: Build
+      if: steps.filter.outputs.is-test-matrix-empty == 'False' || inputs.test-matrix == ''
+      id: build
+      env:
+        BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
+        BRANCH: ${{ steps.parse-ref.outputs.branch }}
+        # TODO duplicated
+        AWS_DEFAULT_REGION: us-east-1
+        PR_NUMBER: ${{ github.event.pull_request.number }}
+        SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
+        SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
+        SCCACHE_S3_KEY_PREFIX: ${{ github.workflow }}
+        XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla
+        PR_LABELS: ${{ toJson(github.event.pull_request.labels.*.name) }}
+        TORCH_CUDA_ARCH_LIST: ${{ inputs.cuda-arch-list }}
+        DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
+        XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }}
+        DEBUG: ${{ inputs.build-with-debug == 'true' && '1' || '0' }}
+        OUR_GITHUB_JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
+        HUGGING_FACE_HUB_TOKEN: ${{ inputs.HUGGING_FACE_HUB_TOKEN }}
+      shell: bash
+      run: |
+        # detached container should get cleaned up by teardown_ec2_linux
+        container_name=$(docker run \
+          -e BUILD_ENVIRONMENT \
+          -e MAX_JOBS="$(nproc --ignore=2)" \
+          -e AWS_DEFAULT_REGION \
+          -e PR_NUMBER \
+          -e SHA1 \
+          -e BRANCH \
+          -e SCCACHE_BUCKET \
+          -e SCCACHE_S3_KEY_PREFIX \
+          -e XLA_CUDA \
+          -e XLA_CLANG_CACHE_S3_BUCKET_NAME \
+          -e SKIP_SCCACHE_INITIALIZATION=1 \
+          -e TORCH_CUDA_ARCH_LIST \
+          -e PR_LABELS \
+          -e OUR_GITHUB_JOB_ID \
+          -e HUGGING_FACE_HUB_TOKEN \
+          --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \
+          --security-opt seccomp=unconfined \
+          --cap-add=SYS_PTRACE \
+          --tty \
+          --detach \
+          --user jenkins \
+          -v "${GITHUB_WORKSPACE}:/var/lib/jenkins/workspace" \
+          -w /var/lib/jenkins/workspace \
+          "${DOCKER_IMAGE}"
+        )
+        docker exec -t "${container_name}" sh -c '.ci/pytorch/build.sh'
+
+    - name: Archive artifacts into zip
+      if: inputs.build-generates-artifacts == 'true' && steps.build.outcome != 'skipped'
+      shell: bash
+      run: |
+        zip -1 -r artifacts.zip dist/ build/custom_test_artifacts build/lib build/bin .additional_ci_files
+
+    - name: Store PyTorch Build Artifacts on S3
+      uses: seemethere/upload-artifact-s3@v5
+      if: inputs.build-generates-artifacts == 'true' && steps.build.outcome != 'skipped'
+      with:
+        name: ${{ inputs.build-environment }}
+        retention-days: 14
+        if-no-files-found: error
+        path: artifacts.zip
+        s3-bucket: ${{ inputs.s3-bucket }}
+
+    - name: Upload sccache stats
+      if: steps.build.outcome != 'skipped'
+      uses: seemethere/upload-artifact-s3@v5
+      with:
+        s3-prefix: |
+          ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact
+        retention-days: 365
+        if-no-files-found: warn
+        path: sccache-stats-*.json
+        s3-bucket: ${{ inputs.s3-bucket }}
+
+    - name: Teardown Linux
+      uses: pytorch/test-infra/.github/actions/teardown-linux@main
+      if: always()

.github/actions/linux-test/action.yml

@@ -0,0 +1,384 @@
+name: linux-test
+
+inputs:
+  build-environment:
+    required: true
+    type: string
+    description: Top-level label for what's being built/tested.
+  test-matrix:
+    required: true
+    type: string
+    description: JSON description of what test configs to run.
+  docker-image:
+    required: true
+    type: string
+    description: Docker image to run in.
+  sync-tag:
+    required: false
+    type: string
+    default: ""
+    description: |
+      If this is set, our linter will use this to make sure that every other
+      job with the same `sync-tag` is identical.
+  use-gha:
+    required: false
+    type: string
+    default: ""
+    description: If set to any value, upload to GHA. Otherwise upload to S3.
+  dashboard-tag:
+    required: false
+    type: string
+    default: ""
+  s3-bucket:
+    description: S3 bucket to download artifact
+    required: false
+    type: string
+    default: "gha-artifacts"
+  aws-role-to-assume:
+    description: role to assume for downloading artifacts
+    required: false
+    type: string
+    default: ""
+  HUGGING_FACE_HUB_TOKEN:
+    description: |
+      HF Auth token to avoid rate limits when downloading models or datasets from hub
+    required: false
+    default: ""
+  GITHUB_TOKEN:
+    description: GitHub token
+    required: true
+
+#env:
+#  GIT_DEFAULT_BRANCH: ${{ inputs.default_branch }}
+
+runs:
+  using: composite
+  steps:
+    - name: Setup Linux
+      uses: ./.github/actions/setup-linux
+
+    - name: configure aws credentials
+      if : ${{ inputs.aws-role-to-assume != '' }}
+      uses: aws-actions/configure-aws-credentials@v3
+      with:
+        role-to-assume: ${{ inputs.aws-role-to-assume }}
+        role-session-name: gha-linux-test
+        aws-region: us-east-1
+
+    - name: Calculate docker image
+      id: calculate-docker-image
+      uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+      with:
+        docker-image-name: ${{ inputs.docker-image }}
+
+    - name: Use following to pull public copy of the image
+      id: print-ghcr-mirror
+      env:
+        ECR_DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
+      shell: bash
+      run: |
+        tag=${ECR_DOCKER_IMAGE##*/}
+        echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}"
+
+    - name: Pull docker image
+      uses: pytorch/test-infra/.github/actions/pull-docker-image@main
+      with:
+        docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
+
+    - name: Check if in a ARC runner
+      shell: bash
+      id: check_arc_runner
+      run: echo "IN_ARC_RUNNER=$([ -f /.inarc ] && echo true || echo false)" >> "$GITHUB_OUTPUT"
+
+    - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
+      id: install-nvidia-driver
+      uses: pytorch/test-infra/.github/actions/setup-nvidia@main
+      if: ${{ contains(inputs.build-environment, 'cuda') && !contains(matrix.config, 'nogpu') && steps.check_arc_runner.outputs.IN_ARC_RUNNER == 'false' }}
+
+    - name: Lock NVIDIA A100 40GB Frequency
+      shell: bash
+      run: |
+        sudo nvidia-smi -pm 1
+        sudo nvidia-smi -ac 1215,1410
+        nvidia-smi
+      if: contains(matrix.runner, 'a100')
+
+    - name: Start monitoring script
+      id: monitor-script
+      shell: bash
+      continue-on-error: true
+      run: |
+        python3 -m pip install psutil==5.9.1 nvidia-ml-py==11.525.84
+        python3 -m tools.stats.monitor > usage_log.txt 2>&1 &
+        echo "monitor-script-pid=${!}" >> "${GITHUB_OUTPUT}"
+
+    - name: Download build artifacts
+      uses: ./.github/actions/download-build-artifacts
+      with:
+        name: ${{ inputs.build-environment }}
+        s3-bucket: ${{ inputs.s3-bucket }}
+
+    - name: Download TD artifacts
+      continue-on-error: true
+      uses: ./.github/actions/download-td-artifacts
+
+    - name: Parse ref
+      id: parse-ref
+      shell: bash
+      run: .github/scripts/parse_ref.py
+
+    - name: Get workflow job id
+      id: get-job-id
+      uses: ./.github/actions/get-workflow-job-id
+      if: always()
+      with:
+        github-token: ${{ inputs.GITHUB_TOKEN }}
+
+    - name: Check for keep-going label and re-enabled test issues
+      # This uses the filter-test-configs action because it conviniently
+      # checks for labels and re-enabled test issues.  It does not actually do
+      # any filtering.  All filtering is done in the build step.
+      id: keep-going
+      uses: ./.github/actions/filter-test-configs
+      with:
+        github-token: ${{ inputs.GITHUB_TOKEN }}
+        test-matrix: ${{ inputs.test-matrix }}
+        job-name: ${{ steps.get-job-id.outputs.job-name }}
+
+    - name: Test
+      id: test
+      env:
+        BUILD_ENVIRONMENT: ${{ inputs.build-environment }}
+        PR_NUMBER: ${{ github.event.pull_request.number }}
+        GITHUB_REPOSITORY: ${{ github.repository }}
+        GITHUB_WORKFLOW: ${{ github.workflow }}
+        GITHUB_JOB: ${{ github.job }}
+        GITHUB_RUN_ID: ${{ github.run_id }}
+        GITHUB_RUN_NUMBER: ${{ github.run_number }}
+        GITHUB_RUN_ATTEMPT: ${{ github.run_attempt }}
+        JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
+        JOB_NAME: ${{ steps.get-job-id.outputs.job-name }}
+        BRANCH: ${{ steps.parse-ref.outputs.branch }}
+        SHA1: ${{ github.event.pull_request.head.sha || github.sha }}
+        BASE_SHA: ${{ github.event.pull_request.base.sha || github.sha }}
+        TEST_CONFIG: ${{ matrix.config }}
+        SHARD_NUMBER: ${{ matrix.shard }}
+        NUM_TEST_SHARDS: ${{ matrix.num_shards }}
+        REENABLED_ISSUES: ${{ steps.keep-going.outputs.reenabled-issues }}
+        CONTINUE_THROUGH_ERROR: ${{ steps.keep-going.outputs.keep-going }}
+        VERBOSE_TEST_LOGS: ${{ steps.keep-going.outputs.ci-verbose-test-logs }}
+        NO_TEST_TIMEOUT: ${{ steps.keep-going.outputs.ci-no-test-timeout }}
+        NO_TD: ${{ steps.keep-going.outputs.ci-no-td }}
+        TD_DISTRIBUTED: ${{ steps.keep-going.outputs.ci-td-distributed }}
+        SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
+        SCCACHE_S3_KEY_PREFIX: ${{ github.workflow }}
+        SHM_SIZE: ${{ contains(inputs.build-environment, 'cuda') && '2g' || '1g' }}
+        DOCKER_IMAGE: ${{ inputs.docker-image }}
+        XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }}
+        XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla
+        PYTORCH_TEST_CUDA_MEM_LEAK_CHECK: ${{ matrix.mem_leak_check && '1' || '0' }}
+        PYTORCH_TEST_RERUN_DISABLED_TESTS: ${{ matrix.rerun_disabled_tests && '1' || '0' }}
+        DASHBOARD_TAG: ${{ inputs.dashboard-tag }}
+        HUGGING_FACE_HUB_TOKEN: ${{ inputs.HUGGING_FACE_HUB_TOKEN }}
+      shell: bash
+      run: |
+        set -x
+
+        if [[ $TEST_CONFIG == 'multigpu' ]]; then
+          TEST_COMMAND=.ci/pytorch/multigpu-test.sh
+        elif [[ $BUILD_ENVIRONMENT == *onnx* ]]; then
+          TEST_COMMAND=.ci/onnx/test.sh
+        else
+          TEST_COMMAND=.ci/pytorch/test.sh
+        fi
+
+        # detached container should get cleaned up by teardown_ec2_linux
+        # TODO: Stop building test binaries as part of the build phase
+        # Used for GPU_FLAG since that doesn't play nice
+        # shellcheck disable=SC2086,SC2090
+        container_name=$(docker run \
+          ${GPU_FLAG:-} \
+          -e BUILD_ENVIRONMENT \
+          -e PR_NUMBER \
+          -e GITHUB_ACTIONS \
+          -e GITHUB_REPOSITORY \
+          -e GITHUB_WORKFLOW \
+          -e GITHUB_JOB \
+          -e GITHUB_RUN_ID \
+          -e GITHUB_RUN_NUMBER \
+          -e GITHUB_RUN_ATTEMPT \
+          -e JOB_ID \
+          -e JOB_NAME \
+          -e BASE_SHA \
+          -e BRANCH \
+          -e SHA1 \
+          -e AWS_DEFAULT_REGION \
+          -e IN_WHEEL_TEST \
+          -e SHARD_NUMBER \
+          -e TEST_CONFIG \
+          -e NUM_TEST_SHARDS \
+          -e REENABLED_ISSUES \
+          -e CONTINUE_THROUGH_ERROR \
+          -e VERBOSE_TEST_LOGS \
+          -e NO_TEST_TIMEOUT \
+          -e NO_TD \
+          -e TD_DISTRIBUTED \
+          -e PR_LABELS \
+          -e MAX_JOBS="$(nproc --ignore=2)" \
+          -e SCCACHE_BUCKET \
+          -e SCCACHE_S3_KEY_PREFIX \
+          -e XLA_CUDA \
+          -e XLA_CLANG_CACHE_S3_BUCKET_NAME \
+          -e PYTORCH_TEST_CUDA_MEM_LEAK_CHECK \
+          -e PYTORCH_TEST_RERUN_DISABLED_TESTS \
+          -e SKIP_SCCACHE_INITIALIZATION=1 \
+          -e HUGGING_FACE_HUB_TOKEN \
+          -e DASHBOARD_TAG \
+          --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \
+          --security-opt seccomp=unconfined \
+          --cap-add=SYS_PTRACE \
+          --ipc=host \
+          --shm-size="${SHM_SIZE}" \
+          --tty \
+          --detach \
+          --name="${container_name}" \
+          --user jenkins \
+          -v "${GITHUB_WORKSPACE}:/var/lib/jenkins/workspace" \
+          -w /var/lib/jenkins/workspace \
+          "${DOCKER_IMAGE}"
+        )
+        # Propagate download.pytorch.org IP to container
+        grep download.pytorch.org /etc/hosts | docker exec -i "${container_name}" sudo bash -c "/bin/cat >> /etc/hosts"
+        echo "DOCKER_CONTAINER_ID=${container_name}" >> "${GITHUB_ENV}"
+        docker exec -t "${container_name}" sh -c "pip install $(echo dist/*.whl)[opt-einsum] && ${TEST_COMMAND}"
+
+    - name: Upload pytest cache if tests failed
+      uses: ./.github/actions/pytest-cache-upload
+      continue-on-error: true
+      if: failure() && steps.test.conclusion && steps.test.conclusion == 'failure'
+      with:
+        cache_dir: .pytest_cache
+        shard: ${{ matrix.shard }}
+        sha: ${{ github.event.pull_request.head.sha || github.sha }}
+        test_config: ${{ matrix.config }}
+        job_identifier: ${{ github.workflow }}_${{ inputs.build-environment }}
+
+    - name: Print remaining test logs
+      shell: bash
+      if: always() && steps.test.conclusion
+      run: |
+        cat test/**/*_toprint.log || true
+
+    - name: Stop monitoring script
+      if: always() && steps.monitor-script.outputs.monitor-script-pid
+      shell: bash
+      continue-on-error: true
+      env:
+        MONITOR_SCRIPT_PID: ${{ steps.monitor-script.outputs.monitor-script-pid }}
+      run: |
+        kill "$MONITOR_SCRIPT_PID"
+
+    - name: Upload test artifacts
+      uses: ./.github/actions/upload-test-artifacts
+      if: always() && steps.test.conclusion && steps.test.conclusion != 'skipped'
+      with:
+        file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.get-job-id.outputs.job-id }}
+        use-gha: ${{ inputs.use-gha }}
+        s3-bucket: ${{ inputs.s3-bucket }}
+
+    - name: Collect backtraces from coredumps (if any)
+      if: always()
+      shell: bash
+      run: |
+        # shellcheck disable=SC2156
+        find . -iname "core.[1-9]*" -exec docker exec "${DOCKER_CONTAINER_ID}" sh -c "gdb python {} -ex 'bt' -ex 'q'" \;
+
+    - name: Store Core dumps on S3
+      uses: seemethere/upload-artifact-s3@v5
+      if: failure()
+      with:
+        name: coredumps-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}
+        retention-days: 14
+        if-no-files-found: ignore
+        path: ./**/core.[1-9]*
+
+    - name: Teardown Linux
+      uses: pytorch/test-infra/.github/actions/teardown-linux@main
+      if: always()
+
+    # NB: We are currently having an intermittent GPU-related issue on G5 runners with
+    # A10G GPU. Once this happens, trying to reset the GPU as done in setup-nvidia does
+    # not seem to help. Here are some symptoms:
+    #   * Calling nvidia-smi timeouts after 60 second
+    #   * Fail to run nvidia-smi with an unable to determine the device handle for GPU
+    #     unknown error
+    #   * Test fails with a missing CUDA GPU error when initializing CUDA in PyTorch
+    #   * Run docker --gpus all fails with error response from daemon
+    #
+    # As both the root cause and recovery path are unclear, let's take the runner out of
+    # service so that it doesn't get any more jobs
+    - name: Check NVIDIA driver installation step
+      if: failure() && steps.install-nvidia-driver.outcome && steps.install-nvidia-driver.outcome != 'skipped'
+      shell: bash
+      env:
+        RUNNER_WORKSPACE: ${{ runner.workspace }}
+      run: |
+        set +e
+        set -x
+
+        nvidia-smi
+        # NB: Surprisingly, nvidia-smi command returns successfully with return code 0 even in
+        # the case where the driver has already crashed as it still can get the driver version
+        # and some basic information like the bus ID.  However, the rest of the information
+        # would be missing (ERR!), for example:
+        #
+        # +-----------------------------------------------------------------------------+
+        # | NVIDIA-SMI 525.89.02    Driver Version: 525.89.02    CUDA Version: 12.0     |
+        # |-------------------------------+----------------------+----------------------+
+        # | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
+        # | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
+        # |                               |                      |               MIG M. |
+        # |===============================+======================+======================|
+        # |   0  ERR!                Off  | 00000000:00:1E.0 Off |                 ERR! |
+        # |ERR!  ERR! ERR!    ERR! / ERR! |   4184MiB / 23028MiB |    ERR!      Default |
+        # |                               |                      |                 ERR! |
+        # +-------------------------------+----------------------+----------------------+
+        #
+        # +-----------------------------------------------------------------------------+
+        # | Processes:                                                                  |
+        # |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
+        # |        ID   ID                                                   Usage      |
+        # |=============================================================================|
+        # +-----------------------------------------------------------------------------+
+        #
+        # This should be reported as a failure instead as it will guarantee to fail when
+        # Docker tries to run with --gpus all
+        #
+        # So, the correct check here is to query one of the missing piece of info like
+        # GPU name, so that the command can fail accordingly
+        nvidia-smi --query-gpu=gpu_name --format=csv,noheader --id=0
+        NVIDIA_SMI_STATUS=$?
+
+        # These are acceptable return code from nvidia-smi as copied from setup-nvidia GitHub action
+        if [ "$NVIDIA_SMI_STATUS" -ne 0 ] && [ "$NVIDIA_SMI_STATUS" -ne 14 ]; then
+          echo "NVIDIA driver installation has failed, shutting down the runner..."
+          .github/scripts/stop_runner_service.sh
+        fi
+
+        # For runner with multiple GPUs, we also want to confirm that the number of GPUs are the
+        # power of 2, i.e. 1, 2, 4, or 8. This is to avoid flaky test issue when one GPU fails
+        # https://github.com/pytorch/test-infra/issues/4000
+        GPU_COUNT=$(nvidia-smi --list-gpus | wc -l)
+        NVIDIA_SMI_STATUS=$?
+
+        # These are acceptable return code from nvidia-smi as copied from setup-nvidia GitHub action
+        if [ "$NVIDIA_SMI_STATUS" -ne 0 ] && [ "$NVIDIA_SMI_STATUS" -ne 14 ]; then
+          echo "NVIDIA driver installation has failed, shutting down the runner..."
+          .github/scripts/stop_runner_service.sh
+        fi
+
+        # Check the GPU count to be a power of 2
+        if [ "$GPU_COUNT" -le 8 ] && [ "$GPU_COUNT" -ne 1 ] && [ "$GPU_COUNT" -ne 2 ] && [ "$GPU_COUNT" -ne 4 ] && [ "$GPU_COUNT" -ne 8 ]; then
+          echo "NVIDIA driver detects $GPU_COUNT GPUs. The runner has a broken GPU, shutting it down..."
+          .github/scripts/stop_runner_service.sh
+        fi

.github/actions/pytest-cache-download/action.yml

@@ -9,6 +9,10 @@ inputs:
   job_identifier:
     description: Text that uniquely identifies a given job type within a workflow. All shards of a job should share the same job identifier.
     required: true
+  s3_bucket:
+    description: S3 bucket to download PyTest cache
+    required: false
+    default: "gha-artifacts"
 
 runs:
   using: composite
@@ -30,6 +34,7 @@ runs:
         CACHE_DIR: ${{ inputs.cache_dir }}
         JOB_IDENTIFIER: ${{ inputs.job_identifier }}
         REPO: ${{ github.repository }}
+        BUCKET: ${{ inputs.s3_bucket }}
       run: |
         python3 .github/scripts/pytest_cache.py \
           --download \
@@ -38,3 +43,4 @@ runs:
           --job_identifier $JOB_IDENTIFIER \
           --temp_dir $RUNNER_TEMP \
           --repo $REPO \
+          --bucket $BUCKET \

.github/actions/setup-linux/action.yml

@@ -15,10 +15,12 @@ runs:
           category=$1
           # If it is GCP runner (runner name contains gcp), do not run this
           runner_name_str=${{ runner.name }}
-          if [[ $runner_name_str != *"gcp"* ]]; then
-            curl -fsSL "http://169.254.169.254/latest/meta-data/${category}"
-          else
+          if [[ -f /.inarc ]]; then
+            echo "ARC Runner, no info on ec2 metadata"
+          elif [[ $runner_name_str == *"gcp"* ]]; then
             echo "Runner is from Google Cloud Platform, No info on ec2 metadata"
+          else
+            curl -fsSL "http://169.254.169.254/latest/meta-data/${category}"
           fi
         }
         echo "ami-id: $(get_ec2_metadata ami-id)"
@@ -26,8 +28,14 @@ runs:
         echo "instance-type: $(get_ec2_metadata instance-type)"
         echo "system info $(uname -a)"
 
+    - name: Check if in a ARC runner
+      shell: bash
+      id: check_arc_runner
+      run: echo "IN_ARC_RUNNER=$([ -f /.inarc ] && echo true || echo false)"  >> $GITHUB_OUTPUT
+
     - name: Start docker if docker deamon is not running
       shell: bash
+      if: ${{ steps.check_arc_runner.outputs.IN_ARC_RUNNER == 'false' }}
       run: |
         if systemctl is-active --quiet docker; then
             echo "Docker daemon is running...";
@@ -58,6 +66,7 @@ runs:
         env | grep '^CI' >> "/tmp/github_env_${GITHUB_RUN_ID}"
 
     - name: Kill any existing containers, clean up images
+      if: ${{ steps.check_arc_runner.outputs.IN_ARC_RUNNER == 'false' }}
       shell: bash
       run: |
         # ignore expansion of "docker ps -q" since it could be empty
@@ -96,3 +105,28 @@ runs:
 
         echo "${RESOLVED_IP} ${PT_DOMAIN}" | sudo tee -a /etc/hosts
         cat /etc/hosts
+
+    - name: Check that the docker daemon is running
+      shell: bash
+      continue-on-error: true
+      if: ${{ steps.check_arc_runner.outputs.IN_ARC_RUNNER == 'true' }}
+      run: |
+        set +x
+
+        max_attempts=30
+        delay=10
+        attempt=1
+
+        for attempt in $(seq 1 $max_attempts); do
+          echo "Attempt $attempt of $max_attempts: Checking if Docker daemon is running..."
+          if docker info > /dev/null 2>&1; then
+            echo "Docker is running. Proceeding with the next steps"
+            exit 0
+          else
+            echo "Docker is not running yet."
+            echo "Retrying in $delay seconds..."
+            sleep $delay
+          fi
+        done
+        echo "Reached maximum attempts to connect to Docker. Exiting."
+        exit 1

.github/actions/upload-test-artifacts/action.yml

@@ -11,6 +11,10 @@ inputs:
       Suffix to add to the filename of the artifacts. This should include the
       workflow job id, see [Job id in artifacts].
     required: true
+  s3-bucket:
+    description: S3 bucket to download builds
+    required: false
+    default: "gha-artifacts"
 
 runs:
   using: composite
@@ -42,7 +46,7 @@ runs:
       env:
         FILE_SUFFIX: ${{ inputs.file-suffix }}
       run: |
-        # Remove any previous test reports if they exist
+        # Remove any previous usage logs if they exist
         rm -f logs-*.zip
         # this workflow is also run in bazel build test, but we dont generate usage reports for it
         # so check to see if the file exists first
@@ -53,6 +57,18 @@ runs:
             zip -r "logs-${FILE_SUFFIX}.zip" test -i '*.log'
         fi
 
+    - name: Zip debugging artifacts for upload
+      if: runner.os != 'Windows' && !inputs.use-gha
+      shell: bash
+      env:
+        FILE_SUFFIX: ${{ inputs.file-suffix }}
+      run: |
+        # Remove any previous debugging artifacts if they exist
+        rm -f debug-*.zip
+        if [ -d 'test/debug' ]; then
+          zip -r "debug-${FILE_SUFFIX}.zip" test/debug
+        fi
+
     # Windows zip
     - name: Zip JSONs for upload
       if: runner.os == 'Windows' && !inputs.use-gha
@@ -87,6 +103,7 @@ runs:
       uses: seemethere/upload-artifact-s3@v5
       if: ${{ !inputs.use-gha }}
       with:
+        s3-bucket: ${{ inputs.s3-bucket }}
         s3-prefix: |
           ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact
         retention-days: 14
@@ -97,6 +114,7 @@ runs:
       uses: seemethere/upload-artifact-s3@v5
       if: ${{ !inputs.use-gha }}
       with:
+        s3-bucket: ${{ inputs.s3-bucket }}
         s3-prefix: |
           ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact
         retention-days: 14
@@ -108,12 +126,25 @@ runs:
       if: ${{ !inputs.use-gha }}
       continue-on-error: true
       with:
+        s3-bucket: ${{ inputs.s3-bucket }}
         s3-prefix: |
           ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact
         retention-days: 14
         if-no-files-found: ignore
         path: logs-*.zip
 
+    - name: Store Debug Artifacts on S3
+      uses: seemethere/upload-artifact-s3@v5
+      if: ${{ !inputs.use-gha }}
+      continue-on-error: true
+      with:
+        s3-bucket: ${{ inputs.s3-bucket }}
+        s3-prefix: |
+          ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact
+        retention-days: 14
+        if-no-files-found: ignore
+        path: debug-*.zip
+
     # GHA upload
     - name: Store Test Downloaded JSONs on Github
       uses: actions/upload-artifact@v3

.github/ci_commit_pins/audio.txt

@@ -1 +1 @@
-87aeb554d3e2f7855b7abe5120c282f59648ed7a
+ea437b31ce316ea3d66fe73768c0dcb94edb79ad

.github/ci_commit_pins/vision.txt

@@ -1 +1 @@
-2c127da8b5e2e8f44b50994c6cb931bcca267cfe
+2c4665ffbb64f03f5d18016d3398af4ac4da5f03

.github/ci_commit_pins/xla.txt

@@ -1 +1 @@
-707a632930bfde19ffb361cdf5c31a7682af4e67
+58a412cb271a3f98ae2e01fd1d24bdbb66645d4e

.github/label_to_label.yml

@@ -0,0 +1,13 @@
+# Use this to auto apply labels based on other labels.  Applies to both PRs and
+# issues. Currently only supports any and all
+- any:
+  - "module: custom operators"
+  - "module: aotdispatch"
+  then:
+  - "module: pt2-dispatcher"
+- any:
+  - "module: dynamo"
+  - "module: pt2-dispatcher"
+  - "module: inductor"
+  then:
+  - "oncall: pt2"

.github/labeler.yml

@@ -35,6 +35,9 @@
 - test/distributed/tensor/parallel/test_fsdp_2d_parallel.py
 - torch/distributed/_tensor/**
 - torch/distributed/fsdp/**
+- torch/csrc/inductor/**
+- test/cpp/aoti_abi_check/**
+- test/cpp/aoti_inference/**
 
 "module: cpu":
 - aten/src/ATen/cpu/**

.github/merge_rules.yaml

@@ -28,7 +28,6 @@
   - caffe2/python/onnx/**
   approved_by:
   - BowenBao
-  - abock
   - justinchuby
   - shubhambhokare1
   - thiagocrepaldi
@@ -236,6 +235,23 @@
   - Lint
   - pull
 
+- name: XPU ATen
+  patterns:
+  - aten/src/ATen/xpu/**
+  - c10/xpu/**
+  - torch/csrc/xpu/**
+  - torch/xpu/**
+  - test/xpu/**
+  - third_party/xpu.txt
+  approved_by:
+  - EikanWang
+  - jgong5
+  - gujinghui
+  mandatory_checks_name:
+  - EasyCLA
+  - Lint
+  - pull
+
 - name: Distributions
   patterns:
   - torch/distributions/**
@@ -357,12 +373,14 @@
 
 - name: CPU inductor
   patterns:
+  - torch/_inductor/mkldnn_lowerings.py
   - torch/_inductor/fx_passes/mkldnn_fusion.py
   - torch/_inductor/fx_passes/quantization.py
   - torch/_inductor/codegen/cpp.py
   - test/inductor/test_mkldnn_pattern_matcher.py
   - test/inductor/test_cpu_repo.py
   - test/inductor/test_cpu_cpp_wrapper.py
+  - aten/src/ATen/cpu/**
   - aten/src/ATen/native/quantized/cpu/**
   - test/quantization/core/test_quantized_op.py
   - torch/ao/quantization/quantizer/x86_inductor_quantizer.py

.github/pytorch-probot.yml

@@ -1,5 +1,6 @@
 tracking_issue: 24422
 ciflow_tracking_issue: 64124
+TD_rollout_issue: 123120
 ciflow_push_tags:
 - ciflow/binaries
 - ciflow/binaries_conda
@@ -15,9 +16,12 @@ ciflow_push_tags:
 - ciflow/trunk
 - ciflow/unstable
 - ciflow/xpu
+- ciflow/torchbench
 retryable_workflows:
 - lint
 - pull
 - trunk
 - linux-binary
 - windows-binary
+labeler_config: labeler.yml
+label_to_label_config: label_to_label.yml

.github/requirements-gha-cache.txt

@@ -5,7 +5,7 @@
 #   functorch/docs/requirements.txt
 #   .ci/docker/requirements-ci.txt
 boto3==1.19.12
-jinja2==3.0.1
+jinja2==3.1.3
 lintrunner==0.10.7
 ninja==1.10.0.post1
 nvidia-ml-py==11.525.84

.github/requirements/pip-requirements-iOS.txt

@@ -1,4 +1,4 @@
 # iOS simulator requirements
 coremltools==5.0b5
 protobuf==3.20.2
-optree==0.9.1
+optree==0.11.0

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

@@ -26,4 +26,7 @@ pytest-cpp==2.3.0
 rockset==1.0.3
 z3-solver==4.12.2.0
 tensorboard==2.13.0
-optree==0.9.1
+optree==0.11.0
+# NB: test_hparams_* from test_tensorboard is failing with protobuf 5.26.0 in
+# which the stringify metadata is wrong when escaping double quote
+protobuf==3.20.2

.github/scripts/amd/package_triton_wheel.sh

@@ -0,0 +1,99 @@
+set -ex
+
+# Set ROCM_HOME isn't available, use ROCM_PATH if set or /opt/rocm
+ROCM_HOME="${ROCM_HOME:-${ROCM_PATH:-/opt/rocm}}"
+
+# Find rocm_version.h header file for ROCm version extract
+rocm_version_h="${ROCM_HOME}/include/rocm-core/rocm_version.h"
+if [ ! -f "$rocm_version_h" ]; then
+    rocm_version_h="${ROCM_HOME}/include/rocm_version.h"
+fi
+
+# Error out if rocm_version.h not found
+if [ ! -f "$rocm_version_h" ]; then
+    echo "Error: rocm_version.h not found in expected locations." >&2
+    exit 1
+fi
+
+# Extract major, minor and patch ROCm version numbers
+MAJOR_VERSION=$(grep 'ROCM_VERSION_MAJOR' "$rocm_version_h" | awk '{print $3}')
+MINOR_VERSION=$(grep 'ROCM_VERSION_MINOR' "$rocm_version_h" | awk '{print $3}')
+PATCH_VERSION=$(grep 'ROCM_VERSION_PATCH' "$rocm_version_h" | awk '{print $3}')
+ROCM_INT=$(($MAJOR_VERSION * 10000 + $MINOR_VERSION * 100 + $PATCH_VERSION))
+echo "ROCm version: $ROCM_INT"
+
+# Check TRITON_ROCM_DIR is set
+if [[ -z "${TRITON_ROCM_DIR}" ]]; then
+    export TRITON_ROCM_DIR=third_party/amd/backend
+fi
+
+# Remove packaged libs and headers
+rm -rf $TRITON_ROCM_DIR/include/*
+
+LIBTINFO_PATH="/usr/lib64/libtinfo.so.5"
+LIBNUMA_PATH="/usr/lib64/libnuma.so.1"
+LIBELF_PATH="/usr/lib64/libelf.so.1"
+
+OS_SO_PATHS=(
+    $LIBELF_PATH
+    $LIBNUMA_PATH
+    $LIBTINFO_PATH
+)
+
+for lib in "${OS_SO_PATHS[@]}"
+do
+    cp $lib $TRITON_ROCM_DIR/lib/
+done
+
+# Required ROCm libraries
+if [[ "${MAJOR_VERSION}" == "6" ]]; then
+    libamdhip="libamdhip64.so.6"
+else
+    libamdhip="libamdhip64.so.5"
+fi
+
+# Required ROCm libraries - ROCm 6.0
+ROCM_SO=(
+    "${libamdhip}"
+    "libhsa-runtime64.so.1"
+    "libamd_comgr.so.2"
+    "libdrm.so.2"
+    "libdrm_amdgpu.so.1"
+)
+
+if [[ $ROCM_INT -ge 60100 ]]; then
+    ROCM_SO+=("librocprofiler-register.so.0")
+fi
+
+for lib in "${ROCM_SO[@]}"
+do
+    file_path=($(find $ROCM_HOME/lib/ -name "$lib")) # First search in lib
+    if [[ -z $file_path ]]; then
+        if [ -d "$ROCM_HOME/lib64/" ]; then
+            file_path=($(find $ROCM_HOME/lib64/ -name "$lib")) # Then search in lib64
+        fi
+    fi
+    if [[ -z $file_path ]]; then
+        file_path=($(find $ROCM_HOME/ -name "$lib")) # Then search in ROCM_HOME
+    fi
+    if [[ -z $file_path ]]; then
+        file_path=($(find /opt/ -name "$lib")) # Then search in /opt
+    fi
+    if [[ -z $file_path ]]; then
+            echo "Error: Library file $lib is not found." >&2
+            exit 1
+    fi
+
+    cp $file_path $TRITON_ROCM_DIR/lib
+    # When running locally, and not building a wheel, we need to satisfy shared objects requests that don't look for versions
+    LINKNAME=$(echo $lib | sed -e 's/\.so.*/.so/g')
+    ln -sf $lib $TRITON_ROCM_DIR/lib/$LINKNAME
+
+done
+
+# Copy Include Files
+cp -r $ROCM_HOME/include/hip $TRITON_ROCM_DIR/include
+
+# Copy linker
+mkdir -p $TRITON_ROCM_DIR/llvm/bin
+cp $ROCM_HOME/llvm/bin/ld.lld $TRITON_ROCM_DIR/llvm/bin/

.github/scripts/amd/patch_triton_wheel.sh

@@ -0,0 +1,99 @@
+#!/bin/bash
+set -x
+
+WHEELHOUSE_DIR=/artifacts
+PATCHELF_BIN=patchelf
+ROCM_LIB=backends/amd/lib
+ROCM_LD=backends/amd/llvm/bin
+PREFIX=triton
+fname_without_so_number() {
+    LINKNAME=$(echo $1 | sed -e 's/\.so.*/.so/g')
+    echo "$LINKNAME"
+}
+
+replace_needed_sofiles() {
+    find $1 -name '*.so*' -o -name 'ld.lld' | while read sofile; do
+        origname=$2
+        patchedname=$3
+        if [[ "$origname" != "$patchedname" ]]; then
+            set +e
+            origname=$($PATCHELF_BIN --print-needed $sofile | grep "$origname.*")
+            ERRCODE=$?
+            set -e
+            if [ "$ERRCODE" -eq "0" ]; then
+                echo "patching $sofile entry $origname to $patchedname"
+                $PATCHELF_BIN --replace-needed $origname $patchedname $sofile
+            fi
+        fi
+    done
+}
+
+mkdir  -p "/tmp_dir"
+pushd /tmp_dir
+for pkg in /$WHEELHOUSE_DIR/*triton*.whl; do
+    echo "Modifying $pkg"
+    rm -rf tmp
+    mkdir -p tmp
+    cd tmp
+    cp $pkg .
+    unzip -q $(basename $pkg)
+    rm -f $(basename $pkg)
+    $PATCHELF_BIN --set-rpath ${LD_SO_RPATH:-'$ORIGIN:$ORIGIN/../../lib'} $PREFIX/$ROCM_LD/ld.lld
+    $PATCHELF_BIN --print-rpath $PREFIX/$ROCM_LD/ld.lld
+    # Modify libtriton.so as it sits in _C directory apart from its dependencies
+    find $PREFIX/_C -type f -name "*.so*" | while read sofile; do
+        echo "Setting rpath of $sofile"
+        $PATCHELF_BIN --set-rpath ${C_SO_RPATH:-'$ORIGIN:$ORIGIN/'../$ROCM_LIB} ${FORCE_RPATH:-} $sofile
+        $PATCHELF_BIN --print-rpath $sofile
+    done
+
+    # All included dependencies are included in a single lib directory
+    deps=()
+    deps_soname=()
+    while read sofile; do
+        echo "Setting rpath of $sofile to ${LIB_SO_RPATH:-'$ORIGIN'}"
+        $PATCHELF_BIN --set-rpath ${LIB_SO_RPATH:-'$ORIGIN'} ${FORCE_RPATH:-} $sofile
+        $PATCHELF_BIN --print-rpath $sofile
+        deps+=("$sofile")
+        deps_soname+=("$(basename $sofile)")
+    done < <(find $PREFIX/$ROCM_LIB -type f -name "*.so*")
+
+    patched=()
+    for filepath in "${deps[@]}"; do
+        filename=$(basename $filepath)
+        destpath=$PREFIX/$ROCM_LIB/$filename
+        if [[ "$filepath" != "$destpath" ]]; then
+            cp $filepath $destpath
+        fi
+        patchedpath=$(fname_without_so_number $destpath)
+        patchedname=$(basename $patchedpath)
+        if [[ "$destpath" != "$patchedpath" ]]; then
+            mv $destpath $patchedpath
+        fi
+        patched+=("$patchedname")
+        echo "Copied $filepath to $patchedpath"
+    done
+
+    # Go through all required shared objects and see if any of our other objects are dependants.  If so, replace so.ver wth so
+    for ((i=0;i<${#deps[@]};++i)); do
+        echo "replacing "${deps_soname[i]} ${patched[i]}
+        replace_needed_sofiles $PREFIX/$ROCM_LIB ${deps_soname[i]} ${patched[i]}
+        replace_needed_sofiles $PREFIX/_C ${deps_soname[i]} ${patched[i]}
+        replace_needed_sofiles $PREFIX/$ROCM_LD ${deps_soname[i]} ${patched[i]}
+    done
+
+    # Re-bundle whl with so adjustments
+    zip -rqy $(basename $pkg) *
+
+    if [[ -z "${MANYLINUX_VERSION}" ]]; then
+        newpkg=$pkg
+    else
+        newpkg=$(echo $pkg | sed -e "s/\linux_x86_64/${MANYLINUX_VERSION}/g")
+    fi
+
+    # Remove original whl
+    rm -f $pkg
+
+    # Move rebuilt whl to original location with new name.
+    mv $(basename $pkg) $newpkg
+done

.github/scripts/build_triton_wheel.py

@@ -10,9 +10,6 @@ from typing import Optional
 SCRIPT_DIR = Path(__file__).parent
 REPO_DIR = SCRIPT_DIR.parent.parent
 
-# TODO: Remove me once Triton version is again in sync for vanilla and ROCm
-ROCM_TRITION_VERSION = "2.1.0"
-
 
 def read_triton_pin(rocm_hash: bool = False) -> str:
     triton_file = "triton.txt" if not rocm_hash else "triton-rocm.txt"
@@ -32,27 +29,6 @@ def check_and_replace(inp: str, src: str, dst: str) -> str:
     return inp.replace(src, dst)
 
 
-def patch_setup_py(
-    path: Path,
-    *,
-    version: str,
-    name: str = "triton",
-    expected_version: Optional[str] = None,
-) -> None:
-    with open(path) as f:
-        orig = f.read()
-    # Replace name
-    orig = check_and_replace(orig, 'name="triton",', f'name="{name}",')
-    # Replace version
-    if not expected_version:
-        expected_version = read_triton_version()
-    orig = check_and_replace(
-        orig, f'version="{expected_version}",', f'version="{version}",'
-    )
-    with open(path, "w") as f:
-        f.write(orig)
-
-
 def patch_init_py(
     path: Path, *, version: str, expected_version: Optional[str] = None
 ) -> None:
@@ -92,14 +68,20 @@ def build_triton(
     with TemporaryDirectory() as tmpdir:
         triton_basedir = Path(tmpdir) / "triton"
         triton_pythondir = triton_basedir / "python"
+        triton_repo = "https://github.com/openai/triton"
         if build_rocm:
-            triton_repo = "https://github.com/ROCmSoftwarePlatform/triton"
             triton_pkg_name = "pytorch-triton-rocm"
         else:
-            triton_repo = "https://github.com/openai/triton"
             triton_pkg_name = "pytorch-triton"
         check_call(["git", "clone", triton_repo], cwd=tmpdir)
-        check_call(["git", "checkout", commit_hash], cwd=triton_basedir)
+        if release:
+            ver, rev, patch = version.split(".")
+            check_call(
+                ["git", "checkout", f"release/{ver}.{rev}.x"], cwd=triton_basedir
+            )
+        else:
+            check_call(["git", "checkout", commit_hash], cwd=triton_basedir)
+
         if build_conda:
             with open(triton_basedir / "meta.yaml", "w") as meta:
                 print(
@@ -155,18 +137,15 @@ def build_triton(
         patch_init_py(
             triton_pythondir / "triton" / "__init__.py",
             version=f"{version}",
-            expected_version=ROCM_TRITION_VERSION if build_rocm else None,
+            expected_version=None,
         )
 
         if build_rocm:
-            # TODO: Remove me when ROCM triton is updated
-            patch_setup_py(
-                triton_pythondir / "setup.py",
-                name=triton_pkg_name,
-                version=f"{version}",
-                expected_version=ROCM_TRITION_VERSION,
+            check_call(
+                [f"{SCRIPT_DIR}/amd/package_triton_wheel.sh"],
+                cwd=triton_basedir,
+                shell=True,
             )
-            check_call("scripts/amd/setup_rocm_libs.sh", cwd=triton_basedir, shell=True)
             print("ROCm libraries setup for triton installation...")
 
         check_call(
@@ -177,8 +156,11 @@ def build_triton(
         shutil.copy(whl_path, Path.cwd())
 
         if build_rocm:
-            check_call("scripts/amd/fix_so.sh", cwd=triton_basedir, shell=True)
-
+            check_call(
+                [f"{SCRIPT_DIR}/amd/patch_triton_wheel.sh"],
+                cwd=triton_basedir,
+                shell=True,
+            )
         return Path.cwd() / whl_path.name
 
 

.github/scripts/cherry_pick.py

@@ -29,7 +29,7 @@ def parse_args() -> Any:
         "--onto-branch", type=str, required=True, help="the target release branch"
     )
     parser.add_argument(
-        "--github-actor", type=str, required=True, help="all the world’s a stage"
+        "--github-actor", type=str, required=True, help="all the world's a stage"
     )
     parser.add_argument(
         "--classification",

.github/scripts/comment_on_pr.py

@@ -23,8 +23,10 @@ def main() -> None:
 
     job_link = f"[job]({run_url})" if run_url is not None else "job"
     msg = (
-        f"The {args.action} {job_link} was canceled. If you believe this is a mistake,"
-        + f" then you can re trigger it through [pytorch-bot]({BOT_COMMANDS_WIKI})."
+        f"The {args.action} {job_link} was canceled or timed out. This most often happen if two merge requests were issued"
+        + " for the same PR, or if merge job was waiting for more than 6 hours for tests to finish."
+        + " In later case, please do not hesitate to reissue the merge command\n"
+        + f" For more information see [pytorch-bot wiki]({BOT_COMMANDS_WIKI})."
     )
 
     gh_post_pr_comment(org, project, args.pr_num, msg)

.github/scripts/delete_old_branches.py

@@ -18,7 +18,7 @@ ESTIMATED_TOKENS = [0]
 
 TOKEN = os.environ["GITHUB_TOKEN"]
 if not TOKEN:
-    raise Exception("GITHUB_TOKEN is not set")
+    raise Exception("GITHUB_TOKEN is not set")  # noqa: TRY002
 
 REPO_ROOT = Path(__file__).parent.parent.parent
 

.github/scripts/filter_test_configs.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import json
+import logging
 import os
 import re
 import subprocess
@@ -8,6 +9,7 @@ import sys
 import warnings
 from enum import Enum
 from functools import lru_cache
+from logging import info
 from typing import Any, Callable, Dict, List, Optional, Set
 from urllib.request import Request, urlopen
 
@@ -17,33 +19,7 @@ REENABLE_TEST_REGEX = "(?i)(Close(d|s)?|Resolve(d|s)?|Fix(ed|es)?) (#|https://gi
 
 PREFIX = "test-config/"
 
-# Same as shard names
-VALID_TEST_CONFIG_LABELS = {
-    f"{PREFIX}{label}"
-    for label in {
-        "backwards_compat",
-        "crossref",
-        "default",
-        "deploy",
-        "distributed",
-        "docs_tests",
-        "dynamo",
-        "force_on_cpu",
-        "functorch",
-        "inductor",
-        "inductor_distributed",
-        "inductor_huggingface",
-        "inductor_timm",
-        "inductor_torchbench",
-        "jit_legacy",
-        "multigpu",
-        "nogpu_AVX512",
-        "nogpu_NO_AVX2",
-        "slow",
-        "tsan",
-        "xla",
-    }
-}
+logging.basicConfig(level=logging.INFO)
 
 
 def is_cuda_or_rocm_job(job_name: Optional[str]) -> bool:
@@ -90,6 +66,12 @@ def parse_args() -> Any:
     parser.add_argument(
         "--test-matrix", type=str, required=True, help="the original test matrix"
     )
+    parser.add_argument(
+        "--selected-test-configs",
+        type=str,
+        default="",
+        help="a comma-separated list of test configurations from the test matrix to keep",
+    )
     parser.add_argument(
         "--workflow", type=str, help="the name of the current workflow, i.e. pull"
     )
@@ -155,19 +137,25 @@ def get_labels(pr_number: int) -> Set[str]:
     }
 
 
+def filter_labels(labels: Set[str], label_regex: Any) -> Set[str]:
+    """
+    Return the list of matching labels
+    """
+    return {l for l in labels if re.match(label_regex, l)}
+
+
 def filter(test_matrix: Dict[str, List[Any]], labels: Set[str]) -> Dict[str, List[Any]]:
     """
     Select the list of test config to run from the test matrix. The logic works
     as follows:
 
-    If the PR has one or more labels as specified in the VALID_TEST_CONFIG_LABELS set, only
-    these test configs will be selected.  This also works with ciflow labels, for example,
-    if a PR has both ciflow/trunk and test-config/functorch, only trunk functorch builds
-    and tests will be run
+    If the PR has one or more test-config labels as specified, only these test configs
+    will be selected.  This also works with ciflow labels, for example, if a PR has both
+    ciflow/trunk and test-config/functorch, only trunk functorch builds and tests will
+    be run.
 
     If the PR has none of the test-config label, all tests are run as usual.
     """
-
     filtered_test_matrix: Dict[str, List[Any]] = {"include": []}
 
     for entry in test_matrix.get("include", []):
@@ -177,23 +165,46 @@ def filter(test_matrix: Dict[str, List[Any]], labels: Set[str]) -> Dict[str, Lis
 
         label = f"{PREFIX}{config_name.strip()}"
         if label in labels:
-            print(
-                f"Select {config_name} because label {label} is presented in the pull request by the time the test starts"
-            )
+            msg = f"Select {config_name} because label {label} is present in the pull request by the time the test starts"
+            info(msg)
             filtered_test_matrix["include"].append(entry)
 
-    valid_test_config_labels = labels.intersection(VALID_TEST_CONFIG_LABELS)
-
-    if not filtered_test_matrix["include"] and not valid_test_config_labels:
-        # Found no valid label and the filtered test matrix is empty, return the same
+    test_config_labels = filter_labels(labels, re.compile(f"{PREFIX}.+"))
+    if not filtered_test_matrix["include"] and not test_config_labels:
+        info("Found no test-config label on the PR, so all test configs are included")
+        # Found no test-config label and the filtered test matrix is empty, return the same
         # test matrix as before so that all tests can be run normally
         return test_matrix
     else:
+        msg = f"Found {test_config_labels} on the PR so only these test configs are run"
+        info(msg)
         # When the filter test matrix contain matches or if a valid test config label
         # is found in the PR, return the filtered test matrix
         return filtered_test_matrix
 
 
+def filter_selected_test_configs(
+    test_matrix: Dict[str, List[Any]], selected_test_configs: Set[str]
+) -> Dict[str, List[Any]]:
+    """
+    Keep only the selected configs if the list if not empty. Otherwise, keep all test configs.
+    This filter is used when the workflow is dispatched manually.
+    """
+    if not selected_test_configs:
+        return test_matrix
+
+    filtered_test_matrix: Dict[str, List[Any]] = {"include": []}
+    for entry in test_matrix.get("include", []):
+        config_name = entry.get("config", "")
+        if not config_name:
+            continue
+
+        if config_name in selected_test_configs:
+            filtered_test_matrix["include"].append(entry)
+
+    return filtered_test_matrix
+
+
 def set_periodic_modes(
     test_matrix: Dict[str, List[Any]], job_name: Optional[str]
 ) -> Dict[str, List[Any]]:
@@ -374,30 +385,33 @@ def process_jobs(
         # - If the target record has the job (config) name, only that test config
         #   will be skipped or marked as unstable
         if not target_job_cfg:
-            print(
+            msg = (
                 f"Issue {target_url} created by {author} has {issue_type.value} "
                 + f"all CI jobs for {workflow} / {job_name}"
             )
+            info(msg)
             return _filter_jobs(
                 test_matrix=test_matrix,
                 issue_type=issue_type,
             )
 
         if target_job_cfg == BUILD_JOB_NAME:
-            print(
+            msg = (
                 f"Issue {target_url} created by {author} has {issue_type.value} "
                 + f"the build job for {workflow} / {job_name}"
             )
+            info(msg)
             return _filter_jobs(
                 test_matrix=test_matrix,
                 issue_type=issue_type,
             )
 
         if target_job_cfg in (TEST_JOB_NAME, BUILD_AND_TEST_JOB_NAME):
-            print(
+            msg = (
                 f"Issue {target_url} created by {author} has {issue_type.value} "
                 + f"all the test jobs for {workflow} / {job_name}"
             )
+            info(msg)
             return _filter_jobs(
                 test_matrix=test_matrix,
                 issue_type=issue_type,
@@ -463,7 +477,7 @@ def parse_reenabled_issues(s: Optional[str]) -> List[str]:
 
 
 def get_reenabled_issues(pr_body: str = "") -> List[str]:
-    default_branch = os.getenv("GIT_DEFAULT_BRANCH", "main")
+    default_branch = f"origin/{os.environ.get('GIT_DEFAULT_BRANCH', 'main')}"
     try:
         commit_messages = subprocess.check_output(
             f"git cherry -v {default_branch}".split(" ")
@@ -494,10 +508,15 @@ def perform_misc_tasks(
         "ci-no-test-timeout", check_for_setting(labels, pr_body, "ci-no-test-timeout")
     )
     set_output("ci-no-td", check_for_setting(labels, pr_body, "ci-no-td"))
+    # Only relevant for the one linux distributed cuda job, delete this when TD
+    # is rolled out completely
+    set_output(
+        "ci-td-distributed", check_for_setting(labels, pr_body, "ci-td-distributed")
+    )
 
     # Obviously, if the job name includes unstable, then this is an unstable job
     is_unstable = job_name and IssueType.UNSTABLE.value in job_name
-    if not is_unstable and test_matrix:
+    if not is_unstable and test_matrix and test_matrix.get("include"):
         # Even when the job name doesn't mention unstable, we will also mark it as
         # unstable when the test matrix only includes unstable jobs. Basically, this
         # logic allows build or build-and-test jobs to be marked as unstable too.
@@ -567,6 +586,16 @@ def main() -> None:
         # No PR number, no tag, we can just return the test matrix as it is
         filtered_test_matrix = test_matrix
 
+    if args.selected_test_configs:
+        selected_test_configs = {
+            v.strip().lower()
+            for v in args.selected_test_configs.split(",")
+            if v.strip()
+        }
+        filtered_test_matrix = filter_selected_test_configs(
+            filtered_test_matrix, selected_test_configs
+        )
+
     if args.event_name == "schedule" and args.schedule == "29 8 * * *":
         # we don't want to run the mem leak check or disabled tests on normal
         # periodically scheduled jobs, only the ones at this time

.github/scripts/generate_ci_workflows.py

@@ -378,7 +378,9 @@ def main() -> None:
     for template, workflows in template_and_workflows:
         # added Iterable check to appease the mypy gods
         if not isinstance(workflows, Iterable):
-            raise Exception(f"How is workflows not iterable? {workflows}")
+            raise Exception(  # noqa: TRY002
+                f"How is workflows not iterable? {workflows}"
+            )  # noqa: TRY002
         for workflow in workflows:
             workflow.generate_workflow_file(workflow_template=template)
 

.github/scripts/get_workflow_job_id.py

@@ -4,6 +4,7 @@
 
 import argparse
 import json
+import operator
 import os
 import re
 import sys
@@ -126,7 +127,7 @@ def find_job_id_name(args: Any) -> Tuple[str, str]:
 
     # Sort the jobs list by start time, in descending order. We want to get the most
     # recently scheduled job on the runner.
-    jobs.sort(key=lambda job: job["started_at"], reverse=True)
+    jobs.sort(key=operator.itemgetter("started_at"), reverse=True)
 
     for job in jobs:
         if job["runner_name"] == args.runner_name:

.github/scripts/lintrunner.sh

@@ -6,6 +6,9 @@ CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
 eval "$(command conda 'shell.bash' 'hook' 2> /dev/null)"
 conda activate "${CONDA_ENV}"
 
+# Use uv to speed up lintrunner init
+python3 -m pip install uv
+
 CACHE_DIRECTORY="/tmp/.lintbin"
 # Try to recover the cached binaries
 if [[ -d "${CACHE_DIRECTORY}" ]]; then

.github/scripts/td_llm_indexer.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+set -euxo pipefail
+
+# Download requirements
+cd llm-target-determinator
+pip install -q -r requirements.txt
+cd ../codellama
+pip install -e .
+
+# Run indexer
+cd ../llm-target-determinator
+
+torchrun \
+    --standalone \
+    --nnodes=1 \
+    --nproc-per-node=1 \
+    indexer.py \
+    --experiment-name indexer-files \
+    --granularity FILE

.github/scripts/test_filter_test_configs.py

@@ -9,6 +9,7 @@ from unittest import main, mock, TestCase
 import yaml
 from filter_test_configs import (
     filter,
+    filter_selected_test_configs,
     get_labels,
     mark_unstable_jobs,
     parse_reenabled_issues,
@@ -17,7 +18,6 @@ from filter_test_configs import (
     remove_disabled_jobs,
     set_periodic_modes,
     SUPPORTED_PERIODICAL_MODES,
-    VALID_TEST_CONFIG_LABELS,
 )
 
 
@@ -273,13 +273,13 @@ class TestConfigFilter(TestCase):
         testcases = [
             {
                 "test_matrix": '{include: [{config: "default", runner: "linux"}]}',
-                "expected": '{"include": [{"config": "default", "runner": "linux"}]}',
-                "description": "No match, keep the same test matrix",
+                "expected": '{"include": []}',
+                "description": "Request test-config/cfg but the test matrix doesn't have it",
             },
             {
                 "test_matrix": '{include: [{config: "default", runner: "linux"}, {config: "plain-cfg"}]}',
-                "expected": '{"include": [{"config": "default", "runner": "linux"}, {"config": "plain-cfg"}]}',
-                "description": "No match because there is no prefix or suffix, keep the same test matrix",
+                "expected": '{"include": []}',
+                "description": "A valid test config label needs to start with test-config/",
             },
             {
                 "test_matrix": '{include: [{config: "default", runner: "linux"}, {config: "cfg", shard: 1}]}',
@@ -294,9 +294,8 @@ class TestConfigFilter(TestCase):
             )
             self.assertEqual(case["expected"], json.dumps(filtered_test_matrix))
 
-    def test_filter_with_valid_label(self) -> None:
+    def test_filter_with_test_config_label(self) -> None:
         mocked_labels = {f"{PREFIX}cfg", "ciflow/trunk"}
-        VALID_TEST_CONFIG_LABELS.add(f"{PREFIX}cfg")
 
         testcases = [
             {
@@ -317,6 +316,51 @@ class TestConfigFilter(TestCase):
             )
             self.assertEqual(case["expected"], json.dumps(filtered_test_matrix))
 
+    def test_filter_selected_test_configs(self) -> None:
+        testcases = [
+            {
+                "test_matrix": '{include: [{config: "default"}]}',
+                "selected_test_configs": "",
+                "expected": '{"include": [{"config": "default"}]}',
+                "description": "No selected test configs",
+            },
+            {
+                "test_matrix": '{include: [{config: "default"}]}',
+                "selected_test_configs": "foo",
+                "expected": '{"include": []}',
+                "description": "A different test config is selected",
+            },
+            {
+                "test_matrix": '{include: [{config: "default"}]}',
+                "selected_test_configs": "foo, bar",
+                "expected": '{"include": []}',
+                "description": "A different set of test configs is selected",
+            },
+            {
+                "test_matrix": '{include: [{config: "default"}]}',
+                "selected_test_configs": "foo, bar,default",
+                "expected": '{"include": [{"config": "default"}]}',
+                "description": "One of the test config is selected",
+            },
+            {
+                "test_matrix": '{include: [{config: "default"}, {config: "bar"}]}',
+                "selected_test_configs": "foo, bar,Default",
+                "expected": '{"include": [{"config": "default"}, {"config": "bar"}]}',
+                "description": "Several test configs are selected",
+            },
+        ]
+
+        for case in testcases:
+            selected_test_configs = {
+                v.strip().lower()
+                for v in case["selected_test_configs"].split(",")
+                if v.strip()
+            }
+            filtered_test_matrix = filter_selected_test_configs(
+                yaml.safe_load(case["test_matrix"]), selected_test_configs
+            )
+            self.assertEqual(case["expected"], json.dumps(filtered_test_matrix))
+
     def test_set_periodic_modes(self) -> None:
         testcases: List[Dict[str, str]] = [
             {
@@ -641,6 +685,7 @@ class TestConfigFilter(TestCase):
             ci_verbose_test_logs: bool = False,
             ci_no_test_timeout: bool = False,
             ci_no_td: bool = False,
+            ci_td_distributed: bool = False,
             is_unstable: bool = False,
             reenabled_issues: str = "",
         ) -> str:
@@ -649,6 +694,7 @@ class TestConfigFilter(TestCase):
                 f"ci-verbose-test-logs={ci_verbose_test_logs}\n"
                 f"ci-no-test-timeout={ci_no_test_timeout}\n"
                 f"ci-no-td={ci_no_td}\n"
+                f"ci-td-distributed={ci_td_distributed}\n"
                 f"is-unstable={is_unstable}\n"
                 f"reenabled-issues={reenabled_issues}\n"
             )

.github/scripts/test_trymerge.py

@@ -205,7 +205,6 @@ def mocked_read_merge_rules(repo: Any, org: str, project: str) -> List[MergeRule
             approved_by=["pytorch/metamates", "ngimel"],
             mandatory_checks_name=[
                 "Lint",
-                "Facebook CLA Check",
                 "pull / linux-xenial-cuda11.3-py3.7-gcc7 / build",
             ],
             ignore_flaky_failures=True,
@@ -398,7 +397,7 @@ class TestTryMerge(TestCase):
     def test_gql_retrieve_checksuites(self, *args: Any) -> None:
         "Fetch comments and conclusions for PR with 60 commits"
         pr = GitHubPR("pytorch", "pytorch", 94787)
-        self.assertEqual(len(pr.get_checkrun_conclusions()), 183)
+        self.assertEqual(len(pr.get_checkrun_conclusions()), 182)
 
     def test_team_members(self, *args: Any) -> None:
         "Test fetching team members works"
@@ -742,6 +741,30 @@ class TestBypassFailures(TestCase):
         self.assertTrue(len(failed) == 0)
         self.assertTrue(len(ignorable["UNSTABLE"]) == 1)
 
+        # Add another test case where there is no unstable keyword in the job name, but
+        # the job has already been marked as unstable
+        pr = GitHubPR("pytorch", "executorch", 3318)
+        checks = pr.get_checkrun_conclusions()
+        checks = get_classifications(
+            pr.pr_num,
+            pr.project,
+            checks,
+            [],
+        )
+        print(checks)
+        workflow_name = "test-llama-app"
+        job_name = "mobile-job (android)"
+        self.assertTrue(
+            checks[f"Android / {workflow_name} / {job_name}"].classification
+            == "UNSTABLE"
+        )
+        pending, failed, ignorable = categorize_checks(
+            checks, list(checks.keys()), ok_failed_checks_threshold=1
+        )
+        self.assertTrue(len(pending) == 0)
+        self.assertTrue(len(failed) == 0)
+        self.assertTrue(len(ignorable["UNSTABLE"]) == 1)
+
     def test_get_classifications_broken_trunk(self, *args: Any) -> None:
         # The mock merge base is the actual value returned by gh_fetch_merge_base
         test_cases = [
@@ -833,6 +856,41 @@ class TestBypassFailures(TestCase):
         self.assertTrue(len(ignorable["FLAKY"]) == 4)
         self.assertTrue(len(ignorable["BROKEN_TRUNK"]) == 2)
 
+    def test_get_classifications_wrong_workflow_name(self, *args: Any) -> None:
+        pr = GitHubPR("pytorch", "pytorch", 123104)
+        checks = pr.get_checkrun_conclusions()
+
+        check_name = "linux-binary-conda / conda-py3_8-cuda11_8-build / build"
+        check_name_workflow_path = ".github/workflows/generated-linux-binary-conda-nightly.yml / conda-py3_8-cuda11_8-build / build"
+
+        # Mock a check where the workflow name uses the full path
+        checks[check_name_workflow_path] = JobCheckState(
+            check_name_workflow_path,
+            checks[check_name].url,
+            checks[check_name].status,
+            checks[check_name].classification,
+            checks[check_name].job_id,
+            checks[check_name].title,
+            checks[check_name].summary,
+        )
+        del checks[check_name]
+
+        checks = get_classifications(
+            pr.pr_num,
+            pr.project,
+            checks,
+            [],
+        )
+        pending, failed, ignorable = categorize_checks(
+            checks,
+            list(checks.keys()),
+        )
+
+        self.assertTrue(len(pending) == 0)
+        self.assertTrue(len(failed) == 0)
+        self.assertTrue(len(ignorable["FLAKY"]) == 1)
+        self.assertTrue(len(ignorable["BROKEN_TRUNK"]) == 0)
+
     @mock.patch("trymerge.read_merge_rules", side_effect=xla_merge_rules)
     def test_dont_ignore_flaky_failures(self, *args: Any) -> None:
         """

.github/scripts/trymerge.py

@@ -123,6 +123,7 @@ fragment PRCheckSuites on CheckSuiteConnection {
         workflow {
           name
         }
+        databaseId
         url
       }
       checkRuns(first: 50) {
@@ -1398,7 +1399,10 @@ def find_matching_merge_rule(
         )
         required_checks = list(
             filter(
-                lambda x: "EasyCLA" in x or not skip_mandatory_checks, mandatory_checks
+                lambda x: ("EasyCLA" in x)
+                or ("Facebook CLA Check" in x)
+                or not skip_mandatory_checks,
+                mandatory_checks,
             )
         )
         pending_checks, failed_checks, _ = categorize_checks(
@@ -1409,6 +1413,13 @@ def find_matching_merge_rule(
             else 0,
         )
 
+        # categorize_checks assumes all tests are required if required_checks is empty.
+        # this is a workaround as we want to keep that behavior for categorize_checks
+        # generally.
+        if not required_checks:
+            pending_checks = []
+            failed_checks = []
+
         hud_link = f"https://hud.pytorch.org/{pr.org}/{pr.project}/commit/{pr.last_commit()['oid']}"
         if len(failed_checks) > 0:
             if reject_reason_score < 30000:
@@ -1608,28 +1619,59 @@ def remove_job_name_suffix(name: str, replacement: str = ")") -> str:
 
 
 def is_broken_trunk(
-    name: str,
+    check: JobCheckState,
     drci_classifications: Any,
 ) -> bool:
-    if not name or not drci_classifications:
+    if not check or not drci_classifications:
         return False
 
+    name = check.name
+    job_id = check.job_id
+
     # Consult the list of broken trunk failures from Dr.CI
     return any(
-        name == broken_trunk["name"]
+        (name == broken_trunk["name"]) or (job_id and job_id == broken_trunk["id"])
         for broken_trunk in drci_classifications.get("BROKEN_TRUNK", [])
     )
 
 
-def is_flaky(
-    name: str,
+def is_unstable(
+    check: JobCheckState,
     drci_classifications: Any,
 ) -> bool:
-    if not name or not drci_classifications:
+    if not check or not drci_classifications:
         return False
 
+    name = check.name
+    job_id = check.job_id
+
+    # The job name has the unstable keyword. This is the original way to mark a job
+    # as unstable on HUD, Dr.CI, and trymerge
+    if "unstable" in name:
+        return True
+
+    # Consult the list of unstable failures from Dr.CI
+    return any(
+        (name == unstable["name"] or (job_id and job_id == unstable["id"]))
+        for unstable in drci_classifications.get("UNSTABLE", [])
+    )
+
+
+def is_flaky(
+    check: JobCheckState,
+    drci_classifications: Any,
+) -> bool:
+    if not check or not drci_classifications:
+        return False
+
+    name = check.name
+    job_id = check.job_id
+
     # Consult the list of flaky failures from Dr.CI
-    return any(name == flaky["name"] for flaky in drci_classifications.get("FLAKY", []))
+    return any(
+        (name == flaky["name"] or (job_id and job_id == flaky["id"]))
+        for flaky in drci_classifications.get("FLAKY", [])
+    )
 
 
 def is_invalid_cancel(
@@ -1702,7 +1744,7 @@ def get_classifications(
         if check.status == "SUCCESS" or check.status == "NEUTRAL":
             continue
 
-        if "unstable" in name:
+        if is_unstable(check, drci_classifications):
             checks_with_classifications[name] = JobCheckState(
                 check.name,
                 check.url,
@@ -1716,7 +1758,7 @@ def get_classifications(
 
         # NB: It's important to note that when it comes to ghstack and broken trunk classification,
         # Dr.CI uses the base of the whole stack
-        if is_broken_trunk(name, drci_classifications):
+        if is_broken_trunk(check, drci_classifications):
             checks_with_classifications[name] = JobCheckState(
                 check.name,
                 check.url,
@@ -1728,7 +1770,7 @@ def get_classifications(
             )
             continue
 
-        elif is_flaky(name, drci_classifications):
+        elif is_flaky(check, drci_classifications):
             checks_with_classifications[name] = JobCheckState(
                 check.name,
                 check.url,

.github/scripts/tryrebase.py

@@ -60,7 +60,7 @@ def rebase_onto(
     repo._run_git("rebase", onto_branch, branch)
 
     if repo.rev_parse(branch) == repo.rev_parse(onto_branch):
-        raise Exception(SAME_SHA_ERROR)
+        raise Exception(SAME_SHA_ERROR)  # noqa: TRY002
 
     if dry_run:
         push_result = repo._run_git("push", "--dry-run", "-f", remote_url, refspec)
@@ -100,7 +100,7 @@ def rebase_ghstack_onto(
     repo._run_git("rebase", onto_branch, orig_ref)
 
     if repo.rev_parse(orig_ref) == repo.rev_parse(onto_branch):
-        raise Exception(SAME_SHA_ERROR)
+        raise Exception(SAME_SHA_ERROR)  # noqa: TRY002
 
     # steal the identity of the committer of the commit on the orig branch
     email = repo._run_git("log", orig_ref, "--pretty=format:%ae", "-1")
@@ -126,7 +126,7 @@ def rebase_ghstack_onto(
         print(push_result)
         if ghstack_result.returncode != 0:
             print(ghstack_result.stderr.decode("utf-8"))
-            raise Exception(f"\n```{push_result}```")
+            raise Exception(f"\n```{push_result}```")  # noqa: TRY002
         # The contents of a successful push result should look like:
         # Summary of changes (ghstack 0.6.0)
 

.github/workflows/_bazel-build-test.yml

@@ -86,9 +86,14 @@ jobs:
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
+      - name: Check if in a ARC runner
+        shell: bash
+        id: check_arc_runner
+        run: echo "IN_ARC_RUNNER=$([ -f /.inarc ] && echo true || echo false)" >> "$GITHUB_OUTPUT"
+
       - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
         uses: pytorch/test-infra/.github/actions/setup-nvidia@main
-        if: ${{ inputs.cuda-version != 'cpu' }}
+        if: ${{ inputs.cuda-version != 'cpu' && steps.check_arc_runner.outputs.IN_ARC_RUNNER == 'false' }}
 
       - name: Output disk space left
         run: |

.github/workflows/_binary-build-linux.yml

@@ -78,7 +78,7 @@ on:
 jobs:
   build:
     runs-on: ${{ inputs.runs_on }}
-    timeout-minutes: 180
+    timeout-minutes: 210
     env:
       PYTORCH_ROOT: ${{ inputs.PYTORCH_ROOT }}
       BUILDER_ROOT: ${{ inputs.BUILDER_ROOT }}

.github/workflows/_docs.yml

@@ -28,7 +28,21 @@ on:
         description: |
           If this is set, our linter will use this to make sure that every other
           job with the same `sync-tag` is identical.
-
+      s3-bucket:
+        description: S3 bucket to download artifact
+        required: false
+        type: string
+        default: "gha-artifacts"
+      aws-role-to-assume:
+        description: role to assume for downloading artifacts
+        required: false
+        type: string
+        default: ""
+      upload-aws-role-to-assume:
+        description: role to assume for downloading artifacts
+        required: false
+        type: string
+        default: ""
     secrets:
       GH_PYTORCHBOT_TOKEN:
         required: false
@@ -82,6 +96,14 @@ jobs:
       - name: Setup Linux
         uses: ./.github/actions/setup-linux
 
+      - name: configure aws credentials
+        if : ${{ inputs.aws-role-to-assume != '' }}
+        uses: aws-actions/configure-aws-credentials@v3
+        with:
+          role-to-assume: ${{ inputs.aws-role-to-assume }}
+          role-session-name: gha-linux-test
+          aws-region: us-east-1
+
       - name: Calculate docker image
         id: calculate-docker-image
         uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
@@ -97,6 +119,7 @@ jobs:
         uses: ./.github/actions/download-build-artifacts
         with:
           name: ${{ inputs.build-environment }}
+          s3-bucket: ${{ inputs.s3-bucket }}
 
       - name: Generate netrc (only for docs-push)
         if: inputs.push
@@ -156,6 +179,14 @@ jobs:
         uses: ./.github/actions/chown-workspace
         if: always()
 
+      - name: configure aws credentials
+        if : ${{ inputs.upload-aws-role-to-assume != '' }}
+        uses: aws-actions/configure-aws-credentials@v3
+        with:
+          role-to-assume: ${{ inputs.upload-aws-role-to-assume }}
+          role-session-name: gha-linux-test
+          aws-region: us-east-1
+
       - name: Upload Python Docs Preview
         uses: seemethere/upload-artifact-s3@v5
         if: ${{ github.event_name == 'pull_request' && matrix.docs_type == 'python' && steps.build-docs.outcome == 'success' }}
@@ -163,7 +194,7 @@ jobs:
           retention-days: 14
           s3-bucket: doc-previews
           if-no-files-found: error
-          path: pytorch.github.io/docs/main/
+          path: pytorch_docs/main/
           s3-prefix: pytorch/pytorch/${{ github.event.pull_request.number }}
 
       - name: Upload C++ Docs Preview

.github/workflows/_linux-build-label.yml

@@ -0,0 +1,109 @@
+name: linux-build
+
+on:
+  workflow_call:
+    inputs:
+      build-environment:
+        required: true
+        type: string
+        description: Top-level label for what's being built/tested.
+      docker-image-name:
+        required: true
+        type: string
+        description: Name of the base docker image to build with.
+      build-generates-artifacts:
+        required: false
+        type: boolean
+        default: true
+        description: If set, upload generated build artifacts.
+      build-with-debug:
+        required: false
+        type: boolean
+        default: false
+        description: If set, build in debug mode.
+      sync-tag:
+        required: false
+        type: string
+        default: ""
+        description: |
+          If this is set, our linter will use this to make sure that every other
+          job with the same `sync-tag` is identical.
+      cuda-arch-list:
+        required: false
+        type: string
+        default: "5.2"
+        description: Runner label to select worker type
+      runner:
+        required: false
+        type: string
+        default: "linux.2xlarge"
+        description: |
+          List of CUDA architectures CI build should target.
+      test-matrix:
+        required: false
+        type: string
+        description: |
+          An option JSON description of what test configs to run later on. This
+          is moved here from the Linux test workflow so that we can apply filter
+          logic using test-config labels earlier and skip unnecessary builds
+      s3-bucket:
+        description: S3 bucket to download artifact
+        required: false
+        type: string
+        default: "gha-artifacts"
+      aws-role-to-assume:
+        description: role to assume for downloading artifacts
+        required: false
+        type: string
+        default: ""
+    secrets:
+      HUGGING_FACE_HUB_TOKEN:
+        required: false
+        description: |
+          HF Auth token to avoid rate limits when downloading models or datasets from hub
+
+    outputs:
+      docker-image:
+        value: ${{ jobs.build.outputs.docker-image }}
+        description: The docker image containing the built PyTorch.
+      test-matrix:
+        value: ${{ jobs.build.outputs.test-matrix }}
+        description: An optional JSON description of what test configs to run later on.
+
+jobs:
+  build:
+    # Don't run on forked repos
+    if: github.repository_owner == 'pytorch'
+    runs-on: ${{ inputs.runner }}
+    timeout-minutes: 240
+    outputs:
+      docker-image: ${{ steps.linux-build.outputs.docker-image }}
+      test-matrix: ${{ steps.linux-build.outputs.test-matrix }}
+    steps:
+      - name: Setup SSH (Click me for login details)
+        uses: pytorch/test-infra/.github/actions/setup-ssh@main
+        with:
+          github-secret: ${{ secrets.GITHUB_TOKEN }}
+
+      # [pytorch repo ref]
+      # Use a pytorch/pytorch reference instead of a reference to the local
+      # checkout because when we run this action we don't *have* a local
+      # checkout. In other cases you should prefer a local checkout.
+      - name: Checkout PyTorch
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+
+      - name: Linux Build
+        id: linux-build
+        uses: ./.github/actions/linux-build
+        with:
+          build-environment: ${{ inputs.build-environment }}
+          docker-image-name: ${{ inputs.docker-image-name }}
+          build-generates-artifacts: ${{ inputs.build-generates-artifacts }}
+          build-with-debug: ${{ inputs.build-with-debug }}
+          sync-tag: ${{ inputs.sync-tag }}
+          cuda-arch-list: ${{ inputs.cuda-arch-list }}
+          test-matrix: ${{ inputs.test-matrix }}
+          s3-bucket: ${{ inputs.s3-bucket }}
+          aws-role-to-assume: ${{ inputs.aws-role-to-assume }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}

.github/workflows/_linux-build-rg.yml

@@ -0,0 +1,105 @@
+name: linux-build-rg
+
+on:
+  workflow_call:
+    inputs:
+      build-environment:
+        required: true
+        type: string
+        description: Top-level label for what's being built/tested.
+      docker-image-name:
+        required: true
+        type: string
+        description: Name of the base docker image to build with.
+      build-generates-artifacts:
+        required: false
+        type: boolean
+        default: true
+        description: If set, upload generated build artifacts.
+      build-with-debug:
+        required: false
+        type: boolean
+        default: false
+        description: If set, build in debug mode.
+      sync-tag:
+        required: false
+        type: string
+        default: ""
+        description: |
+          If this is set, our linter will use this to make sure that every other
+          job with the same `sync-tag` is identical.
+      cuda-arch-list:
+        required: false
+        type: string
+        default: "5.2"
+        description: |
+          List of CUDA architectures CI build should target.
+      runner-group:
+        required: false
+        type: string
+        default: "arc-lf-linux.2xlarge"
+        description: Runner group to select group type
+      test-matrix:
+        required: false
+        type: string
+        description: |
+          An option JSON description of what test configs to run later on. This
+          is moved here from the Linux test workflow so that we can apply filter
+          logic using test-config labels earlier and skip unnecessary builds
+      s3-bucket:
+        description: S3 bucket to download artifact
+        required: false
+        type: string
+        default: "gha-artifacts"
+      aws-role-to-assume:
+        description: role to assume for downloading artifacts
+        required: false
+        type: string
+        default: ""
+    secrets:
+      HUGGING_FACE_HUB_TOKEN:
+        required: false
+        description: |
+          HF Auth token to avoid rate limits when downloading models or datasets from hub
+
+    outputs:
+      docker-image:
+        value: ${{ jobs.build.outputs.docker-image }}
+        description: The docker image containing the built PyTorch.
+      test-matrix:
+        value: ${{ jobs.build.outputs.test-matrix }}
+        description: An optional JSON description of what test configs to run later on.
+
+jobs:
+  build:
+    # Don't run on forked repos
+    if: github.repository_owner == 'pytorch'
+    runs-on:
+      group: ${{ inputs.runner-group }}
+    timeout-minutes: 240
+    outputs:
+      docker-image: ${{ steps.linux-build.outputs.docker-image }}
+      test-matrix: ${{ steps.linux-build.outputs.test-matrix }}
+    steps:
+      # [pytorch repo ref]
+      # Use a pytorch/pytorch reference instead of a reference to the local
+      # checkout because when we run this action we don't *have* a local
+      # checkout. In other cases you should prefer a local checkout.
+      - name: Checkout PyTorch
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+
+      - name: Linux Build
+        id: linux-build
+        uses: ./.github/actions/linux-build
+        with:
+          build-environment: ${{ inputs.build-environment }}
+          docker-image-name: ${{ inputs.docker-image-name }}
+          build-generates-artifacts: ${{ inputs.build-generates-artifacts }}
+          build-with-debug: ${{ inputs.build-with-debug }}
+          sync-tag: ${{ inputs.sync-tag }}
+          cuda-arch-list: ${{ inputs.cuda-arch-list }}
+          test-matrix: ${{ inputs.test-matrix }}
+          s3-bucket: ${{ inputs.s3-bucket }}
+          aws-role-to-assume: ${{ inputs.aws-role-to-assume }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}

.github/workflows/_linux-build.yml

@@ -47,6 +47,23 @@ on:
           An option JSON description of what test configs to run later on. This
           is moved here from the Linux test workflow so that we can apply filter
           logic using test-config labels earlier and skip unnecessary builds
+      selected-test-configs:
+        description: |
+          A comma-separated list of test configurations from the test matrix to keep,
+          The empty list means we are going to keep every configurations by defaults
+        required: false
+        type: string
+        default: ""
+      s3-bucket:
+        description: S3 bucket to download artifact
+        required: false
+        type: string
+        default: "gha-artifacts"
+      aws-role-to-assume:
+        description: Role to assume for downloading artifacts
+        required: false
+        type: string
+        default: ""
     secrets:
       HUGGING_FACE_HUB_TOKEN:
         required: false
@@ -87,6 +104,14 @@ jobs:
       - name: Setup Linux
         uses: ./.github/actions/setup-linux
 
+      - name: configure aws credentials
+        uses: aws-actions/configure-aws-credentials@v3
+        if: ${{ inputs.aws-role-to-assume != '' }}
+        with:
+          role-to-assume: ${{ inputs.aws-role-to-assume }}
+          role-session-name: gha-linux-build
+          aws-region: us-east-1
+
       - name: Calculate docker image
         id: calculate-docker-image
         uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
@@ -125,6 +150,7 @@ jobs:
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
           test-matrix: ${{ inputs.test-matrix }}
+          selected-test-configs: ${{ inputs.selected-test-configs }}
           job-name: ${{ steps.get-job-id.outputs.job-name }}
 
       - name: Download pytest cache
@@ -133,6 +159,7 @@ jobs:
         with:
           cache_dir: .pytest_cache
           job_identifier: ${{ github.workflow }}_${{ inputs.build-environment }}
+          s3_bucket: ${{ inputs.s3-bucket }}
 
       - name: Build
         if: steps.filter.outputs.is-test-matrix-empty == 'False' || inputs.test-matrix == ''
@@ -197,6 +224,7 @@ jobs:
           retention-days: 14
           if-no-files-found: error
           path: artifacts.zip
+          s3-bucket: ${{ inputs.s3-bucket }}
 
       - name: Upload sccache stats
         if: steps.build.outcome != 'skipped'
@@ -207,6 +235,7 @@ jobs:
           retention-days: 365
           if-no-files-found: warn
           path: sccache-stats-*.json
+          s3-bucket: ${{ inputs.s3-bucket }}
 
       - name: Teardown Linux
         uses: pytorch/test-infra/.github/actions/teardown-linux@main

.github/workflows/_linux-test-label.yml

@@ -0,0 +1,85 @@
+name: linux-test-rg
+
+on:
+  workflow_call:
+    inputs:
+      build-environment:
+        required: true
+        type: string
+        description: Top-level label for what's being built/tested.
+      test-matrix:
+        required: true
+        type: string
+        description: JSON description of what test configs to run.
+      docker-image:
+        required: true
+        type: string
+        description: Docker image to run in.
+      sync-tag:
+        required: false
+        type: string
+        default: ""
+        description: |
+          If this is set, our linter will use this to make sure that every other
+          job with the same `sync-tag` is identical.
+      timeout-minutes:
+        required: false
+        type: number
+        default: 240
+        description: |
+          Set the maximum (in minutes) how long the workflow should take to finish
+      use-gha:
+        required: false
+        type: string
+        default: ""
+        description: If set to any value, upload to GHA. Otherwise upload to S3.
+      dashboard-tag:
+        required: false
+        type: string
+        default: ""
+      s3-bucket:
+        description: S3 bucket to download artifact
+        required: false
+        type: string
+        default: "gha-artifacts"
+      aws-role-to-assume:
+        description: role to assume for downloading artifacts
+        required: false
+        type: string
+        default: ""
+    secrets:
+      HUGGING_FACE_HUB_TOKEN:
+        required: false
+        description: |
+          HF Auth token to avoid rate limits when downloading models or datasets from hub
+
+env:
+  GIT_DEFAULT_BRANCH: ${{ github.event.repository.default_branch }}
+
+jobs:
+  test:
+    # Don't run on forked repos or empty test matrix
+    if: github.repository_owner == 'pytorch' && toJSON(fromJSON(inputs.test-matrix).include) != '[]'
+    strategy:
+      matrix: ${{ fromJSON(inputs.test-matrix) }}
+      fail-fast: false
+    runs-on: ${{ matrix.runner }}
+    timeout-minutes: ${{ matrix.mem_leak_check == 'mem_leak_check' && 600 || inputs.timeout-minutes }}
+    steps:
+      - name: Checkout PyTorch
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+
+      - name: Linux Test
+        id: linux-test
+        uses: ./.github/actions/linux-test
+        with:
+          build-environment: ${{ inputs.build-environment }}
+          test-matrix: ${{ inputs.test-matrix }}
+          docker-image: ${{ inputs.docker-image }}
+          sync-tag: ${{ inputs.sync-tag }}
+          use-gha: ${{ inputs.use-gha }}
+          dashboard-tag: ${{ inputs.dashboard-tag }}
+          s3-bucket: ${{ inputs.s3-bucket }}
+          aws-role-to-assume: ${{ inputs.aws-role-to-assume }}
+          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/_linux-test-rg.yml

@@ -0,0 +1,86 @@
+name: linux-test-label
+
+on:
+  workflow_call:
+    inputs:
+      build-environment:
+        required: true
+        type: string
+        description: Top-level label for what's being built/tested.
+      test-matrix:
+        required: true
+        type: string
+        description: JSON description of what test configs to run.
+      docker-image:
+        required: true
+        type: string
+        description: Docker image to run in.
+      sync-tag:
+        required: false
+        type: string
+        default: ""
+        description: |
+          If this is set, our linter will use this to make sure that every other
+          job with the same `sync-tag` is identical.
+      timeout-minutes:
+        required: false
+        type: number
+        default: 240
+        description: |
+          Set the maximum (in minutes) how long the workflow should take to finish
+      use-gha:
+        required: false
+        type: string
+        default: ""
+        description: If set to any value, upload to GHA. Otherwise upload to S3.
+      dashboard-tag:
+        required: false
+        type: string
+        default: ""
+      s3-bucket:
+        description: S3 bucket to download artifact
+        required: false
+        type: string
+        default: "gha-artifacts"
+      aws-role-to-assume:
+        description: role to assume for downloading artifacts
+        required: false
+        type: string
+        default: ""
+    secrets:
+      HUGGING_FACE_HUB_TOKEN:
+        required: false
+        description: |
+          HF Auth token to avoid rate limits when downloading models or datasets from hub
+
+env:
+  GIT_DEFAULT_BRANCH: ${{ github.event.repository.default_branch }}
+
+jobs:
+  test:
+    # Don't run on forked repos or empty test matrix
+    if: github.repository_owner == 'pytorch' && toJSON(fromJSON(inputs.test-matrix).include) != '[]'
+    strategy:
+      matrix: ${{ fromJSON(inputs.test-matrix) }}
+      fail-fast: false
+    runs-on:
+      group: ${{ matrix.runner }}
+    timeout-minutes: ${{ matrix.mem_leak_check == 'mem_leak_check' && 600 || inputs.timeout-minutes }}
+    steps:
+      - name: Checkout PyTorch
+        uses: pytorch/pytorch/.github/actions/checkout-pytorch@main
+
+      - name: Linux Test
+        id: linux-test
+        uses: ./.github/actions/linux-test
+        with:
+          build-environment: ${{ inputs.build-environment }}
+          test-matrix: ${{ inputs.test-matrix }}
+          docker-image: ${{ inputs.docker-image }}
+          sync-tag: ${{ inputs.sync-tag }}
+          use-gha: ${{ inputs.use-gha }}
+          dashboard-tag: ${{ inputs.dashboard-tag }}
+          s3-bucket: ${{ inputs.s3-bucket }}
+          aws-role-to-assume: ${{ inputs.aws-role-to-assume }}
+          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/_linux-test.yml

@@ -37,6 +37,16 @@ on:
         required: false
         type: string
         default: ""
+      s3-bucket:
+        description: S3 bucket to download artifact
+        required: false
+        type: string
+        default: "gha-artifacts"
+      aws-role-to-assume:
+        description: role to assume for downloading artifacts
+        required: false
+        type: string
+        default: ""
     secrets:
       HUGGING_FACE_HUB_TOKEN:
         required: false
@@ -71,6 +81,14 @@ jobs:
       - name: Setup Linux
         uses: ./.github/actions/setup-linux
 
+      - name: configure aws credentials
+        if : ${{ inputs.aws-role-to-assume != '' }}
+        uses: aws-actions/configure-aws-credentials@v3
+        with:
+          role-to-assume: ${{ inputs.aws-role-to-assume }}
+          role-session-name: gha-linux-test
+          aws-region: us-east-1
+
       - name: Calculate docker image
         id: calculate-docker-image
         uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
@@ -91,10 +109,15 @@ jobs:
         with:
           docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }}
 
+      - name: Check if in a ARC runner
+        shell: bash
+        id: check_arc_runner
+        run: echo "IN_ARC_RUNNER=$([ -f /.inarc ] && echo true || echo false)" >> "$GITHUB_OUTPUT"
+
       - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
         id: install-nvidia-driver
         uses: pytorch/test-infra/.github/actions/setup-nvidia@main
-        if: contains(inputs.build-environment, 'cuda') && !contains(matrix.config, 'nogpu')
+        if: ${{ contains(inputs.build-environment, 'cuda') && !contains(matrix.config, 'nogpu') && steps.check_arc_runner.outputs.IN_ARC_RUNNER == 'false' }}
 
       - name: Lock NVIDIA A100 40GB Frequency
         run: |
@@ -116,6 +139,7 @@ jobs:
         uses: ./.github/actions/download-build-artifacts
         with:
           name: ${{ inputs.build-environment }}
+          s3-bucket: ${{ inputs.s3-bucket }}
 
       - name: Download TD artifacts
         continue-on-error: true
@@ -176,6 +200,7 @@ jobs:
           VERBOSE_TEST_LOGS: ${{ steps.keep-going.outputs.ci-verbose-test-logs }}
           NO_TEST_TIMEOUT: ${{ steps.keep-going.outputs.ci-no-test-timeout }}
           NO_TD: ${{ steps.keep-going.outputs.ci-no-td }}
+          TD_DISTRIBUTED: ${{ steps.keep-going.outputs.ci-td-distributed }}
           SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2
           SCCACHE_S3_KEY_PREFIX: ${{ github.workflow }}
           SHM_SIZE: ${{ contains(inputs.build-environment, 'cuda') && '2g' || '1g' }}
@@ -228,6 +253,7 @@ jobs:
             -e VERBOSE_TEST_LOGS \
             -e NO_TEST_TIMEOUT \
             -e NO_TD \
+            -e TD_DISTRIBUTED \
             -e PR_LABELS \
             -e MAX_JOBS="$(nproc --ignore=2)" \
             -e SCCACHE_BUCKET \
@@ -240,7 +266,6 @@ jobs:
             -e HUGGING_FACE_HUB_TOKEN \
             -e DASHBOARD_TAG \
             --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \
-            --ulimit stack=10485760:83886080 \
             --security-opt seccomp=unconfined \
             --cap-add=SYS_PTRACE \
             --ipc=host \
@@ -290,6 +315,7 @@ jobs:
         with:
           file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.get-job-id.outputs.job-id }}
           use-gha: ${{ inputs.use-gha }}
+          s3-bucket: ${{ inputs.s3-bucket }}
 
       - name: Collect backtraces from coredumps (if any)
         if: always()

.github/workflows/_win-test.yml

@@ -92,7 +92,7 @@ jobs:
           retry_wait_seconds: 30
           command: |
             set -eu
-            python3 -m pip install rockset==1.0.3
+            python3 -m pip install rockset==1.0.3 'xdoctest>=1.1.0'
 
       - name: Start monitoring script
         id: monitor-script

.github/workflows/build-triton-wheel.yml

@@ -119,8 +119,7 @@ jobs:
 
       - uses: actions/upload-artifact@v3
         with:
-          # NB: Use the same name here and all wheels can be downloaded by referring to the same artifact
-          name: pytorch-triton-wheel
+          name: pytorch-triton-wheel-${{ matrix.py_vers }}-${{ matrix.device }}
           if-no-files-found: error
           path: ${{ runner.temp }}/artifacts/*
 
@@ -157,8 +156,15 @@ jobs:
       - name: Download Build Artifacts
         uses: actions/download-artifact@v3
         with:
-          name: pytorch-triton-wheel
-          path: ${{ runner.temp }}/artifacts/
+          # Download all available artifacts
+          path: ${{ runner.temp }}/artifacts-all
+
+      - name: Select Wheel Artifacts
+        shell: bash
+        run: |
+          set -x
+          mkdir -p "${RUNNER_TEMP}/artifacts/"
+          mv "${RUNNER_TEMP}"/artifacts-all/pytorch-triton-wheel-*/* "${RUNNER_TEMP}/artifacts/"
 
       - name: Set DRY_RUN (only for tagged pushes)
         if: ${{ github.event_name == 'push' && (github.event.ref == 'refs/heads/main' || startsWith(github.event.ref, 'refs/tags/v')) }}
@@ -246,8 +252,7 @@ jobs:
 
       - uses: actions/upload-artifact@v3
         with:
-          # NB: Use the same name here and all wheels can be downloaded by referring to the same artifact
-          name: pytorch-triton-conda
+          name: pytorch-triton-conda-${{ matrix.py_vers }}
           if-no-files-found: error
           path: ${{ runner.temp }}/artifacts/*
 
@@ -267,8 +272,15 @@ jobs:
       - name: Download Build Artifacts
         uses: actions/download-artifact@v3
         with:
-          name: pytorch-triton-conda
-          path: ${{ runner.temp }}/artifacts/
+          # Download all available artifacts
+          path: ${{ runner.temp }}/artifacts-all
+
+      - name: Select Conda Artifacts
+        shell: bash
+        run: |
+          set -x
+          mkdir -p "${RUNNER_TEMP}/artifacts/"
+          mv "${RUNNER_TEMP}"/artifacts-all/pytorch-triton-conda-*/* "${RUNNER_TEMP}/artifacts/"
 
       - name: Set DRY_RUN (only for tagged pushes)
         if: ${{ github.event_name == 'push' && (github.event.ref == 'refs/heads/main' || startsWith(github.event.ref, 'refs/tags/v')) }}

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

@@ -4,15 +4,17 @@ on:
   schedule:
     - cron: 0 7 * * 1-6
     - cron: 0 7 * * 0
+  # NB: GitHub has an upper limit of 10 inputs here, so before we can sort it
+  # out, let try to run torchao cudagraphs_low_precision as part of cudagraphs
   workflow_dispatch:
     inputs:
       training:
-        description: Run training?
+        description: Run training (on by default)?
         required: false
         type: boolean
         default: true
       inference:
-        description: Run inference?
+        description: Run inference (off by default)?
         required: false
         type: boolean
         default: false
@@ -20,22 +22,17 @@ on:
         description: Run inductor_default?
         required: false
         type: boolean
-        default: true
+        default: false
       dynamic:
         description: Run inductor_dynamic_shapes?
         required: false
         type: boolean
-        default: true
+        default: false
       cudagraphs:
         description: Run inductor_cudagraphs?
         required: false
         type: boolean
         default: true
-      cppwrapper:
-        description: Run inductor_cpp_wrapper for inference?
-        required: false
-        type: boolean
-        default: false
       freezing_cudagraphs:
         description: Run inductor_cudagraphs with freezing for inference?
         required: false
@@ -56,6 +53,11 @@ on:
         required: false
         type: boolean
         default: false
+      benchmark_configs:
+        description: The list of configs used the benchmark
+        required: false
+        type: string
+        default: inductor_huggingface_perf,inductor_timm_perf,inductor_torchbench_perf
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }}
@@ -86,6 +88,7 @@ jobs:
           { config: "inductor_torchbench_perf", shard: 3, num_shards: 4, runner: "linux.gcp.a100.large" },
           { config: "inductor_torchbench_perf", shard: 4, num_shards: 4, runner: "linux.gcp.a100.large" },
         ]}
+      selected-test-configs: ${{ inputs.benchmark_configs }}
     secrets:
       HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
 
@@ -96,7 +99,7 @@ jobs:
     if: github.event.schedule == '0 7 * * 1-6'
     with:
       build-environment: linux-focal-cuda12.1-py3.10-gcc9-sm80
-      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-aotinductor-true-freezing_cudagraphs-true
+      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-aotinductor-true-freezing_cudagraphs-true-cudagraphs_low_precision-true
       docker-image: ${{ needs.linux-focal-cuda12_1-py3_10-gcc9-inductor-build.outputs.docker-image }}
       test-matrix: ${{ needs.linux-focal-cuda12_1-py3_10-gcc9-inductor-build.outputs.test-matrix }}
       use-gha: anything-non-empty-to-use-gha
@@ -111,7 +114,7 @@ jobs:
     if: github.event.schedule == '0 7 * * 0'
     with:
       build-environment: linux-focal-cuda12.1-py3.10-gcc9-sm80
-      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-aotinductor-true-freezing_cudagraphs-true-maxautotune-true-freeze_autotune_cudagraphs-true
+      dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-aotinductor-true-freezing_cudagraphs-true-maxautotune-true-freeze_autotune_cudagraphs-true-cudagraphs_low_precision-true
       docker-image: ${{ needs.linux-focal-cuda12_1-py3_10-gcc9-inductor-build.outputs.docker-image }}
       test-matrix: ${{ needs.linux-focal-cuda12_1-py3_10-gcc9-inductor-build.outputs.test-matrix }}
       use-gha: anything-non-empty-to-use-gha
@@ -126,7 +129,7 @@ jobs:
     if: github.event_name == 'workflow_dispatch'
     with:
       build-environment: linux-focal-cuda12.1-py3.10-gcc9-sm80
-      dashboard-tag: training-${{ inputs.training }}-inference-${{ inputs.inference }}-default-${{ inputs.default }}-dynamic-${{ inputs.dynamic }}-cudagraphs-${{ inputs.cudagraphs }}-cppwrapper-${{ inputs.cppwrapper }}-aotinductor-${{ inputs.aotinductor }}-maxautotune-${{ inputs.maxautotune }}-freezing_cudagraphs-${{ inputs.freezing_cudagraphs }}
+      dashboard-tag: training-${{ inputs.training }}-inference-${{ inputs.inference }}-default-${{ inputs.default }}-dynamic-${{ inputs.dynamic }}-cudagraphs-${{ inputs.cudagraphs }}-cppwrapper-false-aotinductor-${{ inputs.aotinductor }}-maxautotune-${{ inputs.maxautotune }}-freezing_cudagraphs-${{ inputs.freezing_cudagraphs }}-cudagraphs_low_precision-${{ inputs.cudagraphs }}
       docker-image: ${{ needs.linux-focal-cuda12_1-py3_10-gcc9-inductor-build.outputs.docker-image }}
       test-matrix: ${{ needs.linux-focal-cuda12_1-py3_10-gcc9-inductor-build.outputs.test-matrix }}
       use-gha: anything-non-empty-to-use-gha

.github/workflows/inductor.yml

@@ -111,7 +111,7 @@ jobs:
     name: linux-jammy-cpu-py3.8-gcc11-inductor
     uses: ./.github/workflows/_linux-build.yml
     with:
-      build-environment: linux-jammy-py3_8-gcc11-build
+      build-environment: linux-jammy-py3.8-gcc11-build
       docker-image-name: pytorch-linux-jammy-py3.8-gcc11-inductor-benchmarks
       test-matrix: |
         { include: [
@@ -135,7 +135,7 @@ jobs:
     uses: ./.github/workflows/_linux-test.yml
     needs: linux-jammy-cpu-py3_8-gcc11-inductor-build
     with:
-      build-environment: linux-jammy-py3_8-gcc11-build
+      build-environment: linux-jammy-py3.8-gcc11-build
       docker-image: ${{ needs.linux-jammy-cpu-py3_8-gcc11-inductor-build.outputs.docker-image }}
       test-matrix: ${{ needs.linux-jammy-cpu-py3_8-gcc11-inductor-build.outputs.test-matrix }}
     secrets:

.github/workflows/labeler.yml

@@ -1,37 +0,0 @@
-name: Labeler
-
-on:
-  # We need pull_request_target to be able to add labels to PRs from forks.
-  # Only allow pull_request_target when targeting main, not some historical branch.
-  #
-  # Make sure to don't introduce explicit checking out and installing/running
-  # untrusted user code into this workflow!
-  pull_request_target:
-    types: [opened, synchronize, reopened]
-    branches: [main]
-
-  # To add labels on ghstack PRs.
-  # Note: as pull_request doesn't trigger on PRs targeting main,
-  # to test changes to the workflow itself one needs to create
-  # a PR that targets a gh/**/base branch.
-  pull_request:
-    types: [opened, synchronize, reopened]
-    branches: [gh/**/base]
-
-jobs:
-  triage:
-    permissions:
-      contents: read
-      pull-requests: write
-    runs-on: ubuntu-latest
-    # Do not auto-label nightly builds PR
-    if: ${{ github.event.pull_request.number != 26921 && github.repository_owner == 'pytorch' }}
-    steps:
-    - uses: actions/labeler@v4
-      with:
-        repo-token: "${{ secrets.GITHUB_TOKEN }}"
-        sync-labels: ''
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
-  cancel-in-progress: true

.github/workflows/llm_td_retrieval.yml

@@ -0,0 +1,120 @@
+name: Retrieval PyTorch Tests for Target Determination
+
+on:
+  workflow_call:
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  llm-retrieval:
+    runs-on: linux.4xlarge
+    continue-on-error: true
+    steps:
+      - name: Clone PyTorch
+        uses: actions/checkout@v3
+        with:
+          repository: pytorch/pytorch
+          fetch-depth: 0
+          path: pytorch
+
+      - name: Setup Linux
+        uses: ./pytorch/.github/actions/setup-linux
+
+      - name: Clone CodeLlama
+        uses: actions/checkout@v3
+        with:
+          repository: osalpekar/codellama
+          ref: main
+          path: codellama
+
+      - name: Clone Target Determination Code
+        uses: actions/checkout@v3
+        with:
+          repository: osalpekar/llm-target-determinator
+          ref: v0.0.2
+          path: llm-target-determinator
+
+      - name: Setup Conda
+        uses: conda-incubator/setup-miniconda@v2.1.1
+        with:
+          miniconda-version: "py39_4.12.0"
+          python-version: 3.9
+
+      - name: Install Requirements
+        shell: bash -l {0}
+        run: |
+          set -euxo pipefail
+          conda create \
+            --yes \
+            --quiet \
+            --name "tdenv" \
+            "python=3.9"
+          conda activate tdenv
+          cd "${GITHUB_WORKSPACE}/llm-target-determinator"
+          pip install -r requirements.txt
+          cd ../codellama
+          pip install -e .
+
+      - name: Fetch CodeLlama Checkpoint
+        shell: bash -l {0}
+        run: |
+          set -euxo pipefail
+          conda activate tdenv
+          cd codellama/
+          mkdir "CodeLlama-7b-Python"
+          aws s3 cp "s3://target-determinator-assets/CodeLlama-7b-Python" "CodeLlama-7b-Python" --recursive --no-progress
+
+      - name: Fetch indexes
+        uses: nick-fields/retry@v2.8.2
+        with:
+          max_attempts: 3
+          retry_wait_seconds: 10
+          timeout_minutes: 5
+          shell: bash
+          command: |
+            set -euxo pipefail
+            python3 -m pip install awscli==1.29.40
+            cd "${GITHUB_WORKSPACE}"/llm-target-determinator/assets
+            aws s3 cp "s3://target-determinator-assets/indexes/latest" . --recursive
+
+            unzip -o indexer-files\*.zip
+            rm indexer-files*.zip
+
+      - name: Run Retriever
+        id: run_retriever
+        continue-on-error: true  # ghstack not currently supported due to problems getting git diff
+        shell: bash -l {0}
+        run: |
+          set -euxo pipefail
+          conda activate tdenv
+          cd "${GITHUB_WORKSPACE}"/llm-target-determinator
+          torchrun \
+            --standalone \
+            --nnodes=1 \
+            --nproc-per-node=1 \
+            retriever.py \
+            --experiment-name indexer-files \
+            --pr-parse-format GITDIFF
+          cd assets
+          zip -r mappings.zip mappings
+
+      - name: Upload results to s3
+        uses: seemethere/upload-artifact-s3@v5
+        if: ${{ steps.run_retriever.outcome == 'success' }}
+        with:
+          name: llm_results
+          retention-days: 14
+          if-no-files-found: warn
+          path: llm-target-determinator/assets/mappings.zip
+        env:
+          AWS_ACCESS_KEY_ID: ""
+          AWS_SECRET_ACCESS_KEY: ""
+          AWS_SESSION_TOKEN: ""
+          AWS_DEFAULT_REGION: ""
+          AWS_REGION: ""
+
+      - name: Teardown Linux
+        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        if: always()

.github/workflows/nightly.yml

@@ -31,7 +31,7 @@ jobs:
     with:
       build-environment: linux-jammy-py3.8-gcc11
       docker-image: ${{ needs.docs-build.outputs.docker-image }}
-      push: ${{ github.event_name == 'schedule' || startsWith(github.event.ref, 'refs/tags/v') }}
+      push: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' || startsWith(github.event.ref, 'refs/tags/v') }}
       run-doxygen: true
     secrets:
       GH_PYTORCHBOT_TOKEN: ${{ secrets.GH_PYTORCHBOT_TOKEN }}

.github/workflows/periodic.yml

@@ -23,9 +23,17 @@ concurrency:
 permissions: read-all
 
 jobs:
+  llm-td:
+    name: before-test
+    uses: ./.github/workflows/llm_td_retrieval.yml
+    permissions:
+      id-token: write
+      contents: read
+
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
+    needs: llm-td
     permissions:
       id-token: write
       contents: read

.github/workflows/pull.yml

@@ -20,16 +20,24 @@ concurrency:
 permissions: read-all
 
 jobs:
+  llm-td:
+    name: before-test
+    uses: ./.github/workflows/llm_td_retrieval.yml
+    permissions:
+      id-token: write
+      contents: read
+
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
+    needs: llm-td
     permissions:
       id-token: write
       contents: read
 
   linux-jammy-py3_8-gcc11-build:
     name: linux-jammy-py3.8-gcc11
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-jammy-py3.8-gcc11
       docker-image-name: pytorch-linux-jammy-py3.8-gcc11
@@ -66,7 +74,7 @@ jobs:
 
   linux-jammy-py3_8-gcc11-no-ops:
     name: linux-jammy-py3.8-gcc11-no-ops
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-jammy-py3.8-gcc11-no-ops
       docker-image-name: pytorch-linux-jammy-py3.8-gcc11
@@ -77,7 +85,7 @@ jobs:
 
   linux-jammy-py3_8-gcc11-pch:
     name: linux-jammy-py3.8-gcc11-pch
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-jammy-py3.8-gcc11-pch
       docker-image-name: pytorch-linux-jammy-py3.8-gcc11
@@ -86,9 +94,10 @@ jobs:
           { config: "default", shard: 1, num_shards: 1 },
         ]}
 
+
   linux-jammy-py3_10-clang15-asan-build:
     name: linux-jammy-py3.10-clang15-asan
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-jammy-py3.10-clang15-asan
       docker-image-name: pytorch-linux-jammy-py3-clang15-asan
@@ -103,6 +112,7 @@ jobs:
         ]}
       sync-tag: asan-build
 
+
   linux-jammy-py3_10-clang15-asan-test:
     name: linux-jammy-py3.10-clang15-asan
     uses: ./.github/workflows/_linux-test.yml
@@ -117,7 +127,7 @@ jobs:
 
   linux-focal-py3_8-clang10-onnx-build:
     name: linux-focal-py3.8-clang10-onnx
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-py3.8-clang10-onnx
       docker-image-name: pytorch-linux-focal-py3-clang10-onnx
@@ -140,7 +150,7 @@ jobs:
 
   linux-focal-py3_8-clang10-build:
     name: linux-focal-py3.8-clang10
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-py3.8-clang10
       docker-image-name: pytorch-linux-focal-py3.8-clang10
@@ -155,7 +165,6 @@ jobs:
           { config: "dynamo", shard: 2, num_shards: 3, runner: "linux.2xlarge" },
           { config: "dynamo", shard: 3, num_shards: 3, runner: "linux.2xlarge" },
         ]}
-
   linux-focal-py3_8-clang10-test:
     name: linux-focal-py3.8-clang10
     uses: ./.github/workflows/_linux-test.yml
@@ -169,7 +178,7 @@ jobs:
 
   linux-focal-py3_11-clang10-build:
     name: linux-focal-py3.11-clang10
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-py3.11-clang10
       docker-image-name: pytorch-linux-focal-py3.11-clang10
@@ -185,6 +194,7 @@ jobs:
           { config: "dynamo", shard: 3, num_shards: 3, runner: "linux.2xlarge" },
         ]}
 
+
   linux-focal-py3_11-clang10-test:
     name: linux-focal-py3.11-clang10
     uses: ./.github/workflows/_linux-test.yml
@@ -198,7 +208,7 @@ jobs:
 
   linux-focal-py3_12-clang10-build:
     name: linux-focal-py3.12-clang10
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-py3.12-clang10
       docker-image-name: pytorch-linux-focal-py3.12-clang10
@@ -207,6 +217,9 @@ jobs:
           { config: "default", shard: 1, num_shards: 3, runner: "linux.2xlarge" },
           { config: "default", shard: 2, num_shards: 3, runner: "linux.2xlarge" },
           { config: "default", shard: 3, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "dynamo", shard: 1, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "dynamo", shard: 2, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "dynamo", shard: 3, num_shards: 3, runner: "linux.2xlarge" },
         ]}
 
   linux-focal-py3_12-clang10-test:
@@ -221,7 +234,7 @@ jobs:
 
   linux-focal-cuda11_8-py3_10-gcc9-build:
     name: linux-focal-cuda11.8-py3.10-gcc9
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-cuda11.8-py3.10-gcc9
       docker-image-name: pytorch-linux-focal-cuda11.8-cudnn8-py3-gcc9
@@ -246,7 +259,7 @@ jobs:
 
   linux-focal-cuda12_1-py3_10-gcc9-build:
     name: linux-focal-cuda12.1-py3.10-gcc9
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-cuda12.1-py3.10-gcc9
       docker-image-name: pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9
@@ -274,7 +287,7 @@ jobs:
 
   linux-jammy-py3-clang12-mobile-build:
     name: linux-jammy-py3-clang12-mobile-build
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-jammy-py3-clang12-mobile-build
       docker-image-name: pytorch-linux-jammy-py3-clang15-asan
@@ -286,7 +299,7 @@ jobs:
 
   linux-jammy-cuda-11_8-cudnn8-py3_8-clang12-build:
     name: linux-jammy-cuda11.8-cudnn8-py3.8-clang12
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-jammy-cuda11.8-cudnn8-py3.8-clang12
       docker-image-name: pytorch-linux-jammy-cuda11.8-cudnn8-py3.8-clang12
@@ -297,7 +310,7 @@ jobs:
 
   linux-focal-py3-clang9-mobile-custom-build-static:
     name: linux-focal-py3-clang9-mobile-custom-build-static
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-py3-clang9-mobile-custom-build-static
       docker-image-name: pytorch-linux-focal-py3-clang9-android-ndk-r21e
@@ -309,9 +322,9 @@ jobs:
 
   linux-focal-py3_8-clang9-xla-build:
     name: linux-focal-py3_8-clang9-xla
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
-      build-environment: linux-focal-py3_8-clang9-xla
+      build-environment: linux-focal-py3.8-clang9-xla
       docker-image-name: 308535385114.dkr.ecr.us-east-1.amazonaws.com/pytorch/xla_base:v1.1-lite
       test-matrix: |
         { include: [
@@ -323,7 +336,7 @@ jobs:
     uses: ./.github/workflows/_linux-test.yml
     needs: linux-focal-py3_8-clang9-xla-build
     with:
-      build-environment: linux-focal-py3_8-clang9-xla
+      build-environment: linux-focal-py3.8-clang9-xla
       docker-image: ${{ needs.linux-focal-py3_8-clang9-xla-build.outputs.docker-image }}
       test-matrix: ${{ needs.linux-focal-py3_8-clang9-xla-build.outputs.test-matrix }}
 
@@ -391,7 +404,7 @@ jobs:
 
   linux-jammy-py3_8-gcc11-mobile-lightweight-dispatch-build:
     name: linux-jammy-py3.8-gcc11-mobile-lightweight-dispatch-build
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-jammy-py3.8-gcc111-mobile-lightweight-dispatch-build
       docker-image-name: pytorch-linux-jammy-py3.8-gcc11
@@ -405,7 +418,7 @@ jobs:
     # don't run build twice on main
     if: github.event_name == 'pull_request'
     name: linux-focal-rocm6.0-py3.8
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-rocm6.0-py3.8
       docker-image-name: pytorch-linux-focal-rocm-n-py3
@@ -419,7 +432,7 @@ jobs:
 
   linux-focal-cuda12_1-py3_10-gcc9-sm86-build:
     name: linux-focal-cuda12.1-py3.10-gcc9-sm86
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-cuda12.1-py3.10-gcc9-sm86
       docker-image-name: pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9
@@ -446,7 +459,7 @@ jobs:
 
   linux-jammy-py3-clang12-executorch-build:
     name: linux-jammy-py3-clang12-executorch
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-jammy-py3-clang12-executorch
       docker-image-name: pytorch-linux-jammy-py3-clang12-executorch

.github/workflows/rocm.yml

@@ -27,7 +27,7 @@ jobs:
 
   linux-focal-rocm6_0-py3_8-build:
     name: linux-focal-rocm6.0-py3.8
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-rocm6.0-py3.8
       docker-image-name: pytorch-linux-focal-rocm-n-py3

.github/workflows/slow.yml

@@ -21,9 +21,17 @@ concurrency:
 permissions: read-all
 
 jobs:
+  llm-td:
+    name: before-test
+    uses: ./.github/workflows/llm_td_retrieval.yml
+    permissions:
+      id-token: write
+      contents: read
+
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
+    needs: llm-td
     permissions:
       id-token: write
       contents: read
@@ -37,10 +45,12 @@ jobs:
       cuda-arch-list: 8.6
       test-matrix: |
         { include: [
-          { config: "default", shard: 1, num_shards: 4, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "default", shard: 2, num_shards: 4, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "default", shard: 3, num_shards: 4, runner: "linux.g5.4xlarge.nvidia.gpu" },
-          { config: "default", shard: 4, num_shards: 4, runner: "linux.g5.4xlarge.nvidia.gpu" },
+          { config: "default", shard: 1, num_shards: 6, runner: "linux.g5.4xlarge.nvidia.gpu" },
+          { config: "default", shard: 2, num_shards: 6, runner: "linux.g5.4xlarge.nvidia.gpu" },
+          { config: "default", shard: 3, num_shards: 6, runner: "linux.g5.4xlarge.nvidia.gpu" },
+          { config: "default", shard: 4, num_shards: 6, runner: "linux.g5.4xlarge.nvidia.gpu" },
+          { config: "default", shard: 5, num_shards: 6, runner: "linux.g5.4xlarge.nvidia.gpu" },
+          { config: "default", shard: 6, num_shards: 6, runner: "linux.g5.4xlarge.nvidia.gpu" },
         ]}
 
   linux-focal-cuda12_1-py3-gcc9-slow-gradcheck-test:
@@ -128,7 +138,7 @@ jobs:
 
   linux-jammy-py3_10-clang15-asan-build:
     name: linux-jammy-py3.10-clang15-asan
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-jammy-py3.10-clang15-asan
       docker-image-name: pytorch-linux-jammy-py3-clang15-asan

.github/workflows/target-determination-indexer.yml

@@ -2,7 +2,8 @@ name: Index PyTorch Tests for Target Determination
 
 on:
   workflow_dispatch:
-  # TODO: Trigger every few hours
+  schedule:
+    - cron: '0 0 * * *'
 
 permissions:
   id-token: write
@@ -13,14 +14,20 @@ jobs:
     runs-on: linux.g5.4xlarge.nvidia.gpu # 1 GPU A10G 24GB each
     environment: target-determinator-env
     steps:
+      - name: Clone PyTorch
+        uses: actions/checkout@v3
+        with:
+          path: pytorch
+
       - name: Setup Linux
-        uses: ./.github/actions/setup-linux
+        uses: ./pytorch/.github/actions/setup-linux
 
       - name: Calculate docker image
         id: calculate-docker-image
         uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
         with:
           docker-image-name: pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9
+          working-directory: pytorch
 
       - name: Use following to pull public copy of the image
         id: print-ghcr-mirror
@@ -40,112 +47,97 @@ jobs:
         id: install-nvidia-driver
         uses: pytorch/test-infra/.github/actions/setup-nvidia@main
 
-      - name: Clone PyTorch
-        uses: actions/checkout@v3
-        with:
-          path: pytorch
-
       - name: Clone CodeLlama
         uses: actions/checkout@v3
         with:
           repository: osalpekar/codellama
-          ref: main
+          ref: 1ec50e0cfc0fadc3b6ceb146617e2119ab26eb34
           path: codellama
 
       - name: Clone Target Determination Code
         uses: actions/checkout@v3
         with:
           repository: osalpekar/llm-target-determinator
-          ref: v0.0.1
+          ref: v0.0.2
           path: llm-target-determinator
 
-      - name: Install Requirements
-        shell: bash -l {0}
-        run: |
-          set -euxo pipefail
-
-          conda create \
-            --yes \
-            --quiet \
-            --name "tdenv" \
-            "python=3.9"
-
-          conda activate tdenv
-
-          cd "${GITHUB_WORKSPACE}"
-          pwd
-          cd llm-target-determinator
-          pip install -r requirements.txt
-          cd ../codellama
-          pip install -e .
-
       - name: Configure AWS credentials
         uses: aws-actions/configure-aws-credentials@v3
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_target_determinator_s3_read_write
           aws-region: us-east-1
 
-      - name: Fetch CodeLlama Checkpoint
-        shell: bash -l {0}
+      - name: Download checkpoint
+        shell: bash
+        env:
+          AWS_DEFAULT_REGION: us-east-1
         run: |
-          set -euxo pipefail
-
-          conda activate tdenv
-          pip install awscli==1.32.18
-          cd codellama/
+          # Do this outside of docker so I don't have to put env vars in
+          pip3 install awscli==1.29.40
+          cd codellama
           mkdir "CodeLlama-7b-Python"
           aws s3 cp \
             "s3://target-determinator-assets/CodeLlama-7b-Python" \
             "CodeLlama-7b-Python" \
             --recursive
 
-      - name: Run Indexer
-        id: indexer
+      - name: Run indexer
         shell: bash -l {0}
+        env:
+          DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          AWS_DEFAULT_REGION: us-east-1
         run: |
-          set -euxo pipefail
+          # detached container should get cleaned up by teardown_ec2_linux
+          container_name=$(docker run \
+            ${GPU_FLAG:-} \
+            -e MAX_JOBS="$(nproc --ignore=2)" \
+            -e AWS_DEFAULT_REGION \
+            --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \
+            --security-opt seccomp=unconfined \
+            --cap-add=SYS_PTRACE \
+            --tty \
+            --detach \
+            --user jenkins \
+            -v "${GITHUB_WORKSPACE}:/var/lib/jenkins/workspace" \
+            -w /var/lib/jenkins/workspace \
+            "${DOCKER_IMAGE}"
+          )
+          chmod +x pytorch/.github/scripts/td_llm_indexer.sh
+          docker exec -t "${container_name}" sh -c 'pytorch/.github/scripts/td_llm_indexer.sh'
 
-          conda activate tdenv
-          cd "${GITHUB_WORKSPACE}"/llm-target-determinator
-
-          python create_filelist.py
-
-          torchrun \
-            --standalone \
-            --nnodes=1 \
-            --nproc-per-node=1 \
-            indexer.py \
-            --experiment-name indexer-files
-
-      - name: Upload Index to S3
+      - name: Upload to s3
         shell: bash -l {0}
-        if: ${{ steps.indexer.outcome == 'success' }}
+        env:
+          AWS_DEFAULT_REGION: us-east-1
         run: |
-          set -euxo pipefail
-          conda activate tdenv
-          cd "${GITHUB_WORKSPACE}"/llm-target-determinator/assets
+          cd llm-target-determinator/assets
 
           TIMESTAMP=$(date -Iseconds)
-          ZIP_NAME = "indexer-files-${TIMESTAMP}.zip"
+          ZIP_NAME="indexer-files-${TIMESTAMP}.zip"
 
           # Create a zipfile with all the generated indices
           zip -r "${ZIP_NAME}" indexer-files
 
+          # Note that because the below 2 operations are not atomic, there will
+          # be a period of a few seconds between these where there is no index
+          # present in the latest/ folder. To account for this, the retriever
+          # should have some retry logic with backoff to ensure fetching the
+          # index doesn't fail.
           # Move the old index into the archived/ folder
-          aws s3 cp \
-            "s3://target-determinator-assets/indexes/latest/*" \
-            "s3://target-determinator-assets/indexes/archived/"
+          aws s3 mv \
+            "s3://target-determinator-assets/indexes/latest" \
+            "s3://target-determinator-assets/indexes/archived" \
+            --recursive
 
           # Move the new index into the latestl/ folder
           aws s3 cp \
             "${ZIP_NAME}" \
             "s3://target-determinator-assets/indexes/latest/${ZIP_NAME}"
 
-          # Note that because the above 2 operations are not atomic, there will
-          # be a period of a few seconds between these where there is no index
-          # present in the latest/ folder. To account for this, the retriever
-          # should have some retry logic with backoff to ensure fetching the
-          # index doesn't fail.
+      - name: Teardown Linux
+        uses: pytorch/test-infra/.github/actions/teardown-linux@main
+        if: always()
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}

.github/workflows/target_determination.yml

@@ -35,6 +35,13 @@ jobs:
           cache_dir: .pytest_cache
           job_identifier: ${{ github.workflow }}
 
+      - name: Download LLM Artifacts from S3
+        uses: seemethere/download-artifact-s3@v4
+        continue-on-error: true
+        with:
+          name: llm_results
+          path: .additional_ci_files/llm_results
+
       - name: Do TD
         id: td
         continue-on-error: true
@@ -46,10 +53,12 @@ jobs:
           GITHUB_RUN_ID: ${{ github.run_id }}
           GITHUB_RUN_NUMBER: ${{ github.run_number }}
           GITHUB_RUN_ATTEMPT: ${{ github.run_attempt }}
+          GITHUB_REF: ${{ github.ref }}
           JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
           JOB_NAME: ${{ steps.get-job-id.outputs.job-name }}
           PR_NUMBER: ${{ github.event.pull_request.number }}
         run: |
+          unzip -o .additional_ci_files/llm_results/mappings.zip -d .additional_ci_files/llm_results || true
           python3 -m pip install boto3==1.19.12
           python3 tools/testing/do_target_determination_for_s3.py
 

.github/workflows/torchbench.yml

@@ -0,0 +1,38 @@
+name: torchbench
+
+on:
+  push:
+    tags:
+      - ciflow/torchbench/*
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }}
+  cancel-in-progress: true
+
+jobs:
+  linux-focal-cuda12_1-py3_10-gcc9-torchbench-build-gcp:
+    name: cuda12.1-py3.10-gcc9-sm80
+    uses: ./.github/workflows/_linux-build.yml
+    with:
+      build-environment: linux-focal-cuda12.1-py3.10-gcc9-sm80
+      docker-image-name: pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9-inductor-benchmarks
+      cuda-arch-list: '8.0'
+      test-matrix: |
+        { include: [
+          { config: "torchbench_gcp_smoketest", shard: 1, num_shards: 1, runner: "linux.gcp.a100" },
+        ]}
+    secrets:
+      HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
+
+  linux-focal-cuda12_1-py3_10-gcc9-torchbench-test-gcp:
+    name: cuda12.1-py3.10-gcc9-sm80
+    uses: ./.github/workflows/_linux-test.yml
+    needs: linux-focal-cuda12_1-py3_10-gcc9-torchbench-build-gcp
+    with:
+      build-environment: linux-focal-cuda12.1-py3.10-gcc9-sm80
+      docker-image: ${{ needs.linux-focal-cuda12_1-py3_10-gcc9-torchbench-build-gcp.outputs.docker-image }}
+      test-matrix: ${{ needs.linux-focal-cuda12_1-py3_10-gcc9-torchbench-build-gcp.outputs.test-matrix }}
+      use-gha: anything-non-empty-to-use-gha
+    secrets:
+      HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}

.github/workflows/trunk.yml

@@ -19,9 +19,17 @@ concurrency:
 permissions: read-all
 
 jobs:
+  llm-td:
+    name: before-test
+    uses: ./.github/workflows/llm_td_retrieval.yml
+    permissions:
+      id-token: write
+      contents: read
+
   target-determination:
     name: before-test
     uses: ./.github/workflows/target_determination.yml
+    needs: llm-td
     permissions:
       id-token: write
       contents: read
@@ -192,7 +200,7 @@ jobs:
 
   linux-focal-rocm6_0-py3_8-build:
     name: linux-focal-rocm6.0-py3.8
-    uses: ./.github/workflows/_linux-build.yml
+    uses: ./.github/workflows/_linux-build-label.yml
     with:
       build-environment: linux-focal-rocm6.0-py3.8
       docker-image-name: pytorch-linux-focal-rocm-n-py3

.github/workflows/unstable.yml

@@ -32,3 +32,117 @@ jobs:
           echo
           echo "Once the jobs are deemed stable enough (% red signal < 5% and TTS < 3h),"
           echo " they can graduate and move back to pull or trunk."
+
+  #
+  # Experimental ARC jobs
+  #
+
+  linux-jammy-py3_8-gcc11-build:
+    name: linux-jammy-py3.8-gcc11
+    uses: ./.github/workflows/_linux-build-rg.yml
+    with:
+      build-environment: linux-jammy-py3.8-gcc11
+      docker-image-name: pytorch-linux-jammy-py3.8-gcc11
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "default", shard: 2, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "default", shard: 3, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "docs_test", shard: 1, num_shards: 1,  runner: "linux.2xlarge" },
+          { config: "jit_legacy", shard: 1, num_shards: 1, runner: "linux.2xlarge" },
+          { config: "backwards_compat", shard: 1, num_shards: 1, runner: "linux.2xlarge" },
+          { config: "distributed", shard: 1, num_shards: 2, runner: "linux.2xlarge" },
+          { config: "distributed", shard: 2, num_shards: 2, runner: "linux.2xlarge" },
+        ]}
+
+
+  linux-jammy-py3_8-gcc11-no-ops:
+    name: linux-jammy-py3.8-gcc11-no-ops
+    uses: ./.github/workflows/_linux-build-rg.yml
+    with:
+      build-environment: linux-jammy-py3.8-gcc11-no-ops
+      docker-image-name: pytorch-linux-jammy-py3.8-gcc11
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 1 },
+        ]}
+
+  linux-jammy-py3_8-gcc11-pch:
+    name: linux-jammy-py3.8-gcc11-pch
+    uses: ./.github/workflows/_linux-build-rg.yml
+    with:
+      build-environment: linux-jammy-py3.8-gcc11-pch
+      docker-image-name: pytorch-linux-jammy-py3.8-gcc11
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 1 },
+        ]}
+
+  linux-focal-py3_8-clang10-onnx-build:
+    name: linux-focal-py3.8-clang10-onnx
+    uses: ./.github/workflows/_linux-build-rg.yml
+    with:
+      build-environment: linux-focal-py3.8-clang10-onnx
+      docker-image-name: pytorch-linux-focal-py3-clang10-onnx
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 2, runner: "linux.2xlarge" },
+          { config: "default", shard: 2, num_shards: 2, runner: "linux.2xlarge" },
+        ]}
+
+  linux-jammy-py3_10-clang15-asan-build:
+    name: linux-jammy-py3.10-clang15-asan
+    uses: ./.github/workflows/_linux-build-rg.yml
+    with:
+      build-environment: linux-jammy-py3.10-clang15-asan
+      docker-image-name: pytorch-linux-jammy-py3-clang15-asan
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 6, runner: "linux.4xlarge" },
+          { config: "default", shard: 2, num_shards: 6, runner: "linux.4xlarge" },
+          { config: "default", shard: 3, num_shards: 6, runner: "linux.4xlarge" },
+          { config: "default", shard: 4, num_shards: 6, runner: "linux.4xlarge" },
+          { config: "default", shard: 5, num_shards: 6, runner: "linux.4xlarge" },
+          { config: "default", shard: 6, num_shards: 6, runner: "linux.4xlarge" },
+        ]}
+      sync-tag: asan-build-arc
+
+  linux-focal-py3_8-clang10-build:
+    name: linux-focal-py3.8-clang10
+    uses: ./.github/workflows/_linux-build-rg.yml
+    with:
+      build-environment: linux-focal-py3.8-clang10
+      docker-image-name: pytorch-linux-focal-py3.8-clang10
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "default", shard: 2, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "default", shard: 3, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "crossref", shard: 1, num_shards: 2, runner: "linux.2xlarge" },
+          { config: "crossref", shard: 2, num_shards: 2, runner: "linux.2xlarge" },
+          { config: "dynamo", shard: 1, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "dynamo", shard: 2, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "dynamo", shard: 3, num_shards: 3, runner: "linux.2xlarge" },
+        ]}
+
+  linux-focal-py3_11-clang10-build:
+    name: linux-focal-py3.11-clang10
+    uses: ./.github/workflows/_linux-build-rg.yml
+    with:
+      build-environment: linux-focal-py3.11-clang10
+      docker-image-name: pytorch-linux-focal-py3.11-clang10
+      test-matrix: |
+        { include: [
+          { config: "default", shard: 1, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "default", shard: 2, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "default", shard: 3, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "crossref", shard: 1, num_shards: 2, runner: "linux.2xlarge" },
+          { config: "crossref", shard: 2, num_shards: 2, runner: "linux.2xlarge" },
+          { config: "dynamo", shard: 1, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "dynamo", shard: 2, num_shards: 3, runner: "linux.2xlarge" },
+          { config: "dynamo", shard: 3, num_shards: 3, runner: "linux.2xlarge" },
+        ]}
+
+  #
+  # End of Experimental ARC jobs
+  #

.github/workflows/weekly.yml

@@ -25,7 +25,7 @@ jobs:
         with:
           repo-name: xla
           branch: master
-          pin-folder: .ci/docker/ci_commit_pins
+          pin-folder: .github/ci_commit_pins
           test-infra-ref: main
           updatebot-token: ${{ secrets.UPDATEBOT_TOKEN }}
           pytorchbot-token: ${{ secrets.GH_PYTORCHBOT_TOKEN }}

.gitignore

@@ -54,6 +54,7 @@ test/.coverage
 test/.hypothesis/
 test/cpp/api/mnist
 test/custom_operator/model.pt
+test/debug/
 test/jit_hooks/*.pt
 test/data/legacy_modules.t7
 test/data/*.pt

.gitmodules

@@ -149,3 +149,6 @@
 [submodule "third_party/mimalloc"]
 	path = third_party/mimalloc
 	url = https://github.com/microsoft/mimalloc.git
+[submodule "third_party/opentelemetry-cpp"]
+	path = third_party/opentelemetry-cpp
+	url = https://github.com/open-telemetry/opentelemetry-cpp.git

.isort.cfg

@@ -1,6 +0,0 @@
-[settings]
-include_trailing_comma=True
-multi_line_output=3
-skip=third_party
-skip_gitignore=True
-use_parentheses=True

.lintrunner.toml

@@ -135,7 +135,7 @@ init_command = [
     'numpy==1.24.3 ; python_version == "3.8"',
     'numpy==1.26.0 ; python_version >= "3.9"',
     'expecttest==0.1.6',
-    'mypy==1.8.0',
+    'mypy==1.9.0',
     'sympy==1.11.1',
     'types-requests==2.27.25',
     'types-PyYAML==6.0.7',
@@ -147,8 +147,8 @@ init_command = [
     'filelock==3.13.1',
     'junitparser==2.1.1',
     'rich==10.9.0',
-    'pyyaml==6.0',
-    'optree==0.10.0',
+    'pyyaml==6.0.1',
+    'optree==0.11.0',
 ]
 
 [[linter]]
@@ -186,11 +186,14 @@ command = [
 [[linter]]
 code = 'CLANGTIDY'
 include_patterns = [
-    'aten/src/ATen/core/*.cpp',
     # Enable coverage of headers in aten/src/ATen
     # and excluding most sub-directories for now.
     'aten/src/ATen/*.h',
     'aten/src/ATen/*.cpp',
+    'aten/src/ATen/core/*.h',
+    'aten/src/ATen/core/*.cpp',
+    'aten/src/ATen/functorch/*.h',
+    'aten/src/ATen/functorch/*.cpp',
     'c10/**/*.cpp',
     'c10/**/*.h',
     'torch/csrc/*.h',
@@ -204,9 +207,7 @@ exclude_patterns = [
     # CUDA files are also excluded.
     '**/fb/**',
     '**/*pb.h',
-    'c10/**/cuda/*pp',
     'aten/**/cuda/*pp',
-    '**/cuda/*pp',
     'c10/xpu/**/*.h',
     'c10/xpu/**/*.cpp',
     'c10/cuda/CUDAAlgorithm.h',
@@ -225,8 +226,8 @@ exclude_patterns = [
     'third_party/**/*',
     'torch/csrc/api/**',
     'torch/csrc/autograd/generated/**',
-    'torch/csrc/dynamo/*',
     'torch/csrc/distributed/**/*',
+    'torch/csrc/dynamo/eval_frame.h',
     'torch/csrc/inductor/**/*',
     'torch/csrc/jit/**/*',
     'torch/csrc/jit/serialization/import_legacy.cpp',
@@ -979,7 +980,7 @@ init_command = [
     'python3',
     'tools/linter/adapters/pip_init.py',
     '--dry-run={{DRYRUN}}',
-    'PyYAML==6.0',
+    'PyYAML==6.0.1',
 ]
 
 # Black + usort
@@ -1013,425 +1014,9 @@ exclude_patterns = [
     'test/_nvfuser/test_dynamo.py',
     'test/_nvfuser/test_python_frontend.py',
     'test/_nvfuser/test_torchscript.py',
-    'test/_test_bazel.py',
-    'test/ao/sparsity/test_activation_sparsifier.py',
-    'test/ao/sparsity/test_composability.py',
-    'test/ao/sparsity/test_data_scheduler.py',
-    'test/ao/sparsity/test_data_sparsifier.py',
-    'test/ao/sparsity/test_kernels.py',
-    'test/ao/sparsity/test_parametrization.py',
-    'test/ao/sparsity/test_qlinear_packed_params.py',
-    'test/ao/sparsity/test_scheduler.py',
-    'test/ao/sparsity/test_sparsifier.py',
-    'test/ao/sparsity/test_sparsity_utils.py',
-    'test/ao/sparsity/test_structured_sparsifier.py',
-    'test/autograd/test_complex.py',
-    'test/autograd/test_fallback.py',
-    'test/autograd/test_functional.py',
-    'test/backends/xeon/test_launch.py',
-    'test/benchmark_utils/test_benchmark_utils.py',
-    'test/bottleneck_test/test.py',
-    'test/bottleneck_test/test_args.py',
-    'test/bottleneck_test/test_cuda.py',
-    'test/conftest.py',
-    'test/cpp/__init__.py',
-    'test/cpp/aot_inductor/test.py',
-    'test/cpp/api/init_baseline.py',
-    'test/cpp/api/optim_baseline.py',
-    'test/cpp/jit/__init__.py',
-    'test/cpp/jit/tests_setup.py',
-    'test/cpp_api_parity/__init__.py',
-    'test/cpp_api_parity/functional_impl_check.py',
-    'test/cpp_api_parity/module_impl_check.py',
-    'test/cpp_api_parity/parity_table_parser.py',
-    'test/cpp_api_parity/sample_functional.py',
-    'test/cpp_api_parity/sample_module.py',
-    'test/cpp_api_parity/utils.py',
-    'test/cpp_extensions/no_python_abi_suffix_test/setup.py',
-    'test/cpp_extensions/setup.py',
-    'test/cpp_extensions/torch_test_cpp_extension/__init__.py',
-    'test/create_dummy_torchscript_model.py',
-    'test/custom_backend/backend.py',
-    'test/custom_backend/test_custom_backend.py',
-    'test/custom_operator/model.py',
-    'test/custom_operator/test_custom_ops.py',
     'test/delete.py',
-    'test/distributed/_shard/sharded_optim/test_sharded_optim.py',
-    'test/distributed/_shard/sharded_tensor/ops/test_binary_cmp.py',
-    'test/distributed/_shard/sharded_tensor/ops/test_embedding.py',
-    'test/distributed/_shard/sharded_tensor/ops/test_embedding_bag.py',
-    'test/distributed/_shard/sharded_tensor/ops/test_init.py',
-    'test/distributed/_shard/sharded_tensor/ops/test_tensor_ops.py',
-    'test/distributed/_shard/sharded_tensor/test_logger.py',
-    'test/distributed/_shard/sharded_tensor/test_sharded_tensor.py',
-    'test/distributed/_shard/sharded_tensor/test_sharded_tensor_reshard.py',
-    'test/distributed/_shard/sharding_plan/test_sharding_plan.py',
-    'test/distributed/_shard/sharding_spec/test_sharding_spec.py',
-    'test/distributed/_shard/test_sharder.py',
-    'test/distributed/_tools/test_memory_tracker.py',
-    'test/distributed/algorithms/ddp_comm_hooks/test_ddp_hooks.py',
-    'test/distributed/algorithms/quantization/test_quantization.py',
-    'test/distributed/algorithms/test_join.py',
-    'test/distributed/argparse_util_test.py',
-    'test/distributed/bin/test_script.py',
-    'test/distributed/elastic/agent/server/test/__init__.py',
-    'test/distributed/elastic/agent/server/test/api_test.py',
-    'test/distributed/elastic/agent/server/test/local_elastic_agent_test.py',
-    'test/distributed/elastic/events/lib_test.py',
-    'test/distributed/elastic/metrics/__init__.py',
-    'test/distributed/elastic/metrics/api_test.py',
-    'test/distributed/elastic/multiprocessing/api_test.py',
-    'test/distributed/elastic/multiprocessing/bin/echo1.py',
-    'test/distributed/elastic/multiprocessing/bin/echo2.py',
-    'test/distributed/elastic/multiprocessing/bin/echo3.py',
-    'test/distributed/elastic/multiprocessing/bin/test_script.py',
-    'test/distributed/elastic/multiprocessing/bin/zombie_test.py',
-    'test/distributed/elastic/multiprocessing/errors/api_test.py',
-    'test/distributed/elastic/multiprocessing/errors/error_handler_test.py',
-    'test/distributed/elastic/multiprocessing/redirects_test.py',
-    'test/distributed/elastic/multiprocessing/tail_log_test.py',
-    'test/distributed/elastic/rendezvous/__init__.py',
-    'test/distributed/elastic/rendezvous/api_test.py',
-    'test/distributed/elastic/rendezvous/c10d_rendezvous_backend_test.py',
-    'test/distributed/elastic/rendezvous/dynamic_rendezvous_test.py',
-    'test/distributed/elastic/rendezvous/etcd_rendezvous_backend_test.py',
-    'test/distributed/elastic/rendezvous/etcd_rendezvous_test.py',
-    'test/distributed/elastic/rendezvous/etcd_server_test.py',
-    'test/distributed/elastic/rendezvous/rendezvous_backend_test.py',
-    'test/distributed/elastic/rendezvous/static_rendezvous_test.py',
-    'test/distributed/elastic/rendezvous/utils_test.py',
-    'test/distributed/elastic/timer/__init__.py',
-    'test/distributed/elastic/timer/api_test.py',
-    'test/distributed/elastic/timer/file_based_local_timer_test.py',
-    'test/distributed/elastic/timer/local_timer_example.py',
-    'test/distributed/elastic/timer/local_timer_test.py',
-    'test/distributed/elastic/utils/__init__.py',
-    'test/distributed/elastic/utils/data/__init__.py',
-    'test/distributed/elastic/utils/data/cycling_iterator_test.py',
-    'test/distributed/elastic/utils/distributed_test.py',
-    'test/distributed/elastic/utils/logging_test.py',
-    'test/distributed/elastic/utils/util_test.py',
-    'test/distributed/launcher/__init__.py',
-    'test/distributed/launcher/api_test.py',
-    'test/distributed/launcher/bin/test_script.py',
-    'test/distributed/launcher/bin/test_script_init_method.py',
-    'test/distributed/launcher/bin/test_script_is_torchelastic_launched.py',
-    'test/distributed/launcher/bin/test_script_local_rank.py',
-    'test/distributed/launcher/launch_test.py',
-    'test/distributed/launcher/run_test.py',
-    'test/distributed/nn/jit/__init__.py',
-    'test/distributed/nn/jit/test_instantiator.py',
-    'test/distributed/optim/test_apply_optimizer_in_backward.py',
-    'test/distributed/optim/test_named_optimizer.py',
-    'test/distributed/optim/test_zero_redundancy_optimizer.py',
-    'test/distributed/pipeline/sync/__init__.py',
-    'test/distributed/pipeline/sync/conftest.py',
-    'test/distributed/pipeline/sync/skip/__init__.py',
-    'test/distributed/pipeline/sync/skip/test_api.py',
-    'test/distributed/pipeline/sync/skip/test_gpipe.py',
-    'test/distributed/pipeline/sync/skip/test_inspect_skip_layout.py',
-    'test/distributed/pipeline/sync/skip/test_leak.py',
-    'test/distributed/pipeline/sync/skip/test_portal.py',
-    'test/distributed/pipeline/sync/skip/test_stash_pop.py',
-    'test/distributed/pipeline/sync/skip/test_tracker.py',
-    'test/distributed/pipeline/sync/skip/test_verify_skippables.py',
-    'test/distributed/pipeline/sync/test_balance.py',
-    'test/distributed/pipeline/sync/test_bugs.py',
-    'test/distributed/pipeline/sync/test_checkpoint.py',
-    'test/distributed/pipeline/sync/test_copy.py',
-    'test/distributed/pipeline/sync/test_deferred_batch_norm.py',
-    'test/distributed/pipeline/sync/test_dependency.py',
-    'test/distributed/pipeline/sync/test_inplace.py',
-    'test/distributed/pipeline/sync/test_microbatch.py',
-    'test/distributed/pipeline/sync/test_phony.py',
-    'test/distributed/pipeline/sync/test_pipe.py',
-    'test/distributed/pipeline/sync/test_pipeline.py',
-    'test/distributed/pipeline/sync/test_stream.py',
-    'test/distributed/pipeline/sync/test_transparency.py',
-    'test/distributed/pipeline/sync/test_worker.py',
-    'test/distributed/rpc/cuda/test_tensorpipe_agent.py',
-    'test/distributed/rpc/test_faulty_agent.py',
-    'test/distributed/rpc/test_share_memory.py',
-    'test/distributed/rpc/test_tensorpipe_agent.py',
-    'test/distributed/tensor/parallel/__init__.py',
-    'test/distributed/tensor/parallel/test_ddp_2d_parallel.py',
-    'test/distributed/tensor/parallel/test_fsdp_2d_parallel.py',
-    'test/distributed/tensor/parallel/test_parallelize_api.py',
-    'test/distributed/tensor/parallel/test_tp_examples.py',
-    'test/distributed/tensor/parallel/test_tp_random_state.py',
-    'test/distributed/tensor/parallel/test_tp_style.py',
-    'test/distributed/tensor/parallel/test_view_sharding_dim_change.py',
-    'test/distributed/test_c10d_common.py',
-    'test/distributed/test_c10d_gloo.py',
-    'test/distributed/test_c10d_logger.py',
-    'test/distributed/test_c10d_nccl.py',
-    'test/distributed/test_c10d_object_collectives.py',
-    'test/distributed/test_c10d_pypg.py',
-    'test/distributed/test_c10d_spawn.py',
-    'test/distributed/test_c10d_spawn_gloo.py',
-    'test/distributed/test_c10d_spawn_nccl.py',
-    'test/distributed/test_c10d_spawn_ucc.py',
-    'test/distributed/test_c10d_ucc.py',
-    'test/distributed/test_collective_utils.py',
-    'test/distributed/test_data_parallel.py',
-    'test/distributed/test_distributed_spawn.py',
-    'test/distributed/test_dynamo_distributed.py',
-    'test/distributed/test_fake_pg.py',
-    'test/distributed/test_functional_api.py',
-    'test/distributed/test_inductor_collectives.py',
-    'test/distributed/test_launcher.py',
-    'test/distributed/test_multi_threaded_pg.py',
-    'test/distributed/test_nccl.py',
-    'test/distributed/test_pg_wrapper.py',
-    'test/distributed/test_store.py',
-    'test/distributions/test_constraints.py',
-    'test/distributions/test_distributions.py',
-    'test/distributions/test_transforms.py',
-    'test/distributions/test_utils.py',
-    'test/error_messages/storage.py',
     'test/expect/__init__.py',
-    'test/export/test_db.py',
-    'test/export/test_export.py',
-    'test/export/test_funtionalized_assertions.py',
-    'test/export/test_pass_infra.py',
-    'test/export/test_passes.py',
-    'test/export/test_serialize.py',
-    'test/export/test_upgrade.py',
-    'test/export/test_verifier.py',
-    'test/export/test_unflatten.py',
-    'test/forward_backward_compatibility/check_forward_backward_compatibility.py',
-    'test/forward_backward_compatibility/dump_all_function_schemas.py',
-    'test/functorch/attn_ft.py',
-    'test/functorch/attn_positional.py',
-    'test/functorch/common_utils.py',
-    'test/functorch/discover_coverage.py',
-    'test/functorch/functorch_additional_op_db.py',
-    'test/functorch/test_aotdispatch.py',
-    'test/functorch/test_control_flow.py',
-    'test/functorch/test_dims.py',
-    'test/functorch/test_eager_transforms.py',
-    'test/functorch/test_logging.py',
-    'test/functorch/test_memory_efficient_fusion.py',
-    'test/functorch/test_minifier.py',
-    'test/functorch/test_ops.py',
-    'test/functorch/test_parsing.py',
-    'test/functorch/test_rearrange.py',
-    'test/functorch/test_vmap.py',
-    'test/functorch/test_vmap_registrations.py',
-    'test/functorch/xfail_suggester.py',
-    'test/fx/named_tup.py',
-    'test/fx/quantization.py',
-    'test/fx/test_common_passes.py',
-    'test/fx/test_cse_pass.py',
-    'test/fx/test_dce_pass.py',
-    'test/fx/test_future.py',
-    'test/fx/test_fx_const_fold.py',
-    'test/fx/test_fx_param_shape_control_flow.py',
-    'test/fx/test_gradual_type.py',
-    'test/fx/test_matcher_utils.py',
-    'test/fx/test_pass_infra.py',
-    'test/fx/test_source_matcher_utils.py',
-    'test/fx/test_subgraph_rewriter.py',
-    'test/fx/test_z3_gradual_types.py',
-    'test/fx/test_fx_split.py',
-    'test/jit/__init__.py',
-    'test/jit/_imported_class_test/__init__.py',
-    'test/jit/_imported_class_test/bar.py',
-    'test/jit/_imported_class_test/foo.py',
-    'test/jit/_imported_class_test/very/__init__.py',
-    'test/jit/_imported_class_test/very/very/__init__.py',
-    'test/jit/_imported_class_test/very/very/nested.py',
-    'test/jit/fixtures_srcs/__init__.py',
-    'test/jit/fixtures_srcs/fixtures_src.py',
-    'test/jit/fixtures_srcs/generate_models.py',
-    'test/jit/fixtures_srcs/test_upgrader_models_generation.py',
-    'test/jit/myexception.py',
-    'test/jit/test_alias_analysis.py',
-    'test/jit/test_async.py',
-    'test/jit/test_aten_pow.py',
-    'test/jit/test_attr.py',
-    'test/jit/test_autodiff.py',
-    'test/jit/test_autodiff_subgraph_slicing.py',
-    'test/jit/test_await.py',
-    'test/jit/test_backend_nnapi.py',
-    'test/jit/test_backends.py',
-    'test/jit/test_batch_mm.py',
-    'test/jit/test_builtins.py',
-    'test/jit/test_class_type.py',
-    'test/jit/test_complex.py',
-    'test/jit/test_complexity.py',
-    'test/jit/test_convert_activation.py',
-    'test/jit/test_cuda.py',
-    'test/jit/test_custom_operators.py',
-    'test/jit/test_data_parallel.py',
-    'test/jit/test_dataclasses.py',
-    'test/jit/test_dce.py',
-    'test/jit/test_device_analysis.py',
-    'test/jit/test_dtype_analysis.py',
-    'test/jit/test_enum.py',
-    'test/jit/test_exception.py',
-    'test/jit/test_freezing.py',
-    'test/jit/test_functional_blocks.py',
-    'test/jit/test_fuser_common.py',
-    'test/jit/test_graph_rewrite_passes.py',
-    'test/jit/test_hash.py',
-    'test/jit/test_hooks.py',
-    'test/jit/test_hooks_modules.py',
-    'test/jit/test_ignorable_args.py',
-    'test/jit/test_ignore_context_manager.py',
-    'test/jit/test_isinstance.py',
-    'test/jit/test_jit_utils.py',
-    'test/jit/test_list_dict.py',
-    'test/jit/test_logging.py',
-    'test/jit/test_misc.py',
-    'test/jit/test_models.py',
-    'test/jit/test_module_apis.py',
-    'test/jit/test_module_containers.py',
-    'test/jit/test_module_interface.py',
-    'test/jit/test_modules.py',
-    'test/jit/test_op_decompositions.py',
-    'test/jit/test_optimize_for_mobile_preserve_debug_info.py',
-    'test/jit/test_parametrization.py',
-    'test/jit/test_pdt.py',
-    'test/jit/test_peephole.py',
-    'test/jit/test_profiler.py',
-    'test/jit/test_python_bindings.py',
-    'test/jit/test_python_builtins.py',
-    'test/jit/test_python_ir.py',
-    'test/jit/test_recursive_script.py',
-    'test/jit/test_remove_mutation.py',
-    'test/jit/test_save_load.py',
-    'test/jit/test_save_load_for_op_version.py',
-    'test/jit/test_script_profile.py',
-    'test/jit/test_scriptmod_ann.py',
-    'test/jit/test_slice.py',
-    'test/jit/test_sparse.py',
-    'test/jit/test_string_formatting.py',
-    'test/jit/test_symbolic_shape_analysis.py',
-    'test/jit/test_tensor_creation_ops.py',
-    'test/jit/test_tensor_methods.py',
-    'test/jit/test_torchbind.py',
-    'test/jit/test_tracer.py',
-    'test/jit/test_type_sharing.py',
-    'test/jit/test_types.py',
-    'test/jit/test_typing.py',
-    'test/jit/test_union.py',
-    'test/jit/test_unsupported_ops.py',
-    'test/jit/test_upgraders.py',
-    'test/jit/test_warn.py',
-    'test/jit/test_with.py',
-    'test/jit/xnnpack/test_xnnpack_delegate.py',
-    'test/jit_hooks/model.py',
-    'test/lazy/__init__.py',
-    'test/lazy/test_bindings.py',
-    'test/lazy/test_debug_util.py',
-    'test/lazy/test_extract_compiled_graph.py',
-    'test/lazy/test_meta_kernel.py',
-    'test/lazy/test_reuse_ir.py',
-    'test/lazy/test_step_closures.py',
-    'test/lazy/test_ts_opinfo.py',
-    'test/linear.py',
-    'test/load_torchscript_model.py',
-    'test/mkl_verbose.py',
-    'test/mkldnn_verbose.py',
-    'test/mobile/custom_build/prepare_model.py',
-    'test/mobile/lightweight_dispatch/tests_setup.py',
-    'test/mobile/model_test/android_api_module.py',
-    'test/mobile/model_test/builtin_ops.py',
-    'test/mobile/model_test/gen_test_model.py',
-    'test/mobile/model_test/math_ops.py',
-    'test/mobile/model_test/nn_ops.py',
-    'test/mobile/model_test/quantization_ops.py',
-    'test/mobile/model_test/sampling_ops.py',
-    'test/mobile/model_test/tensor_ops.py',
-    'test/mobile/model_test/torchvision_models.py',
-    'test/mobile/model_test/update_production_ops.py',
-    'test/mobile/nnc/aot_test_model.py',
-    'test/mobile/test_bytecode.py',
-    'test/mobile/test_lite_script_module.py',
-    'test/mobile/test_lite_script_type.py',
-    'test/mobile/test_quantize_fx_lite_script_module.py',
-    'test/mobile/test_upgrader_codegen.py',
-    'test/mobile/test_upgraders.py',
-    'test/nn/test_convolution.py',
-    'test/nn/test_dropout.py',
-    'test/nn/test_embedding.py',
-    'test/nn/test_init.py',
-    'test/nn/test_lazy_modules.py',
-    'test/nn/test_load_state_dict.py',
-    'test/nn/test_module_hooks.py',
-    'test/nn/test_multihead_attention.py',
-    'test/nn/test_packed_sequence.py',
-    'test/nn/test_parametrization.py',
-    'test/nn/test_pooling.py',
-    'test/nn/test_pruning.py',
-    'test/onnx_caffe2/export_onnx_tests_filter.py',
-    'test/onnx_caffe2/export_onnx_tests_generator.py',
-    'test/onnx_caffe2/test_caffe2_common.py',
-    'test/onnx_caffe2/test_custom_ops.py',
-    'test/onnx_caffe2/test_pytorch_helper.py',
-    'test/onnx_caffe2/test_pytorch_onnx_caffe2.py',
-    'test/onnx_caffe2/test_pytorch_onnx_caffe2_quantized.py',
-    'test/onnx_caffe2/test_verify.py',
-    'test/optim/test_lrscheduler.py',
-    'test/optim/test_optim.py',
-    'test/optim/test_swa_utils.py',
-    'test/package/__init__.py',
-    'test/package/common.py',
-    'test/package/generate_bc_packages.py',
-    'test/package/module_a.py',
-    'test/package/module_a_remapped_path.py',
-    'test/package/package_a/__init__.py',
-    'test/package/package_a/fake_interface.py',
-    'test/package/package_a/fake_script_class.py',
-    'test/package/package_a/long_name.py',
-    'test/package/package_a/std_sys_module_hacks.py',
-    'test/package/package_a/subpackage.py',
-    'test/package/package_a/test_all_leaf_modules_tracer.py',
-    'test/package/package_a/test_module.py',
-    'test/package/package_a/test_nn_module.py',
-    'test/package/package_a/use_dunder_package.py',
-    'test/package/package_a/use_torch_package_importer.py',
-    'test/package/package_b/__init__.py',
-    'test/package/package_b/subpackage_0/__init__.py',
-    'test/package/package_b/subpackage_0/subsubpackage_0/__init__.py',
-    'test/package/package_b/subpackage_1.py',
-    'test/package/package_b/subpackage_2.py',
-    'test/package/package_c/__init__.py',
-    'test/package/package_c/test_module.py',
-    'test/package/package_d/__init__.py',
-    'test/package/package_d/imports_directly.py',
-    'test/package/package_d/imports_indirectly.py',
-    'test/package/package_d/subpackage_0/__init__.py',
-    'test/package/package_d/subpackage_0/subsubpackage_0/__init__.py',
-    'test/package/test_analyze.py',
-    'test/package/test_dependency_api.py',
-    'test/package/test_dependency_hooks.py',
-    'test/package/test_digraph.py',
-    'test/package/test_directory_reader.py',
-    'test/package/test_glob_group.py',
-    'test/package/test_importer.py',
-    'test/package/test_load_bc_packages.py',
-    'test/package/test_mangling.py',
-    'test/package/test_misc.py',
-    'test/package/test_model.py',
-    'test/package/test_package_fx.py',
-    'test/package/test_package_script.py',
-    'test/package/test_repackage.py',
-    'test/package/test_resources.py',
-    'test/package/test_save_load.py',
-    'test/package/test_trace_dep/__init__.py',
-    'test/profiler/test_memory_profiler.py',
-    'test/profiler/test_profiler.py',
-    'test/profiler/test_profiler_tree.py',
     'test/quantization/__init__.py',
-    'test/quantization/ao_migration/__init__.py',
-    'test/quantization/ao_migration/common.py',
-    'test/quantization/ao_migration/test_ao_migration.py',
-    'test/quantization/ao_migration/test_quantization.py',
-    'test/quantization/ao_migration/test_quantization_fx.py',
-    'test/quantization/bc/__init__.py',
-    'test/quantization/bc/test_backward_compatibility.py',
     'test/quantization/core/__init__.py',
     'test/quantization/core/experimental/apot_fx_graph_mode_ptq.py',
     'test/quantization/core/experimental/apot_fx_graph_mode_qat.py',
@@ -1466,50 +1051,17 @@ exclude_patterns = [
     'test/quantization/fx/test_numeric_suite_fx.py',
     'test/quantization/fx/test_quantize_fx.py',
     'test/quantization/fx/test_subgraph_rewriter.py',
-    'test/quantization/jit/__init__.py',
-    'test/quantization/jit/test_deprecated_jit_quant.py',
-    'test/quantization/jit/test_fusion_passes.py',
-    'test/quantization/jit/test_ondevice_quantization.py',
-    'test/quantization/jit/test_quantize_jit.py',
-    'test/quantization/pt2e/test_graph_utils.py',
-    'test/quantization/pt2e/test_quantize_pt2e.py',
-    'test/quantization/pt2e/test_x86inductor_quantizer.py',
-    'test/scripts/cuda_memcheck_common.py',
-    'test/scripts/run_cuda_memcheck.py',
-    'test/simulate_nccl_errors.py',
-    'test/test_ao_sparsity.py',
-    'test/test_autocast.py',
-    'test/test_autograd.py',
-    'test/test_binary_ufuncs.py',
-    'test/test_bundled_images.py',
-    'test/test_bundled_inputs.py',
-    'test/test_comparison_utils.py',
-    'test/test_compile_benchmark_util.py',
-    'test/test_complex.py',
-    'test/test_cpp_api_parity.py',
-    'test/test_cpp_extensions_aot.py',
-    'test/test_cpp_extensions_jit.py',
-    'test/test_cpp_extensions_open_device_registration.py',
-    'test/test_cuda.py',
-    'test/test_cuda_expandable_segments.py',
-    'test/test_cuda_multigpu.py',
-    'test/test_cuda_nvml_based_avail.py',
-    'test/test_cuda_primary_ctx.py',
-    'test/test_cuda_sanitizer.py',
-    'test/test_cuda_trace.py',
     'test/test_custom_op_testing.py',
     'test/test_dataloader.py',
     'test/test_datapipe.py',
     'test/test_decomp.py',
     'test/test_deploy.py',
     'test/test_determination.py',
-    'test/test_dispatch.py',
     'test/test_dlpack.py',
     'test/test_dynamic_shapes.py',
     'test/test_expanded_weights.py',
     'test/test_fake_tensor.py',
     'test/test_flop_counter.py',
-    'test/test_foreach.py',
     'test/test_function_schema.py',
     'test/test_functional_autograd_benchmark.py',
     'test/test_functional_optim.py',
@@ -1538,7 +1090,6 @@ exclude_patterns = [
     'test/test_jit_string.py',
     'test/test_jiterator.py',
     'test/test_kernel_launch_checks.py',
-    'test/test_legacy_vmap.py',
     'test/test_license.py',
     'test/test_linalg.py',
     'test/test_logging.py',
@@ -1556,7 +1107,6 @@ exclude_patterns = [
     'test/test_modules.py',
     'test/test_monitor.py',
     'test/test_mps.py',
-    'test/test_multiprocessing.py',
     'test/test_multiprocessing_spawn.py',
     'test/test_namedtensor.py',
     'test/test_namedtuple_return_api.py',
@@ -1570,10 +1120,6 @@ exclude_patterns = [
     'test/test_nvfuser_dynamo.py',
     'test/test_nvfuser_frontend.py',
     'test/test_openmp.py',
-    'test/test_ops.py',
-    'test/test_ops_fwd_gradients.py',
-    'test/test_ops_gradients.py',
-    'test/test_ops_jit.py',
     'test/test_optim.py',
     'test/test_out_dtype_op.py',
     'test/test_overrides.py',
@@ -1583,7 +1129,6 @@ exclude_patterns = [
     'test/test_proxy_tensor.py',
     'test/test_pruning_op.py',
     'test/test_public_bindings.py',
-    'test/test_python_dispatch.py',
     'test/test_quantization.py',
     'test/test_reductions.py',
     'test/test_scatter_gather_ops.py',
@@ -1692,24 +1237,6 @@ exclude_patterns = [
     'torch/_export/serde/upgrade.py',
     'torch/_export/trace.py',
     'torch/_export/verifier.py',
-    'torch/_functorch/__init__.py',
-    'torch/_functorch/aot_autograd.py',
-    'torch/_functorch/apis.py',
-    'torch/_functorch/autograd_function.py',
-    'torch/_functorch/batch_norm_replacement.py',
-    'torch/_functorch/benchmark_utils.py',
-    'torch/_functorch/compile_utils.py',
-    'torch/_functorch/compilers.py',
-    'torch/_functorch/config.py',
-    'torch/_functorch/deprecated.py',
-    'torch/_functorch/eager_transforms.py',
-    'torch/_functorch/fx_minifier.py',
-    'torch/_functorch/partitioners.py',
-    'torch/_functorch/pyfunctorch.py',
-    'torch/_functorch/python_key.py',
-    'torch/_functorch/top_operators_github_usage.py',
-    'torch/_functorch/utils.py',
-    'torch/_functorch/vmap.py',
     'torch/_higher_order_ops/__init__.py',
     'torch/_higher_order_ops/out_dtype.py',
     'torch/_higher_order_ops/wrap.py',
@@ -1907,6 +1434,7 @@ exclude_patterns = [
     'torch/compiler/__init__.py',
     'torch/contrib/__init__.py',
     'torch/contrib/_tensorboard_vis.py',
+    "torch/cuda/_gpu_trace.py",
     'torch/cuda/_memory_viz.py',  # mypy: Value of type "object" is not indexable
     'torch/distributed/__init__.py',
     'torch/distributed/_composable_state.py',
@@ -2319,25 +1847,6 @@ exclude_patterns = [
     'torch/nn/utils/rnn.py',
     'torch/nn/utils/spectral_norm.py',
     'torch/nn/utils/weight_norm.py',
-    'torch/optim/__init__.py',
-    'torch/optim/_functional.py',
-    'torch/optim/_multi_tensor/__init__.py',
-    'torch/optim/adadelta.py',
-    'torch/optim/adagrad.py',
-    'torch/optim/adam.py',
-    'torch/optim/adamax.py',
-    'torch/optim/adamw.py',
-    'torch/optim/asgd.py',
-    'torch/optim/lbfgs.py',
-    'torch/optim/lr_scheduler.py',
-    'torch/optim/nadam.py',
-    'torch/optim/optimizer.py',
-    'torch/optim/radam.py',
-    'torch/optim/rmsprop.py',
-    'torch/optim/rprop.py',
-    'torch/optim/sgd.py',
-    'torch/optim/sparse_adam.py',
-    'torch/optim/swa_utils.py',
     'torch/overrides.py',
     'torch/quasirandom.py',
     'torch/random.py',
@@ -2372,7 +1881,7 @@ exclude_patterns = [
     'torch/testing/_internal/common_subclass.py',
     'torch/testing/_internal/common_utils.py',
     'torch/testing/_internal/composite_compliance.py',
-    'torch/testing/_internal/control_flow_opinfo_db.py',
+    'torch/testing/_internal/hop_db.py',
     'torch/testing/_internal/custom_op_db.py',
     'torch/testing/_internal/data/__init__.py',
     'torch/testing/_internal/data/network1.py',
@@ -2428,13 +1937,10 @@ exclude_patterns = [
     'torch/testing/_internal/test_module/__init__.py',
     'torch/testing/_internal/test_module/future_div.py',
     'torch/testing/_internal/test_module/no_future_div.py',
-    'torch/torch_version.py',
-    'torch/types.py',
     'torch/utils/__init__.py',
     'torch/utils/_contextlib.py',
     'torch/utils/_cpp_extension_versioner.py',
     'torch/utils/_crash_handler.py',
-    'torch/utils/_cuda_trace.py',
     'torch/utils/_device.py',
     'torch/utils/_foreach_utils.py',
     'torch/utils/_freeze.py',
@@ -2443,7 +1949,6 @@ exclude_patterns = [
     'torch/utils/_stats.py',
     'torch/utils/_sympy/__init__.py',
     'torch/utils/_sympy/functions.py',
-    'torch/utils/_sympy/value_ranges.py',
     'torch/utils/_traceback.py',
     'torch/utils/_zip.py',
     'torch/utils/backcompat/__init__.py',
@@ -2563,6 +2068,7 @@ exclude_patterns = [
     'torch/utils/viz/__init__.py',
     'torch/utils/viz/_cycles.py',
     'torch/utils/weak.py',
+    'torch/xpu/_gpu_trace.py',
 ]
 init_command = [
     'python3',
@@ -2655,7 +2161,7 @@ init_command = [
     'python3',
     'tools/linter/adapters/pip_init.py',
     '--dry-run={{DRYRUN}}',
-    'ruff==0.3.0',
+    'ruff==0.4.1',
 ]
 is_formatter = true
 
@@ -2740,7 +2246,6 @@ exclude_patterns = [
     "aten/src/ATen/Context.cpp",
     "aten/src/ATen/DLConvertor.cpp",
     "aten/src/ATen/core/Array.h",
-    "aten/src/ATen/native/ConvUtils.h",
     "aten/src/ATen/native/quantized/ConvUtils.h",
     "aten/src/ATen/native/sparse/SparseBlasImpl.cpp",  # triton implementation
     "aten/src/ATen/native/transformers/attention.cpp",

BUILD.bazel

@@ -1762,12 +1762,10 @@ py_library(
     visibility = ["//visibility:public"],
     srcs = glob(["torch/**/*.py"], exclude = ["torch/version.py"]) + [":torch/version.py"] + glob(["functorch/**/*.py"]),
     deps = [
-        rules.requirement("future"),
         rules.requirement("numpy"),
         rules.requirement("pyyaml"),
         rules.requirement("requests"),
         rules.requirement("setuptools"),
-        rules.requirement("six"),
         rules.requirement("sympy"),
         rules.requirement("typing_extensions"),
         "//torchgen",

CITATION.cff

@@ -6,68 +6,111 @@ authors:
 url: https://pytorch.org
 preferred-citation:
   type: conference-paper
-  title: "PyTorch: An Imperative Style, High-Performance Deep Learning Library"
+  title: "PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation"
   authors:
-    - family-names: Paszke
-      given-names: Adam
-    - family-names: Gross
-      given-names: Sam
-    - family-names: Massa
-      given-names: Francisco
-    - family-names: Lerer
-      given-names: Adam
-    - family-names: Bradbury
-      given-names: James
-    - family-names: Chanan
-      given-names: Gregory
-    - family-names: Killeen
-      given-names: Trevor
-    - family-names: Lin
-      given-names: Zeming
-    - family-names: Gimelshein
-      given-names: Natalia
-    - family-names: Antiga
-      given-names: Luca
-    - family-names: Desmaison
-      given-names: Alban
-    - family-names: Kopf
-      given-names: Andreas
+    - family-names: Ansel
+      given-names: Jason
     - family-names: Yang
       given-names: Edward
+    - family-names: He
+      given-names: Horace
+    - family-names: Gimelshein
+      given-names: Natalia
+    - family-names: Jain
+      given-names: Animesh
+    - family-names: Voznesensky
+      given-names: Michael
+    - family-names: Bao
+      given-names: Bin
+    - family-names: Bell
+      given-names: Peter
+    - family-names: Berard
+      given-names: David
+    - family-names: Burovski
+      given-names: Evgeni
+    - family-names: Chauhan
+      given-names: Geeta
+    - family-names: Chourdia
+      given-names: Anjali
+    - family-names: Constable
+      given-names: Will
+    - family-names: Desmaison
+      given-names: Alban
     - family-names: DeVito
       given-names: Zachary
-    - family-names: Raison
-      given-names: Martin
-    - family-names: Tejani
-      given-names: Alykhan
-    - family-names: Chilamkurthy
-      given-names: Sasank
-    - family-names: Steiner
-      given-names: Benoit
-    - family-names: Fang
-      given-names: Lu
-    - family-names: Bai
-      given-names: Junjie
+    - family-names: Ellison
+      given-names: Elias
+    - family-names: Feng
+      given-names: Will
+    - family-names: Gong
+      given-names: Jiong
+    - family-names: Gschwind
+      given-names: Michael
+    - family-names: Hirsh
+      given-names: Brian
+    - family-names: Huang
+      given-names: Sherlock
+    - family-names: Kalambarkar
+      given-names: Kshiteej
+    - family-names: Kirsch
+      given-names: Laurent
+    - family-names: Lazos
+      given-names: Michael
+    - family-names: Lezcano
+      given-names: Mario
+    - family-names: Liang
+      given-names: Yanbo
+    - family-names: Liang
+      given-names: Jason
+    - family-names: Lu
+      given-names: Yinghai
+    - family-names: Luk
+      given-names: CK
+    - family-names: Maher
+      given-names: Bert
+    - family-names: Pan
+      given-names: Yunjie
+    - family-names: Puhrsch
+      given-names: Christian
+    - family-names: Reso
+      given-names: Matthias
+    - family-names: Saroufim
+      given-names: Mark
+    - family-names: Siraichi
+      given-names: Marcos Yukio
+    - family-names: Suk
+      given-names: Helen
+    - family-names: Suo
+      given-names: Michael
+    - family-names: Tillet
+      given-names: Phil
+    - family-names: Wang
+      given-names: Eikan
+    - family-names: Wang
+      given-names: Xiaodong
+    - family-names: Wen
+      given-names: William
+    - family-names: Zhang
+      given-names: Shunting
+    - family-names: Zhao
+      given-names: Xu
+    - family-names: Zhou
+      given-names: Keren
+    - family-names: Zou
+      given-names: Richard
+    - family-names: Mathews
+      given-names: Ajit
+    - family-names: Chanan
+      given-names: Gregory
+    - family-names: Wu
+      given-names: Peng
     - family-names: Chintala
       given-names: Soumith
-  collection-title: Advances in Neural Information Processing Systems 32
+  collection-title: "29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2 (ASPLOS '24)"
   collection-type: proceedings
-  editors:
-    - family-names: Wallach
-      given-names: H.
-    - family-names: Larochelle
-      given-names: H.
-    - family-names: Beygelzimer
-      given-names: A.
-    - family-names: "d'Alché-Buc"
-      given-names: F.
-    - family-names: Fox
-      given-names: E.
-    - family-names: Garnett
-      given-names: R.
-  start: 8024
-  end: 8035
-  year: 2019
+  month: 4
+  year: 2024
   publisher:
-    name: Curran Associates, Inc.
-  url: http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
+    name: ACM
+  doi: "10.1145/3620665.3640366"
+  url: "https://pytorch.org/assets/pytorch2-2.pdf"

CMakeLists.txt

@@ -19,6 +19,12 @@ cmake_policy(SET CMP0069 NEW)
 # nice when it's possible, and it's possible on our Windows configs.
 cmake_policy(SET CMP0092 NEW)
 
+# Prohibit in-source builds
+if(${CMAKE_SOURCE_DIR} STREQUAL ${CMAKE_BINARY_DIR})
+message(FATAL_ERROR "In-source build are not supported")
+endif()
+
+
 # ---[ Project and semantic versioning.
 project(Torch CXX C)
 

CODEOWNERS

@@ -43,12 +43,12 @@ nn/qat/ @jerryzh168
 /torch/csrc/distributed/rpc/tensorpipe_agent.h @jiayisuse @osalpekar @lw
 
 # ONNX Export
-/torch/_dynamo/backends/onnxrt.py @bowenbao @abock @thiagocrepaldi @wschin
-/torch/csrc/jit/passes/onnx.h @bowenbao @abock @thiagocrepaldi
-/torch/csrc/jit/passes/onnx.cpp @bowenbao @abock @thiagocrepaldi
-/torch/csrc/jit/passes/onnx/ @bowenbao @abock @thiagocrepaldi
-/torch/onnx/ @bowenbao @abock @thiagocrepaldi @wschin
-/test/onnx/ @bowenbao @abock @thiagocrepaldi @wschin
+/torch/_dynamo/backends/onnxrt.py @bowenbao @thiagocrepaldi @wschin
+/torch/csrc/jit/passes/onnx.h @bowenbao @thiagocrepaldi
+/torch/csrc/jit/passes/onnx.cpp @bowenbao @thiagocrepaldi
+/torch/csrc/jit/passes/onnx/ @bowenbao @thiagocrepaldi
+/torch/onnx/ @bowenbao @thiagocrepaldi @wschin
+/test/onnx/ @bowenbao @thiagocrepaldi @wschin
 
 # CI
 /.ci  @pytorch/pytorch-dev-infra
@@ -67,6 +67,7 @@ nn/qat/ @jerryzh168
 /test/run_test.py @pytorch/pytorch-dev-infra
 /torch/testing/_internal/common_device_type.py @mruberry
 /torch/testing/_internal/common_utils.py @pytorch/pytorch-dev-infra
+/torch/testing/_internal/hop_db.py @tugsbayasgalan @zou3519 @ydwu4
 
 # Parametrizations
 /torch/nn/utils/parametriz*.py @lezcano
@@ -101,6 +102,10 @@ test/test_mps.py @kulinseth @malfet
 aten/src/ATen/mps/ @kulinseth @malfet
 aten/src/ATen/native/mps/ @kulinseth @malfet
 
+# MTIA
+aten/src/ATen/detail/MTIAHooksInterface.h @egienvalue
+torch/csrc/mtia/ @egienvalue
+
 # Profiler
 torch/csrc/autograd/profiler* @aaronenyeshi
 torch/autograd/profiler* @aaronenyeshi

Dockerfile

@@ -7,8 +7,8 @@
 #
 #       For reference:
 #           https://docs.docker.com/develop/develop-images/build_enhancements/
-ARG BASE_IMAGE=ubuntu:20.04
-ARG PYTHON_VERSION=3.8
+ARG BASE_IMAGE=ubuntu:22.04
+ARG PYTHON_VERSION=3.11
 
 FROM ${BASE_IMAGE} as dev-base
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
@@ -26,7 +26,7 @@ RUN mkdir /opt/ccache && ccache --set-config=cache_dir=/opt/ccache
 ENV PATH /opt/conda/bin:$PATH
 
 FROM dev-base as conda
-ARG PYTHON_VERSION=3.8
+ARG PYTHON_VERSION=3.11
 # Automatically set by buildx
 ARG TARGETPLATFORM
 # translating Docker's TARGETPLATFORM into miniconda arches
@@ -57,12 +57,12 @@ COPY --from=submodule-update /opt/pytorch /opt/pytorch
 RUN make triton
 RUN --mount=type=cache,target=/opt/ccache \
     export eval ${CMAKE_VARS} && \
-    TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX 8.0" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
+    TORCH_CUDA_ARCH_LIST="7.0 7.2 7.5 8.0 8.6 8.7 8.9 9.0 9.0a" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
     CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
     python setup.py install
 
 FROM conda as conda-installs
-ARG PYTHON_VERSION=3.8
+ARG PYTHON_VERSION=3.11
 ARG CUDA_VERSION=12.1
 ARG CUDA_CHANNEL=nvidia
 ARG INSTALL_CHANNEL=pytorch-nightly

RELEASE.md

@@ -50,6 +50,7 @@ Following is the Release Compatibility Matrix for PyTorch releases:
 
 | PyTorch version | Python | Stable CUDA | Experimental CUDA |
 | --- | --- | --- | --- |
+| 2.3 | >=3.8, <=3.11, (3.12 experimental) | CUDA 11.8, CUDNN 8.7.0.84 | CUDA 12.1, CUDNN 8.9.2.26 |
 | 2.2 | >=3.8, <=3.11, (3.12 experimental) | CUDA 11.8, CUDNN 8.7.0.84 | CUDA 12.1, CUDNN 8.9.2.26 |
 | 2.1 | >=3.8, <=3.11 | CUDA 11.8, CUDNN 8.7.0.84 | CUDA 12.1, CUDNN 8.9.2.26 |
 | 2.0 | >=3.8, <=3.11 | CUDA 11.7, CUDNN 8.5.0.96 | CUDA 11.8, CUDNN 8.7.0.84 |
@@ -218,7 +219,7 @@ Validate the release jobs for pytorch and domain libraries should be green. Vali
   * [TorchVision](https://hud.pytorch.org/hud/pytorch/vision/release%2F1.12)
   * [TorchAudio](https://hud.pytorch.org/hud/pytorch/audio/release%2F1.12)
 
-Validate that the documentation build has completed and generated entry corresponding to the release in  [docs folder](https://github.com/pytorch/pytorch.github.io/tree/site/docs/) of pytorch.github.io repository
+Validate that the documentation build has completed and generated entry corresponding to the release in  [docs repository](https://github.com/pytorch/docs/tree/main/).
 
 ### Cherry Picking Fixes
 

SECURITY.md

@@ -1,9 +1,56 @@
-# Reporting Security Issues
+# Security Policy
 
-If you believe you have found a security vulnerability in PyTorch, we encourage you to let us know right away. We will investigate all legitimate reports and do our best to quickly fix the problem.
+ - [**Reporting a Vulnerability**](#reporting-a-vulnerability)
+ - [**Using Pytorch Securely**](#using-pytorch-securely)
+   - [Untrusted models](#untrusted-models)
+   - [Untrusted inputs](#untrusted-inputs)
+   - [Data privacy](#data-privacy)
+
+## Reporting Security Issues
+
+Beware that none of the topics under [Using Pytorch Securely](#using-pytorch-securely) are considered vulnerabilities of Pytorch.
+
+However, if you believe you have found a security vulnerability in PyTorch, we encourage you to let us know right away. We will investigate all legitimate reports and do our best to quickly fix the problem.
 
 Please report security issues using https://github.com/pytorch/pytorch/security/advisories/new
 
 Please refer to the following page for our responsible disclosure policy, reward guidelines, and those things that should not be reported:
 
 https://www.facebook.com/whitehat
+
+
+## Using Pytorch Securely
+**Pytorch models are programs**, so treat its security seriously -- running untrusted models is equivalent to running untrusted code. In general we recommend that model weights and the python code for the model are distributed independently. That said, be careful about where you get the python code from and who wrote it (preferentially check for a provenance or checksums, do not run any pip installed package).
+
+### Untrusted models
+Be careful when running untrusted models. This classification includes models created by unknown developers or utilizing data obtained from unknown sources[^data-poisoning-sources].
+
+**Prefer to execute untrusted models within a secure, isolated environment such as a sandbox** (e.g., containers, virtual machines). This helps protect your system from potentially malicious code. You can find further details and instructions in [this page](https://developers.google.com/code-sandboxing).
+
+**Be mindful of risky model formats**. Give preference to share and load weights with the appropriate format for your use case. [safetensors](https://huggingface.co/docs/safetensors/en/index) gives the most safety but is the most restricted in what it supports. [`torch.load`](https://pytorch.org/docs/stable/generated/torch.load.html#torch.load) with `weights_only=True` is also secure to our knowledge even though it offers significantly larger surface of attack. Loading un-trusted checkpoint with `weights_only=False` MUST never be done.
+
+
+
+Important Note: The trustworthiness of a model is not binary. You must always determine the proper level of caution depending on the specific model and how it matches your use case and risk tolerance.
+
+[^data-poisoning-sources]: To understand risks of utilization of data from unknown sources, read the following Cornell papers on Data poisoning:
+    https://arxiv.org/abs/2312.04748
+    https://arxiv.org/abs/2401.05566
+
+### Untrusted inputs during training and prediction
+
+If you plan to open your model to untrusted inputs, be aware that inputs can also be used as vectors by malicious agents. To minimize risks, make sure to give your model only the permisisons strictly required, and keep your libraries updated with the lates security patches.
+
+If applicable, prepare your model against bad inputs and prompt injections. Some recommendations:
+- Pre-analysis: check how the model performs by default when exposed to prompt injection (e.g. using fuzzing for prompt injection).
+- Input Sanitation: Before feeding data to the model, sanitize inputs rigorously. This involves techniques such as:
+    - Validation: Enforce strict rules on allowed characters and data types.
+    - Filtering: Remove potentially malicious scripts or code fragments.
+    - Encoding: Convert special characters into safe representations.
+    - Verification: Run tooling that identifies potential script injections (e.g. [models that detect prompt injection attempts](https://python.langchain.com/docs/guides/safety/hugging_face_prompt_injection)).
+
+### Data privacy
+
+**Take special security measures if your model if you train models with sensitive data**. Prioritize [sandboxing](https://developers.google.com/code-sandboxing) your models and:
+- Do not feed sensitive data to untrusted model (even if runs in a sandboxed environment)
+- If you consider publishing a model that was partially trained with sensitive data, be aware that data can potentially be recovered from the trained weights (especially if model overfits).

WORKSPACE

@@ -71,6 +71,13 @@ http_archive(
     ],
 )
 
+http_archive(
+    name = "com_github_opentelemetry-cpp",
+    urls = [
+        "https://github.com/open-telemetry/opentelemetry-cpp/archive/refs/tags/v1.14.2.tar.gz",
+    ],
+)
+
 new_local_repository(
     name = "gloo",
     build_file = "//third_party:gloo.BUILD",
@@ -155,6 +162,12 @@ new_local_repository(
     path = "third_party/kineto",
 )
 
+new_local_repository(
+    name = "opentelemetry-cpp",
+    build_file = "//third_party::opentelemetry-cpp.BUILD",
+    path = "third_party/opentelemetry-cpp",
+)
+
 new_patched_local_repository(
     name = "tbb",
     build_file = "//third_party:tbb.BUILD",
@@ -206,11 +219,11 @@ py_repositories()
 load("@rules_python//python:repositories.bzl", "python_register_toolchains")
 
 python_register_toolchains(
-    name = "python3_8",
-    python_version = "3.8",
+    name = "python3_10",
+    python_version = "3.10",
 )
 
-load("@python3_8//:defs.bzl", "interpreter")
+load("@python3_10//:defs.bzl", "interpreter")
 load("@rules_python//python:pip.bzl", "pip_parse")
 
 pip_parse(

aten/CMakeLists.txt

@@ -18,6 +18,8 @@ cmake_policy(SET CMP0012 NEW)
 #############################################
 
 set(ATen_CPU_SRCS)
+set(ATen_XPU_SRCS)
+set(ATen_XPU_INCLUDE)
 set(ATen_CPU_TEST_SRCS)
 set(ATen_CPU_INCLUDE)
 set(ATen_THIRD_PARTY_INCLUDE)
@@ -39,6 +41,7 @@ set(ATen_XPU_INCLUDE)
 set(ATen_XPU_TEST_SRCS)
 set(ATen_VULKAN_TEST_SRCS)
 set(ATen_CPU_DEPENDENCY_LIBS)
+set(ATen_XPU_DEPENDENCY_LIBS)
 set(ATen_CUDA_DEPENDENCY_LIBS)
 set(ATen_HIP_DEPENDENCY_LIBS)
 set(ATen_PUBLIC_CUDA_DEPENDENCY_LIBS)
@@ -105,6 +108,8 @@ add_subdirectory(src/ATen)
 # Pass source, includes, and libs to parent
 set(ATen_CPU_SRCS ${ATen_CPU_SRCS} PARENT_SCOPE)
 set(ATen_CORE_SRCS ${ATen_CORE_SRCS} PARENT_SCOPE)
+set(ATen_XPU_SRCS ${ATen_XPU_SRCS} PARENT_SCOPE)
+set(ATen_XPU_INCLUDE ${ATen_XPU_INCLUDE} PARENT_SCOPE)
 set(ATen_CUDA_CU_SRCS ${ATen_CUDA_CU_SRCS} PARENT_SCOPE)
 set(ATen_CUDA_CPP_SRCS ${ATen_CUDA_CPP_SRCS} PARENT_SCOPE)
 set(ATen_CUDA_LINALG_SRCS ${ATen_CUDA_LINALG_SRCS} PARENT_SCOPE)
@@ -130,6 +135,7 @@ set(ATen_HIP_INCLUDE ${ATen_HIP_INCLUDE} PARENT_SCOPE)
 set(ATen_XPU_INCLUDE ${ATen_XPU_INCLUDE} PARENT_SCOPE)
 set(ATen_THIRD_PARTY_INCLUDE ${ATen_THIRD_PARTY_INCLUDE} PARENT_SCOPE)
 set(ATen_CPU_DEPENDENCY_LIBS ${ATen_CPU_DEPENDENCY_LIBS} PARENT_SCOPE)
+set(ATen_XPU_DEPENDENCY_LIBS ${ATen_XPU_DEPENDENCY_LIBS} PARENT_SCOPE)
 set(ATen_CUDA_DEPENDENCY_LIBS ${ATen_CUDA_DEPENDENCY_LIBS} PARENT_SCOPE)
 set(ATen_HIP_DEPENDENCY_LIBS ${ATen_HIP_DEPENDENCY_LIBS} PARENT_SCOPE)
 set(ATen_CORE_TEST_SRCS ${ATen_CORE_TEST_SRCS} PARENT_SCOPE)

aten/src/ATen/BlasBackend.h

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <c10/util/Exception.h>
+
+#include <ostream>
+#include <string>
+
+namespace at {
+
+enum class BlasBackend : int8_t { Cublas, Cublaslt };
+
+inline std::string BlasBackendToString(at::BlasBackend backend) {
+  switch (backend) {
+    case BlasBackend::Cublas:
+      return "at::BlasBackend::Cublas";
+    case BlasBackend::Cublaslt:
+      return "at::BlasBackend::Cublaslt";
+    default:
+      TORCH_CHECK(false, "Unknown blas backend");
+  }
+}
+
+inline std::ostream& operator<<(std::ostream& stream, at::BlasBackend backend) {
+  return stream << BlasBackendToString(backend);
+}
+
+} // namespace at

aten/src/ATen/CMakeLists.txt

@@ -48,6 +48,7 @@ endif()
 # NB: If you edit these globs, you'll have to update setup.py package_data as well
 file(GLOB_RECURSE ATen_CORE_HEADERS  "core/*.h")
 file(GLOB_RECURSE ATen_CORE_SRCS "core/*.cpp")
+file(GLOB_RECURSE ATen_TRANSFORMER_HEADERS "native/transformers/*.h")
 if(NOT BUILD_LITE_INTERPRETER)
   file(GLOB_RECURSE ATen_CORE_TEST_SRCS "core/*_test.cpp")
 endif()
@@ -84,6 +85,8 @@ file(GLOB miopen_cpp "miopen/*.cpp")
 file(GLOB mkl_cpp "mkl/*.cpp")
 file(GLOB mkldnn_cpp "mkldnn/*.cpp")
 
+file(GLOB mkldnn_xpu_cpp "native/mkldnn/xpu/*.cpp" "native/mkldnn/xpu/detail/*.cpp")
+
 file(GLOB native_cpp "native/*.cpp")
 file(GLOB native_mkl_cpp "native/mkl/*.cpp")
 file(GLOB native_mkldnn_cpp "native/mkldnn/*.cpp")
@@ -237,6 +240,20 @@ else()
   set(all_cpu_cpp ${all_cpu_cpp} ${vulkan_cpp})
 endif()
 
+if(USE_XPU)
+  list(APPEND ATen_XPU_SRCS ${mkldnn_xpu_cpp})
+  list(APPEND ATen_XPU_DEPENDENCY_LIBS xpu_mkldnn)
+
+  list(APPEND ATen_XPU_DEPENDENCY_LIBS ${OCL_LIBRARY})
+  list(APPEND ATen_XPU_INCLUDE ${CMAKE_CURRENT_SOURCE_DIR}/native/mkldnn/xpu)
+  list(APPEND ATen_XPU_INCLUDE ${CMAKE_CURRENT_SOURCE_DIR}/native/mkldnn/xpu/detail)
+  list(APPEND ATen_XPU_INCLUDE ${PROJECT_SOURCE_DIR}/third_party/ideep/mkl-dnn/include)
+  list(APPEND ATen_XPU_INCLUDE ${XPU_MKLDNN_INCLUDE})
+
+  list(APPEND ATen_XPU_INCLUDE ${SYCL_INCLUDE_DIR})
+  list(APPEND ATen_XPU_DEPENDENCY_LIBS ${SYCL_LIBRARY})
+endif()
+
 # Metal
 if(USE_PYTORCH_METAL_EXPORT)
   # Add files needed from exporting metal models(optimized_for_mobile)
@@ -419,32 +436,25 @@ if(NOT CMAKE_SYSTEM_PROCESSOR MATCHES "^(s390x|ppc64le)$")
   list(APPEND ATen_CPU_DEPENDENCY_LIBS cpuinfo)
 endif()
 
-if(NOT MSVC AND NOT EMSCRIPTEN AND NOT INTERN_BUILD_MOBILE)
-  # Preserve values for the main build
-  set(__aten_sleef_build_shared_libs ${BUILD_SHARED_LIBS})
-  set(__aten_sleef_build_tests ${BUILD_TESTS})
-
-  # Unset our restrictive C++ flags here and reset them later.
-  # Remove this once we use proper target_compile_options.
-  set(OLD_CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS})
-  set(CMAKE_CXX_FLAGS)
-
-  # Bump up optimization level for sleef to -O1, since at -O0 the compiler
-  # excessively spills intermediate vector registers to the stack
-  # and makes things run impossibly slowly
-  set(OLD_CMAKE_C_FLAGS_DEBUG ${CMAKE_C_FLAGS_DEBUG})
-  if(${CMAKE_C_FLAGS_DEBUG} MATCHES "-O0")
-    string(REGEX REPLACE "-O0" "-O1" CMAKE_C_FLAGS_DEBUG ${OLD_CMAKE_C_FLAGS_DEBUG})
-  else()
-    set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -O1")
+if(NOT EMSCRIPTEN AND NOT INTERN_BUILD_MOBILE)
+  if(NOT MSVC)
+    # Bump up optimization level for sleef to -O1, since at -O0 the compiler
+    # excessively spills intermediate vector registers to the stack
+    # and makes things run impossibly slowly
+    set(OLD_CMAKE_C_FLAGS_DEBUG ${CMAKE_C_FLAGS_DEBUG})
+    if(${CMAKE_C_FLAGS_DEBUG} MATCHES "-O0")
+      string(REGEX REPLACE "-O0" "-O1" CMAKE_C_FLAGS_DEBUG ${OLD_CMAKE_C_FLAGS_DEBUG})
+    else()
+      set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -O1")
+    endif()
   endif()
 
   if(NOT USE_SYSTEM_SLEEF)
-    set(BUILD_SHARED_LIBS OFF CACHE BOOL "Build sleef static" FORCE)
-    set(BUILD_DFT OFF CACHE BOOL "Don't build sleef DFT lib" FORCE)
-    set(BUILD_GNUABI_LIBS OFF CACHE BOOL "Don't build sleef gnuabi libs" FORCE)
-    set(BUILD_TESTS OFF CACHE BOOL "Don't build sleef tests" FORCE)
-    set(OLD_CMAKE_BUILD_TYPE ${CMAKE_BUILD_TYPE})
+    set(SLEEF_BUILD_SHARED_LIBS OFF CACHE BOOL "Build sleef static" FORCE)
+    set(SLEEF_BUILD_DFT OFF CACHE BOOL "Don't build sleef DFT lib" FORCE)
+    set(SLEEF_BUILD_GNUABI_LIBS OFF CACHE BOOL "Don't build sleef gnuabi libs" FORCE)
+    set(SLEEF_BUILD_TESTS OFF CACHE BOOL "Don't build sleef tests" FORCE)
+    set(SLEEF_BUILD_SCALAR_LIB OFF CACHE BOOL "libsleefscalar will be built." FORCE)
     if(CMAKE_SYSTEM_NAME STREQUAL "Darwin")
       if(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64" OR CMAKE_OSX_ARCHITECTURES MATCHES "arm64")
         set(DISABLE_SVE ON CACHE BOOL "Xcode's clang-12.5 crashes while trying to compile SVE code" FORCE)
@@ -465,12 +475,9 @@ if(NOT MSVC AND NOT EMSCRIPTEN AND NOT INTERN_BUILD_MOBILE)
   endif()
   list(APPEND ATen_CPU_DEPENDENCY_LIBS sleef)
 
-  set(CMAKE_C_FLAGS_DEBUG ${OLD_CMAKE_C_FLAGS_DEBUG})
-  set(CMAKE_CXX_FLAGS ${OLD_CMAKE_CXX_FLAGS})
-
-  # Set these back. TODO: Use SLEEF_ to pass these instead
-  set(BUILD_SHARED_LIBS ${__aten_sleef_build_shared_libs} CACHE BOOL "Build shared libs" FORCE)
-  set(BUILD_TESTS ${__aten_sleef_build_tests} CACHE BOOL "Build tests" FORCE)
+  if(NOT MSVC)
+    set(CMAKE_C_FLAGS_DEBUG ${OLD_CMAKE_C_FLAGS_DEBUG})
+  endif()
 endif()
 
 if(USE_CUDA AND NOT USE_ROCM)
@@ -574,7 +581,7 @@ configure_file(ATenConfig.cmake.in "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/AT
 install(FILES "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/ATenConfig.cmake"
   DESTINATION "${AT_INSTALL_SHARE_DIR}/cmake/ATen")
 
-set(INSTALL_HEADERS ${base_h} ${ATen_CORE_HEADERS} ${native_nested_h})
+set(INSTALL_HEADERS ${base_h} ${ATen_CORE_HEADERS} ${native_nested_h} ${ATen_TRANSFORMER_HEADERS})
 if(NOT INTERN_BUILD_MOBILE)
   list(APPEND INSTALL_HEADERS ${native_h} ${native_cpu_h} ${native_ao_sparse_h} ${native_quantized_h} ${cuda_h} ${native_cuda_h} ${native_hip_h} ${cudnn_h} ${hip_h} ${xpu_h} ${mps_h} ${native_mps_h} ${native_utils_h} ${miopen_h})
   # Metal
@@ -638,6 +645,7 @@ list(APPEND ATen_MOBILE_BENCHMARK_SRCS
 # Pass source, includes, and libs to parent
 set(ATen_CORE_SRCS ${ATen_CORE_SRCS} PARENT_SCOPE)
 set(ATen_CPU_SRCS ${ATen_CPU_SRCS} PARENT_SCOPE)
+set(ATen_XPU_SRCS ${ATen_XPU_SRCS} PARENT_SCOPE)
 set(ATen_CUDA_CU_SRCS ${ATen_CUDA_CU_SRCS} PARENT_SCOPE)
 set(ATen_CUDA_CPP_SRCS ${ATen_CUDA_CPP_SRCS} PARENT_SCOPE)
 set(ATen_CUDA_LINALG_SRCS ${ATen_CUDA_LINALG_SRCS} PARENT_SCOPE)
@@ -667,6 +675,7 @@ set(ATen_XPU_INCLUDE ${ATen_XPU_INCLUDE} PARENT_SCOPE)
 set(ATen_VULKAN_INCLUDE ${ATen_VULKAN_INCLUDE} PARENT_SCOPE)
 set(ATen_CPU_DEPENDENCY_LIBS ${ATen_CPU_DEPENDENCY_LIBS} PARENT_SCOPE)
 set(ATen_CUDA_DEPENDENCY_LIBS ${ATen_CUDA_DEPENDENCY_LIBS} PARENT_SCOPE)
+set(ATen_XPU_DEPENDENCY_LIBS ${ATen_XPU_DEPENDENCY_LIBS} PARENT_SCOPE)
 set(ATen_HIP_DEPENDENCY_LIBS ${ATen_HIP_DEPENDENCY_LIBS} PARENT_SCOPE)
 set(FLASH_ATTENTION_CUDA_SOURCES ${FLASH_ATTENTION_CUDA_SOURCES} PARENT_SCOPE)
 set(MEM_EFF_ATTENTION_CUDA_SOURCES ${MEM_EFF_ATTENTION_CUDA_SOURCES} PARENT_SCOPE)

aten/src/ATen/Context.cpp

@@ -263,6 +263,23 @@ void Context::setLinalgPreferredBackend(at::LinalgBackend b) {
   }
 }
 
+at::BlasBackend Context::blasPreferredBackend() const {
+  return blas_preferred_backend;
+}
+
+void Context::setBlasPreferredBackend(at::BlasBackend b) {
+  TORCH_CHECK((b != at::BlasBackend::Cublaslt) || hasCuBLASLt(),
+      "Cannot set preferred backend to cuBLASLt if PyTorch has not been compiled with cuBLASLt.");
+  if (b != at::BlasBackend::Cublas) {
+    TORCH_WARN_ONCE(
+      "torch.backends.cuda.preferred_blas_library is an experimental feature. "
+      "If you see any error or unexpected behavior when this flag is set "
+      "please file an issue on GitHub."
+    );
+  }
+  blas_preferred_backend = b;
+}
+
 bool Context::allowFP16ReductionCuBLAS() const {
   return allow_fp16_reduction_cublas;
 }