Alternative loop limits for accessing data in multi-dimensional tensors
Methods, systems, and apparatus for accessing a N-dimensional tensor are described. In some implementations, a method includes, for each of one or more first iterations of a first nested loop, performing iterations of a second nested loop that is nested within the first nested loop until a first loop bound for the second nested loop is reached. A number of iterations of the second nested loop for the one or more first iterations of the first nested loop is limited by the first loop bound in response to the second nested loop having a total number of iterations that exceeds a value of a hardware property of the computing system. After a penultimate iteration of the first nested loop has completed, one or more iterations of the second nested loop are performed for a final iteration of the first nested loop until an alternative loop bound is reached.
1. A method performed by a computing system for accessing an N-dimensional tensor, comprising:
determining that a number of tensor elements of a particular dimension of the N-dimensional tensor is not an exact multiple of a number of individual computing units of the computing system that will be used to determine memory addresses for the tensor elements and in response:
generating a first loop for controlling a number of iterations of a second loop used to traverse the particular dimension;
determining a first loop bound for the second loop and an alternative loop bound for the second loop based on the number of tensor elements of the particular dimension and the number of individual computing units of the computing system, wherein the first loop bound controls a number of iterations of the second loop for one or more first iterations of the first loop and the alternative loop bound controls the number of iterations of the second loop for a final iteration of the first loop such that the number of iterations of the second loop does not exceed a number of tensor elements for which a memory address will be determined during the final iteration of the first loop; and
generating code that has the second loop nested within the first loop;
executing the code to access the N-dimensional tensor, including:
for each of one or more first iterations of the first loop, performing iterations of the second loop that is nested within the first loop until the first loop bound for the second loop is reached, wherein:
for each iteration of the second loop, a computing unit of a plurality of computing units determines a memory address for a tensor element of the N-dimensional tensor, each tensor element including a data value;
after a penultimate iteration of the first loop has completed, performing one or more iterations of the second loop for a final iteration of the first loop until the alternative loop bound is reached, wherein the alternative loop bound is less than the first loop bound; and
outputting data indicating each determined memory address to a processing unit that accesses tensor elements of the N-dimensional tensor using each determined memory address.
2. The method of claim 1 , further comprising substituting the alternative loop bound for the first loop bound for the final iteration of the first loop in response to determining that the penultimate iteration of the first loop has completed.
3. The method of claim 1 , wherein each individual computing unit comprises a hardware compute tile, a processor, or a hardware math unit.
4. The method of claim 1 , wherein:
performing iterations of the second loop that is nested within the first loop until the first loop bound for the second loop is reached comprises performing each iteration of the second loop in parallel using the computing units; and
each computing unit performs a respective iteration of the second nested loop.
5. The method of claim 1 , wherein the alternative loop bound is based on a remainder value resulting from dividing the number of tensor elements of the particular dimension by the number of individual computing units.
6. The method of claim 1 , wherein a set of loops including the first loop and the second loop includes one or more loops nested between the first loop and the second loop and the second loop is nested within another loop.
7. The method of claim 1 , wherein the second loop is nested directly within the first loop without any other loops nested between the first loop and the second loop.
8. A system for accessing an N-dimensional tensor, the system comprising:
a plurality of individual computing units;
one or more processors configured to:
determine that a number of tensor elements of a particular dimension of the N-dimensional tensor is not an exact multiple of a number of individual computing units of the computing system that will be used to determine memory addresses for the tensor elements and in response:
generate a first loop for controlling a number of iterations of a second loop used to traverse the particular dimension;
determine a first loop bound for the second loop and an alternative loop bound for the second loop based on the number of tensor elements of the particular dimension and the number of individual computing units of the computing system, wherein the first loop bound controls a number of iterations of the second loop for one or more first iterations of the first loop and the alternative loop bound controls the number of iterations of the second loop for a final iteration of the first loop such that the number of iterations of the second loop does not exceed a number of tensor elements for which a memory address will be determined during the final iteration of the first loop;
generate code that has the second loop nested within the first loop; and
execute the code to access the N-dimensional tensor, including:
for each of one or more first iterations of the first loop, performing iterations of the second loop that is nested within the first loop until the first loop bound for the second loop is reached, wherein:
for each iteration of the second loop, a computing unit of the plurality of individual computing units determines a memory address for a tensor element of the N-dimensional tensor, each tensor element including a data value;
after a penultimate iteration of the first loop has completed, performing one or more iterations of the second loop for a final iteration of the first loop until the alternative loop bound is reached, wherein the alternative loop bound is less than the first loop bound; and
outputting data indicating each determined memory address to a processing unit that accesses tensor elements of the N-dimensional tensor using each determined memory address.
9. The system of claim 8 , wherein the operations further comprise substituting the alternative loop bound for the first loop bound for the final iteration of the first loop in response to determining that the penultimate iteration of the first loop has completed.
10. The system of claim 8 , wherein each individual computing unit comprises a hardware compute tile, a processor, or a hardware math unit.
11. The system of claim 8 , wherein:
performing iterations of the second loop that is nested within the first loop until the first loop bound for the second loop is reached comprises performing each iteration of the second loop in parallel using the individual computing units; and
each computing unit performs a respective iteration of the second loop.
12. The system of claim 8 , wherein the alternative loop bound is based on a remainder value resulting from dividing the number of tensor elements of the particular dimension by the number of individual computing units.
13. The system of claim 8 , wherein a set of nested loops including the first loop and the second loop includes one or more loops nested between the first loop and the second loop and the second loop is nested within another loop.
14. The system of claim 8 , wherein the second loop is nested directly within the first loop without any other loops nested between the first loop and the second loop.
15. An apparatus for accessing an N-dimensional tensor, the apparatus comprising:
a plurality of individual computing units that each compute memory addresses for tensor elements;
one or more processors configured to:
determine that a number of tensor elements of a particular dimension of the N-dimensional tensor is not an exact multiple of a number of individual computing units of the computing system that will be used to determine memory addresses for the tensor elements and in response:
generate a first loop for controlling a number of iterations of a second loop used to traverse the particular dimension;
determine a first loop bound for the second loop and an alternative loop bound for the second loop based on the number of tensor elements of the particular dimension and the number of individual computing units of the computing system, wherein the first loop bound controls a number of iterations of the second loop for one or more first iterations of the first loop and the alternative loop bound controls the number of iterations of the second loop for a final iteration of the first loop such that the number of iterations of the second loop does not exceed a number of tensor elements for which a memory address will be determined during the final iteration of the first loop;
generate code that has the second loop nested within the first loop; and
a controller configured to:
execute the code to access the N-dimensional tensor, including:
assign iterations of nested loops to the individual computing units by performing operations comprising:
for each of one or more first iterations of a first loop, performing iterations of a second loop that is nested within the first loop until a first loop bound for the second loop is reached, wherein:
for each iteration of the second loop, a computing unit of the plurality of individual computing units determines a memory address for a tensor element of the N-dimensional tensor; each tensor element including a data value and
a number of iterations of the second loop for the one or more first iterations of the first loop is limited by the first loop bound in response to the second loop having a total number of iterations that exceeds a number of the individual computing units;
after a penultimate iteration of the first loop has completed, performing one or more iterations of the second loop for a final iteration of the first loop until an alternative loop bound is reached, wherein the alternative loop bound is less than the first loop bound; and
output data indicating each determined memory address to a processing unit that accesses tensor elements of the N-dimensional tensor using each determined memory address.
16. The apparatus of claim 15 , wherein the controller substitutes the alternative loop bound for the first loop bound for the final iteration of the first loop in response to determining that the penultimate iteration of the first loop has completed.
17. The apparatus of claim 15 , wherein each individual computing unit comprises a hardware compute tile, a processor, or a hardware math unit.
18. The apparatus of claim 15 , wherein:
performing iterations of second loop that is nested within the first loop based until the first loop bound for the second loop is reached comprises performing each iteration of the second loop in parallel using the individual computing units; and
each computing unit performs a respective iteration of the second loop.
19. The apparatus of claim 15 , wherein the alternative loop bound is based on a remainder value resulting from dividing the number of tensor elements of the particular dimension by the number of individual computing units.