Broadcast

Applicability

Product	Supported
Atlas 350 Accelerator Card	√
Atlas A3 training product / Atlas A3 inference product	√
Atlas A2 training product / Atlas A2 inference product	√
Atlas 200I/500 A2 inference product	x
Atlas inference product AI Core	√
Atlas inference product Vector Core	x
Atlas training product	x

Function Usage

Broadcasts the input based on the output shape.

For example, if the shape of A is (2, 1) and the target shape after broadcasting is (2, 16), then the original single column will be expanded to 16 identical columns.

      
           Input data:
[[ 1]
 [ 2]]
Output data:
[[ 1  1  1  1  1  1  1  1  1  1  1  1  1  1  1  1]
 [ 2  2  2  2  2  2  2  2  2  2  2  2  2  2  2  2]]

Principles

The figure below illustrates the internal algorithm block diagram of Broadcast high-level APIs, taking the float type, ND format, and broadcasting from [m, 1] to [m, k] as examples.

Figure 1 Broadcast algorithm block diagram

The computation process is divided into the following steps, all of which are performed on vectors:

brcb step: Broadcast each element as a data block.
Copy step: Copy each data block to multiple data blocks. In the k-aligned scenario, the result is y.
In the non-k-aligned scenario, use GatherMask to truncate [m, k] elements, where k' indicates the size obtained by padding k upwards to be 32-byte aligned.

Prototype

Pass the temporary space through the sharedTmpBuffer input parameter.

        
             template <typename T, int32_t dim, int32_t axis, bool isReuseSource = false>
__aicore__ inline void Broadcast(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const uint32_t dstShape[dim], const uint32_t srcShape[dim], LocalTensor<uint8_t>& sharedTmpBuffer)

Allocate the temporary space through the API framework.

        
             template <typename T, int32_t dim, int32_t axis, bool isReuseSource = false>
__aicore__ inline void Broadcast(const LocalTensor<T>& dstLocal, const LocalTensor<T>& srcLocal, const uint32_t dstShape[dim], const uint32_t srcShape[dim])

Dynamic shapes are supported.

It is supported only on Atlas 350 Accelerator Card.

         
              template<class T, int constRank=-1, uint32_t* constDstShape = nullptr, uint32_t* constSrcShape = nullptr, bool constSrcInnerPad = false>
__aicore__ inline void Broadcast(const LocalTensor<T>& dst, const LocalTensor<T>& src, const uint32_t* dstShape, const uint32_t* srcShape, BroadcastTiling* tiling)

This API requires extra temporary space to store intermediate variables during computation. The temporary space can be passed through the sharedTmpBuffer input parameter or allocated through the API framework.

When the sharedTmpBuffer input parameter is used for passing the temporary space, the tensor serves as the temporary space. In this case, the API framework is not required for temporary space allocation. This enables you to manage the sharedTmpBuffer space and reuse the buffer after calling the API, so that the buffer is not repeatedly allocated and deallocated, improving the flexibility and buffer utilization.
When the API framework is used for temporary space allocation, you do not need to allocate the space, but must reserve the required size for the temporary space.

If sharedTmpBuffer is used, you must allocate space for the tensor. If the API framework is used, you must reserve the temporary space. To obtain the size of the temporary space (BufferSize) to be reserved, use the API provided in GetBroadCastMaxMinTmpSize.

In addition, an API on the kernel is provided on the Atlas 350 Accelerator Card to compute tiling based on the broadcast implementation and obtain the tiling result. The template parameters of this API are the same as those of the Broadcast API that supports dynamic shape. For details about other parameters, see Table 5.

Tiling computation API on the kernel

        
             template<class T, int constRank=-1, uint32_t* constDstShape = nullptr, uint32_t* constSrcShape = nullptr>
__aicore__ inline void GetBroadcastTilingInfo(uint32_t rank, const uint32_t* dstShape, const uint32_t* srcShape, bool srcInnerPad, BroadcastTiling& tiling)

Parameters

