warp.tile_extract#

warp.tile_extract(a: Tile[Any, tuple[int]], i: int32) → Any#

Kernel: true
Python: false
Differentiable: true

Extract a single element from a tile.

Each calling thread receives the element at the index it passes, so a block-uniform index gives every thread the same value while a per-thread index gives each thread its own element. The current implementation places a in shared memory before extraction, where it counts against the block’s shared-memory budget (see Shared Memory Budget).

Pass one index per tile dimension to read the element itself. A tile of vectors accepts one extra index, which selects a scalar component. A tile of matrices accepts one extra index, which returns the corresponding row as a vector, or two, which select a scalar element.

Parameters:

a – Tile to extract the element from. All index arguments may be runtime values.
i – Index along the first dimension of a. May be a runtime value. It must be non-negative and less than that dimension’s extent.

Returns:

The element at index i of a 1D tile, with the same data type as the tile’s elements.

Example

@wp.kernel
def shift_right(a: wp.array[float], out: wp.array[float]):
    block, lane = wp.tid()
    t = wp.tile_load(a, shape=5)

    # the index is block-uniform and may differ between blocks
    value = wp.tile_extract(t, block + 1)
    if lane == 0:
        out[block] = value

a = wp.array(np.arange(5, dtype=np.float32) * 10.0, dtype=float)
out = wp.zeros(4, dtype=float)

wp.launch_tiled(shift_right, dim=4, inputs=[a], outputs=[out], block_dim=4)

print(out.numpy())

[10. 20. 30. 40.]

warp.tile_extract( a: Tile[Any, tuple[int, ...]], i: int32, j: int32, ) → Any

Kernel: true
Python: false
Differentiable: true

Extract a single element from a tile.

Two-index overload: both indices select an element of a 2D tile, or i indexes a 1D tile of vectors and j selects a component of that vector, or i indexes a 1D tile of matrices and j selects a row of that matrix. See the one-index overload for the shared contract and a usage example.

Parameters:

a – Tile to extract the element from. All index arguments may be runtime values.
i – Index along the first dimension
j – Index along the second dimension, or vector component index, or matrix row index

Returns:

The element at (i, j) of a 2D tile, with the same data type as the tile’s elements; the scalar component j of the vector element at i; or row j of the matrix element at i, as a vector.

warp.tile_extract( a: Tile[Any, tuple[int, ...]], i: int32, j: int32, k: int32, ) → Any

Kernel: true
Python: false
Differentiable: true

Extract a single element from a tile.

Three-index overload: the indices select an element of a 3D tile, a component from a 2D tile of vectors, a row from a 2D tile of matrices, or an entry from a 1D tile of matrices. See the one-index overload for the shared contract and a usage example.

Parameters:

a – Tile to extract the element from. All index arguments may be runtime values.
i – Index along the first dimension
j – Index along the second dimension, or first matrix index
k – Index along the third dimension, or vector index, or second matrix index

Returns:

The selected tile element, vector component, matrix row, or matrix entry. A matrix row is returned as a vector; a vector component or matrix entry is returned as a scalar.

warp.tile_extract( a: Tile[Any, tuple[int, ...]], i: int32, j: int32, k: int32, l: int32, ) → Any

Kernel: true
Python: false
Differentiable: true

Extract a single element from a tile.

Four-index overload: the indices select an element of a 4D tile, a component from a 3D tile of vectors, a row from a 3D tile of matrices, or an entry from a 2D tile of matrices. See the one-index overload for the shared contract and a usage example.

Parameters:

a – Tile to extract the element from. All index arguments may be runtime values.
i – Index along the first dimension
j – Index along the second dimension
k – Index along the third dimension, or first matrix index
l – Index along the fourth dimension, or vector index, or second matrix index

Returns:

The selected tile element, vector component, matrix row, or matrix entry. A matrix row is returned as a vector; a vector component or matrix entry is returned as a scalar.

warp.tile_extract( a: Tile[Any, tuple[int, ...]], i: int32, j: int32, k: int32, l: int32, m: int32, ) → Any

Kernel: true
Python: false
Differentiable: true

Extract a single element from a tile.

Five-index overload: the indices select an element of a 5D tile, a component from a 4D tile of vectors, a row from a 4D tile of matrices, or an entry from a 3D tile of matrices. See the one-index overload for the shared contract and a usage example.

Parameters:

a – Tile to extract the element from. All index arguments may be runtime values.
i – Index along the first dimension
j – Index along the second dimension
k – Index along the third dimension
l – Index along the fourth dimension, or first matrix index
m – Vector index, or second matrix index

Returns:

The selected tile element, vector component, matrix row, or matrix entry. A matrix row is returned as a vector; a vector component or matrix entry is returned as a scalar.

warp.tile_extract( a: Tile[Any, tuple[int, int, int, int]], i: int32, j: int32, k: int32, l: int32, m: int32, n: int32, ) → Any

Kernel: true
Python: false
Differentiable: true

Extract a single element from a tile.

Six-index overload: the indices select an element of a 4D tile of matrices, followed by a row and a column index. See the one-index overload for the shared contract and a usage example.

Parameters:

a – Tile to extract the element from
i – Index along the first dimension
j – Index along the second dimension
k – Index along the third dimension
l – Index along the fourth dimension
m – First matrix index
n – Second matrix index

Returns:

The scalar component (m, n) of the matrix element at the specified tile indices.