Coverage for src/flag_gems/runtime/backend/_sunrise/ops/scatter_reduce.py: 0%

1"""Triton implementation of torch.scatter_reduce for FlagGems. 

2 

3Supports all reduce modes: sum, prod, mean, amax, amin. 

4Handles 1D-5D tensors with up to 5D coordinate decoding via padding. 

5 

6Vendor compatibility: 

7 - NVIDIA: native atomic_max/min for amax/amin reduce 

8 - Iluvatar: CAS-based fallback for atomic_max/min (no native support) 

9 - Metax: larger BLOCK=256 for better occupancy 

10 

11Performance notes: 

12 - Sum/mean use tl.atomic_add with relaxed semantics for throughput 

13 - Prod uses CAS loop with NaN detection guard (no tl.atomic_mul exists) 

14 - All offset arithmetic uses int64 to avoid overflow for N > 2^31 

15 - LOOP=4: each program processes LOOP*BLOCK elements to amortize launch overhead 

16 - 2D fast path: specialized kernels for 2D tensors avoid 5D coordinate decoding 

17""" 

18 

19import logging 

20 

21import torch 

22import triton 

23import triton.language as tl 

24 

25import flag_gems 

26from flag_gems.runtime import torch_device_fn 

27from flag_gems.utils import libentry 

28 

29logger = logging.getLogger(__name__) 

30 

31 

32def heur_block(args): 

33 """Vendor-aware block size heuristic. 

34 

35 Metax and Iluvatar GPUs benefit from larger blocks (256) for better 

36 occupancy. NVIDIA GPUs default to 128 which balances occupancy and 

37 register pressure. 

38 """ 

39 if flag_gems.vendor_name in ["metax", "iluvatar"]: 

40 return 256 

41 return 128 

42 

43 

44def heur_loop(args): 

45 """Loop unrolling factor. 

46 

47 Each program processes LOOP*BLOCK elements to amortize kernel launch 

48 overhead. LOOP=4 is optimal for Iluvatar BI-V150. 

49 """ 

50 return 4 

51 

52 

53def heur_scan_block(args): 

54 """Source-dimension tile size for deterministic product scan.""" 

55 return 128 

56 

57 

58# --------------------------------------------------------------------------- 

59# Helpers 

60# --------------------------------------------------------------------------- 

61 

62 

63def _pad5(lst, fill): 

64 """Pad a list to exactly 5 elements from the left with `fill`. 

65 

66 This enables uniform 5D coordinate decoding in kernels regardless 

67 of the actual tensor dimensionality (1D-5D). Shapes are padded with 1, 

68 strides with 0. 

69 """ 

70 return [fill] * (5 - len(lst)) + lst if len(lst) < 5 else lst 

71 

72 

73def _needs_cas_fallback(): 

74 """Check if the current vendor needs CAS-based fallback for atomic_max/min. 

75 

76 Iluvatar GPUs lack native tl.atomic_max/min, so we fall back to a 

77 CAS (Compare-And-Swap) loop for amax/amin reduce modes. 

78 """ 

79 return flag_gems.vendor_name in ["iluvatar"] 

80 

81 

82@libentry() 

83@triton.heuristics({"BLOCK": heur_scan_block}) 

84@triton.jit(do_not_specialize=["out_numel"]) 

85def scatter_reduce_prod_scan_kernel( 

86 index_ptr, 

87 src_ptr, 

88 out_ptr, 

89 mask_ptr, 

90 out_numel, 

91 DIM: tl.constexpr, 

92 USE_MASK: tl.constexpr, 

93 src_shape_dim: tl.constexpr, 

94 src_stride_0, 

95 src_stride_1, 

96 src_stride_2, 

97 src_stride_3, 

98 src_stride_4, 

99 idx_shape_0, 

100 idx_shape_1, 

101 idx_shape_2, 

102 idx_shape_3, 

103 idx_shape_4, 

104 src_shape_0, 

105 src_shape_1, 

106 src_shape_2, 

107 src_shape_3, 

108 src_shape_4, 

109 idx_stride_0, 

110 idx_stride_1, 

111 idx_stride_2, 

112 idx_stride_3, 

113 idx_stride_4, 

114 out_shape_0, 

115 out_shape_1, 

116 out_shape_2, 

117 out_shape_3, 

118 out_shape_4, 

119 out_stride_0, 

120 out_stride_1, 

121 out_stride_2, 

122 out_stride_3, 

123 out_stride_4, 

124 BLOCK: tl.constexpr, 

125): 

