Indexing Analysis

HLO indexing analysis is a dataflow analysis that describes how elements of one tensor relate to another via "indexing maps". For example, how indices of an HLO instruction output map to indices of HLO instruction operands.

Example

For a broadcast from tensor<20xf32> to tensor<10x20x30xf32>

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

the indexing map from the output to input is (i, j, k) -> (j) for i in [0, 10], j in [0, 20] and k in [0, 30].

Motivation

XLA uses several bespoke solutions to reason about coalescing, operand utilization, and tiling schemes (more details below). The goal of indexing analysis is providing a reusable component for such use cases. Indexing analysis is built on XLA's custom SymbolicExpr and SymbolicMap infrastructure and adds HLO semantics.

Coalescing

Reasoning about memory coalescing becomes feasible for non-trivial cases, when we know what elements/slices of the inputs are read to compute an element of the output.

Operand Utilization

Operand utilization in XLA indicates how much each input of the instruction is used assuming its output is fully used. Currently, utilization is also not computed for a generic case. Indexing analysis allows us to compute utilization precisely.

Tiling

A tile/slice is hyper-rectangular subset of a tensor parameterized by offsets, sizes and strides. Tile propagation is a way to compute tile parameters of the producer/consumer of the op using the tiling parameters of the op itself. There is already a library that does it for softmax and dot. Tile propagation can be made more generic and robust if it is expressed via indexing maps.

Indexing map

An indexing map is a combination of

  • a symbolically expressed function that maps every element of one tensor A to ranges of elements in tensor B;
  • constraints on valid function arguments, including function's domain.

Function arguments are split into 3 categories to better communicate their nature:

  • dimension variables of the tensor A or a GPU grid we are mapping from; values are known statically. Index elements are also called dimension variables.

  • range variables. They define a one-to-many mapping and specify a set of elements in B used to compute a single value of A; values are known statically. The contracting dimension of a matrix multiplication is an example of a range variable.

  • runtime variables that are only known at during execution. For example, indices argument of gather operation.

Result of the function is an index of the target B tensor.

In short, an indexing function from tensor A to tensor B for operation x is

map_ab(index in A, range variables, runtime variables) -> index in B.

To better separate the types of mapping arguments we write them as:

map_ab(index in A)[range variables]{runtime variables} -> (index in B)

For example, let's look at the indexing maps for the reduce operation f32[4, 8] out = reduce(f32[2, 4, 8, 16] in, 0), dimensions={0,3}:

  • to map elements of in to out our function can be expressed as (d0, d1, d2, d3) -> (d1, d2). The constraints of the variables d0 in [0, 1], d1 in [0, 3], d2 in [0, 7], d3 in [0, 15] are defined by the shape of in.

  • To map elements of out to in: out has only two dimensions, and reduction introduces two range variables that cover reducing dimensions. Thus the mapping function is (d0, d1)[s0, s1] -> (s0, d0, d1, s1), where (d0, d1) is index of out. s0, s1 are ranges defined by operation's semantics and span dimension 0 and 3 of the in tensor. The constraints are d0 in [0, 3], d1 in [0, 7], s0 in [0,1], s1 in [0, 15].

It's important to note that in most scenarios we are interested in mapping from the elements of the output. For computation

C = op1(A, B)
E = op2(C, D)

we can talk about "indexing of B" meaning "mapping of elements of E into the elements of B". This might be counter-intuitive compared to other types of data-flow analysis that work from input toward outputs.

Constraints on variables enable optimization opportunities and help with code generation. In the documentation and implementation constraints are also referred to as domain as they define all valid combinations or argument values of the mapping function. For many operation, constraints simply describe the dimensions of tensors but for some operations they might be more complicated; see examples below.

By having functions and argument constraints expressed symbolically and being able to combine functions and constraints we can compute a compact indexing mapping for an arbitrary large computation (fusion).

Expressiveness of symbolic function and constraints is a balance between implementation complexity and optimization gains we get from having a more precise representation. For some HLO operations we capture access patterns only approximately.

Implementation

To minimize recomputation, we need a framework for symbolic computation. This is implemented as SymbolicExpr and SymbolicMap.

A typical SymbolicMap looks like:

(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)

SymbolicMap has two types of parameters: dimensions and symbols. Dimensions correspond to the dimension variables d; symbols correspond to the range variables r and runtime variables rt. SymbolicMap does not contain any metadata about constraints of the parameters, so we have to provide them separately.

class IndexingMap {
   // Variable represents dimension, range or runtime variable.
  struct Variable {
    // struct Interval represents a closed interval [lower_bound, upper_bound].
    Interval bounds;
    // Name of the variable is used for nicer printing.
    std::string name = "";
  };

  SymbolicMap symbolic_map_;

  // A dimension variable represents a dimension of a tensor or a GPU grid.
  // Dimension variables correspond to the dimensions of the `symbolic_map_`.
  std::vector<Variable> dim_vars_;

  // A range variable represents a range of values, e.g. to compute a single
  // element of the reduction's result we need a range of values from the input
  // tensor. Range variables correspond to the front portion of the
  // symbols in `symbolic_map_`.
  std::vector<Variable> range_vars_;

  // A runtime variable represents a runtime symbol, e.g. a dynamic offset in of
  // a HLO dynamic-update-slice op. Runtime variables correspond to the back
  // portion of the symbols in `symbolic_map_`.
  std::vector<Variable> rt_vars_;

   // Inequality constraints for symbolic expressions. They restrict the feasible
  // set for the domain of the indexing map. It contains symbolic expressions
  // other than SymbolicDimExpr and SymbolicSymbolExpr.
  llvm::MapVector<SymbolicExpr, Interval> constraints_;
};

Code reference: indexing_map.h#L114

dim_vars_ encode the inclusive box constraints for the dimension variables d of the indexing map, which usually coincide with the shape of the output tensor for ops like transpose, reduce, elementwise, dot, but there are some exceptions like HloConcatenateInstruction.

range_vars_ all values that range variables s take. The range variables are needed when multiple values are necessary to compute a single element of the tensor we are mapping from, e.g. for output->input indexing map of reductions or input->output map for broadcasts.

rt_vars_ encode the feasible values in runtime. For example, the offset is dynamic for a 1D HloDynamicSliceInstruction. The corresponding RTVar will have feasible values between 0 and tensor_size - slice_size - 1.

constraints_ capture relations between values in form <expression> in <range>, e.g. d0 + s0 in [0, 20]. Together with Variable.bounds they define the "domain" of indexing function.

Let's study-by-example to understand what all of the above actually means.

Indexing Maps for Unfused Ops

Elementwise

For elementwise ops the indexing map is an identity.

  p0 = f32[10, 20] parameter(0)
  p1 = f32[10, 20] parameter(1)
  output = f32[10, 20] add(p0, p1)

The output to input map output -> p0:

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
d1 in [0, 19]

The input to output map p0 -> output:

(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
d1 in [0, 19]

Broadcast

Broadcasting means that some of the dimensions will be removed when we map output to input and added when we map input to output.

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

The output to input map:

(d0, d1, d2) -> (d1),
domain:
d0 in [0, 9],
d1 in [0, 19],
d2 in [0, 29]

The input to output map:

(d0)[s0, s1] -> (s0, d0, s1),
domain:
d0 in [0, 19],
s0 in [0, 9],
s1 in [0, 29]

Note that now we have range variables s on the right side for the input-to-output mapping. Those are the symbols that represent ranges of values. For example, in this particular case every element of input with index d0 is mapped to a 10x1x30 slice of the output.

Iota

Iota has no input tensor operand, so there is no input index arguments.