#include "mlir/Analysis/LoopAnalysis.h"
#include "mlir/Analysis/NestedMatcher.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Analysis/Utils.h"
#include "mlir/Dialect/AffineOps/AffineOps.h"
#include "mlir/Dialect/StandardOps/Ops.h"
#include "mlir/Dialect/VectorOps/Utils.h"
#include "mlir/Dialect/VectorOps/VectorOps.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/Types.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/Functional.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/Passes.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"

Include dependency graph for Vectorize.cpp:

Macros
#define	DEBUG_TYPE "early-vect"

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE "early-vect"

Implements a high-level vectorization strategy on a Function. The abstraction used is that of super-vectors, which provide a single, compact, representation in the vector types, information that is expected to reduce the impact of the phase ordering problem

Vector granularity:

This pass is designed to perform vectorization at a super-vector granularity. A super-vector is loosely defined as a vector type that is a multiple of a "good" vector size so the HW can efficiently implement a set of high-level primitives. Multiple is understood along any dimension; e.g. both vector<16xf32> and vector<2x8xf32> are valid super-vectors for a vector<8xf32> HW vector. Note that a "good vector size so the HW can efficiently implement a set of high-level primitives" is not necessarily an integer multiple of actual hardware registers. We leave details of this distinction unspecified for now.

Some may prefer the terminology a "tile of HW vectors". In this case, one should note that super-vectors implement an "always full tile" abstraction. They guarantee no partial-tile separation is necessary by relying on a high-level copy-reshape abstraction that we call vector.transfer. This copy-reshape operations is also responsible for performing layout transposition if necessary. In the general case this will require a scoped allocation in some notional local memory.

Whatever the mental model one prefers to use for this abstraction, the key point is that we burn into a single, compact, representation in the vector types, information that is expected to reduce the impact of the phase ordering problem. Indeed, a vector type conveys information that:

the associated loops have dependency semantics that do not prevent vectorization;
the associate loops have been sliced in chunks of static sizes that are compatible with vector sizes (i.e. similar to unroll-and-jam);
the inner loops, in the unroll-and-jam analogy of 2, are captured by the vector type and no vectorization hampering transformations can be applied to them anymore;
the underlying memrefs are accessed in some notional contiguous way that allows loading into vectors with some amount of spatial locality; In other words, super-vectorization provides a level of separation of concern by way of opacity to subsequent passes. This has the effect of encapsulating and propagating vectorization constraints down the list of passes until we are ready to lower further.

For a particular target, a notion of minimal n-d vector size will be specified and vectorization targets a multiple of those. In the following paragraph, let "k ." represent "a multiple of", to be understood as a multiple in the same dimension (e.g. vector<16 x k . 128> summarizes vector<16 x 128>, vector<16 x 256>, vector<16 x 1024>, etc).

Some non-exhaustive notable super-vector sizes of interest include:

CPU: vector<k . HW_vector_size>, vector<k' . core_count x k . HW_vector_size>, vector<socket_count x k' . core_count x k . HW_vector_size>;
GPU: vector<k . warp_size>, vector<k . warp_size x float2>, vector<k . warp_size x float4>, vector<k . warp_size x 4 x 4x 4> (for tensor_core sizes).

Loops and operations are emitted that operate on those super-vector shapes. Subsequent lowering passes will materialize to actual HW vector sizes. These passes are expected to be (gradually) more target-specific.

At a high level, a vectorized load in a loop will resemble:

affine.for %i = ? to ? step ? {
  %v_a = vector.transfer_read A[%i] : memref<?xf32>, vector<128xf32>
}

It is the responsibility of the implementation of vector.transfer_read to materialize vector registers from the original scalar memrefs. A later (more target-dependent) lowering pass will materialize to actual HW vector sizes. This lowering may be occur at different times:

at the MLIR level into a combination of loops, unrolling, DmaStartOp + DmaWaitOp + vectorized operations for data transformations and shuffle; thus opening opportunities for unrolling and pipelining. This is an instance of library call "whiteboxing"; or
later in the a target-specific lowering pass or hand-written library call; achieving full separation of concerns. This is an instance of library call; or
a mix of both, e.g. based on a model. In the future, these operations will expose a contract to constrain the search on vectorization patterns and sizes.

Occurrence of super-vectorization in the compiler flow:

This is an active area of investigation. We start with 2 remarks to position super-vectorization in the context of existing ongoing work: LLVM VPLAN and LLVM SLP Vectorizer.

LLVM VPLAN:

The astute reader may have noticed that in the limit, super-vectorization can be applied at a similar time and with similar objectives than VPLAN. For instance, in the case of a traditional, polyhedral compilation-flow (for instance, the PPCG project uses ISL to provide dependence analysis, multi-level(scheduling + tiling), lifting footprint to fast memory, communication synthesis, mapping, register optimizations) and before unrolling. When vectorization is applied at this late level in a typical polyhedral flow, and is instantiated with actual hardware vector sizes, super-vectorization is expected to match (or subsume) the type of patterns that LLVM's VPLAN aims at targeting. The main difference here is that MLIR is higher level and our implementation should be significantly simpler. Also note that in this mode, recursive patterns are probably a bit of an overkill although it is reasonable to expect that mixing a bit of outer loop and inner loop vectorization + unrolling will provide interesting choices to MLIR.

LLVM SLP Vectorizer:

Super-vectorization however is not meant to be usable in a similar fashion to the SLP vectorizer. The main difference lies in the information that both vectorizers use: super-vectorization examines contiguity of memory references along fastest varying dimensions and loops with recursive nested patterns capturing imperfectly-nested loop nests; the SLP vectorizer, on the other hand, performs flat pattern matching inside a single unrolled loop body and stitches together pieces of load and store operations into full 1-D vectors. We envision that the SLP vectorizer is a good way to capture innermost loop, control-flow dependent patterns that super-vectorization may not be able to capture easily. In other words, super-vectorization does not aim at replacing the SLP vectorizer and the two solutions are complementary.

Ongoing investigations:

We discuss the following early places where super-vectorization is applicable and touch on the expected benefits and risks . We list the opportunities in the context of the traditional polyhedral compiler flow described in PPCG. There are essentially 6 places in the MLIR pass pipeline we expect to experiment with super-vectorization:

Right after language lowering to MLIR: this is the earliest time where super-vectorization is expected to be applied. At this level, all the language/user/library-level annotations are available and can be fully exploited. Examples include loop-type annotations (such as parallel, reduction, scan, dependence distance vector, vectorizable) as well as memory access annotations (such as non-aliasing writes guaranteed, indirect accesses that are permutations by construction) accesses or that a particular operation is prescribed atomic by the user. At this level, anything that enriches what dependence analysis can do should be aggressively exploited. At this level we are close to having explicit vector types in the language, except we do not impose that burden on the programmer/library: we derive information from scalar code + annotations.
After dependence analysis and before polyhedral scheduling: the information that supports vectorization does not need to be supplied by a higher level of abstraction. Traditional dependence analysis is available in MLIR and will be used to drive vectorization and cost models.

Let's pause here and remark that applying super-vectorization as described in 1. and 2. presents clear opportunities and risks:

the opportunity is that vectorization is burned in the type system and is protected from the adverse effect of loop scheduling, tiling, loop interchange and all passes downstream. Provided that subsequent passes are able to operate on vector types; the vector shapes, associated loop iterator properties, alignment, and contiguity of fastest varying dimensions are preserved until we lower the super-vector types. We expect this to significantly rein in on the adverse effects of phase ordering.
the risks are that a. all passes after super-vectorization have to work on elemental vector types (not that this is always true, wherever vectorization is applied) and b. that imposing vectorization constraints too early may be overall detrimental to loop fusion, tiling and other transformations because the dependence distances are coarsened when operating on elemental vector types. For this reason, the pattern profitability analysis should include a component that also captures the maximal amount of fusion available under a particular pattern. This is still at the stage of rough ideas but in this context, search is our friend as the Tensor Comprehensions and auto-TVM contributions demonstrated previously. Bottom-line is we do not yet have good answers for the above but aim at making it easy to answer such questions.

Back to our listing, the last places where early super-vectorization makes sense are:

right after polyhedral-style scheduling: PLUTO-style algorithms are known to improve locality, parallelism and be configurable (e.g. max-fuse, smart-fuse etc). They can also have adverse effects on contiguity properties that are required for vectorization but the vector.transfer copy-reshape-pad-transpose abstraction is expected to help recapture these properties.
right after polyhedral-style scheduling+tiling;
right after scheduling+tiling+rescheduling: points 4 and 5 represent probably the most promising places because applying tiling achieves a separation of concerns that allows rescheduling to worry less about locality and more about parallelism and distribution (e.g. min-fuse).

At these levels the risk-reward looks different: on one hand we probably lost a good deal of language/user/library-level annotation; on the other hand we gained parallelism and locality through scheduling and tiling. However we probably want to ensure tiling is compatible with the full-tile-only abstraction used in super-vectorization or suffer the consequences. It is too early to place bets on what will win but we expect super-vectorization to be the right abstraction to allow exploring at all these levels. And again, search is our friend.

