1 //===- Builders.cpp - MLIR Declarative Linalg Builders --------------------===//
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 #include "mlir/IR/Builders.h"
10 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
11 #include "mlir/Dialect/Linalg/EDSC/Builders.h"
12 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
13 #include "mlir/Dialect/SCF/EDSC/Builders.h"
14 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
15 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
16 #include "mlir/IR/AffineExpr.h"
17
18 using namespace mlir;
19 using namespace mlir::edsc;
20 using namespace mlir::edsc::intrinsics;
21 using namespace mlir::linalg;
22 using namespace mlir::scf;
23
makeGenericLinalgOp(ArrayRef<IteratorType> iteratorTypes,ArrayRef<StructuredIndexed> inputs,ArrayRef<StructuredIndexed> outputBuffers,ArrayRef<Value> initTensors,ArrayRef<StructuredIndexed> resultTensorTypes,function_ref<void (ValueRange)> regionBuilder,ArrayRef<Value> otherValues,ArrayRef<Attribute> otherAttributes)24 Operation *mlir::edsc::makeGenericLinalgOp(
25 ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
26 ArrayRef<StructuredIndexed> outputBuffers, ArrayRef<Value> initTensors,
27 ArrayRef<StructuredIndexed> resultTensorTypes,
28 function_ref<void(ValueRange)> regionBuilder, ArrayRef<Value> otherValues,
29 ArrayRef<Attribute> otherAttributes) {
30 OpBuilder &builder = edsc::ScopedContext::getBuilderRef();
31
32 // Build maps
33 SmallVector<SmallVector<AffineExpr, 4>, 4> exprsList;
34 exprsList.reserve(inputs.size() + outputBuffers.size() + initTensors.size());
35 for (auto container : {inputs, outputBuffers, resultTensorTypes})
36 for (const StructuredIndexed &s : container)
37 exprsList.emplace_back(s.getExprs().begin(), s.getExprs().end());
38 auto maps = AffineMap::inferFromExprList(exprsList);
39
40 SmallVector<Type, 4> types;
41 assert(llvm::all_of(resultTensorTypes, [](const StructuredIndexed &s) {
42 return !s.hasValue();
43 }));
44 std::copy(resultTensorTypes.begin(), resultTensorTypes.end(),
45 std::back_inserter(types));
46
47 SmallVector<Value, 4> inputValues, outputBufferValues, initTensorValues;
48 inputValues.reserve(inputs.size());
49 outputBufferValues.reserve(outputBuffers.size());
50 initTensorValues.reserve(initTensors.size());
51 std::copy(inputs.begin(), inputs.end(), std::back_inserter(inputValues));
52 std::copy(outputBuffers.begin(), outputBuffers.end(),
53 std::back_inserter(outputBufferValues));
54 std::copy(initTensors.begin(), initTensors.end(),
55 std::back_inserter(initTensorValues));
56
57 auto iteratorStrTypes =
58 llvm::to_vector<8>(llvm::map_range(iteratorTypes, toString));
59 // clang-format off
60 auto *op =
61 edsc::ScopedContext::getBuilderRef()
62 .create<linalg::GenericOp>(
63 edsc::ScopedContext::getLocation(),
64 types,
65 inputValues,
66 outputBufferValues,
67 initTensorValues,
68 builder.getAffineMapArrayAttr(maps),
69 builder.getStrArrayAttr(iteratorStrTypes),
70 StringAttr() /*doc*/,
71 StringAttr() /*library_call*/,
72 ArrayAttr() /*sparse*/
73 /* TODO: other attributes in op */
74 )
75 .getOperation();
76 // clang-format on
77
78 using namespace edsc;
79 SmallVector<Type, 4> blockTypes;
80 blockTypes.reserve(inputs.size() + outputBuffers.size() + initTensors.size());
81 for (auto container : {inputs, outputBuffers})
82 for (const StructuredIndexed &s : container)
83 blockTypes.push_back(getElementTypeOrSelf(s.getType()));
84 for (Value v : initTensors)
85 blockTypes.push_back(getElementTypeOrSelf(v.getType()));
86
87 assert(op->getNumRegions() == 1);
88 assert(op->getRegion(0).empty());
89 OpBuilder opBuilder(op);
90 ScopedContext scope(opBuilder, op->getLoc());
91 buildInNewBlock(op->getRegion(0), blockTypes, regionBuilder);
92 assert(llvm::hasSingleElement(op->getRegion(0)));
93 return op;
94 }
95
mulRegionBuilder(ValueRange args)96 void mlir::edsc::ops::mulRegionBuilder(ValueRange args) {
97 using edsc::op::operator+;
98 using edsc::op::operator*;
99 assert(args.size() == 2 && "expected 2 block arguments");
100 Value a(args[0]), b(args[1]);
101 linalg_yield(a * b);
102 }
103
macRegionBuilder(ValueRange args)104 void mlir::edsc::ops::macRegionBuilder(ValueRange args) {
105 using edsc::op::operator+;
106 using edsc::op::operator*;
107 assert(args.size() == 3 && "expected 3 block arguments");
108 Value a(args[0]), b(args[1]), c(args[2]);
109 linalg_yield(c + a * b);
110 }
111
linalg_generic_pointwise(UnaryPointwiseOpBuilder unaryOp,StructuredIndexed I,StructuredIndexed O)112 Operation *mlir::edsc::ops::linalg_generic_pointwise(
113 UnaryPointwiseOpBuilder unaryOp, StructuredIndexed I, StructuredIndexed O) {
114 SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
115 IteratorType::Parallel);
116 auto fun = [&unaryOp](ValueRange args) {
117 assert(!args.empty() && "expected >= 1 block arguments");
118 Value a(args[0]);
119 linalg_yield(unaryOp(a));
120 };
121 if (O.getType().isa<RankedTensorType>())
122 return makeGenericLinalgOp(iterTypes, /*inputs=*/{I}, /*outputBuffers=*/{},
123 /*initTensors=*/{}, /*resultTensorTypes=*/{O},
124 fun);
125 return makeGenericLinalgOp(iterTypes, /*inputs=*/{I}, /*outputBuffers=*/{O},
126 /*initTensors=*/{}, /*resultTensorTypes=*/{}, fun);
127 }
128
linalg_generic_pointwise_tanh(StructuredIndexed I,StructuredIndexed O)129 Operation *mlir::edsc::ops::linalg_generic_pointwise_tanh(StructuredIndexed I,
130 StructuredIndexed O) {
131 UnaryPointwiseOpBuilder unOp([](Value a) -> Value { return std_tanh(a); });
132 return linalg_generic_pointwise(unOp, I, O);
133 }
134
135 /// Binary pointwise operation (with broadcast) entry point.
linalg_generic_pointwise(BinaryPointwiseOpBuilder binaryOp,StructuredIndexed I1,StructuredIndexed I2,StructuredIndexed O)136 Operation *mlir::edsc::ops::linalg_generic_pointwise(
137 BinaryPointwiseOpBuilder binaryOp, StructuredIndexed I1,
138 StructuredIndexed I2, StructuredIndexed O) {
139 SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
140 IteratorType::Parallel);
141 auto fun = [&binaryOp](ValueRange args) {
142 assert(args.size() >= 2 && "expected >= 2 block arguments");
143 Value a(args[0]), b(args[1]);
144 linalg_yield(binaryOp(a, b));
145 };
146 if (O.getType().isa<RankedTensorType>())
147 return makeGenericLinalgOp(
148 iterTypes, /*inputs=*/{I1, I2}, /*outputBuffers=*/{},
149 /*initTensors=*/{}, /*resultTensorTypes=*/{O}, fun);
150 return makeGenericLinalgOp(iterTypes, /*inputs=*/{I1, I2},
151 /*outputBuffers=*/{O},
152 /*initTensors=*/{}, /*resultTensorTypes=*/{}, fun);
153 }
154
linalg_generic_pointwise_add(StructuredIndexed I1,StructuredIndexed I2,StructuredIndexed O)155 Operation *mlir::edsc::ops::linalg_generic_pointwise_add(StructuredIndexed I1,
156 StructuredIndexed I2,
157 StructuredIndexed O) {
158 using edsc::op::operator+;
159 BinaryPointwiseOpBuilder binOp(
160 [](Value a, Value b) -> Value { return a + b; });
161 return linalg_generic_pointwise(binOp, I1, I2, O);
162 }
163
linalg_generic_pointwise_max(StructuredIndexed I1,StructuredIndexed I2,StructuredIndexed O)164 Operation *mlir::edsc::ops::linalg_generic_pointwise_max(StructuredIndexed I1,
165 StructuredIndexed I2,
166 StructuredIndexed O) {
167 BinaryPointwiseOpBuilder binOp([](Value a, Value b) -> Value {
168 using edsc::op::sgt;
169 return std_select(sgt(a, b), a, b);
170 });
171 return linalg_generic_pointwise(binOp, I1, I2, O);
172 }
173
174 Operation *
linalg_generic_matmul(Value vA,Value vB,Value vC,MatmulRegionBuilder regionBuilder)175 mlir::edsc::ops::linalg_generic_matmul(Value vA, Value vB, Value vC,
176 MatmulRegionBuilder regionBuilder) {
177 // clang-format off
178 AffineExpr m, n, k;
179 bindDims(ScopedContext::getContext(), m, n, k);
180 StructuredIndexed A(vA), B(vB), C(vC);
181 return makeGenericLinalgOp(
182 {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
183 /*inputs=*/{A({m, k}), B({k, n})},
184 /*outputBuffers=*/{C({m, n})},
185 /*initTensors=*/{},
186 /*resultTensorTypes=*/{},
187 regionBuilder);
188 // clang-format on
189 }
190
191 Operation *
linalg_generic_matmul(Value vA,Value vB,Value vC,RankedTensorType tD,MatmulRegionBuilder regionBuilder)192 mlir::edsc::ops::linalg_generic_matmul(Value vA, Value vB, Value vC,
193 RankedTensorType tD,
194 MatmulRegionBuilder regionBuilder) {
195 // clang-format off
196 AffineExpr m, n, k;
197 bindDims(ScopedContext::getContext(), m, n, k);
198 StructuredIndexed A(vA), B(vB), C(vC), D(tD);
199 return makeGenericLinalgOp(
200 {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
201 /*inputs=*/{A({m, k}), B({k, n})},
202 /*outputBuffers=*/{},
203 /*initTensors=*/{C({m, n})},
204 /*resultTensorTypes=*/{D({m, n})},
205 regionBuilder);
206 // clang-format on
207 }
208
linalg_generic_conv_nhwc(Value vI,Value vW,Value vO,ArrayRef<int> strides,ArrayRef<int> dilations)209 Operation *mlir::edsc::ops::linalg_generic_conv_nhwc(Value vI, Value vW,
210 Value vO,
211 ArrayRef<int> strides,
212 ArrayRef<int> dilations) {
213 MLIRContext *ctx = ScopedContext::getContext();
214 // TODO: some template magic to make everything rank-polymorphic.
215 assert((dilations.empty() || dilations.size() == 2) && "only 2-D conv atm");
216 assert((strides.empty() || strides.size() == 2) && "only 2-D conv atm");
217
218 // Some short names.
219 auto par = IteratorType::Parallel;
220 auto red = IteratorType::Reduction;
221 auto s = strides;
222 auto d = dilations;
223
224 AffineExpr b, f, h, w, kh, kw, c;
225 bindDims(ctx, b, f, h, w, kh, kw, c);
226 unsigned numDims = c.cast<AffineDimExpr>().getPosition() + 1;
227 StructuredIndexed I(vI), W(vW), O(vO);
228 // clang-format off
229 return makeGenericLinalgOp(
230 {par, par, par, par, red, red, red},
231 /*inputs=*/{
232 I({b,
233 // Roundtrip to flattened form to serve as canonicalization and ensure
234 // consistent ordering of subexpressions.
235 simplifyAffineExpr(s[0] * h + d[0] * kh, numDims, 0),
236 simplifyAffineExpr(s[1] * w + d[1] * kw, numDims, 0),
237 c}),
238 W({kh, kw, c, f}) },
239 /*outputBuffers=*/{ O({b, h, w, f}) },
240 /*initTensors=*/{},
241 /*resultTensorTypes=*/{},
242 macRegionBuilder);
243 // clang-format on
244 }
245
linalg_generic_dilated_conv_nhwc(Value vI,Value vW,Value vO,int depth_multiplier,ArrayRef<int> strides,ArrayRef<int> dilations)246 Operation *mlir::edsc::ops::linalg_generic_dilated_conv_nhwc(
247 Value vI, Value vW, Value vO, int depth_multiplier, ArrayRef<int> strides,
248 ArrayRef<int> dilations) {
249 MLIRContext *ctx = ScopedContext::getContext();
250 // TODO: some template magic to make everything rank-polymorphic.
251 assert((dilations.empty() || dilations.size() == 2) && "only 2-D conv atm");
252 assert((strides.empty() || strides.size() == 2) && "only 2-D conv atm");
253
254 // Some short names.
255 auto par = IteratorType::Parallel;
256 auto red = IteratorType::Reduction;
257 auto s = strides;
258 auto d = dilations;
259
260 // clang-format off
261 AffineExpr b, dm, c, h, w, kh, kw;
262 bindDims(ctx, b, dm, c, h, w, kh, kw);
263 unsigned numDims = kw.cast<AffineDimExpr>().getPosition() + 1;
264 StructuredIndexed I(vI), W(vW), O(vO);
265 return makeGenericLinalgOp(
266 {par, par, par, par, par, red, red},
267 /*inputs=*/{
268 I({b,
269 // Roundtrip to flattened form to serve as canonicalization and ensure
270 // consistent ordering of subexpressions.
271 simplifyAffineExpr(s[0] * h + d[0] * kh, numDims, 0),
272 simplifyAffineExpr(s[1] * w + d[1] * kw, numDims, 0),
273 c}),
274 W({kh, kw, c, dm})},
275 /*outputBuffers=*/{
276 O({b, h, w, simplifyAffineExpr(c * depth_multiplier + dm, numDims, 0)})},
277 /*initTensors=*/{},
278 /*resultTensorTypes=*/{},
279 macRegionBuilder);
280 // clang-format on
281 }
282