**Table 1** Template parameters
Parameter	Function
T	Operand data type. For the Atlas 350 Accelerator Card, the supported data types are those corresponding to the 8-, 16-,32-, and 64-bit widths. For details about the data types, see Data Types Corresponding to Different Bit Widths. For the Atlas A3 training product / Atlas A3 inference product , the supported data types are int8_t, uint8_t, half, and float. For the Atlas A2 training product / Atlas A2 inference product , the supported data types are int8_t, uint8_t, half, and float. For the Atlas inference product AI Core, the supported data types are int8_t, uint8_t, half, and float.
dim	Dimension of the input/output tensor. Currently, only 1-dimensional and 2-dimensional tensors are supported.
axis	Dimension to be broadcasted. Currently, only dimensions 0 and 1 are supported. The parameter value 0 indicates that the first dimension is to be broadcasted, and the parameter value 1 indicates that the second dimension is to be broadcasted.
isReuseSource	Whether the source operand can be modified. This parameter is reserved. Pass the default value false.

**Table 2** Template parameters of the API that supports dynamic shape
Parameter	Function
T	Data type of the operand. Currently, int8_t, uint8_t, int16_t, uint16_t, half, bfloat16_t, int32_t, uint32_t, float, int64_t, and uint64_t are supported.
constRank	Number of dimensions of the input/output tensor. If the default value -1 indicates the dynamic shape scenario, and the rank parameter in the GetBroadcastTilingInfo API is used for computing. If the value of constRank is greater than 0, the value must be the same as that of rank in the GetBroadcastTilingInfo API. The rank value range is [1, 9].
constDstShape	Shape of the output tensor. It is an array of the uint32_t type. If any dimension value in the array is 0, the dimension is dynamic, and the actual value of the dimension is determined by the value of the corresponding dimension of dstShape. If any dimension value in the array is greater than 0, the dimension is static, and the dimension value is the same as that of the corresponding dimension in dstShape. This parameter is reserved. Pass the default value nullptr.
constSrcShape	Shape of the input tensor. It is an array of the uint32_t type. If any dimension value in the array is 0, the dimension is dynamic, and the actual shape is determined by srcShape. If any dimension value in the array is greater than 0, the dimension is static, and the dimension value is the same as that in srcShape. This parameter is reserved. Pass the default value nullptr.
constSrcInnerPad	Whether the last dimension of the input, srcShape[rank-1], is 32-byte aligned. rank indicates the number of dimensions of the input/output tensor. This parameter is reserved. Pass the default value false.

**Table 3** API parameters
Parameter	Input/Output	Description
dstLocal	Output	Destination operand. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
srcLocal	Input	Source operand. The source operand must have the same data type as the destination operand. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
dstShape	Input	Shape of the output tensor. It is an array of the uint32_t type with a length of 1 or 2. The number of input shape dimensions must be the same as that of output shape dimensions.
srcShape	Input	Shape of the input tensor. It is an array of the uint32_t type with a length of 1 or 2. The number of input shape dimensions must be the same as that of output shape dimensions.
sharedTmpBuffer	Input	Temporary buffer. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT. This parameter is used to store intermediate variables during complex computation in Broadcast and is provided by developers. For details about how to obtain the temporary space size (BufferSize), see GetBroadCastMaxMinTmpSize.

**Table 4** Parameters of the API that supports dynamic shape
Parameter	Input/Output	Description
dst	Output	Destination operand. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
src	Input	Source operand. The source operand must have the same data type as the destination operand. The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
dstShape	Input	Shape of the output tensor. It is an array of the uint32_t type. The length range is [1, 9]. The number of input shape dimensions must be the same as that of output shape dimensions, and the condition dstShape[i] ≥ srcShape[i] must be met.
srcShape	Input	Shape of the input tensor. It is an array of the uint32_t type. The length range is [1, 9]. The number of input shape dimensions must be the same as that of output shape dimensions, and the condition dstShape[i] ≥ srcShape[i] must be met. If the value of srcShape[i] is 1 and dstShape[i] is not equal to srcShape[i], the axis is the broadcast axis.
tiling	Input	Tiling information required by the Broadcast API. BroadcastTiling* type, which is obtained by calling the tiling computation API GetBroadcastTilingInfo on the kernel.

