Emit two unconditional SliceUpdates for ring buffer wrapping. write_pos = start_pos % ring_size first_len = ring_size - write_pos first_chunk = update[:, :, :first_len, :] (Slice clamps to seq_len) actual_first = first_chunk.shape[2]
(
self, P: "MLXProgramBuilder", cache_slot, update_slot, start_slot
)
| 260 | ) |
| 261 | |
| 262 | def _emit_ring_buffer( |
| 263 | self, P: "MLXProgramBuilder", cache_slot, update_slot, start_slot |
| 264 | ): |
| 265 | """ |
| 266 | Emit two unconditional SliceUpdates for ring buffer wrapping. |
| 267 | |
| 268 | write_pos = start_pos % ring_size |
| 269 | first_len = ring_size - write_pos |
| 270 | first_chunk = update[:, :, :first_len, :] (Slice clamps to seq_len) |
| 271 | actual_first = first_chunk.shape[2] (min(first_len, seq_len)) |
| 272 | rest_chunk = update[:, :, actual_first:seq_len, :] |
| 273 | overflow = seq_len - actual_first |
| 274 | SliceUpdate(cache, first_chunk, write_pos, write_pos + actual_first) |
| 275 | SliceUpdate(cache, rest_chunk, 0, overflow) |
| 276 | |
| 277 | When no wrap: actual_first == seq_len, rest_chunk is zero-length, |
| 278 | second SliceUpdate is a no-op (guarded in exec_slice_update). |
| 279 | """ |
| 280 | ring_size = self.ring_size |
| 281 | |
| 282 | # write_pos = start_pos % ring_size |
| 283 | _, write_pos_slot = P.slot_manager.make_tmp_value_slot() |
| 284 | P.emit( |
| 285 | ModIntNode( |
| 286 | a=P.to_int_or_vid(start_slot), |
| 287 | b=IntOrVid.from_literal(ring_size), |
| 288 | out=P.slot_to_vid(write_pos_slot), |
| 289 | ) |
| 290 | ) |
| 291 | |
| 292 | # seq_len = update.shape[2] |
| 293 | _, seq_len_slot = P.slot_manager.make_tmp_value_slot() |
| 294 | P.emit( |
| 295 | SymSizeNode( |
| 296 | a=P.slot_to_tid(update_slot), |
| 297 | dim=2, |
| 298 | out=P.slot_to_vid(seq_len_slot), |
| 299 | ) |
| 300 | ) |
| 301 | |
| 302 | # first_len = ring_size - write_pos (may be > seq_len) |
| 303 | _, first_len_slot = P.slot_manager.make_tmp_value_slot() |
| 304 | P.emit( |
| 305 | SubtractIntNode( |
| 306 | a=IntOrVid.from_literal(ring_size), |
| 307 | b=P.to_int_or_vid(write_pos_slot), |
| 308 | out=P.slot_to_vid(first_len_slot), |
| 309 | ) |
| 310 | ) |
| 311 | |
| 312 | # first_chunk = update[:, :, :first_len, :] (Slice clamps to seq_len) |
| 313 | _, first_chunk_slot = P.make_tmp_slot() |
| 314 | P.emit( |
| 315 | SliceNode( |
| 316 | x=P.slot_to_tid(update_slot), |
| 317 | out=P.slot_to_tid(first_chunk_slot), |
| 318 | axis=IntOrVid.from_literal(2), |
| 319 | start=IntOrVid.from_literal(0), |
no test coverage detected