1 //===- Utils.h - Affine dialect utilities -----------------------*- C++ -*-===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This header file declares a set of utilities for the affine dialect ops. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #ifndef MLIR_DIALECT_AFFINE_UTILS_H 14 #define MLIR_DIALECT_AFFINE_UTILS_H 15 16 #include "mlir/Support/LLVM.h" 17 #include "llvm/ADT/DenseMap.h" 18 #include "llvm/ADT/SmallVector.h" 19 20 namespace mlir { 21 22 class AffineForOp; 23 class AffineIfOp; 24 class AffineParallelOp; 25 struct LogicalResult; 26 class Operation; 27 28 /// Replaces parallel affine.for op with 1-d affine.parallel op. 29 /// mlir::isLoopParallel detect the parallel affine.for ops. 30 /// There is no cost model currently used to drive this parallelization. 31 void affineParallelize(AffineForOp forOp); 32 33 /// Hoists out affine.if/else to as high as possible, i.e., past all invariant 34 /// affine.fors/parallel's. Returns success if any hoisting happened; folded` is 35 /// set to true if the op was folded or erased. This hoisting could lead to 36 /// significant code expansion in some cases. 37 LogicalResult hoistAffineIfOp(AffineIfOp ifOp, bool *folded = nullptr); 38 39 /// Holds parameters to perform n-D vectorization on a single loop nest. 40 /// For example, for the following loop nest: 41 /// 42 /// func @vec2d(%in: memref<64x128x512xf32>, %out: memref<64x128x512xf32>) { 43 /// affine.for %i0 = 0 to 64 { 44 /// affine.for %i1 = 0 to 128 { 45 /// affine.for %i2 = 0 to 512 { 46 /// %ld = affine.load %in[%i0, %i1, %i2] : memref<64x128x512xf32> 47 /// affine.store %ld, %out[%i0, %i1, %i2] : memref<64x128x512xf32> 48 /// } 49 /// } 50 /// } 51 /// return 52 /// } 53 /// 54 /// and VectorizationStrategy = 'vectorSizes = {8, 4}', 'loopToVectorDim = 55 /// {{i1->0}, {i2->1}}', SuperVectorizer will generate: 56 /// 57 /// func @vec2d(%arg0: memref<64x128x512xf32>, %arg1: memref<64x128x512xf32>) { 58 /// affine.for %arg2 = 0 to 64 { 59 /// affine.for %arg3 = 0 to 128 step 8 { 60 /// affine.for %arg4 = 0 to 512 step 4 { 61 /// %cst = constant 0.000000e+00 : f32 62 /// %0 = vector.transfer_read %arg0[%arg2, %arg3, %arg4], %cst : ... 63 /// vector.transfer_write %0, %arg1[%arg2, %arg3, %arg4] : ... 64 /// } 65 /// } 66 /// } 67 /// return 68 /// } 69 // TODO: Hoist to a VectorizationStrategy.cpp when appropriate. 70 struct VectorizationStrategy { 71 // Vectorization factors to apply to each target vector dimension. 72 // Each factor will be applied to a different loop. 73 SmallVector<int64_t, 8> vectorSizes; 74 // Maps each AffineForOp vectorization candidate with its vector dimension. 75 // The candidate will be vectorized using the vectorization factor in 76 // 'vectorSizes' for that dimension. 77 DenseMap<Operation *, unsigned> loopToVectorDim; 78 }; 79 80 /// Vectorizes affine loops in 'loops' using the n-D vectorization factors in 81 /// 'vectorSizes'. By default, each vectorization factor is applied 82 /// inner-to-outer to the loops of each loop nest. 'fastestVaryingPattern' can 83 /// be optionally used to provide a different loop vectorization order. 84 void vectorizeAffineLoops( 85 Operation *parentOp, 86 llvm::DenseSet<Operation *, DenseMapInfo<Operation *>> &loops, 87 ArrayRef<int64_t> vectorSizes, ArrayRef<int64_t> fastestVaryingPattern); 88 89 /// External utility to vectorize affine loops from a single loop nest using an 90 /// n-D vectorization strategy (see doc in VectorizationStrategy definition). 91 /// Loops are provided in a 2D vector container. The first dimension represents 92 /// the nesting level relative to the loops to be vectorized. The second 93 /// dimension contains the loops. This means that: 94 /// a) every loop in 'loops[i]' must have a parent loop in 'loops[i-1]', 95 /// b) a loop in 'loops[i]' may or may not have a child loop in 'loops[i+1]'. 96 /// 97 /// For example, for the following loop nest: 98 /// 99 /// func @vec2d(%in0: memref<64x128x512xf32>, %in1: memref<64x128x128xf32>, 100 /// %out0: memref<64x128x512xf32>, 101 /// %out1: memref<64x128x128xf32>) { 102 /// affine.for %i0 = 0 to 64 { 103 /// affine.for %i1 = 0 to 128 { 104 /// affine.for %i2 = 0 to 512 { 105 /// %ld = affine.load %in0[%i0, %i1, %i2] : memref<64x128x512xf32> 106 /// affine.store %ld, %out0[%i0, %i1, %i2] : memref<64x128x512xf32> 107 /// } 108 /// affine.for %i3 = 0 to 128 { 109 /// %ld = affine.load %in1[%i0, %i1, %i3] : memref<64x128x128xf32> 110 /// affine.store %ld, %out1[%i0, %i1, %i3] : memref<64x128x128xf32> 111 /// } 112 /// } 113 /// } 114 /// return 115 /// } 116 /// 117 /// loops = {{%i0}, {%i2, %i3}}, to vectorize the outermost and the two 118 /// innermost loops; 119 /// loops = {{%i1}, {%i2, %i3}}, to vectorize the middle and the two innermost 120 /// loops; 121 /// loops = {{%i2}}, to vectorize only the first innermost loop; 122 /// loops = {{%i3}}, to vectorize only the second innermost loop; 123 /// loops = {{%i1}}, to vectorize only the middle loop. 124 LogicalResult 125 vectorizeAffineLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops, 126 const VectorizationStrategy &strategy); 127 128 /// Normalize a affine.parallel op so that lower bounds are 0 and steps are 1. 129 /// As currently implemented, this transformation cannot fail and will return 130 /// early if the op is already in a normalized form. 131 void normalizeAffineParallel(AffineParallelOp op); 132 133 } // namespace mlir 134 135 #endif // MLIR_DIALECT_AFFINE_UTILS_H 136