Lastly, we mention it again here:

as a MLIR-based alternative to VPLAN.

Lowering, unrolling, pipelining:

TODO(ntv): point to the proper places.

Algorithm:

The algorithm proceeds in a few steps:

defining super-vectorization patterns and matching them on the tree of AffineForOp. A super-vectorization pattern is defined as a recursive data structures that matches and captures nested, imperfectly-nested loops that have a. conformable loop annotations attached (e.g. parallel, reduction, vectorizable, ...) as well as b. all contiguous load/store operations along a specified minor dimension (not necessarily the fastest varying) ;
analyzing those patterns for profitability (TODO(ntv): and interference);
Then, for each pattern in order: a. applying iterative rewriting of the loop and the load operations in DFS postorder. Rewriting is implemented by coarsening the loops and turning load operations into opaque vector.transfer_read ops; b. keeping track of the load operations encountered as "roots" and the store operations as "terminals"; c. traversing the use-def chains starting from the roots and iteratively propagating vectorized values. Scalar values that are encountered during this process must come from outside the scope of the current pattern (TODO(ntv): enforce this and generalize). Such a scalar value is vectorized only if it is a constant (into a vector splat). The non-constant case is not supported for now and results in the pattern failing to vectorize; d. performing a second traversal on the terminals (store ops) to rewriting the scalar value they write to memory into vector form. If the scalar value has been vectorized previously, we simply replace it by its vector form. Otherwise, if the scalar value is a constant, it is vectorized into a splat. In all other cases, vectorization for the pattern currently fails. e. if everything under the root AffineForOp in the current pattern vectorizes properly, we commit that loop to the IR. Otherwise we discard it and restore a previously cloned version of the loop. Thanks to the recursive scoping nature of matchers and captured patterns, this is transparently achieved by a simple RAII implementation. f. vectorization is applied on the next pattern in the list. Because pattern interference avoidance is not yet implemented and that we do not support further vectorizing an already vector load we need to re-verify that the pattern is still vectorizable. This is expected to make cost models more difficult to write and is subject to improvement in the future.

Points c. and d. above are worth additional comment. In most passes that do not change the type of operands, it is usually preferred to eagerly replaceAllUsesWith. Unfortunately this does not work for vectorization because during the use-def chain traversal, all the operands of an operation must be available in vector form. Trying to propagate eagerly makes the IR temporarily invalid and results in errors such as: `vectorize.mlir:308:13: error: 'addf' op requires the same type for all operands and results s5 = addf a5, b5 : f32`

Lastly, we show a minimal example for which use-def chains rooted in load / vector.transfer_read are not enough. This is what motivated splitting terminal processing out of the use-def chains starting from loads. In the following snippet, there is simply no load::

func @fill(%A : memref<128xf32>) -> () {
  %f1 = constant 1.0 : f32
  affine.for %i0 = 0 to 32 {
    affine.store %f1, %A[%i0] : memref<128xf32, 0>
  }
  return
}

Choice of loop transformation to support the algorithm:

The choice of loop transformation to apply for coarsening vectorized loops is still subject to exploratory tradeoffs. In particular, say we want to vectorize by a factor 128, we want to transform the following input:

affine.for %i = %M to %N {
  %a = affine.load %A[%i] : memref<?xf32>
}

