From a990da22252c5a46d6cd18e1870f4987bac7e670 Mon Sep 17 00:00:00 2001
From: yisonzhu <107918054+yisonzhu@users.noreply.github.com>
Date: Thu, 25 Jul 2024 16:05:38 +0800
Subject: [PATCH] [FIX] Updating SYCL Shuffle API (#2737)

---
 itex/core/kernels/gpu/full_reduction_kernels.h | 2 +-
 itex/core/kernels/gpu/softmax_op_functor.h     | 3 ++-
 itex/core/kernels/gpu/sparse_xent_op.h         | 3 ++-
 itex/core/kernels/gpu/xent_op.h                | 3 ++-
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/itex/core/kernels/gpu/full_reduction_kernels.h b/itex/core/kernels/gpu/full_reduction_kernels.h
index 50a7fdb76..e1c083d7d 100644
--- a/itex/core/kernels/gpu/full_reduction_kernels.h
+++ b/itex/core/kernels/gpu/full_reduction_kernels.h
@@ -141,7 +141,7 @@ struct GroupReduceKernel<itex::int64, OutputT, InputFunctor, OutputFunctor,
     InitValueT result = value;
 #pragma unroll
     for (int i = SubGroupSize / 2; i > 0; i >>= 1) {
-      InitValueT new_value = sg.shuffle_down(result, i);
+      InitValueT new_value = sycl::shift_group_left(sg, result, i);
       result = op_(result, new_value);
     }
     return result;
diff --git a/itex/core/kernels/gpu/softmax_op_functor.h b/itex/core/kernels/gpu/softmax_op_functor.h
index eba173c35..607808dc3 100644
--- a/itex/core/kernels/gpu/softmax_op_functor.h
+++ b/itex/core/kernels/gpu/softmax_op_functor.h
@@ -87,7 +87,8 @@ template <template <typename> typename ReductionOp, typename T,
           int workitem_group_width = kSubGroupSize>
 T SubGroupAllReduce(const sycl::sub_group& sg, T val) {
   for (int mask = workitem_group_width / 2; mask > 0; mask /= 2) {
-    val = ReductionOp<T>()(val, sg.shuffle_xor(val, sycl::id<1>(mask)));
+    val = ReductionOp<T>()(
+        val, sycl::permute_group_by_xor(sg, val, sycl::id<1>(mask)));
   }
   return val;
 }
diff --git a/itex/core/kernels/gpu/sparse_xent_op.h b/itex/core/kernels/gpu/sparse_xent_op.h
index 667a17f2a..499636195 100644
--- a/itex/core/kernels/gpu/sparse_xent_op.h
+++ b/itex/core/kernels/gpu/sparse_xent_op.h
@@ -46,7 +46,8 @@ template <template <typename> typename ReductionOp, typename T,
           int workitem_group_width = kSubGroupSize>
 T SubGroupAllReduce(const sycl::sub_group& sg, T val) {
   for (int mask = workitem_group_width / 2; mask > 0; mask /= 2) {
-    val = ReductionOp<T>()(val, sg.shuffle_xor(val, sycl::id<1>(mask)));
+    val = ReductionOp<T>()(
+        val, sycl::permute_group_by_xor(sg, val, sycl::id<1>(mask)));
   }
   return val;
 }
diff --git a/itex/core/kernels/gpu/xent_op.h b/itex/core/kernels/gpu/xent_op.h
index d1a529c9d..4713c55f8 100644
--- a/itex/core/kernels/gpu/xent_op.h
+++ b/itex/core/kernels/gpu/xent_op.h
@@ -45,7 +45,8 @@ template <template <typename> typename ReductionOp, typename T,
           int workitem_group_width = kSubGroupSize>
 T SubGroupAllReduce(const sycl::sub_group& sg, T val) {
   for (int mask = workitem_group_width / 2; mask > 0; mask /= 2) {
-    val = ReductionOp<T>()(val, sg.shuffle_xor(val, sycl::id<1>(mask)));
+    val = ReductionOp<T>()(
+        val, sycl::permute_group_by_xor(sg, val, sycl::id<1>(mask)));
   }
   return val;
 }