@@ -56,7 +56,7 @@ static void soft_max_f32(const float * x,
5656 : block_size_template;
5757 const int nthreads = block_size;
5858 const int nwarps = nthreads / WARP_SIZE ;
59- size_t nreduce = nwarps / WARP_SIZE ;
59+ const size_t nreduce = nwarps / WARP_SIZE ;
6060
6161 const int tid = item_ct1.get_local_id (2 );
6262
@@ -105,17 +105,15 @@ static void soft_max_f32(const float * x,
105105 max_val = warp_reduce_max<WARP_SIZE >(max_val);
106106
107107 if (block_size > WARP_SIZE ) {
108- if (warp_id == 0 ) {
109- buf_iw[lane_id] = -INFINITY ;
110- }
111- item_ct1.barrier ();
112-
113108 if (lane_id == 0 ) {
114109 buf_iw[warp_id] = max_val;
115110 }
116111 item_ct1.barrier ();
117112
118- max_val = buf_iw[lane_id];
113+ max_val = -INFINITY ;
114+ for (int i = lane_id; i < nwarps; i += WARP_SIZE ) {
115+ max_val = sycl::max (max_val, buf_iw[i]);
116+ }
119117 max_val = warp_reduce_max<WARP_SIZE >(max_val);
120118 }
121119 float tmp = 0 .0f ; // partial sum
@@ -290,7 +288,8 @@ static void soft_max_f32_sycl(const float *x, const T *mask,
290288
291289 cgh.parallel_for (
292290 sycl::nd_range<3 >(block_nums * block_dims, block_dims),
293- [=](sycl::nd_item<3 > item_ct1) {
291+ [=](sycl::nd_item<3 > item_ct1)
292+ [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
294293 soft_max_f32<false , 0 , 0 >(
295294 x, mask, sinks, dst, params,
296295 dpct_local_acc_ct1
0 commit comments