126 pid = tl.program_id(axis=0).to(tl.int64) 

127 in_bounds = pid < out_numel 

128 

129 remaining = pid 

130 coord0 = remaining // (out_shape_1 * out_shape_2 * out_shape_3 * out_shape_4) 

131 remaining = remaining % (out_shape_1 * out_shape_2 * out_shape_3 * out_shape_4) 

132 coord1 = remaining // (out_shape_2 * out_shape_3 * out_shape_4) 

133 remaining = remaining % (out_shape_2 * out_shape_3 * out_shape_4) 

134 coord2 = remaining // (out_shape_3 * out_shape_4) 

135 remaining = remaining % (out_shape_3 * out_shape_4) 

136 coord3 = remaining // out_shape_4 

137 coord4 = remaining % out_shape_4 

138 

139 out_offset = ( 

140 coord0 * out_stride_0 

141 + coord1 * out_stride_1 

142 + coord2 * out_stride_2 

143 + coord3 * out_stride_3 

144 + coord4 * out_stride_4 

145 ) 

146 idx_full_offset = ( 

147 coord0 * idx_stride_0 

148 + coord1 * idx_stride_1 

149 + coord2 * idx_stride_2 

150 + coord3 * idx_stride_3 

151 + coord4 * idx_stride_4 

152 ) 

153 src_full_offset = ( 

154 coord0 * src_stride_0 

155 + coord1 * src_stride_1 

156 + coord2 * src_stride_2 

157 + coord3 * src_stride_3 

158 + coord4 * src_stride_4 

159 ) 

160 

161 if DIM == 0: 

162 target = coord0 

163 idx_base = idx_full_offset - coord0 * idx_stride_0 

164 src_base = src_full_offset - coord0 * src_stride_0 

165 idx_scan_stride = idx_stride_0 

166 src_scan_stride = src_stride_0 

167 idx_scan_shape = idx_shape_0 

168 valid_other = ( 

169 (coord1 < idx_shape_1) 

170 & (coord2 < idx_shape_2) 

171 & (coord3 < idx_shape_3) 

172 & (coord4 < idx_shape_4) 

173 & (coord1 < src_shape_1) 

174 & (coord2 < src_shape_2) 

175 & (coord3 < src_shape_3) 

176 & (coord4 < src_shape_4) 

177 ) 

178 elif DIM == 1: 

179 target = coord1 

180 idx_base = idx_full_offset - coord1 * idx_stride_1 

181 src_base = src_full_offset - coord1 * src_stride_1 

182 idx_scan_stride = idx_stride_1 

183 src_scan_stride = src_stride_1 

184 idx_scan_shape = idx_shape_1 

185 valid_other = ( 

186 (coord0 < idx_shape_0) 

187 & (coord2 < idx_shape_2) 

188 & (coord3 < idx_shape_3) 

189 & (coord4 < idx_shape_4) 

190 & (coord0 < src_shape_0) 

191 & (coord2 < src_shape_2) 

192 & (coord3 < src_shape_3) 

193 & (coord4 < src_shape_4) 

194 ) 

195 elif DIM == 2: 

196 target = coord2 

197 idx_base = idx_full_offset - coord2 * idx_stride_2 

198 src_base = src_full_offset - coord2 * src_stride_2 

199 idx_scan_stride = idx_stride_2 

200 src_scan_stride = src_stride_2 

201 idx_scan_shape = idx_shape_2 

202 valid_other = ( 

203 (coord0 < idx_shape_0) 

204 & (coord1 < idx_shape_1) 

205 & (coord3 < idx_shape_3) 

206 & (coord4 < idx_shape_4) 

207 & (coord0 < src_shape_0) 

208 & (coord1 < src_shape_1) 

209 & (coord3 < src_shape_3) 

210 & (coord4 < src_shape_4) 

211 ) 

