update

Author: yucai-intel

src/ATen/native/xpu/sycl/GroupNormKernels.cpp

@@ -63,13 +63,12 @@ template <typename T, int SIMD>
 struct GNRowwiseMomentsFunctor : public __SYCL_KER_CONFIG_CONVENTION__ {
   using T_ACC = acc_type_device<T, kXPU>;
   using WelfordType = WelfordData<T_ACC, int64_t>;
-  using WelfordOp =
-      WelfordOpsXPU<T_ACC, T_ACC, int64_t, std::pair<T_ACC, T_ACC>>;
+  using WelfordOp = WelfordOps<T_ACC, T_ACC, int64_t, std::pair<T_ACC, T_ACC>>;
 
   [[intel::reqd_sub_group_size(SIMD)]] void operator()(
       sycl::nd_item<1> item) const {
     const int64_t i = item.get_group(0);
-    WelfordOp welford_op = {/*correction=*/0, /*take_sqrt=*/false, item};
+    WelfordOp welford_op = {/*correction=*/0, /*take_sqrt=*/false};
     WelfordType val(0, 0, 0, 0);
     WelfordType identity_element(0, 0, 0, 0);
     for (int64_t j = item.get_local_id(0); j < N_;
@@ -78,8 +77,13 @@ struct GNRowwiseMomentsFunctor : public __SYCL_KER_CONFIG_CONVENTION__ {
       val = welford_op.reduce(val, static_cast<T_ACC>(X_[index]), index);
     }
 
-    val = GroupReduceWithoutBroadcast<WelfordType, WelfordOp, SIMD>(
-        item, val, welford_op, identity_element, shared_);
+    if (item.get_local_range(0) <= SIMD) {
+      val = SubgroupReduceWithoutBroadcast<WelfordType, WelfordOp, SIMD>(
+          item, val, welford_op);
+    } else {
+      val = GroupReduceWithoutBroadcast<WelfordType, WelfordOp, SIMD>(
+          item, val, welford_op, identity_element, shared_);
+    }
 
     if (item.get_local_id(0) == 0) {
       T_ACC m1;
@@ -111,15 +115,14 @@ struct GNRowwiseMomentsVectorizedFunctor
     : public __SYCL_KER_CONFIG_CONVENTION__ {
   using T_ACC = acc_type_device<T, kXPU>;
   using WelfordType = WelfordData<T_ACC, int64_t>;
-  using WelfordOp =
-      WelfordOpsXPU<T_ACC, T_ACC, int64_t, std::pair<T_ACC, T_ACC>>;
+  using WelfordOp = WelfordOps<T_ACC, T_ACC, int64_t, std::pair<T_ACC, T_ACC>>;
   using vec_t = memory::aligned_vector<T, VEC_SIZE>;
 
   [[intel::reqd_sub_group_size(SIMD)]] void operator()(
       sycl::nd_item<1> item) const {
     WelfordType val[VEC_SIZE];
     WelfordType identity_element(0, 0, 0, 0);
-    WelfordOp welford_op = {/*correction=*/0, /*take_sqrt=*/false, item};
+    WelfordOp welford_op = {/*correction=*/0, /*take_sqrt=*/false};
     auto g_start = item.get_group(0) * VEC_SIZE;
 
 #pragma unroll
@@ -140,8 +143,15 @@ struct GNRowwiseMomentsVectorizedFunctor
 
 #pragma unroll
     for (int v = 0; v < VEC_SIZE; ++v) {
-      val[v] = GroupReduceWithoutBroadcast<WelfordType, WelfordOp, SIMD>(
-          item, val[v], welford_op, identity_element, shared_);
+      // val[v] = GroupReduceWithoutBroadcast<WelfordType, WelfordOp, SIMD>(
+      //     item, val[v], welford_op, identity_element, shared_);
+      if (item.get_local_range(0) <= SIMD) {
+        val[v] = SubgroupReduceWithoutBroadcast<WelfordType, WelfordOp, SIMD>(
+            item, val[v], welford_op);
+      } else {
+        val[v] = GroupReduceWithoutBroadcast<WelfordType, WelfordOp, SIMD>(
+            item, val[v], welford_op, identity_element, shared_);
+      }
     }
 
     if (item.get_local_id(0) == 0) {

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

@@ -119,10 +119,12 @@ inline T& SubgroupReduceWithoutBroadcast(
   auto sg = item.get_sub_group();
   auto sg_tid = sg.get_local_linear_id();
 #pragma unroll
-  for (int offset = 1; offset < SIMD; offset <<= 1) {
-    if (sg_tid < SIMD - offset) {
-      val = op.combine(val, sycl::shift_group_left(sg, val, offset));
-    }
+  for (int offset = (SIMD >> 1); offset > 0; offset >>= 1) {
+    // for (int offset = 1; offset < SIMD; offset <<= 1) {
+    T temp = sycl::shift_group_left(sg, val, offset);
+    // if (sg_tid < SIMD - offset) {
+    val = op.combine(val, temp);
+    // }
   }
   return val;
 }
@@ -135,6 +137,7 @@ inline T& GroupReduceWithoutBroadcast(
     const T& identity_element,
     shared_t shared) {
   auto sg = item.get_sub_group();
+  int g_tid = item.get_local_linear_id();
   int sg_tid = sg.get_local_linear_id();
   int sg_id = sg.get_group_linear_id();
   int n_sg = get_local_linear_range<DIM>(item) / SIMD;
@@ -148,9 +151,17 @@ inline T& GroupReduceWithoutBroadcast(
     shared[sg_id] = val;
   }
   item.barrier(sycl_local_fence);
-  val = (sg_id < n_sg) ? shared[sg_id] : identity_element;
+  // val = (g_tid < n_sg) ? shared[sg_id] : identity_element;
+  val = identity_element;
+
   if (sg_id == 0) {
+    for (int i = sg_tid; i < n_sg; i += SIMD) {
+      val = op.combine(val, shared[i]);
+    }
     val = SubgroupReduceWithoutBroadcast<T, ReduceOp, SIMD, DIM>(item, val, op);
+    // for (int i = 1; i < n_sg; i++) {
+    //   val = op.combine(val, shared[i]);
+    // }
   }
   return val;
 }