1 //===- LinalgOps.h - Linalg Operations --------------------------*- 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 #ifndef MLIR_DIALECT_LINALG_LINALGOPS_H_ 10 #define MLIR_DIALECT_LINALG_LINALGOPS_H_ 11 12 #include "mlir/Dialect/Linalg/IR/LinalgTraits.h" 13 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" 14 #include "mlir/Dialect/StandardOps/IR/Ops.h" 15 #include "mlir/Dialect/Utils/StructuredOpsUtils.h" 16 #include "mlir/IR/AffineExpr.h" 17 #include "mlir/IR/AffineMap.h" 18 #include "mlir/IR/BlockAndValueMapping.h" 19 #include "mlir/IR/Builders.h" 20 #include "mlir/IR/BuiltinDialect.h" 21 #include "mlir/IR/BuiltinTypes.h" 22 #include "mlir/IR/OpDefinition.h" 23 #include "mlir/IR/TypeUtilities.h" 24 #include "mlir/IR/Types.h" 25 #include "mlir/Interfaces/CopyOpInterface.h" 26 #include "mlir/Interfaces/SideEffectInterfaces.h" 27 #include "mlir/Interfaces/ViewLikeInterface.h" 28 #include "mlir/Support/LLVM.h" 29 30 #include "llvm/ADT/STLExtras.h" 31 32 namespace mlir { 33 namespace linalg { 34 35 class ConvOp; 36 class LinalgOp; 37 class PoolingMaxOp; 38 class PoolingMinOp; 39 class PoolingSumOp; 40 41 // TOFO: allow an extra ValueRange to specify an indexing and allow 42 // non-hyperrectangular shapes. 43 using LoopRangeBuilder = 44 std::function<SmallVector<Range, 4>(OpBuilder &, Location)>; 45 46 /// Returns the values obtained by applying `map` to the list of values. 47 SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc, 48 AffineMap map, ValueRange values); 49 50 /// Provide a very simple inference procedure to build the loop ranges from the 51 /// op and its operands. This only works with permutation affine maps and 52 /// patterns of the form `(m, n)[s] -> (m + n - s floordiv 2)`. 53 /// A more advanced Tensor-Comprehension like inference is possible but has 54 /// proven to be ambiguous in unfavorable case. 55 /// As a consequence, we relax the default behavior very conservatively and 56 /// provide an op-specified hook so that Linalg ops may override the behavior. 57 LoopRangeBuilder defaultLoopRangesBuilder(LinalgOp op); 58 59 using ReassociationIndices = SmallVector<int64_t, 2>; 60 using ReassociationExprs = SmallVector<AffineExpr, 2>; 61 62 /// Returns the name mangled library call name to disambiguate between different 63 /// overloads at the C level. The name mangling scheme is basic and uses MLIR 64 /// type names: 65 /// 1. form a string which is the concatenation of the linalg op name with all 66 /// the operand type names, separate by underscores; 67 /// 2. drop the `linalg.` prefix, and the `<`, `>`, `?` symbols from the type. 68 /// Assumes `op` is a LinalgOp. 69 /// 70 /// Examples: 71 /// 72 /// 1. linalg.fill(%A, %f) : memref<f32>, f32 73 /// name mangles into `linalg_fill_viewf32_f32_impl` 74 /// 75 /// 2. linalg.dot %A, %B, %C : 76 /// (memref<?xf32, stride_specification>, 77 /// memref<?xf32, stride_specification>, memref<f32>) 78 /// name mangles into `linalg_dot_viewxf32_viewxf32_viewf32_impl` 79 /// 80 /// 3. linalg.matmul(...) : 81 /// memref<?x?xf32, stride_specification>, 82 /// memref<?x?xf32, stride_specification>, 83 /// memref<?x?xf32, stride_specification> 84 /// name mangles into `linalg_matmul_viewxxf32_viewxxf32_viewxxf32_impl` 85 std::string generateLibraryCallName(Operation *op); 86 87 /// Returns `num` AffineDimExpr dimensions at positions 88 /// [startIdx, startIdx + num) and increments `startIdx` to `startIdx + num`. 89 SmallVector<AffineExpr, 4> makeAffineDimExprs(unsigned num, unsigned &startIdx, 90 MLIRContext *context); 91 92 /// Builds the indexing expressions for a ConvOp/PoolingOp `op`. Returns the 93 /// vector of AffineMaps representing: 94 /// `stride[i] * outputDims[i] + dilation[i] * windowDims[i] - pad_low[i]` 95 template <typename PoolingOp> 96 extern SmallVector<AffineExpr, 4> 97 weightedPoolingInputIndex(PoolingOp op, ArrayRef<AffineExpr> outputDims, 98 ArrayRef<AffineExpr> windowDims); 99 100 /// Returns `maybeMap.get()` if `maybeMap` is set, otherwise returns the 101 /// symbol-less identity map of `rank`. 102 AffineMap extractOrIdentityMap(Optional<AffineMap> maybeMap, unsigned rank, 103 MLIRContext *context); 104 105 /// Return the vector that is the concatenation of `a` and `b`. 106 SmallVector<AffineExpr, 4> concat(ArrayRef<AffineExpr> a, 107 ArrayRef<AffineExpr> b); 108 109 /// Return the dims that are `iteratorTypeName` loops in the LinalgOp `op`. 110 /// Assumes `op` is a LinalgOp. 111 void getDimsOfType(Operation *op, StringRef iteratorTypeName, 112 SmallVectorImpl<AffineExpr> &res); 113 114 } // namespace linalg 115 } // namespace mlir 116 117 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterfaces.h.inc" 118 119 #define GET_OP_CLASSES 120 #include "mlir/Dialect/Linalg/IR/LinalgOps.h.inc" 121 122 #define GET_OP_CLASSES 123 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.h.inc" 124 125 #endif // MLIR_DIALECT_LINALG_LINALGOPS_H_ 126