Reshape and broadcast optimizations to avoid unnecessary data movement

Abstract

Methods, systems, and apparatus, including a computer program encoded on a computer storage medium, for transforming patterns of operations on tensors in a computation graph to reduce the storage burden incurred when performing reshape operations, particularly when deployed to hardware platforms with vector instructions or vector storage that require aligned operators.

Description

Background Art

[0001] The TensorFlow library can be used to create machine learning models, such as recurrent neural network ("RNN") models, convolutional neural network ("CNN") models, feedforward neural network models, and random forest models. (TensorFlow is described in Abadi et al., TensorFlow: A System for Large-Scale Design and Implementation (OSDI'16), pp. 265-283, November 2-4, 2016. The software is available at https: / / tensorflow.org)

[0002] The TensorFlow library can be used to represent machine learning models as TensorFlow graphs. Each node in a TensorFlow graph represents an operation. Each edge in a TensorFlow graph is directed and represents the flow of data into or out of the node connected by the edge. Data is in the form of zero-dimensional or multi-dimensional tensors, where each element has the same data type, for example, a 32-bit integer, a double-length floating point number, or a string. Tensors are represented by vectors with pairs of brackets "[]" on the outside. For example, a one-dimensional (1D) tensor (also called a vector) of 3 elements would be represented as [1,2,3]. Zero-dimensional tensors are scalars. Two-dimensional (2D) tensors would be represented as [[1,2,3],[4,5,6]]. The rank of the tensor, that is, the number of dimensions or indices required to uniquely select each element of the tensor, is 2. The shape of the tensor is [2,3]. The number of elements in the zeroth dimension is 2, that is, the two vectors (1D tensors) [1,2,3] and [4,5,6]; and the number of elements in the first dimension is 3; that is, each of the vectors [1,2,3] and [4,5,6] has three elements. The shape of the tensor itself is a 1D tensor. As is customary in many programming environments, the numbering of dimensions starts at zero.

[0003] In this description, examples will be expressed using the Python API for building and executing TensorFlow graphs. TensorFlow modules can be loaded like this:

[0004] import tensorflow as tf

[0005] TensorFlow operations include shape, reshape, broadcast, and reduce operations. These will be described below, omitting parameters and aspects that are not important to this specification from the description.

[0006] When executed, a shape operation returns the shape (i.e., dimensions) of the input tensor as a 1D tensor. In the following example:

[0007] X=tf.constant([[[1,1,1],[2,2,2]],[[3,3,3],[4,4,4]]])

[0008] tf.shape(X)

[0009] The shape operation returns a tensor of [2,2,3] which represents the dimensions of the tensor X.

[0010] When executed, the reshape operation returns a tensor with the same element values ​​in the same order as the input tensor but with a shape defined by the shape tensor input. In the following example,

[0011] X=tf.constant([[[1,1],[2,2]],[[3,3],[4,4]]])

[0012] tf.reshape(X,[2,4])

[0013] The reshape operation receives as input a tensor X and a one-dimensional tensor [2,4] representing the desired shape. The reshape operation returns a tensor [[1,1,2,2],[3,3,4,4]] with the same elements as the input tensor X and with the desired shape, i.e., [2,4]. The desired shape input to the reshape operation can have more or fewer dimensions than the input tensor has.

[0014] Broadcasting operations include broadcast_to. Broadcasting is the process of making arrays with compatible shapes for arithmetic operations. Two shapes are compatible if, for their respective pairs of dimensions, the dimensions are equal or one of them is one. When a tensor is broadcasted to a shape, the operation starts with the trailing dimension and proceeds forward.

[0015] Therefore, when executed, the tensor returned by the broadcast_to operation is the input tensor replicated as many times as necessary to reach the specified shape requested. In the following example:

[0016] V = tf.constant([7,8])

[0017] tf.broadcast_to(V,[2,3])

[0018] The broadcast_to operation receives as input the tensor V and a tensor of the desired shape [2,3]. It returns a tensor of the desired shape [[7,7,7],[8,8,8]].

[0019] Reduction operations include reduce_all, reduce_any, reduce_sum, and reduce_mean. The output tensor returned by a reduction operation usually has a smaller rank and fewer elements than the input tensor.

[0020] The reduce operation receives an input tensor and an axis tensor. The elements of the axis tensor identify the dimensions of the shape of the input tensor. The reduce operation reduces the input tensor along the dimensions specified by the axis tensor. For example,

[0021] X=tf.constant([[1,1,1],[1,1,1]])

[0022] In the above formula, the shape of X is [2,3], that is, X is a tensor with two rows and three columns: 1 1 1 1 1 1

[0025] Taking the specific reduction operation reduce_sum as an example, the axis tensor [0] is used to identify the row, that is, the zeroth dimension of X.

[0026] tf.reduce_sum(x,[0])

[0027] When this operation is performed, the tensor is reduced along the zeroth dimension (rows) and the rows [1,1,1]+[1,1,1] are summed to return [2,2,2].

[0028] tf.reduce_sum(x,[1])

[0029] When executed, it reduces along the first dimension (columns) and sums the columns [1,1]+[1,,1]+[1,1] to return [3,3].

[0030] tf.reduce_sum(x,[0,1])

[0031] When executed, the reduction is performed along both dimensions and summed to return a scalar (0D tensor) 6.

[0032] The shape of the tensor returned by the reduction operation has the dimensions of the input tensors without the indices specified by the axis tensor.

[0033] Other reduction operations return tensors whose element values ​​are computed by other operations. For example, the reduce_all operation computes logical AND, the reduce_any operation computes logical OR, the reduce_mean operation computes the mean, and so on.

[0034] In some scenarios, users use a compiler, such as a just-in-time (JIT) compiler, to compile the TensorFlow graph into a graph for input to the XLA compiler. (The JIT compiler is in The input language of XLA is called "HLO IR", or just HLO (High Level Optimizer). The XLA compiler takes graphs (i.e. computations) defined in HLO and compiles them into machine instructions for various architectures, performing target-dependent optimizations and generating target-dependent code.

[0035] The nodes in the HLO graph represent operations. Each edge in the graph is directed and represents the flow of data into or out of the node connected by the edge. This data is in the form of tensors. The operations represented in the HLO graph correspond to the operations in the TensorFlow process that generated the HLO graph. Specifically, the HLO graph can include reshape, reduce, and broadcast operations.

[0036] The binaries generated by the XLA compiler are deployed to the hardware and executed by the specific processor of the hardware. Some processors implement instructions that operate on vectors. In order for a processor to execute vector instructions that operate on tensor data, the tensors must be stored so that each of the tensor vectors to be operated on by the vector instructions is aligned on a vector boundary, as specified for the processor.

[0037] For example, if a reshape operation receives a tensor [1,2,3,4,5,6,7,8,9] and a tensor of the specified shape [3,3] as input parameters, then the resulting tensor [[1,2,3],[4,5,6],[7,8,9]] has three vectors [1,2,3], [4,5,6], and [7,8,9] that must be moved so that they are aligned on vector boundaries if the vectors are not aligned exactly as required by the hardware's specific processor. This vector boundary alignment operation can be computationally expensive. Summary of the invention

[0038] This specification describes optimization techniques that can be implemented in the XLA compiler to reduce the storage burden of certain sequences of operations that include reshape operations.

[0039] These optimizations are particularly helpful in implementing sequences of operations for machine learning techniques, such as group normalization (https: / / arxiv.org / pdf / 1803.08494.pdf) and ghost batch normalization (https: / / arxiv.org / pdf / 1705.08741.pdf). In a direct implementation of such machine learning techniques, the input tensor is reshaped to a higher dimension. Then, a reduction is performed between some dimensions changed by the reshaping and other dimensions that are not changed by the reshaping. The derivative of the reduction and reshaping is broadcasting, and it is reshaped back to the original shape. Broadcasting is also made into some dimensions changed by the reshaping and some dimensions that are not changed by the reshaping. The reshaping operation has no effect on data on platforms with linear address spaces. However, on platforms with vector storage, reshaping will generally change the shape alignment relative to the vector storage, requiring data to be moved. This specification describes optimizations that reduce the size of the tensor that needs to be moved to achieve alignment relative to the vector storage. Optimizations may be used to reduce the number of reshape operations and / or the number of vectors used to represent intermediate tensors in a sequence of operations, thereby reducing the number of vector boundary alignment operations that need to be performed by the processor. BRIEF DESCRIPTION OF THE DRAWINGS

[0040] Figure 1 is a flow chart illustrating an example process for implementing compiler transformations to detect an operating mode including a reshape operation that may be transformed to minimize the size of the reshape operation.

[0041] Figure 2 is a flow chart illustrating an example process of a transformation.

[0042] The same reference numbers and names in different drawings indicate the same elements. DETAILED DESCRIPTION

[0043] Figure 1 1 is a flow chart illustrating an example process 100 for implementing a compiler transformation that detects a pattern of operations including a reshape operation that can be transformed to minimize the size of the reshape operation. The compiler transformation will be described with reference to the XLA compiler and patterns of reshape, reduce, and broadcast operations in a computation graph. The compiler and process can be implemented on and executed by a system of one or more computers in one or more locations.

[0044] The process detects an operation mode including a reshape operation that can be transformed (102). The mode can be an operation mode in an XLA graph. One such mode is reduce(reshape(X)), when one or more dimensions are reduced and not changed in the reshape. In this mode, the reshape operation is performed on the input tensor X and returns a tensor as input to the reduce operation.

[0045] The process transforms the operation into an operation pattern with a smaller reshape size (104). For example, an operation that fits the reduce(reshape(X)) pattern described above is transformed into an operation of the reduce(reshape(reduce(X))) pattern. This transformation increases the amount of computation by adding additional reductions; but it has the important advantage of reducing the overall size of the reshape, because both computations have a reshape, and the latter reshape has the dimensions that have been reduced and, as a result, is strictly smaller in number of elements than the original. The process continues to detect and transform until no additional patterns are detected, at which point the compiler generates code that is specific to the target hardware and implements the computation including the transformed pattern (106). The compiler, or other element of the TensorFlow infrastructure, then deploys the generated code to the target hardware for execution (108).

[0046] An example broadcasted version of another pattern that can be optimized is reshape(add(reshape(X),broadcast(Y)) , where the input shape of tensor X is the same as the output shape of the tensor returned by the outermost reshape. This operation pattern is transformed into add(X,broadcast(reshape(broadcast(Y)))) , which has fewer reshapes overall, and has smaller size reshapes because the output of the reshape is broadcasted to a larger shape due to the extra dimension. In this way, the number of vector boundary alignment operations required to implement the overall operation is reduced.

[0047] This transformation can be applied to any subcomputation matching the following patterns:

[0048]

[0049] It is transformed into

[0050]

[0051]

[0052] Lowercase letters f, g, h, a, and b represent mathematical operations in the graph. The transformed pattern can be found using a depth-first search of the graph. In some implementations, for simplicity, the search is in post order, i.e., a topological sort with producers before consumers, and the graph is transformed in place. When the search finds a matching subtree, it copies the subtree and replaces the user of the original subtree root with the new subtree root. Other compilers fix duplication and invalid code. Other search methods for searching for pattern graphs in the computational graph can also be used.

[0053] Using these two schema transformations — from the forms reduce(reshape(X)) and reshape(f(reshape(X), broadcast(Y)) — group normalization and virtual batch normalization and their derivatives can be accomplished with smaller reshapes and resulting smaller storage requirements.

[0054] For example, the implementation of group specification is naturally expressed as

[0055] reduce(reshape(image,[B,H,W,C / G,G]),[1,2,3]).

[0056] The image input above has shape [B,H,W,C] with a batch size for the batch of images in the input and dimensions of height, width, and channels of the images. The number of groups is G groups. If this expression is executed in this form, it creates large intermediate tensors and slow reshapes on some hardware platforms. The above transformation improves the computation by transforming it into the following form:

[0057] reduce(reshape(reduce(image,[1,2]),[B,C / G,G]),[1])

[0058] For the purpose of describing the transformation process, it is expressed as:

[0059] Y = reshape(X,[B,H,W,C / G,G])

[0060] Z = reduce(Y,[1,2,3])

[0061] The reduce operation reduces a tensor Y of shape [B,H,W,C / G,G] along the dimensions specified by the axis tensor [1,2,3] and returns a tensor Z. The axis tensor [1,2,3] represents the dimensions [H,W,C / G] of tensor Y along which the reduce operation is performed. The reduce operation returns a tensor Z of shape [B,G].

[0062] The following describes how the parameters of the original mode are mapped to the appropriate inputs of the transformed mode.

[0063] For the purpose of discussion, the reduce(reshape(reduce(X,[1,2]),[B,C / G,G]),[1]) transformation will be represented as:

[0064] W = reduce(X,[1,2])

[0065] Y2=reshape(W,[B,C / G,G])

[0066] Z2=reduce(Y2,[1])

[0067] Figure 2 is a flow chart illustrating an example process 200 for transformation. This will be described with reference to the example mode just described. Process 200 is a flow chart of the example process 200 described above with reference to Figure 1 An example implementation of the transformation (104) is described.

[0068] The process determines the final output dimensions of the tensor returned by the raw mode of the operation (202). The final output dimensions are determined by comparing the shape of X to the final output shape. In the example, the raw mode example receives a tensor X of shape [B, H, W, C] and returns a tensor of shape [B, G].

[0069] The process reduces (204) along the dimensions of the reshaped input tensor that are neither in the final output nor affect the original pattern of the reshaped operation. In the example, according to the original pattern

[0070] Y = reshape(X,[B,H,W,C / G,G])

[0071] Z = reduce(Y,[1,2,3])

[0072] The compiler determines from the axis tensors input to the reshape operation that the zeroth and third indices of tensor X (i.e. B and C, respectively) are in and affect the final output. Therefore, reduce X along its first and second dimensions (i.e. H and W dimensions):

[0073] W = reduce(X,[1,2])

[0074] This process reshapes the output tensor of the reduce operation (206). The output tensor of the reduce operation (W in the example) is reshaped to the shape of the original mode, but without the dimensions that are not in the final output or are transformed. In the original mode, the third dimension (i.e., C) is divided by G, and the fourth dimension (i.e., G) is added. The zeroth dimension (i.e., B) is in the final output. Therefore, the reshape operation in the transformation reshapes the reduced tensor into [B, C / G, G]:

[0075] Y2=reshape(W,[B,C / G,G])

[0076] The process reduces the output tensor of the reshape operation along any dimension that was not in the output tensor of the original mode (208). In the example, the original mode outputs a tensor of shape [B, G]. Therefore, the first index of the output of the reshape operation in the transform is reduced, and the reduce operation returns a tensor of shape [B, G]:

[0077] Z2=reduce(Y2,[1])

[0078] The same rules apply to transforming an original pattern of the form reshape(operator(reshape(X),broadcast(Y))) into the form operator(X,broadcast(reshape(broadcast(Y))))).

[0079] The subject matter described in this specification and the embodiments of the actions and operations can be implemented in digital electronic circuits, in tangibly embodied computer software or firmware, in computer hardware (including the structures disclosed in this specification and their structural equivalents), or in a combination of one or more of them. The embodiments of the subject matter described in this specification can be implemented as one or more computer programs, for example, one or more modules of computer program instructions encoded on a computer program carrier, for execution by a data processing device or for controlling the operation of the data processing device. The carrier can be a tangible non-transitory computer storage medium. Alternatively or additionally, the carrier can be an artificially generated propagation signal (for example, a machine-generated electrical, optical or electromagnetic signal) that is generated to encode information for transmission to a suitable receiver device for execution by a data processing device. The computer storage medium can be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them, or a part of them. The computer storage medium is not a propagation signal.

[0080] The term "data processing apparatus" includes various devices, equipment, and machines for processing data, including, for example, a programmable processor, a computer, or multiple processors or computers. The data processing apparatus may include special-purpose logic circuits, such as an FPGA (field programmable gate array), an ASIC (application-specific integrated circuit), or a GPU (graphics processing unit). In addition to hardware, the apparatus may also include code that creates an execution environment for a computer program, such as code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0081] A computer program (which may also be referred to or described as a program, software, software application, applet, module, software module, engine, script or code) may be written in any form of programming language, including compiled or interpreted languages ​​or declarative or procedural languages; and it may be deployed in any form, including as a stand-alone program or as a module, component, engine, subroutine or other unit suitable for execution in a computing environment, which may include one or more computers interconnected at one or more locations via a data communications network.

[0082] A computer program may, but need not, correspond to a file in a file system. A computer program may be stored as part of a file that holds other programs or data (e.g., one or more scripts stored in a markup language document), in a single file dedicated to the program in question, or in multiple coordinated files (e.g., files that store one or more modules, subroutines, or portions of code).

[0083] The processes and logic flows described in this specification can be performed by one or more computers executing one or more computer programs to perform operations by operating on input data and generating output. The processes and logic flows can also be performed by a special purpose logic circuit (e.g., an FPGA, ASIC, or GPU), or by a combination of a special purpose logic circuit and one or more programmed computers.

[0084] Computers suitable for executing computer programs may be based on general or special purpose microprocessors, or both, or any other type of central processing unit. In general, the central processing unit will receive instructions and data from a read-only memory or a random access memory, or both. The essential elements of a computer are a central processing unit for executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory may be supplemented by, or incorporated in, special purpose logic circuits.

[0085] Typically, a computer will also include one or more mass storage devices, or be operably coupled thereto, to receive data from or send data to. The mass storage device may be, for example, a magnetic disk, a magneto-optical disk or an optical disk, or a solid-state drive. However, a computer need not have such a device. In addition, a computer may be embedded in another device, such as a mobile phone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a global positioning system (GPS) receiver, or a portable storage device (e.g., a universal serial bus (USB) flash drive), to name a few.

[0086] To provide interaction with a user, embodiments of the subject matter described in this specification may be implemented on a computer, or configured to communicate with a computer, having a display device (e.g., an LCD (liquid crystal display) monitor) for displaying information to a user and an input device, such as a keyboard and a pointing device (e.g., a mouse, trackball, or touchpad) through which a user can provide input to the computer. Other types of devices may also be used to provide interaction with a user; for example, the feedback provided to the user may be any form of sensory feedback, such as visual feedback, auditory feedback, or tactile feedback; and the input from the user may be received in any form, including sound, voice, or tactile input. In addition, a computer may interact with a user by sending documents to and receiving documents from a device used by a user; for example, by sending a web page to a web browser on a user's device in response to a request received from the web browser, or by interacting with an application running on a user's device (e.g., a smartphone or electronic tablet). In addition, a computer may interact with a user by sending a text message or other form of message to a personal device (e.g., a smartphone running a messaging application) and receiving a response message from the user in response.

[0087] This specification uses the term "configured to" in connection with systems, devices, and computer program components. A system of one or more computers configured to perform a particular operation or action means that the system has software, firmware, hardware, or a combination thereof installed thereon that, in operation, causes the system to perform the operation or action. One or more computer programs configured to perform a particular operation or action means that the one or more programs include instructions that, when executed by a data processing device, cause the device to perform the operation or action. A dedicated logic circuit configured to perform a particular operation or action means that the circuit has electronic logic to perform the operation or action.

[0088] Although this specification contains many specific implementation details, these should not be interpreted as limitations on the scope of the claims defined by the claims themselves, but as descriptions of features for specific embodiments that may be specific to a particular invention. Certain features described in this specification in the context of separate embodiments may also be implemented in combination in a single embodiment. On the contrary, the various features described in the context of a single embodiment may also be implemented separately in multiple embodiments or in any suitable sub-combination. In addition, although features may be described above as working in certain combinations, and even initially claimed as such, in some cases, one or more features from the claimed combination may be deleted from the combination, and the claimed may be directed to a sub-combination or a variation of the sub-combination.

[0089] Similarly, although the operations depicted in the drawings and recited in the claims are in a particular order, this should not be understood as requiring that the operations be performed in the particular order shown or in a sequential order, or that all of the operations shown be performed, to obtain the desired results. In some cases, multitasking and parallel processing may be advantageous. In addition, the separation of various system modules and components in the above-described embodiments should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated into a single software product or packaged into multiple software products.

[0090] Specific embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve the desired results. As an example, the processes described in the accompanying drawings do not necessarily require the particular order shown or sequential order to achieve the desired results. In some cases, multitasking and parallel processing may be advantageous.

Claims

1. A method performed by one or more computers, the method comprising: detecting a primitive pattern of operations in a graph of computational operations on a tensor, wherein the primitive pattern of operations returns a final output tensor and takes an input tensor as input, wherein the primitive pattern includes a primitive reshape operation that (i) returns the original tensor and (ii) can be transformed to use less storage; Transforming the original mode of operation into a new mode of operation using one or more reshape operations that return a tensor that is smaller than the original tensor, comprising: determining the final output dimensions of the final output tensor returned by the original mode of operation, reducing along a dimension of the input tensor that is neither in the final output tensor nor affects the reshape of the original reshape operation to return a first intermediate result tensor, reshape the first intermediate result tensor to return the second intermediate result tensor, and Along the dimensions of the final output tensor that are not in the original mode from the operation to reduce the second intermediate result tensor by any dimension of ; and Executable code is generated that is specific to the target hardware platform and implements the computation represented by the new mode of operation.

2. The method according to claim 1, further comprising: The generated code is deployed to the target hardware platform for execution.

3. The method according to claim 1 or 2, wherein: The original reshape mode of operation requires data to be moved to meet the alignment requirements for vector instructions or vector storage on the target hardware platform.

4. One or more non-transitory computer-readable storage media encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform actions comprising: detecting a primitive pattern of operations in a graph of computational operations on a tensor, wherein the primitive pattern of operations returns a final output tensor and takes an input tensor as input, wherein the primitive pattern includes a primitive reshape operation that (i) returns the original tensor and (ii) can be transformed to use less storage; Transforming the original mode of operation into a new mode of operation using one or more reshape operations that return a tensor that is smaller than the original tensor, comprising: determining the final output dimensions of the final output tensor returned by the original mode of operation, reducing along a dimension of the input tensor that is neither in the final output tensor nor affects the reshape of the original reshape operation to return a first intermediate result tensor, reshape the first intermediate result tensor to return the second intermediate result tensor, and reducing the second intermediate result tensor along any dimension that was not among the dimensions of the final output tensor from the original mode of operation; and Executable code is generated that is specific to the target hardware platform and implements the computation represented by the new mode of operation.

5. The non-transitory computer-readable storage medium of claim 4, wherein: The actions also include: The generated code is deployed to the target hardware platform for execution.

6. The non-transitory computer-readable storage medium according to claim 4 or 5, wherein: The original reshape mode of operation requires data to be moved to meet the alignment requirements for vector instructions or vector storage on the target hardware platform.

7. A system for transforming a pattern of operations on tensors in a computation graph, comprising: One or more computers and one or more storage devices having instructions stored thereon, which, when executed by the one or more computers, are operable to cause the one or more computers to perform actions, the actions comprising: detecting a primitive pattern of operations in a graph of computational operations on a tensor, wherein the primitive pattern of operations returns a final output tensor and takes an input tensor as input, wherein the primitive pattern includes a primitive reshape operation that (i) returns the original tensor and (ii) can be transformed to use less storage; Transforming the original mode of operation into a new mode of operation using one or more reshape operations that return a tensor that is smaller than the original tensor, comprising: determining the final output dimensions of the final output tensor returned by the original mode of operation, reducing along a dimension of the input tensor that is neither in the final output tensor nor affects the reshape of the original reshape operation to return a first intermediate result tensor, reshape the first intermediate result tensor to return the second intermediate result tensor, and reducing the second intermediate result tensor along any dimension that was not among the dimensions of the final output tensor from the original mode of operation; and Executable code is generated that is specific to the target hardware platform and implements the computation represented by the new mode of operation.

8. The system according to claim 7, wherein: The actions also include: The generated code is deployed to the target hardware platform for execution.

9. A system according to claim 7 or 8, wherein: The original reshape mode of operation requires data to be moved to meet the alignment requirements for vector instructions or vector storage on the target hardware platform.

Images