**Table 5** Parameters of the tiling computing API on the kernel
Parameter	Input/Output	Function
rank	Input	Number of dimensions of the input/output tensor. Currently, the supported value range is [1, 9].
dstShape	Input	Shape of the output tensor. It is an array of the uint32_t type. The length range is [1, 9]. The number of input shape dimensions must be the same as that of output shape dimensions, and the condition dstShape[i] ≥ srcShape[i] must be met.
srcShape	Input	Shape of the input tensor. It is an array of the uint32_t type. The length range is [1, 9]. The number of input shape dimensions must be the same as that of output shape dimensions, and the condition dstShape[i] ≥ srcShape[i] must be met. If the value of srcShape[i] is 1 and dstShape[i] is not equal to srcShape[i], the axis is the broadcast axis.
srcInnerPad	Input	Whether the last dimension of the input, srcShape[rank-1], is 32-byte aligned. Currently, the value can only be false.
tiling	Output	Tiling computing information, which is of the BroadcastTiling& type.

Returns

None

Constraints

For details about the operand address alignment requirements, see General Address Alignment Restrictions.
The source operand address must not overlap the destination operand address.
Currently, only the ND format is supported.
Currently, dim can only be set to 1 or 2, and axis can only be set to 0 or 1.
The dynamic shape API and the tiling API of the kernel are supported only on the Atlas 350 Accelerator Card.
For the Atlas inference product AI Core, when dim is 2 and axis is 1, srcShape[0] must be 32-byte aligned. That is, when the input/output tensor has two dimensions and the broadcast dimension is 1, the data in dimension 0 of the input tensor must be a multiple of 32 bytes.
When dim is 2 and axis is 0, srcShape[1] must be 32-byte aligned.
For the Atlas 350 Accelerator Card, the number of dimensions supported by the input/output tensor, that is, the value range supported by rank is [1, 9].

Examples

For the example calls, see broadcast operator sample.

      
           // dstLocal: output tensor
// srcLocal: input tensor
const uint32_t srcShape[2] = {1, 16}; // Source data shape
const uint32_t dstShape[2] = {16, 16}; // Broadcast data shape
The AscendC::Broadcast<float, 2, 1>(dstLocal, srcLocal, dstShape, srcShape); // The broadcast data type is float, the source data is two-dimensional, and broadcast is applied along dimension 1.

Result example:

       
        
          
          
            Input (srcLocal):
[[ 1]
 [ 2]
 [ 3]
 [ 4]
 [ 5]
 [ 6]
 [ 7]
 [ 8]
 [ 9]
 [10]
 [11]
 [12]
 [13]
 [14]
 [15]
 [16]]
dim: 2
axis: 1
Output (dstLocal):
[[ 1  1  1  1  1  1  1  1  1  1  1  1  1  1  1  1]
 [ 2  2  2  2  2  2  2  2  2  2  2  2  2  2  2  2]
 [ 3  3  3  3  3  3  3  3  3  3  3  3  3  3  3  3]
 [ 4  4  4  4  4  4  4  4  4  4  4  4  4  4  4  4]
 [ 5  5  5  5  5  5  5  5  5  5  5  5  5  5  5  5]
 [ 6  6  6  6  6  6  6  6  6  6  6  6  6  6  6  6]
 [ 7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7]
 [ 8  8  8  8  8  8  8  8  8  8  8  8  8  8  8  8]
 [ 9  9  9  9  9  9  9  9  9  9  9  9  9  9  9  9]
 [10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10]
 [11 11 11 11 11 11 11 11 11 11 11 11 11 11 11 11]
 [12 12 12 12 12 12 12 12 12 12 12 12 12 12 12 12]
 [13 13 13 13 13 13 13 13 13 13 13 13 13 13 13 13]
 [14 14 14 14 14 14 14 14 14 14 14 14 14 14 14 14]
 [15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15]
 [16 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16]]

           

         

       
      

Parent topic: Tensor Transformation