212 elif DIM == 3: 

213 target = coord3 

214 idx_base = idx_full_offset - coord3 * idx_stride_3 

215 src_base = src_full_offset - coord3 * src_stride_3 

216 idx_scan_stride = idx_stride_3 

217 src_scan_stride = src_stride_3 

218 idx_scan_shape = idx_shape_3 

219 valid_other = ( 

220 (coord0 < idx_shape_0) 

221 & (coord1 < idx_shape_1) 

222 & (coord2 < idx_shape_2) 

223 & (coord4 < idx_shape_4) 

224 & (coord0 < src_shape_0) 

225 & (coord1 < src_shape_1) 

226 & (coord2 < src_shape_2) 

227 & (coord4 < src_shape_4) 

228 ) 

229 else: 

230 target = coord4 

231 idx_base = idx_full_offset - coord4 * idx_stride_4 

232 src_base = src_full_offset - coord4 * src_stride_4 

233 idx_scan_stride = idx_stride_4 

234 src_scan_stride = src_stride_4 

235 idx_scan_shape = idx_shape_4 

236 valid_other = ( 

237 (coord0 < idx_shape_0) 

238 & (coord1 < idx_shape_1) 

239 & (coord2 < idx_shape_2) 

240 & (coord3 < idx_shape_3) 

241 & (coord0 < src_shape_0) 

242 & (coord1 < src_shape_1) 

243 & (coord2 < src_shape_2) 

244 & (coord3 < src_shape_3) 

245 ) 

246 

247 lanes = tl.arange(0, BLOCK) 

248 acc = tl.load(out_ptr + out_offset, mask=in_bounds, other=1.0).to(tl.float32) 

249 has_contrib = False 

250 

251 for start in range(0, src_shape_dim, BLOCK): 

252 scan = start + lanes 

253 valid = ( 

254 in_bounds & valid_other & (scan < src_shape_dim) & (scan < idx_scan_shape) 

255 ) 

256 idx_val = tl.load( 

257 index_ptr + idx_base + scan * idx_scan_stride, 

258 mask=valid, 

259 other=-1, 

260 ).to(tl.int64) 

261 match = valid & (idx_val == target) 

262 src_val = tl.load( 

263 src_ptr + src_base + scan * src_scan_stride, 

264 mask=valid, 

265 other=1.0, 

266 ).to(tl.float32) 

267 factors = tl.where(match, src_val, 1.0) 

268 prefix = tl.cumprod(factors, 0) 

269 tile_prod = tl.sum(tl.where(lanes == (BLOCK - 1), prefix, 0.0)) 

270 acc *= tile_prod 

271 has_contrib |= tl.sum(match.to(tl.int32)) > 0 

272 

273 tl.store(out_ptr + out_offset, acc, mask=in_bounds) 

274 if USE_MASK: 

275 tl.store(mask_ptr + out_offset, has_contrib.to(tl.int32), mask=in_bounds) 

276 

277 

278# --------------------------------------------------------------------------- 

279# 2D Fast Path Kernels with LOOP 

280# Specialized for 2D tensors to avoid 5D coordinate decoding overhead. 

281# Uses 1D grid with LOOP=4 to amortize kernel launch overhead. 

282# --------------------------------------------------------------------------- 

283 

284 

285@libentry() 

286@triton.heuristics({"BLOCK": heur_block, "LOOP": heur_loop}) 

287@triton.jit(do_not_specialize=["N"]) 

288def scatter_reduce_sum_2d_kernel( 

289 index_ptr, 

290 src_ptr, 

291 out_ptr, 

292 mask_ptr, 

293 N, 

294 idx_ncols, 

295 src_ncols, 

296 out_ncols, 

297 DIM: tl.constexpr, 

298 USE_MASK: tl.constexpr, 

299 BLOCK: tl.constexpr, 

300 LOOP: tl.constexpr, 

301): 

302 pid = tl.program_id(axis=0) 

303 base_offsets = pid * BLOCK * LOOP + tl.arange(0, BLOCK) 

304 

305 for i in range(LOOP): 

306 offsets = (base_offsets + i * BLOCK).to(tl.int64) 

307 mask = offsets < N 

308 

