Merge branch 'main' into penghuic/add_hook_random_failure

Author: PenghuiCheng

.github/scripts/apply_torch_pr.py

@@ -13,6 +13,8 @@
         "https://github.com/pytorch/pytorch/pull/126516",
         # Modify the tolerance level in TIMM benchmark
         "https://github.com/pytorch/pytorch/pull/143739",
+        # Fix build error caused by incorrect namespace change by #144014
+        "https://github.com/pytorch/pytorch/pull/144450",
     ]
 )
 parser.add_argument('--extra-pr-list', '-e', nargs='+',default=[])

.github/workflows/_linux_transformers.yml

@@ -46,6 +46,8 @@ jobs:
   Torch-XPU-Transformers-Tests:
     runs-on: ${{ inputs.runner != '' && inputs.runner || 'linux.idc.xpu' }}
     env:
+      HF_HOME: ${{ github.workspace }}/.hf_home
+      HF_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
       NEOReadDebugKeys: ${{ inputs.driver == 'rolling' && '1' || '0' }}
       DisableScratchPages: ${{ inputs.driver == 'rolling' && '1' || '0' }}
       python: ${{ inputs.python != '' && inputs.python || '3.10' }}
@@ -115,7 +117,7 @@ jobs:
           cat /sys/class/drm/render*/device/device | tee ${{ github.workspace }}/transformers/tests_log/device_IDs.txt
           echo "xpu-smi output:"
           xpu-smi discovery -y --json --dump -1
-      - name: Sanitry check installed packages
+      - name: Sanity check installed packages
         run: |
           source activate huggingface_transformers_test
           # These checks are to exit earlier if for any reason Transformers
@@ -124,6 +126,9 @@ jobs:
           pip show torchaudio | grep Version | grep xpu
           pip show torchvision | grep Version | grep xpu
           python -c 'import torch; exit(not torch.xpu.is_available())'
+      - name: Clean HF home directory and cache
+        run: |
+          rm -rf ${{ env.HF_HOME }}
       - name: Run -k backbone tests
         env:
           TEST_CASE: 'tests_backbone'
@@ -212,6 +217,11 @@ jobs:
           FAILED_CASES=$(echo $FAILED_CASES | sed 's/^,//')
           echo "Failed cases: [$(echo $FAILED_CASES | sed 's/,/, /g')]"
           test -z "$FAILED_CASES"
+      - name: Clean HF home directory and cache
+        if: ${{ always() }}
+        run: |
+          du -sh ${{ env.HF_HOME }} || true
+          rm -rf ${{ env.HF_HOME }}
       - name: Print results table
         if: ${{ ! cancelled() }}
         run: |

cmake/Codegen.cmake

@@ -8,6 +8,7 @@ file(MAKE_DIRECTORY ${BUILD_TORCH_XPU_ATEN_GENERATED})
 
 set(RegisterXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU.cpp)
 set(RegisterSparseXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU.cpp)
+set(RegisterNestedTensorXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterNestedTensorXPU.cpp)
 set(XPUFallback_PATH ${TORCH_XPU_OPS_ROOT}/src/ATen/native/xpu/XPUFallback.template)
 
 if(WIN32)
@@ -47,6 +48,7 @@ endfunction(GEN_BACKEND)
 
 set(RegisterXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU.cpp)
 set(RegisterSparseXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU.cpp)
+set(RegisterNestedTensorXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterNestedTensorXPU.cpp)
 set(XPUFallback_PATH ${TORCH_XPU_OPS_ROOT}/src/ATen/native/xpu/XPUFallback.template)
 set(XPU_AOTI_INSTALL_DIR ${TORCH_ROOT}/torch/csrc/inductor/aoti_torch/generated/extend)
 function(GEN_XPU file_yaml)
@@ -77,7 +79,7 @@ function(GEN_XPU file_yaml)
     --install-dir ${BUILD_TORCH_XPU_ATEN_GENERATED}
     --per-operator-headers
     --static-dispatch-backend
-    --backend-whitelist XPU SparseXPU
+    --backend-whitelist XPU SparseXPU NestedTensorXPU
     # --xpu: generate in-tree RegisterXPU.cpp for in-tree OPs
     --xpu
     # --update-aoti-c-shim: generate extend/c_shim_xpu.h
