Update GroupReduceUtils.h

yucai-intel · web-flow · commit e799a63c2d61 · 2025-05-26T10:54:27.000+08:00
diff --git a/src/ATen/native/xpu/sycl/GroupReduceUtils.h b/src/ATen/native/xpu/sycl/GroupReduceUtils.h
@@ -149,8 +149,11 @@ inline T& GroupReduceWithoutBroadcast(
     shared[sg_id] = val;
   }
   item.barrier(sycl_local_fence);
-  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);
   }
   return val;

Original file line number	Diff line number	Diff line change
`@@ -149,8 +149,11 @@ inline T& GroupReduceWithoutBroadcast(`
`149`	`149`	`shared[sg_id] = val;`
`150`	`150`	`}`
`151`	`151`	`item.barrier(sycl_local_fence);`
`152`		`- val = (g_tid < n_sg) ? shared[sg_id] : identity_element;`
	`152`	`+ val = identity_element;`
`153`	`153`	`if (sg_id == 0) {`
	`154`	`+ for (int i = sg_tid; i < n_sg; i += SIMD) {`
	`155`	`+ val = op.combine(val, shared[i]);`
	`156`	`+ }`
`154`	`157`	`val = SubgroupReduceWithoutBroadcast<T, ReduceOp, SIMD, DIM>(item, val, op);`
`155`	`158`	`}`
`156`	`159`	`return val;`