MLIR 24.0.0git
TosaToLinalgNamed.cpp
Go to the documentation of this file.
1//===- TosaToLinalgNamed.cpp - Lowering Tosa to Linalg Named Ops ----------===//
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// These rewriters lower from the Tosa to the Linalg named ops.
10//
11//===----------------------------------------------------------------------===//
12
23#include "llvm/ADT/SmallVectorExtras.h"
24
25#include <type_traits>
26
27using namespace mlir;
28using namespace mlir::tosa;
29
31 TypedAttr padAttr, OpBuilder &rewriter) {
32 // Input should be padded only if necessary.
33 if (llvm::all_of(pad, [](int64_t p) { return p == 0; }))
34 return input;
35
36 ShapedType inputTy = cast<ShapedType>(input.getType());
37 Type inputETy = inputTy.getElementType();
38 auto inputShape = inputTy.getShape();
39
40 assert((inputShape.size() * 2) == pad.size());
41
42 SmallVector<int64_t, 4> paddedShape;
45 for (size_t i : llvm::seq(inputShape.size())) {
46 auto lowPad = pad[i * 2];
47 auto highPad = pad[i * 2 + 1];
48 if (ShapedType::isDynamic(inputShape[i]))
49 paddedShape.push_back(inputShape[i]);
50 else
51 paddedShape.push_back(inputShape[i] + highPad + lowPad);
52 lowIndices.push_back(rewriter.getIndexAttr(lowPad));
53 highIndices.push_back(rewriter.getIndexAttr(highPad));
54 }
55
56 Value padValue = arith::ConstantOp::create(rewriter, loc, padAttr);
57
58 return tensor::PadOp::create(rewriter, loc,
59 RankedTensorType::get(paddedShape, inputETy),
60 input, lowIndices, highIndices, padValue);
61}
62
63static mlir::Value
65 Value conv, Value result,
66 ArrayRef<AffineMap> indexingMaps) {
67 ShapedType resultTy = cast<ShapedType>(conv.getType());
68 return linalg::GenericOp::create(
69 rewriter, loc, resultTy, ValueRange({bias, conv}), result,
70 indexingMaps, getNParallelLoopsAttrs(resultTy.getRank()),
71 [](OpBuilder &builder, Location loc, ValueRange args) {
72 Value biasVal = args[0];
73 Type resType = args[1].getType();
74 if (resType != biasVal.getType()) {
75 biasVal =
76 arith::ExtSIOp::create(builder, loc, resType, biasVal);
77 }
78 Value added =
79 arith::AddIOp::create(builder, loc, biasVal, args[1]);
80 linalg::YieldOp::create(builder, loc, added);
81 })
82 .getResult(0);
83}
84
85// Construct the affine map that a linalg generic would use to broadcast the
86// source tensor into the shape of the result tensor.
88 Value result) {
89 ShapedType resultTy = cast<ShapedType>(result.getType());
90 ShapedType sourceTy = cast<ShapedType>(source.getType());
91 const int64_t resultRank = resultTy.getRank();
92 const int64_t sourceRank = sourceTy.getRank();
93
94 // The source tensor is broadcast to all the outer dimensions of the
95 // result tensor.
96 SmallVector<AffineExpr> sourceDims;
97 // In the case of a rank one source tensor with a single element TOSA
98 // specifies that the value be broadcast meaning we need an edge case for a
99 // constant map.
100 assert(sourceTy.hasStaticShape() &&
101 "Dynamic broadcasting shapes not supported!");
102 if (sourceRank == 1 && sourceTy.getDimSize(0) == 1) {
103 sourceDims.push_back(rewriter.getAffineConstantExpr(0));
104 } else {
105 for (auto dim : llvm::seq<int64_t>(0, sourceRank)) {
106 auto expr = rewriter.getAffineDimExpr(dim + resultRank - sourceRank);
107 sourceDims.push_back(expr);
108 }
109 }
110
111 return AffineMap::get(/*dimCount=*/resultRank,
112 /*symbolCount=*/0, sourceDims, rewriter.getContext());
113}
114
115// Broadcast the source value to all the outer dimensions of the result value.
116// If required, the element type is expanded using an arith.extsi or arith.extf
117// operation as appropriate.
119 Location loc, Value source,
120 Value result) {
121 ShapedType resultTy = cast<ShapedType>(result.getType());
122 const int64_t resultRank = resultTy.getRank();
123 // Creating maps for the input and output of the broacast-like generic op.
124 SmallVector<AffineMap, 2> indexingMaps;
125 indexingMaps.push_back(getBroadcastingMap(rewriter, source, result));
126 indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank));
127
128 // Build the broadcast-like operation as a linalg.generic.
129 return linalg::GenericOp::create(
130 rewriter, loc, resultTy, ValueRange({source}), result,
131 indexingMaps, getNParallelLoopsAttrs(resultTy.getRank()),
132 [&resultTy](OpBuilder &builder, Location loc, ValueRange args) {
133 Value biasVal = args[0];
134 Type resType = args[1].getType();
135 if (resType != biasVal.getType()) {
136 biasVal =
137 resultTy.getElementType().isFloat()
138 ? arith::ExtFOp::create(
139 builder, loc, TypeRange{resType},
140 ValueRange{biasVal}, arith::ExtFOp::Properties{})
141 .getResult()
142 : arith::ExtSIOp::create(builder, loc, resType,
143 biasVal)
144 .getResult();
145 }
146 linalg::YieldOp::create(builder, loc, biasVal);
147 })
148 .getResult(0);
149}
150
152 ImplicitLocOpBuilder &builder) {
153 return arith::ConstantIndexOp::create(builder, attr);
154}
155
156// Calculating the output width/height using the formula:
157// H = ((IH+pad_top+pad_bottom-(dilation_y*(KH-1)+1))/stride_y)+1
158// W = ((IW+pad_left+pad_right-(dilation_x*(KW-1)+1))/stride_x)+1
159
161 int64_t padBeforeAttr,
162 int64_t padAfterAttr, Value kernelDim,
163 int64_t strideAttr,
164 int64_t dilationAttr,
165 OpBuilder &rewriter) {
166 ImplicitLocOpBuilder builder(loc, rewriter);
167 auto one = arith::ConstantOp::create(rewriter, loc,
168 IntegerAttr::get(inputDim.getType(), 1));
169 Value padBefore = reifyConstantDim(padBeforeAttr, builder);
170 Value paddedBefore = arith::AddIOp::create(builder, inputDim, padBefore);
171 Value padAfter = reifyConstantDim(padAfterAttr, builder);
172 Value paddedAfter = arith::AddIOp::create(builder, paddedBefore, padAfter);
173
174 Value subOne = arith::SubIOp::create(builder, kernelDim, one);
175 Value dilation = reifyConstantDim(dilationAttr, builder);
176 Value dilated = arith::MulIOp::create(builder, dilation, subOne);
177 Value addOne = arith::AddIOp::create(builder, dilated, one);
178
179 Value subtract = arith::SubIOp::create(builder, paddedAfter, addOne);
180 Value stride = reifyConstantDim(strideAttr, builder);
181 Value divide = arith::DivUIOp::create(builder, subtract, stride);
182 return arith::AddIOp::create(builder, divide, one);
183}
184
185// Creates a vector of the dynamic output dims for Conv2D and Depthwise_Conv2D
187 Location loc, Value input, Value weight, ShapedType resultTy,
188 ArrayRef<int64_t> padAttr, ArrayRef<int64_t> strideAttr,
189 ArrayRef<int64_t> dilationAttr, ArrayRef<int64_t> inputSizeDims,
190 ArrayRef<int64_t> kernelSizeDims, OpBuilder &rewriter) {
191 ShapedType inputTy = cast<ShapedType>(input.getType());
192 int64_t inputRank = inputTy.getRank();
193
194 SmallVector<Value> dynDims;
195 dynDims.resize(resultTy.getRank());
196
197 for (uint32_t i = 0, s = inputSizeDims.size(); i < s; ++i) {
198 int64_t inputDim = inputSizeDims[i];
199 int64_t kernelDim = kernelSizeDims[i];
200 if (resultTy.isDynamicDim(inputDim)) {
201 auto padTop = padAttr[i * 2];
202 auto padBottom = padAttr[i * 2 + 1];
203 auto stride = strideAttr[i];
204 auto dilation = dilationAttr[i];
205 Value initDynDim = tensor::DimOp::create(rewriter, loc, input, inputDim);
206 Value kernelDynDim =
207 tensor::DimOp::create(rewriter, loc, weight, kernelDim);
208 // H = F(IH, pad_top, pad_bottom, dilation_y, KH, stride_y)
209 dynDims[inputDim] =
210 getConvOrPoolOutputDim(loc, initDynDim, padTop, padBottom,
211 kernelDynDim, stride, dilation, rewriter);
212 }
213 }
214
215 // Get the batch/channels dimensions.
216 for (int i = 0; i < inputRank; i++) {
217 if (resultTy.isDynamicDim(i) && !dynDims[i])
218 dynDims[i] = tensor::DimOp::create(rewriter, loc, input, i);
219 }
220
221 SmallVector<Value> filteredDims = condenseValues(dynDims);
222 return filteredDims;
223}
224
225// Creates a map to collapse the last dimension of the Depthwise convolution op
226// due to a shape mismatch
228 int64_t outputRank, SmallVector<ReassociationExprs, 4> &reassociationMap,
229 OpBuilder &rewriter) {
230 reassociationMap.resize(outputRank);
231 for (int i = 0; i < outputRank; i++) {
232 reassociationMap[i].push_back(rewriter.getAffineDimExpr(i));
233 }
234 reassociationMap[outputRank - 1].push_back(
235 rewriter.getAffineDimExpr(outputRank));
236}
237
238namespace {
239
240template <typename TosaConvOp, typename LinalgConvOp, typename LinalgConvQOp>
241class ConvConverter : public OpConversionPattern<TosaConvOp> {
242public:
243 using OpConversionPattern<TosaConvOp>::OpConversionPattern;
244 LogicalResult
245 matchAndRewrite(TosaConvOp op, typename TosaConvOp::Adaptor adaptor,
246 ConversionPatternRewriter &rewriter) const final {
247 Location loc = op->getLoc();
248 Value input = op->getOperand(0);
249 Value weight = op->getOperand(1);
250 Value bias = op->getOperand(2);
251
252 ShapedType inputTy = cast<ShapedType>(input.getType());
253 ShapedType weightTy = cast<ShapedType>(weight.getType());
254 ShapedType biasTy = cast<ShapedType>(bias.getType());
255 ShapedType resultTy = cast<ShapedType>(op->getResult(0).getType());
256
257 Type inputETy = inputTy.getElementType();
258
259 DenseI64ArrayAttr padAttr = op.getPadAttr();
260 DenseI64ArrayAttr strideTosaAttr = op.getStrideAttr();
261 DenseI64ArrayAttr dilationTosaAttr = op.getDilationAttr();
262
263 Type accETy = op.getAccType();
264 Type accTy = RankedTensorType::get(resultTy.getShape(), accETy);
265
266 // Get and verify zero points.
267 FailureOr<int64_t> maybeIZp = op.getInputZeroPoint();
268 if (failed(maybeIZp))
269 return rewriter.notifyMatchFailure(
270 op, "input zero point cannot be statically determined");
271
272 FailureOr<int64_t> maybeWZp = op.getWeightZeroPoint();
273 if (failed(maybeWZp))
274 return rewriter.notifyMatchFailure(
275 op, "weight zero point cannot be statically determined");
276
277 const int64_t inputZpVal = *maybeIZp;
278 const int64_t weightZpVal = *maybeWZp;
279
280 if (op.verifyInputZeroPoint(inputZpVal).failed())
281 return rewriter.notifyMatchFailure(
282 op, "input zero point must be zero for non-int8 integer types");
283
284 if (op.verifyWeightZeroPoint(weightZpVal).failed())
285 return rewriter.notifyMatchFailure(
286 op, "weight zero point must be zero for non-int8 integer types");
287
288 bool hasZp = (inputZpVal != 0) || (weightZpVal != 0);
289
290 if (!weightTy.hasStaticShape() || !biasTy.hasStaticShape())
291 return rewriter.notifyMatchFailure(
292 op, "tosa.conv ops require static shapes for weight and bias");
293
294 if (inputETy.isUnsignedInteger())
295 return rewriter.notifyMatchFailure(
296 op, "tosa.conv ops does not support unsigned integer input");
297
298 llvm::SmallVector<int64_t> inputSizeDims;
299 llvm::SmallVector<int64_t> kernelSizeDims;
300 for (int i = 1; i < resultTy.getRank() - 1; i++) {
301 inputSizeDims.push_back(i);
302 kernelSizeDims.push_back(i);
303 }
304
305 SmallVector<Value> filteredDims = inferDynamicDimsForConv(
306 loc, input, weight, resultTy, padAttr.asArrayRef(),
307 strideTosaAttr.asArrayRef(), dilationTosaAttr.asArrayRef(),
308 inputSizeDims, kernelSizeDims, rewriter);
309
310 auto weightShape = weightTy.getShape();
311
312 // Apply padding as necessary.
313 TypedAttr zeroAttr = rewriter.getZeroAttr(inputETy);
314 if (hasZp) {
315 int64_t intMin =
316 APInt::getSignedMinValue(inputETy.getIntOrFloatBitWidth())
317 .getSExtValue();
318 int64_t intMax =
319 APInt::getSignedMaxValue(inputETy.getIntOrFloatBitWidth())
320 .getSExtValue();
321
322 if (inputZpVal < intMin || inputZpVal > intMax)
323 return rewriter.notifyMatchFailure(
324 op, "tosa.conv op quantization has zp outside of input range");
325
326 zeroAttr = rewriter.getIntegerAttr(inputETy, inputZpVal);
327 }
328
329 llvm::SmallVector<int64_t> pad;
330 pad.resize(2, 0);
331 llvm::append_range(pad, padAttr.asArrayRef());
332 pad.resize(pad.size() + 2, 0);
333 input = applyPad(loc, input, pad, zeroAttr, rewriter);
334
335 if (4 == inputTy.getRank()) {
336 // For 2D convolutions, we need to check if the target convolution op
337 // wants a HWCF kernel layout.
338 bool wantHwcf =
339 hasZp ? std::is_same_v<LinalgConvQOp, linalg::Conv2DNhwcHwcfQOp>
340 : std::is_same_v<LinalgConvOp, linalg::Conv2DNhwcHwcfOp>;
341 if (wantHwcf) {
342 // Transpose the kernel to match dimension ordering of the linalg
343 // convolution operation.
344 // TODO(suderman): See if this can be efficiently folded - check whether
345 // the input is used anywhere else, if not fold the constant.
346 SmallVector<int32_t> weightPerm;
347 for (int i = 1; i < resultTy.getRank(); i++)
348 weightPerm.push_back(i);
349 weightPerm.push_back(0);
350
351 SmallVector<int64_t> newWeightShape;
352 for (auto dim : weightPerm)
353 newWeightShape.push_back(weightShape[dim]);
354 auto weightPermAttr = rewriter.getDenseI32ArrayAttr(weightPerm);
355 Type newWeightTy =
356 RankedTensorType::get(newWeightShape, weightTy.getElementType());
357 weight = tosa::TransposeOp::create(rewriter, loc, newWeightTy, weight,
358 weightPermAttr);
359 }
360 }
361
362 // For Conv3D transpose the kernel to match dimension ordering of the linalg
363 // convolution operation. Conv2D has a 1-1 mapping in linalg so better to
364 // map directly and then transpose later if desired.
365 if (5 == inputTy.getRank()) {
366 // TODO(suderman): See if this can be efficiently folded - check whether
367 // the input is used anywhere else, if not fold the constant.
368 SmallVector<int32_t> weightPerm;
369 for (int i = 1; i < resultTy.getRank(); i++)
370 weightPerm.push_back(i);
371 weightPerm.push_back(0);
372
373 SmallVector<int64_t> newWeightShape;
374 for (auto dim : weightPerm)
375 newWeightShape.push_back(weightShape[dim]);
376 auto weightPermAttr = rewriter.getDenseI32ArrayAttr(weightPerm);
377 Type newWeightTy =
378 RankedTensorType::get(newWeightShape, weightTy.getElementType());
379 weight = tosa::TransposeOp::create(rewriter, loc, newWeightTy, weight,
380 weightPermAttr);
381 }
382
383 // Extract the attributes for convolution.
384 ArrayRef<int64_t> stride = strideTosaAttr;
385 ArrayRef<int64_t> dilation = dilationTosaAttr;
386
387 // Create the convolution op.
388 auto strideAttr = rewriter.getI64TensorAttr(stride);
389 auto dilationAttr = rewriter.getI64TensorAttr(dilation);
390
391 Value biasEmptyTensor = tensor::EmptyOp::create(
392 rewriter, loc, resultTy.getShape(), accETy, filteredDims);
393
394 Value broadcastBias =
395 linalgBroadcastAndMaybeExt(rewriter, loc, bias, biasEmptyTensor);
396
397 if (hasZp) {
398 auto iZp = rewriter.getI32IntegerAttr(inputZpVal);
399 auto kZp = rewriter.getI32IntegerAttr(weightZpVal);
400
401 auto iZpVal = arith::ConstantOp::create(rewriter, loc, iZp);
402 auto kZpVal = arith::ConstantOp::create(rewriter, loc, kZp);
403
404 Value conv = LinalgConvQOp::create(
405 rewriter, loc, resultTy,
406 ValueRange{input, weight, iZpVal, kZpVal},
407 ValueRange{broadcastBias}, strideAttr, dilationAttr)
408 ->getResult(0);
409
410 rewriter.replaceOp(op, conv);
411 return success();
412 }
413
414 Value conv = LinalgConvOp::create(
415 rewriter, loc, accTy, ValueRange{input, weight},
416 ValueRange{broadcastBias}, strideAttr, dilationAttr)
417 ->getResult(0);
418
419 // We may need to truncate back to the result type if the accumulator was
420 // wider than the result.
421 if (resultTy != accTy)
422 conv = tosa::CastOp::create(rewriter, loc, resultTy, conv);
423
424 rewriter.replaceOp(op, conv);
425 return success();
426 }
427};
428
429class DepthwiseConvConverter
430 : public OpConversionPattern<tosa::DepthwiseConv2DOp> {
431public:
432 using OpConversionPattern<tosa::DepthwiseConv2DOp>::OpConversionPattern;
433 LogicalResult
434 matchAndRewrite(tosa::DepthwiseConv2DOp op, OpAdaptor adaptor,
435 ConversionPatternRewriter &rewriter) const final {
436 Location loc = op->getLoc();
437 Value input = op->getOperand(0);
438 Value weight = op->getOperand(1);
439 Value bias = op->getOperand(2);
440
441 ShapedType inputTy = cast<ShapedType>(input.getType());
442 ShapedType weightTy = cast<ShapedType>(weight.getType());
443 ShapedType biasTy = cast<ShapedType>(bias.getType());
444 ShapedType resultTy = cast<ShapedType>(op->getResult(0).getType());
445 int64_t resultRank = resultTy.getRank();
446
447 Type inputETy = inputTy.getElementType();
448 Type resultETy = resultTy.getElementType();
449
450 auto padAttr = op.getPadAttr();
451 auto strideTosaAttr = op.getStrideAttr();
452 auto dilationTosaAttr = op.getDilationAttr();
453
454 Type accETy = op.getAccType();
455
456 if (!weightTy.hasStaticShape() || !biasTy.hasStaticShape())
457 return rewriter.notifyMatchFailure(
458 op, "tosa.depthwise_conv ops require static shapes");
459
460 // Compute output dynamic dims
461 SmallVector<Value> filteredDims = inferDynamicDimsForConv(
462 loc, input, weight, resultTy, padAttr.asArrayRef(),
463 strideTosaAttr.asArrayRef(), dilationTosaAttr.asArrayRef(),
464 /*inputSizeDims=*/{1, 2},
465 /*kernelSizeDims=*/{0, 1}, rewriter);
466
467 // Get and verify zero points.
468
469 FailureOr<int64_t> maybeIZp = op.getInputZeroPoint();
470 FailureOr<int64_t> maybeWZp = op.getWeightZeroPoint();
471 if (failed(maybeIZp))
472 return rewriter.notifyMatchFailure(
473 op, "input zero point cannot be statically determined");
474 if (failed(maybeWZp))
475 return rewriter.notifyMatchFailure(
476 op, "weight zero point cannot be statically determined");
477
478 const int64_t inputZpVal = *maybeIZp;
479 const int64_t weightZpVal = *maybeWZp;
480
481 if (op.verifyInputZeroPoint(inputZpVal).failed())
482 return rewriter.notifyMatchFailure(
483 op, "input zero point must be zero for non-int8 integer types");
484
485 if (op.verifyWeightZeroPoint(weightZpVal).failed())
486 return rewriter.notifyMatchFailure(
487 op, "weight zero point must be zero for non-int8 integer types");
488
489 bool hasNullZps = (inputZpVal == 0) && (weightZpVal == 0);
490 auto weightShape = weightTy.getShape();
491 auto resultShape = resultTy.getShape();
492
493 // Apply padding as necessary.
494 TypedAttr zeroAttr = rewriter.getZeroAttr(inputETy);
495 if (!hasNullZps) {
496 int64_t intMin =
497 APInt::getSignedMinValue(inputETy.getIntOrFloatBitWidth())
498 .getSExtValue();
499 int64_t intMax =
500 APInt::getSignedMaxValue(inputETy.getIntOrFloatBitWidth())
501 .getSExtValue();
502
503 if (inputZpVal < intMin || inputZpVal > intMax)
504 return rewriter.notifyMatchFailure(
505 op, "tosa.depthwise_conv op quantization has zp outside of input "
506 "range");
507
508 zeroAttr = rewriter.getIntegerAttr(inputETy, inputZpVal);
509 }
510
511 llvm::SmallVector<int64_t> pad;
512 pad.resize(2, 0);
513 llvm::append_range(pad, padAttr.asArrayRef());
514 pad.resize(pad.size() + 2, 0);
515
516 input = applyPad(loc, input, pad, zeroAttr, rewriter);
517
518 // Extract the attributes for convolution.
519 ArrayRef<int64_t> stride = strideTosaAttr;
520 ArrayRef<int64_t> dilation = dilationTosaAttr;
521
522 // Create the convolution op.
523 auto strideAttr = rewriter.getI64TensorAttr(stride);
524 auto dilationAttr = rewriter.getI64TensorAttr(dilation);
525 ShapedType linalgConvTy =
526 RankedTensorType::get({resultShape[0], resultShape[1], resultShape[2],
527 weightShape[2], weightShape[3]},
528 accETy);
529
530 auto resultZeroAttr = rewriter.getZeroAttr(accETy);
531 Value emptyTensor = tensor::EmptyOp::create(
532 rewriter, loc, linalgConvTy.getShape(), accETy, filteredDims);
533 Value zero = arith::ConstantOp::create(rewriter, loc, resultZeroAttr);
534 Value zeroTensor = linalg::FillOp::create(rewriter, loc, ValueRange{zero},
535 ValueRange{emptyTensor})
536 .result();
537
538 Value biasEmptyTensor = tensor::EmptyOp::create(
539 rewriter, loc, resultTy.getShape(), resultETy, filteredDims);
540
541 // Broadcast the initial value to the output tensor before convolving.
542 SmallVector<AffineMap, 4> indexingMaps;
543 indexingMaps.push_back(getBroadcastingMap(rewriter, bias, biasEmptyTensor));
544 indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank));
545 indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank));
546
547 if (hasNullZps) {
548 Value conv = linalg::DepthwiseConv2DNhwcHwcmOp::create(
549 rewriter, loc, linalgConvTy, ValueRange{input, weight},
550 ValueRange{zeroTensor}, strideAttr, dilationAttr)
551 .getResult(0);
552
553 // We may need to truncate back to the result type if the accumulator was
554 // wider than the result.
555 if (accETy != resultETy)
556 conv = tosa::CastOp::create(
557 rewriter, loc,
558 RankedTensorType::get(cast<ShapedType>(conv.getType()).getShape(),
559 resultETy),
560 conv);
561
562 SmallVector<ReassociationExprs, 4> reassociationMap;
563 createDepthwiseConvCollapseMap(resultRank, reassociationMap, rewriter);
564 Value convReshape = tensor::CollapseShapeOp::create(
565 rewriter, loc, resultTy, conv, reassociationMap);
566
567 Value result =
568 linalg::GenericOp::create(
569 rewriter, loc, resultTy, ValueRange({bias, convReshape}),
570 biasEmptyTensor, indexingMaps, getNParallelLoopsAttrs(resultRank),
571 [&](OpBuilder &nestedBuilder, Location nestedLoc,
572 ValueRange args) {
573 Value added;
574 if (llvm::isa<FloatType>(inputETy))
575 added = arith::AddFOp::create(nestedBuilder, loc, args[0],
576 args[1]);
577 else
578 added = arith::AddIOp::create(nestedBuilder, loc, args[0],
579 args[1]);
580 linalg::YieldOp::create(nestedBuilder, nestedLoc, added);
581 })
582 .getResult(0);
583 rewriter.replaceOp(op, result);
584 } else {
585 IntegerAttr iZp = rewriter.getI32IntegerAttr(inputZpVal);
586 IntegerAttr wZp = rewriter.getI32IntegerAttr(weightZpVal);
587 auto iZpVal = arith::ConstantOp::create(rewriter, loc, iZp);
588 auto kZpVal = arith::ConstantOp::create(rewriter, loc, wZp);
589 Value conv = linalg::DepthwiseConv2DNhwcHwcmQOp::create(
590 rewriter, loc, linalgConvTy,
591 ValueRange{input, weight, iZpVal, kZpVal},
592 ValueRange{zeroTensor}, strideAttr, dilationAttr)
593 .getResult(0);
594 SmallVector<ReassociationExprs, 4> reassociationMap;
595 createDepthwiseConvCollapseMap(resultRank, reassociationMap, rewriter);
596 Value convReshape = tensor::CollapseShapeOp::create(
597 rewriter, loc, resultTy, conv, reassociationMap);
599 rewriter, loc, bias, convReshape, biasEmptyTensor, indexingMaps);
600 rewriter.replaceOp(op, result);
601 }
602 return success();
603 }
604};
605
606class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
607public:
608 using OpConversionPattern<tosa::MatMulOp>::OpConversionPattern;
609 LogicalResult
610 matchAndRewrite(tosa::MatMulOp op, OpAdaptor adaptor,
611 ConversionPatternRewriter &rewriter) const final {
612 Location loc = op.getLoc();
613
614 auto outputTy = cast<ShapedType>(op.getType());
615 auto outputElementTy = outputTy.getElementType();
616
617 SmallVector<Value> dynDims;
618 dynDims.resize(cast<ShapedType>(op->getResult(0).getType()).getRank());
619
620 if (!outputTy.hasRank() || outputTy.isDynamicDim(0)) {
621 dynDims[0] = tensor::DimOp::create(rewriter, loc, op->getOperand(0), 0);
622 }
623
624 if (!outputTy.hasRank() || outputTy.isDynamicDim(1)) {
625 dynDims[1] = tensor::DimOp::create(rewriter, loc, op->getOperand(0), 1);
626 }
627
628 if (!outputTy.hasRank() || outputTy.isDynamicDim(2)) {
629 dynDims[2] = tensor::DimOp::create(rewriter, loc, op->getOperand(1), 2);
630 }
631
632 SmallVector<Value> filteredDims = condenseValues(dynDims);
633
634 auto zeroAttr = rewriter.getZeroAttr(outputElementTy);
635 Value zero = arith::ConstantOp::create(rewriter, loc, zeroAttr);
636 auto emptyTensor =
637 tensor::EmptyOp::create(rewriter, loc, outputTy.getShape(),
638 outputTy.getElementType(), filteredDims);
639 Value zeroTensor = linalg::FillOp::create(rewriter, loc, ValueRange{zero},
640 ValueRange{emptyTensor})
641 .result();
642
643 FailureOr<int64_t> maybeAZp = op.getAZeroPoint();
644 FailureOr<int64_t> maybeBZp = op.getBZeroPoint();
645 if (failed(maybeAZp))
646 return rewriter.notifyMatchFailure(
647 op, "input a zero point cannot be statically determined");
648 if (failed(maybeBZp))
649 return rewriter.notifyMatchFailure(
650 op, "input b zero point cannot be statically determined");
651
652 const int64_t aZpVal = *maybeAZp;
653 const int64_t bZpVal = *maybeBZp;
654
655 if (op.verifyAZeroPoint(aZpVal).failed())
656 return rewriter.notifyMatchFailure(
657 op, "input a zero point must be zero for non-int8 integer types");
658
659 if (op.verifyBZeroPoint(bZpVal).failed())
660 return rewriter.notifyMatchFailure(
661 op, "input b zero point must be zero for non-int8 integer types");
662
663 if (aZpVal == 0 && bZpVal == 0) {
664 rewriter.replaceOpWithNewOp<linalg::BatchMatmulOp>(
665 op, TypeRange{op.getType()},
666 ValueRange{adaptor.getA(), adaptor.getB()}, ValueRange{zeroTensor});
667 return success();
668 }
669
670 auto aZp = arith::ConstantOp::create(rewriter, loc,
671 rewriter.getI32IntegerAttr(aZpVal));
672 auto bZp = arith::ConstantOp::create(rewriter, loc,
673 rewriter.getI32IntegerAttr(bZpVal));
674 rewriter.replaceOpWithNewOp<linalg::QuantizedBatchMatmulOp>(
675 op, TypeRange{op.getType()},
676 ValueRange{adaptor.getA(), adaptor.getB(), aZp, bZp}, zeroTensor);
677
678 return success();
679 }
680};
681
682class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
683public:
684 using OpConversionPattern::OpConversionPattern;
685
686 // Compute the dynamic output sizes of the maxpool operation.
687 static SmallVector<Value>
688 computeDynamicOutputSizes(tosa::MaxPool2dOp op, OpAdaptor adaptor,
689 ConversionPatternRewriter &rewriter) {
690 TensorType resultTy = op.getType();
691 Location loc = op.getLoc();
692
693 Value input = adaptor.getInput();
694 ArrayRef<int64_t> kernel = op.getKernel();
695 ArrayRef<int64_t> pad = op.getPad();
696 ArrayRef<int64_t> stride = op.getStride();
697
698 SmallVector<Value> dynamicDims;
699
700 // Batch dimension
701 if (resultTy.isDynamicDim(0))
702 dynamicDims.push_back(tensor::DimOp::create(rewriter, loc, input, 0));
703
704 // Height/width dimensions
705 for (int64_t dim : {1, 2}) {
706 if (!resultTy.isDynamicDim(dim))
707 continue;
708
709 // Index into the attribute arrays
710 int64_t index = dim - 1;
711
712 // Input height/width
713 Value ihw = tensor::DimOp::create(rewriter, loc, input, dim);
714
715 // Kernel height/width
716 Value khw = arith::ConstantIndexOp::create(rewriter, loc, kernel[index]);
717
718 // Output height/width
719 Value ohw = getConvOrPoolOutputDim(loc, ihw, pad[index * 2],
720 pad[index * 2 + 1], khw, stride[index],
721 /*dilationAttr=*/1, rewriter);
722 dynamicDims.push_back(ohw);
723 }
724
725 // Channel dimension
726 if (resultTy.isDynamicDim(3))
727 dynamicDims.push_back(tensor::DimOp::create(rewriter, loc, input, 3));
728
729 return dynamicDims;
730 }
731
732 LogicalResult
733 matchAndRewrite(tosa::MaxPool2dOp op, OpAdaptor adaptor,
734 ConversionPatternRewriter &rewriter) const final {
735 Location loc = op.getLoc();
736 Value input = adaptor.getInput();
737 ShapedType inputTy = cast<ShapedType>(input.getType());
738
739 bool isUnsigned = op.getType().getElementType().isUnsignedInteger();
740 ShapedType resultTy =
741 getTypeConverter()->convertType<ShapedType>(op.getType());
742 if (!resultTy)
743 return rewriter.notifyMatchFailure(op, "failed to convert type");
744 Type resultETy = inputTy.getElementType();
745
746 SmallVector<Value> dynamicDims =
747 computeDynamicOutputSizes(op, adaptor, rewriter);
748
749 // Determine what the initial value needs to be for the max pool op.
750 TypedAttr initialAttr;
751 if (resultETy.isF32() || resultETy.isBF16() || resultETy.isF16())
752 initialAttr = rewriter.getFloatAttr(
753 resultETy, APFloat::getLargest(
754 cast<FloatType>(resultETy).getFloatSemantics(), true));
755
756 else if (isUnsigned)
757 initialAttr = rewriter.getIntegerAttr(
758 resultETy, APInt::getZero(resultETy.getIntOrFloatBitWidth()));
759 else if (isa<IntegerType>(resultETy))
760 initialAttr = rewriter.getIntegerAttr(
761 resultETy,
762 APInt::getSignedMinValue(resultETy.getIntOrFloatBitWidth()));
763
764 if (!initialAttr)
765 return rewriter.notifyMatchFailure(
766 op, "Unsupported initial value for tosa.maxpool_2d op");
767
768 // Apply padding as necessary.
769 llvm::SmallVector<int64_t> pad;
770 pad.resize(2, 0);
771 llvm::append_range(pad, op.getPad());
772 pad.resize(pad.size() + 2, 0);
773
774 Value paddedInput = applyPad(loc, input, pad, initialAttr, rewriter);
775
776 Value initialValue = arith::ConstantOp::create(rewriter, loc, initialAttr);
777
778 ArrayRef<int64_t> kernel = op.getKernel();
779 ArrayRef<int64_t> stride = op.getStride();
780
781 Attribute strideAttr = rewriter.getI64VectorAttr(stride);
782 Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
783
784 // Create the linalg op that performs pooling.
785 Value emptyTensor =
786 tensor::EmptyOp::create(rewriter, loc, resultTy.getShape(),
787 resultTy.getElementType(), dynamicDims);
788
789 Value filledEmptyTensor =
790 linalg::FillOp::create(rewriter, loc, initialValue, emptyTensor)
791 .result();
792
793 Value fakeWindowDims =
794 tensor::EmptyOp::create(rewriter, loc, kernel, resultETy);
795
796 if (isUnsigned) {
797 rewriter.replaceOpWithNewOp<linalg::PoolingNhwcMaxUnsignedOp>(
798 op, ArrayRef<Type>{resultTy}, ValueRange{paddedInput, fakeWindowDims},
799 filledEmptyTensor, strideAttr, dilationAttr);
800 return llvm::success();
801 }
802
803 auto resultOp = linalg::PoolingNhwcMaxOp::create(
804 rewriter, op->getLoc(), ArrayRef<Type>{resultTy},
805 ValueRange{paddedInput, fakeWindowDims}, filledEmptyTensor, strideAttr,
806 dilationAttr);
807
808 NanPropagationMode nanMode = op.getNanMode();
809 rewriter.replaceOp(op, resultOp);
810
811 // NaN propagation has no meaning for non floating point types.
812 if (!isa<FloatType>(getElementTypeOrSelf(inputTy)))
813 return success();
814
815 // "PROPAGATE" mode matches the behaviour of the LinAlg named op, so no
816 // compare and select materialization is required.
817 //
818 // In the case of "IGNORE" we need to insert a compare and select. Since
819 // we've already produced a named op we will just take its body and modify
820 // it to include the appropriate checks. If the current value is NaN the
821 // old value of pool will be taken otherwise we use the result.
822 if (nanMode == NanPropagationMode::IGNORE) {
823 auto genericOp = linalg::GenericOp::create(
824 rewriter, loc, resultOp.getType(0), resultOp.getInputs(),
825 resultOp.getOutputs(), resultOp.getIndexingMapsArray(),
826 resultOp.getIteratorTypesArray(),
827 [&](OpBuilder &opBuilder, Location loc, ValueRange blockArgs) {
828 IRMapping map;
829 auto oldBlock = resultOp.getRegion().begin();
830 auto oldArgs = oldBlock->getArguments();
831 auto &oldMaxOp = *resultOp.getBlock()->begin();
832 map.map(oldArgs, blockArgs);
833 auto *newOp = opBuilder.clone(oldMaxOp, map);
834 Value isNaN =
835 arith::CmpFOp::create(opBuilder, loc, arith::CmpFPredicate::UNO,
836 blockArgs.front(), blockArgs.front());
837 auto selectOp = arith::SelectOp::create(
838 opBuilder, loc, isNaN, blockArgs.back(), newOp->getResult(0));
839 linalg::YieldOp::create(opBuilder, loc, selectOp.getResult());
840 });
841 rewriter.replaceOp(resultOp, genericOp);
842 }
843
844 return success();
845 }
846};
847
848class AvgPool2dConverter : public OpRewritePattern<tosa::AvgPool2dOp> {
849public:
850 using OpRewritePattern<tosa::AvgPool2dOp>::OpRewritePattern;
851
852 LogicalResult matchAndRewrite(tosa::AvgPool2dOp op,
853 PatternRewriter &rewriter) const final {
854 Location loc = op.getLoc();
855 Value input = op.getInput();
856 ShapedType inputTy = cast<ShapedType>(input.getType());
857 Type inElementTy = inputTy.getElementType();
858
859 ShapedType resultTy = cast<ShapedType>(op.getType());
860 Type resultETy = cast<ShapedType>(op.getType()).getElementType();
861
862 Type accETy = op.getAccType();
863 ShapedType accTy = resultTy.clone(accETy);
864
865 auto dynamicDimsOr =
866 checkHasDynamicBatchDims(rewriter, op, {input, op.getOutput()});
867 if (!dynamicDimsOr.has_value())
868 return failure();
869 SmallVector<Value> dynamicDims = *dynamicDimsOr;
870
871 FailureOr<int64_t> maybeIZp = op.getInputZeroPoint();
872 FailureOr<int64_t> maybeOZp = op.getOutputZeroPoint();
873 if (failed(maybeIZp))
874 return rewriter.notifyMatchFailure(
875 op, "input zero point could not be statically determined");
876 if (failed(maybeOZp))
877 return rewriter.notifyMatchFailure(
878 op, "output zero point could not be statically determined");
879
880 const int64_t inputZpVal = *maybeIZp;
881 const int64_t outputZpVal = *maybeOZp;
882
883 // Apply padding as necessary.
884 llvm::SmallVector<int64_t> pad;
885 pad.resize(2, 0);
886 llvm::append_range(pad, op.getPad());
887 pad.resize(pad.size() + 2, 0);
888 TypedAttr padAttr = rewriter.getZeroAttr(inElementTy);
889 // Unsupported element type
890 if (!padAttr)
891 return failure();
892 Value paddedInput = applyPad(loc, input, pad, padAttr, rewriter);
893
894 auto initialAttr = rewriter.getZeroAttr(accETy);
895 Value initialValue = arith::ConstantOp::create(rewriter, loc, initialAttr);
896
897 ArrayRef<int64_t> kernel = op.getKernel();
898 ArrayRef<int64_t> stride = op.getStride();
899
900 Attribute strideAttr = rewriter.getI64VectorAttr(stride);
901 Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
902
903 // Create the linalg op that performs pooling.
904 Value poolEmptyTensor = tensor::EmptyOp::create(
905 rewriter, loc, accTy.getShape(), accETy, dynamicDims);
906
907 Value filledEmptyTensor =
908 linalg::FillOp::create(rewriter, loc, ValueRange{initialValue},
909 ValueRange{poolEmptyTensor})
910 .result();
911
912 Value fakeWindowDims =
913 tensor::EmptyOp::create(rewriter, loc, kernel, accETy);
914
915 // Sum across the pooled region.
916 Value poolingOp = linalg::PoolingNhwcSumOp::create(
917 rewriter, loc, ArrayRef<Type>{accTy},
918 ValueRange{paddedInput, fakeWindowDims},
919 filledEmptyTensor, strideAttr, dilationAttr)
920 .getResult(0);
921
922 // Normalize the summed value by the number of elements grouped in each
923 // pool.
924 Value iH = tensor::DimOp::create(rewriter, loc, poolingOp, 1);
925 Value iW = tensor::DimOp::create(rewriter, loc, poolingOp, 2);
926
927 auto one = arith::ConstantIndexOp::create(rewriter, loc, 1);
928 iH = arith::SubIOp::create(rewriter, loc, iH, one);
929 iW = arith::SubIOp::create(rewriter, loc, iW, one);
930
931 Value genericEmptyTensor = tensor::EmptyOp::create(
932 rewriter, loc, resultTy.getShape(), resultETy, dynamicDims);
933
934 auto affineMap = rewriter.getMultiDimIdentityMap(resultTy.getRank());
935 auto genericOp = linalg::GenericOp::create(
936 rewriter, loc, ArrayRef<Type>({resultTy}), ValueRange{poolingOp},
937 ValueRange{genericEmptyTensor},
938 ArrayRef<AffineMap>({affineMap, affineMap}),
939 getNParallelLoopsAttrs(resultTy.getRank()),
940 [&](OpBuilder &b, Location loc, ValueRange args) {
941 auto zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
942
943 // Determines what the portion of valid input is covered by the
944 // kernel.
945 auto padFn = [&](Value valid, Value pos, int64_t pad) -> Value {
946 if (pad == 0)
947 return valid;
948
949 auto padVal = arith::ConstantIndexOp::create(rewriter, loc, pad);
950 Value dpos = arith::SubIOp::create(rewriter, loc, pos, padVal);
951
952 Value offset = arith::MinSIOp::create(rewriter, loc, dpos, zero);
953 return arith::AddIOp::create(rewriter, loc, valid, offset)
954 ->getResult(0);
955 };
956
957 auto coverageFn = [&](int64_t i, Value isize) -> Value {
958 Value strideVal =
959 arith::ConstantIndexOp::create(rewriter, loc, stride[i - 1]);
960 Value val =
961 arith::ConstantIndexOp::create(rewriter, loc, kernel[i - 1]);
962
963 // Find the position relative to the input tensor's ends.
964 Value left = linalg::IndexOp::create(rewriter, loc, i);
965 Value right = arith::SubIOp::create(rewriter, loc, isize, left);
966 left = arith::MulIOp::create(rewriter, loc, left, strideVal);
967 right = arith::MulIOp::create(rewriter, loc, right, strideVal);
968
969 // Determine how much padding was included.
970 val = padFn(val, left, pad[i * 2]);
971 val = padFn(val, right, pad[i * 2 + 1]);
972 return arith::MaxSIOp::create(rewriter, loc, one, val);
973 };
974
975 // Compute the indices from either end.
976 Value kH3 = coverageFn(1, iH);
977 Value kW3 = coverageFn(2, iW);
978
979 // Compute the total number of elements and normalize.
980 auto count = arith::IndexCastOp::create(
981 rewriter, loc, rewriter.getI32Type(),
982 arith::MulIOp::create(rewriter, loc, kH3, kW3));
983
984 // Divide by the number of summed values. For floats this is just
985 // a div however for quantized values input normalization had
986 // to be applied.
987 Value poolVal = args[0];
988 if (isa<FloatType>(accETy)) {
989 auto countF = arith::SIToFPOp::create(rewriter, loc, accETy, count);
990 poolVal = arith::DivFOp::create(rewriter, loc, poolVal, countF)
991 ->getResult(0);
992 if (accETy.getIntOrFloatBitWidth() >
993 resultETy.getIntOrFloatBitWidth())
994 poolVal =
995 arith::TruncFOp::create(rewriter, loc, resultETy, poolVal);
996 } else {
997
998 // If we have quantization information we need to apply an offset
999 // for the input zp value.
1000 if (inputZpVal != 0) {
1001 auto inputZp = arith::ConstantOp::create(
1002 rewriter, loc, b.getIntegerAttr(accETy, inputZpVal));
1003 Value offset =
1004 arith::MulIOp::create(rewriter, loc, accETy, count, inputZp);
1005 poolVal =
1006 arith::SubIOp::create(rewriter, loc, accETy, poolVal, offset);
1007 }
1008
1009 // Compute: k = 32 - count_leading_zeros(value - 1)
1010 Value one32 = arith::ConstantOp::create(
1011 rewriter, loc, rewriter.getI32IntegerAttr(1));
1012 Value thirtyTwo32 = arith::ConstantOp::create(
1013 rewriter, loc, rewriter.getI32IntegerAttr(32));
1014
1015 Value countSubOne =
1016 arith::SubIOp::create(rewriter, loc, count, one32);
1017 Value leadingZeros =
1018 math::CountLeadingZerosOp::create(rewriter, loc, countSubOne);
1019 Value k =
1020 arith::SubIOp::create(rewriter, loc, thirtyTwo32, leadingZeros);
1021
1022 // Compute: numerator = ((1 << 30) + 1) << k
1023 Value k64 =
1024 arith::ExtUIOp::create(rewriter, loc, rewriter.getI64Type(), k);
1025 Value thirtyShiftPlusOne = arith::ConstantOp::create(
1026 rewriter, loc, rewriter.getI64IntegerAttr((1 << 30) + 1));
1027 Value numerator =
1028 arith::ShLIOp::create(rewriter, loc, thirtyShiftPlusOne, k64);
1029
1030 // Compute: scale.multiplier = numerator / value;
1031 Value count64 = arith::ExtUIOp::create(
1032 rewriter, loc, rewriter.getI64Type(), count);
1033 Value multiplier =
1034 arith::DivUIOp::create(rewriter, loc, numerator, count64);
1035 multiplier = arith::TruncIOp::create(
1036 rewriter, loc, rewriter.getI32Type(), multiplier);
1037
1038 // Compute: scale.shift = 30 + k
1039 Value k8 =
1040 arith::TruncIOp::create(rewriter, loc, rewriter.getI8Type(), k);
1041 Value thirty8 = arith::ConstantOp::create(
1042 rewriter, loc, rewriter.getI8IntegerAttr(30));
1043 Value shift = arith::AddIOp::create(rewriter, loc, k8, thirty8);
1044
1045 auto roundingAttr = RoundingModeAttr::get(
1046 rewriter.getContext(), RoundingMode::SINGLE_ROUND);
1047
1048 auto scaled = tosa::ApplyScaleOp::create(
1049 rewriter, loc, rewriter.getI32Type(), poolVal,
1050 multiplier, shift, roundingAttr)
1051 .getResult();
1052
1053 // If we have quantization information we need to apply output
1054 // zeropoint.
1055 if (outputZpVal != 0) {
1056 auto outputZp = arith::ConstantOp::create(
1057 rewriter, loc,
1058 b.getIntegerAttr(scaled.getType(), outputZpVal));
1059 scaled = arith::AddIOp::create(rewriter, loc, scaled, outputZp)
1060 .getResult();
1061 }
1062
1063 // Apply Clip.
1064 int64_t outBitwidth = resultETy.getIntOrFloatBitWidth();
1065
1067 rewriter, loc, accETy,
1068 APInt::getSignedMinValue(outBitwidth).getSExtValue());
1070 rewriter, loc, accETy,
1071 APInt::getSignedMaxValue(outBitwidth).getSExtValue());
1072 auto clamp = clampIntHelper(loc, scaled, min, max, rewriter,
1073 /*isUnsigned=*/false);
1074
1075 poolVal = clamp;
1076 // Convert type.
1077 if (resultETy != clamp.getType()) {
1078 poolVal =
1079 arith::TruncIOp::create(rewriter, loc, resultETy, poolVal);
1080 }
1081 }
1082
1083 linalg::YieldOp::create(rewriter, loc, poolVal);
1084 });
1085
1086 rewriter.replaceOp(op, genericOp.getResult(0));
1087 return success();
1088 }
1089};
1090
1091class TransposeConverter : public OpRewritePattern<tosa::TransposeOp> {
1092public:
1093 using OpRewritePattern<tosa::TransposeOp>::OpRewritePattern;
1094
1095 LogicalResult matchAndRewrite(tosa::TransposeOp op,
1096 PatternRewriter &rewriter) const final {
1097 const llvm::ArrayRef<int32_t> constantPerms = op.getPerms();
1098
1099 Location loc = op.getLoc();
1100 // The verifier should have made sure we have a valid TOSA permutation
1101 // tensor. isPermutationVector doesn't actually check the TOSA perms we
1102 // expect.
1103 SmallVector<OpFoldResult> inputSizes =
1104 tensor::getMixedSizes(rewriter, loc, op.getInput1());
1105 auto permutedSizes =
1106 applyTOSAPermutation<OpFoldResult>(inputSizes, constantPerms);
1107
1108 auto permutedInit =
1109 tensor::EmptyOp::create(rewriter, loc, permutedSizes,
1110 op.getInput1().getType().getElementType());
1111 rewriter.replaceOpWithNewOp<linalg::TransposeOp>(
1112 op, op.getInput1(), permutedInit,
1113 llvm::map_to_vector(constantPerms,
1114 [](int32_t v) -> int64_t { return v; }));
1115 return success();
1116 }
1117};
1118} // namespace
1119
1121 const TypeConverter &converter, RewritePatternSet *patterns,
1122 const TosaToLinalgNamedOptions &options) {
1123 if (options.preferConv2DKernelLayoutHWCF) {
1124 patterns->add<ConvConverter<tosa::Conv2DOp, linalg::Conv2DNhwcHwcfOp,
1125 linalg::Conv2DNhwcHwcfQOp>>(
1126 patterns->getContext());
1127 } else {
1128 patterns->add<ConvConverter<tosa::Conv2DOp, linalg::Conv2DNhwcFhwcOp,
1129 linalg::Conv2DNhwcFhwcQOp>>(
1130 patterns->getContext());
1131 }
1132 patterns->add<
1133 // clang-format off
1134 ConvConverter<tosa::Conv3DOp, linalg::Conv3DNdhwcDhwcfOp, linalg::Conv3DNdhwcDhwcfQOp>,
1135 DepthwiseConvConverter,
1136 MatMulConverter,
1137 AvgPool2dConverter,
1138 TransposeConverter
1139 >(patterns->getContext());
1140
1141 patterns->add<
1142 MaxPool2dConverter
1143 >(converter, patterns->getContext());
1144 // clang-format on
1145}
return success()
b
Return true if permutation is a valid permutation of the outer_dims_perm (case OuterOrInnerPerm::Oute...
static llvm::ManagedStatic< PassManagerOptions > options
static Value clamp(ImplicitLocOpBuilder &builder, Value value, Value lowerBound, Value upperBound)
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
static AffineMap getBroadcastingMap(PatternRewriter &rewriter, Value source, Value result)
static mlir::Value applyPad(Location loc, Value input, ArrayRef< int64_t > pad, TypedAttr padAttr, OpBuilder &rewriter)
static void createDepthwiseConvCollapseMap(int64_t outputRank, SmallVector< ReassociationExprs, 4 > &reassociationMap, OpBuilder &rewriter)
static mlir::Value linalgIntBroadcastExtSIAdd(PatternRewriter &rewriter, Location loc, Value bias, Value conv, Value result, ArrayRef< AffineMap > indexingMaps)
static mlir::Value getConvOrPoolOutputDim(Location loc, Value inputDim, int64_t padBeforeAttr, int64_t padAfterAttr, Value kernelDim, int64_t strideAttr, int64_t dilationAttr, OpBuilder &rewriter)
static mlir::Value linalgBroadcastAndMaybeExt(PatternRewriter &rewriter, Location loc, Value source, Value result)
static mlir::Value reifyConstantDim(int64_t attr, ImplicitLocOpBuilder &builder)
static SmallVector< Value > inferDynamicDimsForConv(Location loc, Value input, Value weight, ShapedType resultTy, ArrayRef< int64_t > padAttr, ArrayRef< int64_t > strideAttr, ArrayRef< int64_t > dilationAttr, ArrayRef< int64_t > inputSizeDims, ArrayRef< int64_t > kernelSizeDims, OpBuilder &rewriter)
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition AffineMap.h:46
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
IntegerAttr getIndexAttr(int64_t value)
Definition Builders.cpp:116
AffineMap getMultiDimIdentityMap(unsigned rank)
Definition Builders.cpp:396
AffineExpr getAffineConstantExpr(int64_t constant)
Definition Builders.cpp:381
AffineExpr getAffineDimExpr(unsigned position)
Definition Builders.cpp:373
MLIRContext * getContext() const
Definition Builders.h:56
ImplicitLocOpBuilder maintains a 'current location', allowing use of the create<> method without spec...
Definition Builders.h:632
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition Location.h:76
This class helps build Operations.
Definition Builders.h:210
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
MLIRContext * getContext() const
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
This class provides an abstraction over the various different ranges of value types.
Definition TypeRange.h:40
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition Types.h:74
bool isF32() const
Definition Types.cpp:40
bool isUnsignedInteger() const
Return true if this is an unsigned integer type (with the specified width).
Definition Types.cpp:90
bool isF16() const
Definition Types.cpp:38
unsigned getIntOrFloatBitWidth() const
Return the bit width of an integer or a float type, assert failure on other types.
Definition Types.cpp:124
bool isBF16() const
Definition Types.cpp:37
This class provides an abstraction over the different types of ranges over Values.
Definition ValueRange.h:389
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition Value.h:96
Type getType() const
Return the type of this value.
Definition Value.h:105
static ConstantIndexOp create(OpBuilder &builder, Location location, int64_t value)
Definition ArithOps.cpp:398
static ConstantIntOp create(OpBuilder &builder, Location location, int64_t value, unsigned width)
Definition ArithOps.cpp:297
detail::InFlightRemark failed(Location loc, RemarkOpts opts)
Report an optimization remark that failed.
Definition Remarks.h:732
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given tensor value.
Definition TensorOps.cpp:91
SmallVector< T > applyTOSAPermutation(ArrayRef< T > input, ArrayRef< int32_t > perms)
SmallVector< utils::IteratorType > getNParallelLoopsAttrs(unsigned nParallelLoops)
SmallVector< Value > condenseValues(const SmallVector< Value > &values)
std::optional< SmallVector< Value > > checkHasDynamicBatchDims(PatternRewriter &rewriter, Op op, ArrayRef< Value > params)
Value clampIntHelper(Location loc, Value arg, Value min, Value max, OpBuilder &rewriter, bool isUnsigned)
void populateTosaToLinalgNamedConversionPatterns(const TypeConverter &converter, RewritePatternSet *patterns, const TosaToLinalgNamedOptions &options)
Populates conversion passes from TOSA dialect to Linalg named operations.
Include the generated interface declarations.
detail::DenseArrayAttrImpl< int64_t > DenseI64ArrayAttr
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...