@@ -93,6 +95,7 @@ function(GEN_XPU file_yaml)
     # Codegen post-process
     COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterXPU_PATH}
     COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseXPU_PATH}
+    COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterNestedTensorXPU_PATH}
     ${SIMPLE_TRACE} 
     WORKING_DIRECTORY ${TORCH_ROOT}
     DEPENDS
@@ -122,6 +125,7 @@ GEN_XPU(
   ${BUILD_TORCH_XPU_ATEN_GENERATED}/XPUFunctions.h
   ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU.cpp
   ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU.cpp
+  ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterNestedTensorXPU.cpp
   ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.h
   ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.cpp
 )
@@ -133,7 +137,7 @@ GEN_XPU(
 # $TORCH_XPU_OPS_INCLUDE_DIRS, so that "#include <ATen/ops/*.h>" works.
 list(APPEND TORCH_XPU_OPS_INCLUDE_DIRS ${CMAKE_BINARY_DIR}/xpu)
 
-list(APPEND xpu_generated_src ${RegisterXPU_PATH} ${RegisterSparseXPU_PATH})
+list(APPEND xpu_generated_src ${RegisterXPU_PATH} ${RegisterSparseXPU_PATH} ${RegisterNestedTensorXPU_PATH})
 list(APPEND xpu_generated_src ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.cpp)
 add_custom_target(TORCH_XPU_GEN_TARGET DEPENDS ${xpu_generated_src})
 set(ATen_XPU_GEN_SRCS ${xpu_generated_src})

src/ATen/native/xpu/TensorAdvancedIndexing.cpp

@@ -7,17 +7,24 @@
 #include <ATen/core/op_registration/adaption.h>
 #include <ATen/native/DispatchStub.h>
 #include <ATen/native/IndexKernel.h>
+#include <ATen/native/ReductionType.h>
 #include <ATen/native/TensorAdvancedIndexing.h>
 #include <ATen/native/TensorAdvancedIndexingUtils.h>
 #include <ATen/native/TensorIterator.h>
+//#include <ATen/native/TensorFactories.cpp>
 #include <ATen/native/xpu/sycl/IndexingKernels.h>
 #include <ATen/native/xpu/sycl/ScatterGatherKernels.h>
+#include <ATen/ops/ones_like.h>
+#include <ATen/ops/zeros_like.h>
 #include <comm/RegisterUtils.h>
 #include <comm/xpu_aten.h>
 #include <torch/library.h>
 
 #include <ATen/ops/index_add_meta.h>
+#include <ATen/ops/index_reduce_meta.h>
 #include <xpu/ATen/ops/index_add_native.h>
+#include <xpu/ATen/ops/index_reduce_native.h> //generated
+//#include <xpu/ATen/ops/index_reduce_prod_native.h> //generated
 
 namespace at {
 
@@ -42,6 +49,7 @@ REGISTER_XPU_DISPATCH(index_fill_stub, &xpu::index_fill_kernel);
 REGISTER_XPU_DISPATCH(index_copy_stub, &xpu::index_copy_kernel);
 REGISTER_XPU_DISPATCH(put_stub, &xpu::put_kernel);
 REGISTER_XPU_DISPATCH(take_stub, &xpu::take_kernel);
+// REGISTER_XPU_DISPATCH(index_reduce_stub, &xpu::index_reduce_kernel);
 
 TORCH_IMPL_FUNC(index_add_xpu_out)
 (const Tensor& self,
@@ -126,5 +134,44 @@ Tensor count_nonzero_xpu(const Tensor& self, IntArrayRef dims) {
   return (self != 0).sum(dims);
 }
 
+TORCH_IMPL_FUNC(index_reduce_xpu_out)
+(const Tensor& self,
+ int64_t dim,
+ const Tensor& index,
+ const Tensor& source,
+ const c10::string_view reduce,
+ bool include_self,
+ const Tensor& result) {
+  TORCH_WARN_ONCE(
+      "index_reduce() is in beta and the API may change at any time.");
+  if (reduce == "prod") {
+    xpu::index_reduce_prod_kernel(
+        self, dim, index, source, include_self, ReductionType::PROD, result);
+  } else if (reduce == "mean") {
+    xpu::index_reduce_mean_kernel(
+        self, dim, index, source, include_self, ReductionType::MEAN, result);
+    auto counts = include_self ? ones_like(result) : zeros_like(result);
+    counts.index_add_(dim, index, ones_like(source));
+    counts.masked_fill_(counts == 0, 1);
+    if (result.is_floating_point() || result.is_complex()) {
+      result.div_(counts);
+    } else {
+      result.div_(counts, "floor");
+    }
+  } else if (reduce == "amax") {
+    xpu::index_reduce_amax_kernel(
+        self, dim, index, source, include_self, ReductionType::MAX, result);
+  } else if (reduce == "amin") {
+    xpu::index_reduce_amin_kernel(
+        self, dim, index, source, include_self, ReductionType::MIN, result);
+  } else {
+    TORCH_CHECK(
+        false,
+        "Only support prod, mean, amax or amin reduce operator. Input was ",
+        reduce,
+        ".");
+  }
+}
+
 } // namespace native
 } // namespace at

src/ATen/native/xpu/XPUFallback.template

@@ -163,7 +163,6 @@ TORCH_LIBRARY_IMPL(aten, XPU, m) {
     "_fft_r2c",
     "_flash_attention_forward",
     "geqrf",
-    "index_reduce.out",
     "linalg_cholesky_ex.L",
     "_linalg_det.result",
     "linalg_eig",

src/ATen/native/xpu/sycl/Atomics.h

@@ -360,6 +360,8 @@ SYCL_ATOMIC_INTEGER(Mul, std::multiplies<int8_t>()(a, b), int8_t)
 SYCL_ATOMIC_INTEGER(Mul, std::multiplies<int16_t>()(a, b), int16_t)
 SYCL_ATOMIC_INTEGER(Mul, std::multiplies<int32_t>()(a, b), int32_t)
 SYCL_ATOMIC_INTEGER(Mul, std::multiplies<int64_t>()(a, b), int64_t)
+SYCL_ATOMIC_INTEGER(Mul, std::multiplies<uint32_t>()(a, b), uint32_t)
+SYCL_ATOMIC_INTEGER(Mul, std::multiplies<uint64_t>()(a, b), uint64_t)
 
 SYCL_ATOMIC_FP(Mul, std::multiplies<float>()(a, b), float)
 SYCL_ATOMIC_FP(Mul, std::multiplies<double>()(a, b), double)
@@ -391,6 +393,8 @@ SYCL_ATOMIC_INTEGER(Max, safe_max<int8_t>(a, b), int8_t)
 SYCL_ATOMIC_INTEGER(Max, safe_max<int16_t>(a, b), int16_t)
 SYCL_ATOMIC_INTEGER(Max, safe_max<int32_t>(a, b), int32_t)
 SYCL_ATOMIC_INTEGER(Max, safe_max<int64_t>(a, b), int64_t)
+SYCL_ATOMIC_INTEGER(Max, safe_max<uint32_t>(a, b), uint32_t)
+SYCL_ATOMIC_INTEGER(Max, safe_max<uint64_t>(a, b), uint64_t)
 
 SYCL_ATOMIC_FP(Max, safe_max<float>(a, b), float)
 SYCL_ATOMIC_FP(Max, safe_max<double>(a, b), double)
@@ -403,6 +407,8 @@ SYCL_ATOMIC_INTEGER(Min, safe_min<int8_t>(a, b), int8_t)
 SYCL_ATOMIC_INTEGER(Min, safe_min<int16_t>(a, b), int16_t)
 SYCL_ATOMIC_INTEGER(Min, safe_min<int32_t>(a, b), int32_t)
 SYCL_ATOMIC_INTEGER(Min, safe_min<int64_t>(a, b), int64_t)
+SYCL_ATOMIC_INTEGER(Min, safe_min<uint32_t>(a, b), uint32_t)
+SYCL_ATOMIC_INTEGER(Min, safe_min<uint64_t>(a, b), uint64_t)
 
 SYCL_ATOMIC_FP(Min, safe_min<float>(a, b), float)
 SYCL_ATOMIC_FP(Min, safe_min<double>(a, b), double)