309 row = offsets // idx_ncols 

310 col = offsets % idx_ncols 

311 

312 if DIM == 0: 

313 src_offsets = row * src_ncols + col 

314 idx = tl.load(index_ptr + src_offsets, mask=mask, other=0).to(tl.int64) 

315 out_offsets = idx * out_ncols + col 

316 else: 

317 src_offsets = row * src_ncols + col 

318 idx = tl.load(index_ptr + src_offsets, mask=mask, other=0).to(tl.int64) 

319 out_offsets = row * out_ncols + idx 

320 

321 src_val = tl.load(src_ptr + src_offsets, mask=mask, other=0).to(tl.float32) 

322 tl.atomic_add(out_ptr + out_offsets, src_val, mask=mask, sem="relaxed") 

323 

324 if USE_MASK: 

325 ones = tl.full((BLOCK,), 1, dtype=tl.int32) 

326 tl.atomic_add(mask_ptr + out_offsets, ones, mask=mask, sem="relaxed") 

327 

328 

329@libentry() 

330@triton.heuristics({"BLOCK": heur_block, "LOOP": heur_loop}) 

331@triton.jit(do_not_specialize=["N"]) 

332def scatter_reduce_prod_2d_kernel( 

333 index_ptr, 

334 src_ptr, 

335 out_ptr, 

336 mask_ptr, 

337 N, 

338 idx_ncols, 

339 src_ncols, 

340 out_ncols, 

341 DIM: tl.constexpr, 

342 USE_MASK: tl.constexpr, 

343 BLOCK: tl.constexpr, 

344 LOOP: tl.constexpr, 

345): 

346 pid = tl.program_id(axis=0) 

347 base_offsets = pid * BLOCK * LOOP + tl.arange(0, BLOCK) 

348 

349 for i in range(LOOP): 

350 offsets = (base_offsets + i * BLOCK).to(tl.int64) 

351 mask = offsets < N 

352 

353 row = offsets // idx_ncols 

354 col = offsets % idx_ncols 

355 

356 if DIM == 0: 

357 src_offsets = row * src_ncols + col 

358 idx = tl.load(index_ptr + src_offsets, mask=mask, other=0).to(tl.int64) 

359 out_offsets = idx * out_ncols + col 

360 else: 

361 src_offsets = row * src_ncols + col 

362 idx = tl.load(index_ptr + src_offsets, mask=mask, other=0).to(tl.int64) 

363 out_offsets = row * out_ncols + idx 

364 

365 src_val = tl.load(src_ptr + src_offsets, mask=mask, other=0).to(tl.float32) 

366 

367 # CAS loop for product 

368 stop = tl.where(mask, 0, 1).to(tl.int1) 

369 block_stop = False 

370 out_ptr_u32 = (out_ptr + out_offsets).to( 

371 tl.pointer_type(tl.uint32, 1), bitcast=True 

372 ) 

373 while not block_stop: 

374 cur_bits = tl.load(out_ptr_u32, mask=mask, other=0) 

375 cur_val = cur_bits.to(tl.float32, bitcast=True) 

376 new_val = tl.where(stop, cur_val, cur_val * src_val) 

377 is_nan = new_val != new_val 

378 new_val = tl.where(is_nan, src_val, new_val) 

379 new_bits = new_val.to(tl.uint32, bitcast=True) 

380 # Sunrise/PTPU is more stable when product CAS operates on the raw 

381 # float32 bit pattern instead of a floating-pointer CAS. 

382 cas_res = tl.atomic_cas(out_ptr_u32, cur_bits, new_bits, sem="acq_rel") 

383 stop |= (cur_bits == cas_res) | is_nan 

384 block_stop = tl.sum(stop.to(tl.int32)) == BLOCK 

385 

386 if USE_MASK: 

387 ones = tl.full((BLOCK,), 1, dtype=tl.int32) 

388 tl.atomic_add(mask_ptr + out_offsets, ones, mask=mask, sem="relaxed") 

389 

390 

391@libentry() 

392@triton.heuristics({"BLOCK": heur_block, "LOOP": heur_loop}) 

393@triton.jit(do_not_specialize=["N"]) 

