tilus.Script.load_global¶
- Script.load_global(src, /, *, offsets, shape=None, layout=None, dims=None, out=None)[source]¶
Load a slice of global tensor into a register tensor.
This instruction loads a slice of the global tensor x into a register tensor, given the offsets for each dimension of the global tensor and the shape of the slice to be loaded.
When we only slice over a subset of the dimensions of the global tensor, we can specify the dims parameter to indicate which dimensions are being sliced.
The optional layout parameter can be used to specify the layout of the register tensor.
When out is provided, the loaded data will be stored in the out register tensor, otherwise a new register tensor will be allocated.
- Parameters:
src (GlobalTensor) – The global tensor to load from.
offsets (Sequence[Expr | int]) – The offsets for each dimension of the global tensor. The length of this sequence must match the number of dimensions of the global tensor.
shape (Sequence[int], optional) – The shape of the slice to be loaded. If not provided, the shape of the global tensor will be used.
layout (RegisterLayout, optional) – The layout of the register tensor. If not provided, the layout will be inferred based on the operations performed on it. When provided, its shape must match the shape parameter.
dims (Sequence[int], optional) – The dimensions of the global tensor that are being sliced. If not provided, it is assumed that all dimensions are being sliced. The length of this sequence must match the number of dimensions of the register tensor being loaded into.
out (RegisterTensor, optional) – The register tensor to store the loaded data into. If not provided, a new register tensor will be allocated.
- Returns:
ret – The register tensor containing the loaded data from the global tensor.
- Return type: