TENSOR TRANSPOSE PROCESSOR
The present invention relates to a processor designed to optimize memory bandwidth utilization for tensor transpositions in machine learning. An example processor includes an input tensor shift buffer, a staging buffer, and an output tensor shift buffer. The input tensor shift buffer reads an input tensor from input memory and performs multiple cycles of input tensor shifting. The shifted tensor data is then written into the staging buffer. The output tensor shift buffer reads the shifted tensor data from the staging buffer and performs multiple cycles of output tensor shifting. Finally, the result is written to the output memory. This configuration facilitates efficient handling and transformation of tensor data, optimizing the computational processes required in machine learning tasks.
1 . A processor for accelerating tensor transposition operations in machine learning, comprising:
an input tensor shift buffer;
a staging buffer; and
an output tensor shift buffer;
wherein:
the input tensor shift buffer is configured to:
read an input tensor from an input memory,
perform multiple cycles of input tensor shifting on the input tensor read from the input memory; and
write a result of the multiple cycles of input tensor shifting into the staging buffer; and
the output tensor shift buffer is configured to:
read the result of the multiple cycles of input tensor shifting from the staging buffer;
perform multiple cycles of output tensor shifting on the result read from the staging buffer; and
write a result of the multiple cycles of output tensor shifting to an output memory.
2 . The processor of claim 1 , wherein the staging buffer comprises a plurality of memory banks.
3 . The processor of claim 2 , wherein the input tensor shift buffer is further configured to:
during each cycle of the multiple cycles of input tensor shifting:
read a sub-tensor of the input tensor from the input memory, and
write the sub-tensor of the input tensor in a first shifting direction across the plurality of memory banks in the staging buffer.
4 . The processor of claim 3 , wherein to read the sub-tensor of the input tensor, the input tensor shift buffer is further configured to:
read a row of data of the input tensor from the input memory, and
to write the sub-tensor of the input tensor in the first shifting direction across the plurality of memory banks in the staging buffer, the input tensor shift buffer is further configured to:
write the row of data read from the input tensor in a column direction across the plurality of memory banks in the staging buffer.
5 . The processor of claim 2 , wherein the output tensor shift buffer is further configured to:
during each cycle of the multiple cycles of output tensor shifting:
read a sub-tensor from the staging buffer, and
write the sub-tensor in a second shifting direction into the output memory.
6 . The processor of claim 5 , wherein to read the sub-tensor from the staging buffer, the output tensor shift buffer is further configured to:
read a row of data from the staging buffer; and
write the row of data from the staging buffer by shifting the row of data in the second shifting direction.
7 . The processor of claim 1 , wherein a direction of the input tensor shifting is opposite to a direction of the output tensor shifting.
8 . The processor of claim 1 , wherein the staging buffer comprises a pair of buffers of a same size.
9 . The processor of claim 8 , wherein the pair of buffers comprise a first buffer and a second buffer that support parallel processing, and
when the input tensor shift buffer is writing to the first buffer, the output tensor shift buffer is reading from the second buffer.
10 . The processor of claim 8 , wherein the pair of buffers comprise a first buffer and a second buffer, and
after the input tensor shift buffer writes a first result of the multiple cycles of input tensor shifting on a first input tensor into the first buffer, the output tensor shift buffer starts reading the first result from the first buffer to perform the multiple cycles of output tensor shifting; and
while the output tensor is reading from the first buffer, the input tensor shift buffer starts writing a second result of the multiple cycles of input tensor shifting on a second input tensor into the second buffer, such that the input tensor shift buffer and the output tensor shift buffer avoid idle time.
11 . A method for accelerating tensor transposition operations in machine learning, comprising:
performing, using an input tensor shift buffer, multiple cycles of input tensor shifting on an input tensor read from an input memory;
during each cycle of the multiple cycles of input tensor shifting, writing, using the input tensor shift buffer, a result of the cycle of input tensor shifting into a staging buffer;
performing, using an output tensor shift buffer, multiple cycles of output tensor shifting on the results read from the staging buffer; and
during each cycle of the multiple cycles of output tensor shifting, writing, using the output tensor shift buffer, a result of the cycle of output tensor shifting to an output memory.
12 . The method of claim 11 , wherein the staging buffer comprises a plurality of memory banks.
13 . The method of claim 12 , wherein the performing the multiple cycles of input tensor shifting on the input tensor read from the input memory comprises:
during each cycle of the multiple cycles of input tensor shifting:
reading a sub-tensor of the input tensor from the input memory, and
shifting the sub-tensor of the input tensor in a first shifting direction across the plurality of memory banks when writing to the staging buffer.
14 . The method of claim 13 , wherein the sub-tensor read from the input memory comprises a row of data of the input tensor, and
the writing the result of the cycle of input tensor shifting into the staging buffer comprises:
writing the row of data read from the input tensor in a column direction across the plurality of memory banks in the staging buffer.
15 . The method of claim 12 , wherein the performing the multiple cycles of output tensor shifting on the results read from the staging buffer comprises:
during each cycle of the multiple cycles of output tensor shifting:
reading a sub-tensor from the staging buffer, and
shifting the sub-tensor in a second shifting direction when writing to the output memory.
16 . The method of claim 15 , wherein the sub-tensor read from the staging buffer comprises a row of data read from the staging buffer, and
the writing the result of the cycle of output tensor shifting to an output memory comprises:
write the row of data from the staging buffer by shifting the row of data in the second shifting direction.
17 . The method of claim 11 , wherein a direction of the input tensor shifting is opposite to a direction of the output tensor shifting.
18 . The method of claim 11 , wherein the staging buffer comprises a pair of buffers of a same size.
19 . The method of claim 18 , wherein the pair of buffers comprise a first buffer and a second buffer that support parallel processing, and
when the input tensor shift buffer is writing to the first buffer, the output tensor shift buffer is reading from the second buffer.
20 . The method of claim 18 , wherein the pair of buffers comprise a first buffer and a second buffer, and
after the input tensor shift buffer writes a first result of the multiple cycles of input tensor shifting on a first input tensor into the first buffer, the output tensor shift buffer starts reading the first result from the first buffer to perform the multiple cycles of output tensor shifting; and
while the output tensor is reading from the first buffer, the input tensor shift buffer starts writing a second result of the multiple cycles of input tensor shifting on a second input tensor into the second buffer, such that the input tensor shift buffer and the output tensor shift buffer avoid idle time.