Non-rectangular matrix computations and data pattern processing using tensor cores
Matrix multiplication operations can be implemented, at least in part, on one or more tensor cores of a parallel processing unit. An efficiency of the matrix multiplication operations can be improved in cases where one of the input operands or the output operand of the matrix multiplication operation is a square matrix having a triangular data pattern. In such cases, the number of computations performed by the tensor cores of the parallel processing unit can be reduced by dropping computations and/or masking out elements of the square matrix input operand on one side of the main diagonal of the square matrix. In other cases where the output operand exhibits the triangular data pattern, computations can be dropped or masked out for the invalid side of the main diagonal of the square matrix. In an embodiment, a library implementing the matrix multiplication operations is provided.
1 . A system, comprising:
a host processor;
a parallel processing unit including at least one tensor core; and
a memory storing instructions that, when executed by the host processor, cause the host processor to perform a matrix multiplication operation by:
partitioning an output operand of the matrix multiplication operation into a number of tiles;
responsive to determining that the output operand of the matrix multiplication operation is a square matrix having a triangular data pattern:
for at least one tile of the output operand:
determining whether the at least one tile is valid based on a location of the tile relative to a main diagonal of the square matrix, and
responsive to determining that the at least one tile is valid, launching an operation on the parallel processing unit to generate values for the elements of the at least one tile using the at least one tensor core, or
responsive to determining that the at least one tile is invalid, preventing the elements of the at least one tile from being processed by the parallel processing unit; or
responsive to determining that an input operand of the matrix multiplication operation is a square matrix having a triangular data pattern:
for at least one tile of the output operand:
partitioning a portion of the input operand into a number of sub-tiles associated with the at least one tile of the output operand,
for at least one sub-tile of the input operand, determining whether the at least one sub-tile is valid based on a location of the at least one sub-tile relative to a main diagonal of the square matrix, and
responsive to determining that the at least one sub-tile is valid, launching an operation on the parallel processing unit to generate values for the elements of the at least one tile using the at least one tensor core to process the elements of the at least one sub-tile, or
responsive to determining that the at least one sub-tile is invalid, preventing the elements of the at least one sub-tile from being processed by the parallel processing unit.
2 . The system of claim 1 , wherein the parallel processing unit includes a plurality of tensor cores, and the operation on the parallel processing unit is configured to:
generate elements for a first tile of the output operand via a first tensor core, and
generate elements for a second tile of the output operand via a second tensor core,
wherein the first tensor core and the second tensor core operate in parallel.
3 . The system of claim 1 , wherein, responsive to determining that the output operand of the matrix multiplication operation is the square matrix having the triangular data pattern, for at least one tile of the output operand, the host processor is further configured to:
determine whether the at least one tile intersects the main diagonal of the square matrix; and
responsive to determining that the at least one tile intersects the main diagonal of the square matrix, mask out at least a portion of the elements of the at least one tile on one side of the main diagonal.
4 . The system of claim 1 , wherein, responsive to determining that the input operand of the matrix multiplication operation is the square matrix having the triangular data pattern, for at least one tile of the output operand, the host processor is further configured to:
determine whether at least one sub-tile of the input operand intersects the main diagonal of the square matrix; and
responsive to determining that the at least one sub-tile intersects the main diagonal of the square matrix, mask out at least a portion of the elements of the at least one sub-tile on one side of the main diagonal.
5 . The system of claim 1 , wherein the square matrix is a triangular matrix.
6 . The system of claim 1 , wherein the square matrix is a symmetric matrix.
7 . The system of claim 6 , wherein the matrix multiplication operation comprises multiplying a rectangular matrix of dimension N×K with a transpose of the rectangular matrix.
8 . The system of claim 1 , wherein the square matrix is a Hermitian matrix.
9 . The system of claim 8 , wherein the matrix multiplication operation comprises multiplying a rectangular matrix of dimension N×K with a conjugate transpose of the rectangular matrix.
10 . The system of claim 1 , wherein the memory stores a library that implements an application programming interface (API) that, responsive to being called by an application, implements at least a portion of the matrix multiplication operation using the parallel processing unit.
11 . The system of claim 8 , wherein the memory further stores a driver configured to launch a kernel on the parallel processing unit.
12 . The system of claim 1 , wherein elements of the output operand of the matrix multiplication operation are processed by a single kernel configured to perform the matrix multiplication operation using the at least one tensor core, and wherein none of the elements are generated by a floating-point arithmetic logic unit of the parallel processing unit.
13 . A computer-implemented method for performing a matrix multiplication operation, the method comprising:
partitioning, by a host processor, an output operand of the matrix multiplication operation into a number of tiles;
responsive to determining that the output operand of the matrix multiplication operation is a square matrix having a triangular data pattern:
for at least one tile of the output operand:
determining whether the at least one tile is valid based on a location of the tile relative to a main diagonal of the square matrix, and
responsive to determining that the at least one tile is valid, launching an operation on a parallel processing unit to generate values for the elements of the at least one tile using a at least one tensor core, or
responsive to determining that the at least one tile is invalid, preventing the elements of the at least one tile from being processed by the parallel processing unit; or
responsive to determining that an input operand of the matrix multiplication operation is a square matrix having a triangular data pattern:
for at least one tile of the output operand:
partitioning a portion of the input operand into a number of sub-tiles associated with the at least one tile of the output operand,
for at least one sub-tile of the input operand, determining whether the at least one sub-tile is valid based on a location of the at least one sub-tile relative to a main diagonal of the square matrix, and
responsive to determining that the at least one sub-tile is valid, launching an operation on the parallel processing unit to generate values for the elements of the at least one tile using the at least one tensor core to process the elements of the at least one sub-tile, or
responsive to determining that the at least one sub-tile is invalid,
preventing the elements of the at least one sub-tile from being processed by the parallel processing unit.
14 . The method of claim 13 , wherein elements of a first tile of the number of tiles are generated by a first tensor core in parallel with elements of a second tile of the number of tiles generated by a second tensor core.
15 . The method of claim 13 , wherein, responsive to determining that the output operand of the matrix multiplication operation is the square matrix having the triangular data pattern, for at least one tile of the output operand, the method further comprises:
determining whether the at least one tile intersects the main diagonal of the square matrix; and
responsive to determining that the at least one tile intersects the main diagonal of the square matrix, masking out at least a portion of the elements of the at least one tile on one side of the main diagonal.
16 . The method of claim 13 , wherein, responsive to determining that the input operand of the matrix multiplication operation is the square matrix having the triangular data pattern, for at least one tile of the output operand, the method further comprises:
determining whether at least one sub-tile of the input operand intersects the main diagonal of the square matrix; and
responsive to determining that the at least one sub-tile intersects the main diagonal of the square matrix, masking out at least a portion of the elements of the at least one sub-tile on one side of the main diagonal.
17 . The method of claim 13 , wherein the square matrix is a triangular matrix, a symmetric matrix, or a Hermitian matrix.
18 . The method of claim 13 , wherein an application executed by the host processor is configured to call an application programming interface of a library configured to implement the matrix multiplication operation, and wherein a driver is configured to launch a kernel on the parallel processing unit in response to a call from one of the library or a runtime API.
19 . The method of claim 13 , wherein elements of the output operand of the matrix multiplication operation are processed by a single kernel configured to perform the multiplication operation using the at least one tensor core, and wherein at least one of the elements is not generated by a floating-point arithmetic logic unit of the parallel processing unit.
20 . A non-transitory computer-readable media storing computer instructions for performing a matrix multiplication operation that, when executed by a host processor, cause the host processor to perform the steps of:
partitioning, by a host processor, an output operand of the matrix multiplication operation into a number of tiles;
responsive to determining that the output operand of the matrix multiplication operation is a square matrix having a triangular data pattern:
for at least one tile of the output operand:
determining whether the at least one tile is valid based on a location of the at least one tile relative to a main diagonal of the square matrix, and
responsive to determining that the at least one tile is valid, launching an operation on the parallel processing unit to generate values for the elements of the at least one tile using the at least one tensor core, or
responsive to determining that the at least one tile is invalid, preventing the elements of the at least one tile from being processed by the parallel processing unit; or
responsive to determining that an input operand of the matrix multiplication operation is a square matrix having a triangular data pattern:
for at least one tile of the output operand:
partitioning a portion of the input operand into a number of sub-tiles associated with the at least one tile of the output operand,
for at least one sub-tile of the input operand, determining whether the at least one sub-tile is valid based on a location of the at least one sub-tile relative to a main diagonal of the square matrix, and
responsive to determining that the at least one sub-tile is valid, launching an operation on a parallel processing unit to generate values for the elements of the at least one tile using at least one tensor core to process the elements of the at least one sub-tile, or
responsive to determining that the at least one sub-tile is invalid, preventing the elements of the at least one sub-tile from being processed by the parallel processing unit.