asc2.copy_in

asc2.copy_in(src: GlobalTensor, offsets: Iterable[RuntimeInt], shape: Iterable[int], location: TensorLocation = TensorLocation.UB, *, real_shape: Iterable[RuntimeInt] | None = None, pad_value: RuntimeNumeric = 0) → LocalTensor
asc2.copy_in(src: GlobalTensor, offsets: Iterable[RuntimeInt]) → PlainValue

Copy data from a global tensor into a local tensor or scalar value.

This function supports two modes of operation:

  1. Load a local tensor with offsets: Load values to a local tensor of the given shape from the global tensor at the specified offsets.

  2. Load a scalar: When shape is not provided, load a single scalar value at the specified offsets.

Parameters:

  • src – The source global tensor.

  • offsets – The offsets into the global tensor for each dimension.

  • shape – The shape of the local tensor to load. If None, loads a scalar value. Must contain static values (e.g., ConstExpr or compile-time constants). For 1D tensors, any shape is supported. For 2D+ tensors in UB, the last dimension must be aligned to 32 bytes (e.g., 8 elements for float32, 16 elements for float16).

  • location – The memory location for the local tensor. Default is TensorLocation.UB. Available locations: UB, L1, L0A, L0B, BT.

  • real_shape – Explicitly specify how many elements to load from the global tensor. The local tensor will have the given shape, but only real_shape elements are loaded; remaining elements are filled with pad_value. Must match the rank of shape and each dimension must not exceed the corresponding tensor dimension. Not supported for TensorLocation.L1, L0A, or L0B.

  • pad_value – The value to use for padding when real_shape is provided. Default is 0.

Returns:

A local tensor loaded from the global tensor (when shape is provided) PlainValue: A scalar value loaded from the global tensor (when shape is None)

Return type:

LocalTensor

Raises:

  • TypeError – If src is not a GlobalTensor or location is not a TensorLocation

  • RuntimeError – If shape is invalid, data alignment check fails, offsets rank mismatch, real_shape exceeds tensor shape, or real_shape is used with TensorLocation.L1, L0A, or L0B

Note

Only 1D and 2D tensors are fully supported and stable; higher-dimensional support is experimental.

Examples

Copy a 1D tensor using explicit offsets:

x_gm = asc2.global_tensor(x_ptr, [1024])
result = asc2.copy_in(x_gm, [256], [128])

Copy a 2D tensor from a 2D global tensor:

x_gm = asc2.global_tensor(x_ptr, [64, 128])
result = asc2.copy_in(x_gm, [8, 16], [16, 32])

Copy a scalar value:

x_gm = asc2.global_tensor(x_ptr, [1024])
scalar = asc2.copy_in(x_gm, [42])

Copy a 1D tensor with padding (load fewer elements than tensor shape):

x_gm = asc2.global_tensor(x_ptr, [256])
result = asc2.copy_in(x_gm, [200], [128], pad_value=2.0)

Copy a 2D tensor with real_shape and padding (load fewer elements than tensor shape):

x_gm = asc2.global_tensor(x_ptr, [100, 100])
result = asc2.copy_in(x_gm, [0, 0], [16, 16], real_shape=[12, 12], pad_value=-1.0)
# result has shape [16, 16], but only 12x12 elements loaded from global tensor, rest padded with -1.0