394def scatter_reduce_mean_2d_kernel( 

395 index_ptr, 

396 src_ptr, 

397 out_ptr, 

398 count_ptr, 

399 mask_ptr, 

400 N, 

401 idx_ncols, 

402 src_ncols, 

403 out_ncols, 

404 DIM: tl.constexpr, 

405 USE_MASK: tl.constexpr, 

406 BLOCK: tl.constexpr, 

407 LOOP: tl.constexpr, 

408): 

409 pid = tl.program_id(axis=0) 

410 base_offsets = pid * BLOCK * LOOP + tl.arange(0, BLOCK) 

411 

412 for i in range(LOOP): 

413 offsets = (base_offsets + i * BLOCK).to(tl.int64) 

414 mask = offsets < N 

415 

416 row = offsets // idx_ncols 

417 col = offsets % idx_ncols 

418 

419 if DIM == 0: 

420 src_offsets = row * src_ncols + col 

421 idx = tl.load(index_ptr + src_offsets, mask=mask, other=0).to(tl.int64) 

422 out_offsets = idx * out_ncols + col 

423 else: 

424 src_offsets = row * src_ncols + col 

425 idx = tl.load(index_ptr + src_offsets, mask=mask, other=0).to(tl.int64) 

426 out_offsets = row * out_ncols + idx 

427 

428 src_val = tl.load(src_ptr + src_offsets, mask=mask, other=0).to(tl.float32) 

429 

430 tl.atomic_add(out_ptr + out_offsets, src_val, mask=mask, sem="relaxed") 

431 ones_f = tl.full((BLOCK,), 1.0, dtype=tl.float32) 

432 tl.atomic_add(count_ptr + out_offsets, ones_f, mask=mask, sem="relaxed") 

433 

434 if USE_MASK: 

435 ones_i = tl.full((BLOCK,), 1, dtype=tl.int32) 

436 tl.atomic_add(mask_ptr + out_offsets, ones_i, mask=mask, sem="relaxed") 

437 

438 

439@libentry() 

440@triton.heuristics({"BLOCK": heur_block, "LOOP": heur_loop}) 

441@triton.jit(do_not_specialize=["N"]) 

442def scatter_reduce_amax_2d_kernel( 

443 index_ptr, 

444 src_ptr, 

445 out_ptr, 

446 mask_ptr, 

447 N, 

448 idx_ncols, 

449 src_ncols, 

450 out_ncols, 

451 DIM: tl.constexpr, 

452 IS_AMAX: tl.constexpr, 

453 USE_MASK: tl.constexpr, 

454 USE_CAS: tl.constexpr, 

455 BLOCK: tl.constexpr, 

456 LOOP: tl.constexpr, 

457): 

458 pid = tl.program_id(axis=0) 

459 base_offsets = pid * BLOCK * LOOP + tl.arange(0, BLOCK) 

460 

461 for i in range(LOOP): 

462 offsets = (base_offsets + i * BLOCK).to(tl.int64) 

463 mask = offsets < N 

464 

465 row = offsets // idx_ncols 

466 col = offsets % idx_ncols 

467 

468 if DIM == 0: 

469 src_offsets = row * src_ncols + col 

470 idx = tl.load(index_ptr + src_offsets, mask=mask, other=0).to(tl.int64) 

471 out_offsets = idx * out_ncols + col 

472 else: 

473 src_offsets = row * src_ncols + col 

474 idx = tl.load(index_ptr + src_offsets, mask=mask, other=0).to(tl.int64) 

475 out_offsets = row * out_ncols + idx 

476 

477 src_val = tl.load(src_ptr + src_offsets, mask=mask, other=0).to(tl.float32) 

478 

479 if USE_CAS: 

480 stop = tl.where(mask, 0, 1).to(tl.int1) 

481 block_stop = False 

482 while not block_stop: 

483 cur_val = tl.load(out_ptr + out_offsets, mask=mask, other=0.0) 

484 if IS_AMAX: 

485 new_val = tl.maximum(cur_val, src_val) 

486 else: 

487 new_val = tl.minimum(cur_val, src_val) 

488 cas_res = tl.atomic_cas( 

489 out_ptr + out_offsets, cur_val, new_val, sem="relaxed" 

490 ) 

491 stop |= cur_val == cas_res 

492 block_stop = tl.sum(stop.to(tl.int32)) == BLOCK 

493 else: 

494 if IS_AMAX: 

495 tl.atomic_max(out_ptr + out_offsets, src_val, mask=mask, sem="relaxed") 

496 else: 

497 tl.atomic_min(out_ptr + out_offsets, src_val, mask=mask, sem="relaxed") 

498 

499 if USE_MASK: 

500 ones = tl.full((BLOCK,), 1, dtype=tl.int32) 

501 tl.atomic_add(mask_ptr + out_offsets, ones, mask=mask, sem="relaxed") 

502 

503 

504# --------------------------------------------------------------------------- 

505# Generic 5D Kernels with LOOP optimization 

506# For tensors with ndim != 2. 

507# --------------------------------------------------------------------------- 

508 

509 

510@libentry() 

511@triton.heuristics({"BLOCK": heur_block, "LOOP": heur_loop}) 

512@triton.jit(do_not_specialize=["N"]) 

513def scatter_reduce_sum_kernel( 

514 index_ptr, 

515 src_ptr, 

516 out_ptr, 

517 mask_ptr, 

518 N, 

519 out_stride_dim, 

520 src_stride_dim, 

521 src_shape_dim, 

522 out_shape_dim, 

523 DIM: tl.constexpr, 

524 USE_MASK: tl.constexpr, 

525 src_stride_0, 

526 src_stride_1, 

527 src_stride_2, 

528 src_stride_3, 

529 src_stride_4, 

530 idx_shape_0, 

531 idx_shape_1, 

532 idx_shape_2, 

533 idx_shape_3, 

534 idx_shape_4, 

535 src_shape_0, 

536 src_shape_1, 

537 src_shape_2, 

538 src_shape_3, 

539 src_shape_4, 

540 idx_stride_0, 

541 idx_stride_1, 

542 idx_stride_2, 

543 idx_stride_3, 

544 idx_stride_4, 

545 out_stride_0, 

546 out_stride_1, 

547 out_stride_2, 

548 out_stride_3, 

549 out_stride_4, 

550 BLOCK: tl.constexpr, 

551 LOOP: tl.constexpr, 

552): 

553 pid = tl.program_id(axis=0) 

554 base_offsets = pid * BLOCK * LOOP + tl.arange(0, BLOCK) 

555 

556 for i in range(LOOP): 

557 offsets = (base_offsets + i * BLOCK).to(tl.int64) 

558 mask = offsets < N 

559 

560 remaining = offsets 

561 coord0 = remaining // (idx_shape_1 * idx_shape_2 * idx_shape_3 * idx_shape_4) 

562 remaining = remaining % (idx_shape_1 * idx_shape_2 * idx_shape_3 * idx_shape_4) 

563 coord1 = remaining // (idx_shape_2 * idx_shape_3 * idx_shape_4) 

564 remaining = remaining % (idx_shape_2 * idx_shape_3 * idx_shape_4) 

565 coord2 = remaining // (idx_shape_3 * idx_shape_4) 

566 remaining = remaining % (idx_shape_3 * idx_shape_4) 

567 coord3 = remaining // idx_shape_4 

568 coord4 = remaining % idx_shape_4 

569 

570 idx_offsets = ( 

571 coord0 * idx_stride_0 

572 + coord1 * idx_stride_1 

573 + coord2 * idx_stride_2 

574 + coord3 * idx_stride_3 

575 + coord4 * idx_stride_4 

576 ) 

577 src_offsets = ( 

578 coord0 * src_stride_0 

579 + coord1 * src_stride_1 

580 + coord2 * src_stride_2 

581 + coord3 * src_stride_3 

582 + coord4 * src_stride_4 

583 ) 

584 

585 idx = tl.load(index_ptr + idx_offsets, mask=mask, other=0).to(tl.int64) 

586 

587 if DIM == 0: 

588 out_offsets = ( 

589 idx * out_stride_0 

590 + coord1 * out_stride_1 

591 + coord2 * out_stride_2 

592 + coord3 * out_stride_3 

593 + coord4 * out_stride_4 

594 ) 

