Storing tensors in memory based on depth

1. A system comprising:

a set of processing units, each processing unit in the set of processing units comprising:

a memory manager; and

memory;

wherein the memory manager is configured to:

for each of a plurality of sets of elements in a three-dimensional (3D) matrix:

determine a position along a height axis and width axis of the 3D matrix;

at the determined position, identify a set of elements along a depth axis of the 3D matrix; and

store the set of elements in a contiguous block of memory in the memory.

2. The system of claim 1 , wherein the set of elements identified along the depth axis of the 3D matrix at a particular determined position is a first set of elements, wherein the first set of elements stored in the contiguous block of memory is a first contiguous block of memory, wherein the memory manager is further configured to:

at the particular determined position, identify a second set of elements along the third axis of the 3D matrix; and

store the second set of elements in a second contiguous block of memory.

3. The system of claim 1 , wherein each contiguous block of memory forms a tile of memory.

4. The system of claim 1 , wherein the 3D matrix is a first 3D matrix, wherein the memory manager is further configured to:

determine a first position of a second 3D matrix along the first and second axes of the first 3D matrix;

perform a first matrix operation on the first 3D matrix and the second 3D matrix based on the first position of the second 3D matrix;

determine a second position of the second 3D matrix along the first and second axes of the first 3D matrix; and

perform a second matrix operation on the first 3D matrix and the second 3D matrix based on the second position of the second 3D matrix.

5. The system of claim 4 , wherein the memory manager is further configured to:

for each of a plurality of sets of elements in the second 3D matrix:

determine a position along a height axis and width axis of the second 3D matrix;

at the determined position, identify a set of elements along a depth axis of the second 3D matrix; and

store the set of elements in a contiguous block of memory.

6. The system of claim 5 , wherein performing the first matrix operation on the first 3D matrix and the second 3D matrix comprises:

reading a first set of elements in the first 3D matrix from a first contiguous block of memory;

reading a second set of elements in the first 3D matrix from a second contiguous block of memory;

reading a third set of elements in the second 3D matrix from a third contiguous block of memory; and

reading a fourth set of elements in the second 3D matrix from a fourth contiguous block of memory,

wherein performing the second matrix operation on the first 3D matrix and the second 3D matrix comprises:

using the first set of elements in the first 3D matrix from the first contiguous block of memory but not the second set of elements in the first 3D matrix from the second contiguous block of memory; and

using the third set of elements in the second 3D matrix from the third contiguous block of memory, but not the fourth set of elements in the second 3D matrix from the fourth contiguous block of memory.

7. The system of claim 6 , wherein the first and second matrix operations are part of a convolution operation between the first 3D matrix and the second 3D matrix.

8. The system of claim 7 , wherein the first matrix operation further comprises:

generating a transpose of the third set of elements; and

performing a matrix multiplication operation on the first set of elements and the transposed third set of elements.

9. The system of claim 1 , wherein the first tile of memory is a multiple of a defined size; wherein a total size of the first set of elements is a multiple of the defined size.

10. The system of claim 1 , wherein the first tile of memory is a multiple of a defined size; wherein the instructions further cause the at least one processing unit to pad a total size of the first set of elements with a defined value so that the padded first set of elements is a multiple of the defined size.

11. A method comprising:

for each of a plurality of sets of elements in a three-dimensional (3D) matrix:

determining a position along a height axis and width axis of the 3D matrix;

at the determined position, identifying a set of elements along a depth axis of the 3D matrix; and

storing the set of elements in a contiguous block of memory.

12. The method of claim 11 , wherein the set of elements identified along the depth axis of the 3D matrix at a particular determined position is a first set of elements, wherein the first set of elements stored in the contiguous block of memory is a first contiguous block of memory, the method further comprising:

at the particular determined position, identifying a second set of elements along the third axis of the 3D matrix; and

storing the second set of elements in a second contiguous block of memory.

13. The method of claim 11 , wherein each contiguous block of memory forms a tile of memory.

14. The method of claim 11 , wherein the 3D matrix is a first 3D matrix, the method further comprising:

determining a first position of a second 3D matrix along the first and second axes of the first 3D matrix;

performing a first matrix operation on the first 3D matrix and the second 3D matrix based on the first position of the second 3D matrix;

determining a second position of the second 3D matrix along the first and second axes of the first 3D matrix; and

performing a second matrix operation on the first 3D matrix and the second 3D matrix based on the second position of the second 3D matrix.

15. The method of claim 14 further comprising:

for each of a plurality of sets of elements in the second 3D matrix:

determine a position along a height axis and width axis of the second 3D matrix;

at the determined position, identify a set of elements along a depth axis of the second 3D matrix; and

store the set of elements in a contiguous block of memory.

16. The method of claim 15 , wherein performing the first matrix operation on the first 3D matrix and the second 3D matrix comprises:

reading a first set of elements in the first 3D matrix from a first contiguous block of memory;

reading a second set of elements in the first 3D matrix from a second contiguous block of memory;

reading a third set of elements in the second 3D matrix from a third contiguous block of memory; and

reading a fourth set of elements in the second 3D matrix from a fourth contiguous block of memory,

wherein performing the second matrix operation on the first 3D matrix and the second 3D matrix comprises:

using the first set of elements in the first 3D matrix from the first contiguous block of memory but not the second set of elements in the first 3D matrix from the second contiguous block of memory; and

using the third set of elements in the second 3D matrix from the third contiguous block of memory, but not the fourth set of elements in the second 3D matrix from the fourth contiguous block of memory.

17. The method of claim 16 , wherein the first and second matrix operations are part of a convolution operation between the first 3D matrix and the second 3D matrix.

18. The method of claim 17 , wherein the first matrix operation further comprises:

generating a transpose of the third set of elements; and

performing a matrix multiplication operation on the first set of elements and the transposed third set of elements.

19. The method of claim 11 , wherein the first tile of memory is a multiple of a defined size; wherein a total size of the first set of elements is a multiple of the defined size.

20. The method of claim 11 , wherein the first tile of memory is a multiple of a defined size; wherein the instructions further cause the at least one processing unit to pad a total size of the first set of elements with a defined value so that the padded first set of elements is a multiple of the defined size.