Traditionally, one would vectorize late (after scheduling, tiling, memory promotion etc) say after stripmining (and potentially unrolling in the case of LLVM's SLP vectorizer):

affine.for %i = floor(%M, 128) to ceil(%N, 128) {
  affine.for %ii = max(%M, 128 * %i) to min(%N, 128*%i + 127) {
    %a = affine.load %A[%ii] : memref<?xf32>
  }
}

Instead, we seek to vectorize early and freeze vector types before scheduling, so we want to generate a pattern that resembles:

affine.for %i = ? to ? step ? {
  %v_a = vector.transfer_read %A[%i] : memref<?xf32>, vector<128xf32>
}

i. simply dividing the lower / upper bounds by 128 creates issues when representing expressions such as ii + 1 because now we only have access to original values that have been divided. Additional information is needed to specify accesses at below-128 granularity; ii. another alternative is to coarsen the loop step but this may have consequences on dependence analysis and fusability of loops: fusable loops probably need to have the same step (because we don't want to stripmine/unroll to enable fusion). As a consequence, we choose to represent the coarsening using the loop step for now and reevaluate in the future. Note that we can renormalize loop steps later if/when we have evidence that they are problematic.

For the simple strawman example above, vectorizing for a 1-D vector abstraction of size 128 returns code similar to:

affine.for %i = %M to %N step 128 {
  %v_a = vector.transfer_read %A[%i] : memref<?xf32>, vector<128xf32>
}

Unsupported cases, extensions, and work in progress (help welcome :-) ):

lowering to concrete vector types for various HW;
reduction support;
non-effecting padding during vector.transfer_read and filter during vector.transfer_write;
misalignment support vector.transfer_read / vector.transfer_write (hopefully without read-modify-writes);
control-flow support;
cost-models, heuristics and search;
Op implementation, extensions and implication on memref views;
many TODOs left around.

Examples:

Consider the following Function:

func @vector_add_2d(%M : index, %N : index) -> f32 {
  %A = alloc (%M, %N) : memref<?x?xf32, 0>
  %B = alloc (%M, %N) : memref<?x?xf32, 0>
  %C = alloc (%M, %N) : memref<?x?xf32, 0>
  %f1 = constant 1.0 : f32
  %f2 = constant 2.0 : f32
  affine.for %i0 = 0 to %M {
    affine.for %i1 = 0 to %N {
      // non-scoped %f1
      affine.store %f1, %A[%i0, %i1] : memref<?x?xf32, 0>
    }
  }
  affine.for %i2 = 0 to %M {
    affine.for %i3 = 0 to %N {
      // non-scoped %f2
      affine.store %f2, %B[%i2, %i3] : memref<?x?xf32, 0>
    }
  }
  affine.for %i4 = 0 to %M {
    affine.for %i5 = 0 to %N {
      %a5 = affine.load %A[%i4, %i5] : memref<?x?xf32, 0>
      %b5 = affine.load %B[%i4, %i5] : memref<?x?xf32, 0>
      %s5 = addf %a5, %b5 : f32
      // non-scoped %f1
      %s6 = addf %s5, %f1 : f32
      // non-scoped %f2
      %s7 = addf %s5, %f2 : f32
      // diamond dependency.
      %s8 = addf %s7, %s6 : f32
      affine.store %s8, %C[%i4, %i5] : memref<?x?xf32, 0>
    }
  }
  %c7 = constant 7 : index
  %c42 = constant 42 : index
  %res = load %C[%c7, %c42] : memref<?x?xf32, 0>
  return %res : f32
}

The -affine-vectorize pass with the following arguments:

-affine-vectorize -virtual-vector-size 256 --test-fastest-varying=0

produces this standard innermost-loop vectorized code:

func @vector_add_2d(%arg0 : index, %arg1 : index) -> f32 {
  %0 = alloc(%arg0, %arg1) : memref<?x?xf32>
  %1 = alloc(%arg0, %arg1) : memref<?x?xf32>
  %2 = alloc(%arg0, %arg1) : memref<?x?xf32>
  %cst = constant 1.0 : f32
  %cst_0 = constant 2.0 : f32
  affine.for %i0 = 0 to %arg0 {
    affine.for %i1 = 0 to %arg1 step 256 {
      %cst_1 = constant dense<vector<256xf32>, 1.0> :
               vector<256xf32>
      vector.transfer_write %cst_1, %0[%i0, %i1] :
               vector<256xf32>, memref<?x?xf32>
    }
  }
  affine.for %i2 = 0 to %arg0 {
    affine.for %i3 = 0 to %arg1 step 256 {
      %cst_2 = constant dense<vector<256xf32>, 2.0> :
               vector<256xf32>
      vector.transfer_write %cst_2, %1[%i2, %i3] :
               vector<256xf32>, memref<?x?xf32>
    }
  }
  affine.for %i4 = 0 to %arg0 {
    affine.for %i5 = 0 to %arg1 step 256 {
      %3 = vector.transfer_read %0[%i4, %i5] :
           memref<?x?xf32>, vector<256xf32>
      %4 = vector.transfer_read %1[%i4, %i5] :
           memref<?x?xf32>, vector<256xf32>
      %5 = addf %3, %4 : vector<256xf32>
      %cst_3 = constant dense<vector<256xf32>, 1.0> :
               vector<256xf32>
      %6 = addf %5, %cst_3 : vector<256xf32>
      %cst_4 = constant dense<vector<256xf32>, 2.0> :
               vector<256xf32>
      %7 = addf %5, %cst_4 : vector<256xf32>
      %8 = addf %7, %6 : vector<256xf32>
      vector.transfer_write %8, %2[%i4, %i5] :
               vector<256xf32>, memref<?x?xf32>
    }
  }
  %c7 = constant 7 : index
  %c42 = constant 42 : index
  %9 = load %2[%c7, %c42] : memref<?x?xf32>
  return %9 : f32
}