jax.experimental.pallas.tpu.BufferedRef

class jax.experimental.pallas.tpu.BufferedRef(_spec, _buffer_type, _in_buffer_count, _out_buffer_count, _grid_rank, window_ref, copy_in_slot, wait_in_slot, copy_out_slot, wait_out_slot, next_fetch, sem_recvs, sem_sends, tiling, is_trivial_windowing=False, has_allocated_buffer=False, prefetched_count=0, await_prefetch=False, output_wait_read_only=False)

A helper class to automate VMEM double buffering in pallas pipelines.

Parameters:

• _spec (pallas_core.BlockSpec)

• _buffer_type (BufferType)

• _in_buffer_count (int)

• _out_buffer_count (int)

• _grid_rank (int | None)

• window_ref (ArrayRef | None)

• copy_in_slot (int | jax.Array | None)

• wait_in_slot (int | jax.Array | None)

• copy_out_slot (int | jax.Array | None)

• wait_out_slot (int | jax.Array | None)

• next_fetch (Sequence[jax.Array | int] | None)

• sem_recvs (SemaphoreTuple | None)

• sem_sends (SemaphoreTuple | None)

• tiling (Tiling | None)

• is_trivial_windowing (bool)

• has_allocated_buffer (bool)

• prefetched_count (int)

• await_prefetch (bool)

• output_wait_read_only (bool)

spec

pallas blockspec.

buffer_type

enum indicating whether this is an input, output, or in/out buffered reference.

window_ref

a multiple-buffer to hold the working and dirty buffers used to copy into and out of. In the case of a BufferedRef targeting a VMEM reference, this simply points to the existing ref.

Type:

ArrayRef | None

copy_in_slot

current slot to copy in for the working buffer.

Type:

int | jax.Array | None

copy_out_slot

current slot to copy out for the working buffer.

Type:

int | jax.Array | None

wait_in_slot

current slot to wait in for the working buffer.

Type:

int | jax.Array | None

wait_out_slot

current slot to wait out for the working buffer.

Type:

int | jax.Array | None

next_fetch

Holds the next grid indices to fetch for lookahead. This is the register state used to track the indices within the pipeline loop.

Type:

Sequence[jax.Array | int] | None

sem_recvs

Multiple buffered semaphores for input DMAs.

Type:

SemaphoreTuple | None

sem_sends

Multiple buffered semaphores for output DMAs.

Type:

SemaphoreTuple | None

tiling

The tiling to assume for the buffers.

Type:

Tiling | None

is_trivial_windowing

Whether the reference uses trivial windowing.

Type:

bool

has_allocated_buffer

Whether the reference has an allocated buffer due to being in a different memory space than the source ref.

Type:

bool

__init__(_spec, _buffer_type, _in_buffer_count, _out_buffer_count, _grid_rank, window_ref, copy_in_slot, wait_in_slot, copy_out_slot, wait_out_slot, next_fetch, sem_recvs, sem_sends, tiling, is_trivial_windowing=False, has_allocated_buffer=False, prefetched_count=0, await_prefetch=False, output_wait_read_only=False)

Parameters:

• _spec (pallas_core.BlockSpec)

• _buffer_type (BufferType)

• _in_buffer_count (int)

• _out_buffer_count (int)

• _grid_rank (int | None)

• window_ref (ArrayRef | None)

• copy_in_slot (int | jax.Array | None)

• wait_in_slot (int | jax.Array | None)

• copy_out_slot (int | jax.Array | None)

• wait_out_slot (int | jax.Array | None)

• next_fetch (Sequence[jax.Array | int] | None)

• sem_recvs (SemaphoreTuple | None)

• sem_sends (SemaphoreTuple | None)

• tiling (Tiling | None)

• is_trivial_windowing (bool)

• has_allocated_buffer (bool)

• prefetched_count (int)

• await_prefetch (bool)

• output_wait_read_only (bool)

Return type:

None

Methods

__init__(_spec, _buffer_type, ...[, ...])
advance_copy_in_slot([predicate])	Switch to the next copy slot.
advance_copy_out_slot([predicate])	Switch to the next copy slot.
advance_next_fetch(grid)
advance_wait_in_slot([predicate])	Switch to the next wait slot.
advance_wait_out_slot([predicate])	Switch to the next wait slot.
bind_existing_ref(window_ref, indices)	For handling VMEM references, the pipeline aliases the existing ref.
compute_slice(grid_indices)	Compute the indexers for the window at given grid indices.
copy_in(src_ref, grid_indices)	Starts copy of HBM dma slice into the current slot.
copy_out(dst_ref, grid_indices)	Starts copy of HBM dma slice from the current slot.
create(spec, dtype_or_type, buffer_type, ...)	Create a BufferedRef.
get_dma_slice(src_ty, grid_indices)
initialize_slots()	Initializes slots to 0.
input(spec, dtype_or_type[, buffer_count])
input_output(spec, dtype_or_type[, buffer_count])	Creates an input_output BufferedRef, which uses in + out - 1 slots.
output(spec, dtype_or_type[, buffer_count])
unbind_refs()
wait_in(src_ref, grid_indices)	Waits for input copy to finish.
wait_out(dst_ref, grid_indices)	Waits for output copy to finish.
wait_out_write(dst_ref, grid_indices)	Waits until the output copies have been written to dst_ref.
with_next_fetch([next_fetch])
with_slot_index([copy_in_slot, ...])	Returns a new BufferedRef with the given slot index.
with_spec(spec)	Returns a new BufferedRef with the given block spec.
with_window_ref(window_ref)

Attributes

await_prefetch
block_shape
buffer_count	Returns the number of buffers used for multiple buffering.
buffer_type
compute_index
cumulative_copy_in	The cumulative number of copy_ins issued on this buffer.
cumulative_copy_out	The cumulative number of copy_outs issued on this buffer.
cumulative_wait_in	The cumulative number of wait_ins issued on this buffer.
cumulative_wait_out	The cumulative number of wait_outs issued on this buffer.
current_copy_in_slot	Index in multiple buffer corresponding to the current slot.
current_copy_out_slot	Index in multiple buffer corresponding to the current copy slot.
current_ref	Returns the current working slice of the double-buffer.
current_wait_in_slot	Index in multiple buffer corresponding to the current wait slot.
current_wait_out_slot	Index in multiple buffer corresponding to the current wait slot.
has_allocated_buffer
has_indirect	Whether any block dimension uses indirect indexing.
in_buffer_count	Number of slots the pipeline prefetches inputs into.
is_buffered	Whether this buffer is multiple-buffered.
is_input
is_input_output
is_manual
is_output
is_trivial_windowing
next_fetch_indices	Returns the next grid indices to fetch from if using lookahead.
out_buffer_count	Number of slots the pipeline writes back from.
output_wait_read_only
prefetched_count
spec
use_lookahead	Whether this buffer allows lookahead for fetching blocks.
window_ref
copy_in_slot
wait_in_slot
copy_out_slot
wait_out_slot
next_fetch
sem_recvs
sem_sends
tiling