595 elif DIM == 1: 

596 out_offsets = ( 

597 coord0 * out_stride_0 

598 + idx * out_stride_1 

599 + coord2 * out_stride_2 

600 + coord3 * out_stride_3 

601 + coord4 * out_stride_4 

602 ) 

603 elif DIM == 2: 

604 out_offsets = ( 

605 coord0 * out_stride_0 

606 + coord1 * out_stride_1 

607 + idx * out_stride_2 

608 + coord3 * out_stride_3 

609 + coord4 * out_stride_4 

610 ) 

611 elif DIM == 3: 

612 out_offsets = ( 

613 coord0 * out_stride_0 

614 + coord1 * out_stride_1 

615 + coord2 * out_stride_2 

616 + idx * out_stride_3 

617 + coord4 * out_stride_4 

618 ) 

619 else: 

620 out_offsets = ( 

621 coord0 * out_stride_0 

622 + coord1 * out_stride_1 

623 + coord2 * out_stride_2 

624 + coord3 * out_stride_3 

625 + idx * out_stride_4 

626 ) 

627 

628 src_val = tl.load(src_ptr + src_offsets, mask=mask, other=0).to(tl.float32) 

629 tl.atomic_add(out_ptr + out_offsets, src_val, mask=mask, sem="relaxed") 

630 

631 if USE_MASK: 

632 ones = tl.full((BLOCK,), 1, dtype=tl.int32) 

633 tl.atomic_add(mask_ptr + out_offsets, ones, mask=mask, sem="relaxed") 

634 

635 

636@libentry() 

637@triton.heuristics({"BLOCK": heur_block, "LOOP": heur_loop}) 

638@triton.jit(do_not_specialize=["N"]) 

639def scatter_reduce_prod_kernel( 

640 index_ptr, 

641 src_ptr, 

642 out_ptr, 

643 mask_ptr, 

644 N, 

645 out_stride_dim, 

646 src_stride_dim, 

647 src_shape_dim, 

648 out_shape_dim, 

649 DIM: tl.constexpr, 

650 USE_MASK: tl.constexpr, 

651 src_stride_0, 

652 src_stride_1, 

653 src_stride_2, 

654 src_stride_3, 

655 src_stride_4, 

656 idx_shape_0, 

657 idx_shape_1, 

658 idx_shape_2, 

659 idx_shape_3, 

660 idx_shape_4, 

661 src_shape_0, 

662 src_shape_1, 

663 src_shape_2, 

664 src_shape_3, 

665 src_shape_4, 

666 idx_stride_0, 

667 idx_stride_1, 

668 idx_stride_2, 

669 idx_stride_3, 

670 idx_stride_4, 

671 out_stride_0, 

672 out_stride_1, 

673 out_stride_2, 

674 out_stride_3, 

675 out_stride_4, 

676 BLOCK: tl.constexpr, 

677 LOOP: tl.constexpr, 

678): 

679 pid = tl.program_id(axis=0) 

680 base_offsets = pid * BLOCK * LOOP + tl.arange(0, BLOCK) 

681 

682 for i in range(LOOP): 

683 offsets = (base_offsets + i * BLOCK).to(tl.int64) 

684 mask = offsets < N 

685 

686 remaining = offsets 

687 coord0 = remaining // (idx_shape_1 * idx_shape_2 * idx_shape_3 * idx_shape_4) 

688 remaining = remaining % (idx_shape_1 * idx_shape_2 * idx_shape_3 * idx_shape_4) 

689 coord1 = remaining // (idx_shape_2 * idx_shape_3 * idx_shape_4) 

690 remaining = remaining % (idx_shape_2 * idx_shape_3 * idx_shape_4) 

691 coord2 = remaining // (idx_shape_3 * idx_shape_4) 

692 remaining = remaining % (idx_shape_3 * idx_shape_4) 

693 coord3 = remaining // idx_shape_4 

694 coord4 = remaining % idx_shape_4 

695 

696 idx_offsets = ( 

697 coord0 * idx_stride_0 

698 + coord1 * idx_stride_1 

699 + coord2 * idx_stride_2 

700 + coord3 * idx_stride_3 

701 + coord4 * idx_stride_4 

702 ) 

