Processor, method, and system for accelerating tensor transpose for machine learning
The present invention relates to a processor designed to optimize memory bandwidth utilization for tensor transpositions in machine learning. It features an instruction decoder that decodes tensor transpose instructions for an input tensor, distinguishing between transposing and stationary axes. An inner transpose engine, equipped with multiple tensor buffer units of varying sizes, transposes the tensor by reading data by rows and writing by columns, efficiently managing memory when the buffer is full. Additionally, an address scheduler determines the target memory addresses in the output tensor memory for the tensor data, improving the handling and transformation of tensors in machine learning environments. This processor significantly enhances the efficiency and speed of tensor operations critical to advanced machine learning applications.
1 . A processor for optimizing memory bandwidth utilization during tensor transpositions in machine learning, comprising:
an instruction decoder configured to decode a tensor transpose instruction for transposing an input tensor stored in an input tensor memory, wherein:
the input tensor comprises a plurality of axes, and
the tensor transpose instruction indicates transposing axes and stationary axes of the plurality of axes, wherein the transposing axes comprise two or more axes being transposed, and the stationary axes comprise one or more axes that are not being transposed;
an inner transpose engine configured to:
provide a plurality of tensor buffer units of different sizes;
activate a tensor buffer unit of the plurality of tensor buffer units;
read tensor data in the transposing axes of the input tensor from the input tensor memory by rows;
write the tensor data into the activated tensor buffer unit by columns; and
in response to the activated tensor buffer unit being full, copy the tensor data from the activated tensor buffer unit into an output tensor memory; and
an address scheduler configured to:
determine target memory addresses in the output tensor memory for writing tensor data from the stationary axes of the input tensor into the output tensor memory.
2 . The processor of claim 1 , wherein the input tensor is segmented into a plurality of sub-tensors based on the activated tensor buffer unit, and each of the sub-tensors has a same size as the activated tensor buffer unit.
3 . The processor of claim 1 , wherein:
the input tensor comprises a channel dimension divided into a plurality of channel groups, each channel group comprising a plurality of channels, wherein a number of the plurality of channels in the channel group is referred to as a size of the channel group of the input tensor, and tensor data in each channel group is stored sequentially in the input tensor memory.
4 . The processor of claim 3 , wherein to activate the tensor buffer unit of the plurality of tensor buffer units, the inner transpose engine is further configured to:
select one tensor buffer unit based on the size of the channel group of the input tensor, wherein the selected tensor buffer unit has a most-inner dimension of a same size as the size of the channel group of the input tensor; and
activate the selected tensor buffer unit.
5 . The processor of claim 4 , wherein to activate the tensor buffer unit of the plurality of tensor buffer units, the inner transpose engine is further configured to:
in response to none of the plurality of tensor buffer units having an inner-most dimension matching the size of the channel group of the input tensor, select an approximate tensor buffer unit that has a smallest inner-most dimension that is greater than the size of the channel group of the input tensor.
6 . The processor of claim 5 , wherein to activate the tensor buffer unit of the plurality of tensor buffer units, the inner transpose engine is further configured to:
apply tensor masks to the approximate tensor buffer unit to obtain a masked tensor buffer unit, such that the masked tensor buffer unit has the inner-most dimension of the same size as the size of the channel group of the input tensor; and
activate the masked tensor buffer unit.
7 . The processor of claim 1 , wherein the input tensor memory has a maximum memory bandwidth, and
each of the plurality of tensor buffer units comprises a plurality of tensor buffer axes, and a product of two most-inner axes of the plurality of tensor buffer axes is equal to the maximum memory bandwidth of the input tensor memory.
8 . The processor of claim 1 , wherein:
each of the plurality of tensor buffer units comprises a first axis, a second axis, and a third axis, wherein the first axis is equal to the second axis, and the third axis represents a channel group dimension that is smaller than a maximum memory bandwidth of the input tensor memory.
9 . The processor of claim 1 , wherein to read the tensor data from the transposing axes of the input tensor from the input tensor memory by rows, the inner transpose engine is further configured to:
determine a continuous unit in the input tensor for transposing, wherein the continuous unit comprises a longest vector of tensor data that is sequentially stored in the input tensor memory.
10 . The processor of claim 9 , wherein the input tensor memory has a maximum memory bandwidth, and
to read the tensor data from the transposing axes of the input tensor from the input tensor memory by rows and write the tensor data into the activated tensor buffer unit by columns, the inner transpose engine is further configured to:
read a first plurality of continuous units in a row of the input tensor to fulfill the maximum memory bandwidth of the input tensor memory;
write the first plurality of continuous units from the row of the input tensor into a column of the activated tensor buffer unit;
move to a next row of the input tensor and read a second plurality of continuous units in the next row to fulfill the maximum memory bandwidth of the input tensor memory; and
write the second plurality of continuous units from the next row of the input tensor into a next column of the activated tensor buffer unit.
11 . The processor of claim 10 , wherein to copy the tensor data from the activated tensor buffer unit into the output tensor memory, the inner transpose engine is further configured to:
in response to the activated tensor buffer unit being full, start copying the tensor data from the activated tensor buffer unit into the output tensor memory without transposing.
12 . A method for optimizing memory bandwidth utilization during tensor transpositions in machine learning, comprising:
decoding a tensor transpose instruction for transposing an input tensor stored in an input tensor memory, wherein the input tensor comprises a plurality of axes, and the tensor transpose instruction indicates transposing axes and stationary axes of the plurality of axes, wherein the transposing axes comprise two or more axes being transposed, and the stationary axes comprise one or more axes that are not being transposed;
activating a tensor buffer unit for transposing the input tensor, wherein the activated tensor buffer unit is selected from a plurality of tensor buffer units of different sizes;
reading tensor data in the transposing axes of the input tensor from the input tensor memory by rows;
writing the tensor data into the activated tensor buffer unit by columns; and
in response to the activated tensor buffer unit being full, copying the tensor data from the activated tensor buffer unit into an output tensor memory.
13 . The method of claim 12 , further comprising:
determining target memory addresses in the output tensor memory for writing tensor data from the stationary axes of the input tensor into the output tensor memory.
14 . The method of claim 12 , wherein the input tensor is segmented into a plurality of sub-tensors based on the activated tensor buffer unit, and each of the sub-tensors has a same size as the activated tensor buffer unit.
15 . The method of claim 12 , wherein the input tensor comprises a channel dimension divided into a plurality of channel groups, each channel group comprising a plurality of channels, wherein a number of the plurality of channels in the channel group is referred to as a size of the channel group of the input tensor, and tensor data within each channel group is stored sequentially in the input tensor memory.
16 . The method of claim 15 , wherein the activating the tensor buffer unit for transposing the input tensor comprises:
selecting one tensor buffer unit based on the size of the channel group of the input tensor, wherein the selected tensor buffer unit has a most-inner dimension of a same size as the size of the channel group of the input tensor; and
activating the selected tensor buffer unit.
17 . The method of claim 16 , wherein the selecting the one tensor buffer unit based on the size of the channel group of the input tensor comprises:
in response to none of the plurality of tensor buffer units having an inner-most dimension matching the size of the channel group of the input tensor, selecting an approximate tensor buffer unit that has a smallest inner-most dimension that is greater than the size of the channel group of the input tensor.
18 . The method of claim 17 , wherein the activating the tensor buffer unit for transposing the input tensor comprises:
applying tensor masks to the approximate tensor buffer unit to obtain a masked tensor buffer unit, such that the masked tensor buffer unit and the input tensor have a same channel group dimension.
19 . The method of claim 12 , wherein the input tensor memory has a maximum memory bandwidth, and
each of the plurality of tensor buffer units comprises a plurality of tensor buffer axes, and a product of two most-inner axes of the plurality of tensor buffer axes equals to the maximum memory bandwidth of the input tensor memory.
20 . The method of claim 12 , wherein each of the plurality of tensor buffer units comprises a first axis, a second axis, and a third axis, wherein the first axis is equal to the second axis, and the third axis represents a channel group dimension that is smaller than a maximum memory bandwidth of the input tensor memory.