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