703 src_offsets = ( 

704 coord0 * src_stride_0 

705 + coord1 * src_stride_1 

706 + coord2 * src_stride_2 

707 + coord3 * src_stride_3 

708 + coord4 * src_stride_4 

709 ) 

710 

711 idx = tl.load(index_ptr + idx_offsets, mask=mask, other=0).to(tl.int64) 

712 

713 if DIM == 0: 

714 out_offsets = ( 

715 idx * out_stride_0 

716 + coord1 * out_stride_1 

717 + coord2 * out_stride_2 

718 + coord3 * out_stride_3 

719 + coord4 * out_stride_4 

720 ) 

721 elif DIM == 1: 

722 out_offsets = ( 

723 coord0 * out_stride_0 

724 + idx * out_stride_1 

725 + coord2 * out_stride_2 

726 + coord3 * out_stride_3 

727 + coord4 * out_stride_4 

728 ) 

729 elif DIM == 2: 

730 out_offsets = ( 

731 coord0 * out_stride_0 

732 + coord1 * out_stride_1 

733 + idx * out_stride_2 

734 + coord3 * out_stride_3 

735 + coord4 * out_stride_4 

736 ) 

737 elif DIM == 3: 

738 out_offsets = ( 

739 coord0 * out_stride_0 

740 + coord1 * out_stride_1 

741 + coord2 * out_stride_2 

742 + idx * out_stride_3 

743 + coord4 * out_stride_4 

744 ) 

745 else: 

746 out_offsets = ( 

747 coord0 * out_stride_0 

748 + coord1 * out_stride_1 

749 + coord2 * out_stride_2 

750 + coord3 * out_stride_3 

751 + idx * out_stride_4 

752 ) 

753 

754 src_val = tl.load(src_ptr + src_offsets, mask=mask, other=0).to(tl.float32) 

755 

756 # CAS loop for product. NaN/Inf guard: if cur_val is NaN, mark as done 

757 # to prevent infinite spin (NaN != NaN causes CAS to always fail). 

758 stop = tl.where(mask, 0, 1).to(tl.int1) 

759 block_stop = False 

760 out_ptr_u32 = (out_ptr + out_offsets).to( 

761 tl.pointer_type(tl.uint32, 1), bitcast=True 

762 ) 

763 while not block_stop: 

764 cur_bits = tl.load(out_ptr_u32, mask=mask, other=0) 

765 cur_val = cur_bits.to(tl.float32, bitcast=True) 

766 new_val = tl.where(stop, cur_val, cur_val * src_val) 

767 # Detect NaN: if new_val != new_val (NaN check), use src_val directly 

768 is_nan = new_val != new_val 

769 new_val = tl.where(is_nan, src_val, new_val) 

770 new_bits = new_val.to(tl.uint32, bitcast=True) 

771 # Sunrise/PTPU is more stable when product CAS operates on the raw 

772 # float32 bit pattern instead of a floating-pointer CAS. 

773 cas_res = tl.atomic_cas(out_ptr_u32, cur_bits, new_bits, sem="acq_rel") 

774 # Mark done if CAS succeeded OR if value is NaN (can't recover) 

775 stop |= (cur_bits == cas_res) | is_nan 

776 block_stop = tl.sum(stop.to(tl.int32)) == BLOCK 

777 

778 if USE_MASK: 

779 ones = tl.full((BLOCK,), 1, dtype=tl.int32) 

780 tl.atomic_add(mask_ptr + out_offsets, ones, mask=mask, sem="relaxed") 

781 

782 

783@libentry() 

784@triton.heuristics({"BLOCK": heur_block, "LOOP": heur_loop}) 

785@triton.jit(do_not_specialize=["N"]) 

786def scatter_reduce_mean_kernel( 

787 index_ptr, 

788 src_ptr, 

789 out_ptr, 

790 count_ptr, 

791 mask_ptr, 

792 N, 

793 out_stride_dim, 

794 src_stride_dim, 

795 src_shape_dim, 

796 out_shape_dim, 

797 DIM: tl.constexpr, 

798 USE_MASK: tl.constexpr, 

799 src_stride_0,