23#include "llvm/ADT/SmallVectorExtras.h"
33 if (llvm::all_of(pad, [](
int64_t p) {
return p == 0; }))
36 ShapedType inputTy = cast<ShapedType>(input.
getType());
37 Type inputETy = inputTy.getElementType();
38 auto inputShape = inputTy.getShape();
40 assert((inputShape.size() * 2) == pad.size());
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]);
51 paddedShape.push_back(inputShape[i] + highPad + lowPad);
56 Value padValue = arith::ConstantOp::create(rewriter, loc, padAttr);
58 return tensor::PadOp::create(rewriter, loc,
59 RankedTensorType::get(paddedShape, inputETy),
60 input, lowIndices, highIndices, padValue);
67 ShapedType resultTy = cast<ShapedType>(conv.
getType());
68 return linalg::GenericOp::create(
72 Value biasVal = args[0];
73 Type resType = args[1].getType();
74 if (resType != biasVal.
getType()) {
76 arith::ExtSIOp::create(builder, loc, resType, biasVal);
79 arith::AddIOp::create(builder, loc, biasVal, args[1]);
80 linalg::YieldOp::create(builder, loc, added);
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();
100 assert(sourceTy.hasStaticShape() &&
101 "Dynamic broadcasting shapes not supported!");
102 if (sourceRank == 1 && sourceTy.getDimSize(0) == 1) {
105 for (
auto dim : llvm::seq<int64_t>(0, sourceRank)) {
107 sourceDims.push_back(expr);
121 ShapedType resultTy = cast<ShapedType>(
result.getType());
122 const int64_t resultRank = resultTy.getRank();
129 return linalg::GenericOp::create(
133 Value biasVal = args[0];
134 Type resType = args[1].getType();
135 if (resType != biasVal.
getType()) {
137 resultTy.getElementType().isFloat()
138 ? arith::ExtFOp::create(
140 ValueRange{biasVal}, arith::ExtFOp::Properties{})
142 : arith::ExtSIOp::create(builder, loc, resType,
146 linalg::YieldOp::create(builder, loc, biasVal);
167 auto one = arith::ConstantOp::create(rewriter, loc,
168 IntegerAttr::get(inputDim.
getType(), 1));
170 Value paddedBefore = arith::AddIOp::create(builder, inputDim, padBefore);
172 Value paddedAfter = arith::AddIOp::create(builder, paddedBefore, padAfter);
174 Value subOne = arith::SubIOp::create(builder, kernelDim, one);
176 Value dilated = arith::MulIOp::create(builder, dilation, subOne);
177 Value addOne = arith::AddIOp::create(builder, dilated, one);
179 Value subtract = arith::SubIOp::create(builder, paddedAfter, addOne);
181 Value divide = arith::DivUIOp::create(builder, subtract, stride);
182 return arith::AddIOp::create(builder, divide, one);
191 ShapedType inputTy = cast<ShapedType>(input.
getType());
192 int64_t inputRank = inputTy.getRank();
195 dynDims.resize(resultTy.getRank());
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);
207 tensor::DimOp::create(rewriter, loc, weight, kernelDim);
211 kernelDynDim, stride, dilation, rewriter);
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);
230 reassociationMap.resize(outputRank);
231 for (
int i = 0; i < outputRank; i++) {
234 reassociationMap[outputRank - 1].push_back(
240template <
typename TosaConvOp,
typename LinalgConvOp,
typename LinalgConvQOp>
241class ConvConverter :
public OpConversionPattern<TosaConvOp> {
243 using OpConversionPattern<TosaConvOp>::OpConversionPattern;
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);
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());
257 Type inputETy = inputTy.getElementType();
263 Type accETy = op.getAccType();
264 Type accTy = RankedTensorType::get(resultTy.getShape(), accETy);
267 FailureOr<int64_t> maybeIZp = op.getInputZeroPoint();
269 return rewriter.notifyMatchFailure(
270 op,
"input zero point cannot be statically determined");
272 FailureOr<int64_t> maybeWZp = op.getWeightZeroPoint();
274 return rewriter.notifyMatchFailure(
275 op,
"weight zero point cannot be statically determined");
277 const int64_t inputZpVal = *maybeIZp;
278 const int64_t weightZpVal = *maybeWZp;
280 if (op.verifyInputZeroPoint(inputZpVal).failed())
281 return rewriter.notifyMatchFailure(
282 op,
"input zero point must be zero for non-int8 integer types");
284 if (op.verifyWeightZeroPoint(weightZpVal).failed())
285 return rewriter.notifyMatchFailure(
286 op,
"weight zero point must be zero for non-int8 integer types");
288 bool hasZp = (inputZpVal != 0) || (weightZpVal != 0);
290 if (!weightTy.hasStaticShape() || !biasTy.hasStaticShape())
291 return rewriter.notifyMatchFailure(
292 op,
"tosa.conv ops require static shapes for weight and bias");
295 return rewriter.notifyMatchFailure(
296 op,
"tosa.conv ops does not support unsigned integer input");
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);
306 loc, input, weight, resultTy, padAttr.
asArrayRef(),
308 inputSizeDims, kernelSizeDims, rewriter);
310 auto weightShape = weightTy.getShape();
313 TypedAttr zeroAttr = rewriter.getZeroAttr(inputETy);
322 if (inputZpVal < intMin || inputZpVal > intMax)
323 return rewriter.notifyMatchFailure(
324 op,
"tosa.conv op quantization has zp outside of input range");
326 zeroAttr = rewriter.getIntegerAttr(inputETy, inputZpVal);
329 llvm::SmallVector<int64_t> pad;
331 llvm::append_range(pad, padAttr.
asArrayRef());
332 pad.resize(pad.size() + 2, 0);
333 input =
applyPad(loc, input, pad, zeroAttr, rewriter);
335 if (4 == inputTy.getRank()) {
339 hasZp ? std::is_same_v<LinalgConvQOp, linalg::Conv2DNhwcHwcfQOp>
340 : std::is_same_v<LinalgConvOp, linalg::Conv2DNhwcHwcfOp>;
346 SmallVector<int32_t> weightPerm;
347 for (
int i = 1; i < resultTy.getRank(); i++)
348 weightPerm.push_back(i);
349 weightPerm.push_back(0);
351 SmallVector<int64_t> newWeightShape;
352 for (
auto dim : weightPerm)
353 newWeightShape.push_back(weightShape[dim]);
354 auto weightPermAttr = rewriter.getDenseI32ArrayAttr(weightPerm);
356 RankedTensorType::get(newWeightShape, weightTy.getElementType());
357 weight = tosa::TransposeOp::create(rewriter, loc, newWeightTy, weight,
365 if (5 == inputTy.getRank()) {
368 SmallVector<int32_t> weightPerm;
369 for (
int i = 1; i < resultTy.getRank(); i++)
370 weightPerm.push_back(i);
371 weightPerm.push_back(0);
373 SmallVector<int64_t> newWeightShape;
374 for (
auto dim : weightPerm)
375 newWeightShape.push_back(weightShape[dim]);
376 auto weightPermAttr = rewriter.getDenseI32ArrayAttr(weightPerm);
378 RankedTensorType::get(newWeightShape, weightTy.getElementType());
379 weight = tosa::TransposeOp::create(rewriter, loc, newWeightTy, weight,
384 ArrayRef<int64_t> stride = strideTosaAttr;
385 ArrayRef<int64_t> dilation = dilationTosaAttr;
388 auto strideAttr = rewriter.getI64TensorAttr(stride);
389 auto dilationAttr = rewriter.getI64TensorAttr(dilation);
391 Value biasEmptyTensor = tensor::EmptyOp::create(
392 rewriter, loc, resultTy.getShape(), accETy, filteredDims);
394 Value broadcastBias =
398 auto iZp = rewriter.getI32IntegerAttr(inputZpVal);
399 auto kZp = rewriter.getI32IntegerAttr(weightZpVal);
401 auto iZpVal = arith::ConstantOp::create(rewriter, loc, iZp);
402 auto kZpVal = arith::ConstantOp::create(rewriter, loc, kZp);
404 Value conv = LinalgConvQOp::create(
405 rewriter, loc, resultTy,
407 ValueRange{broadcastBias}, strideAttr, dilationAttr)
410 rewriter.replaceOp(op, conv);
414 Value conv = LinalgConvOp::create(
415 rewriter, loc, accTy,
ValueRange{input, weight},
416 ValueRange{broadcastBias}, strideAttr, dilationAttr)
421 if (resultTy != accTy)
422 conv = tosa::CastOp::create(rewriter, loc, resultTy, conv);
424 rewriter.replaceOp(op, conv);
429class DepthwiseConvConverter
430 :
public OpConversionPattern<tosa::DepthwiseConv2DOp> {
432 using OpConversionPattern<tosa::DepthwiseConv2DOp>::OpConversionPattern;
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);
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();
447 Type inputETy = inputTy.getElementType();
448 Type resultETy = resultTy.getElementType();
450 auto padAttr = op.getPadAttr();
451 auto strideTosaAttr = op.getStrideAttr();
452 auto dilationTosaAttr = op.getDilationAttr();
454 Type accETy = op.getAccType();
456 if (!weightTy.hasStaticShape() || !biasTy.hasStaticShape())
457 return rewriter.notifyMatchFailure(
458 op,
"tosa.depthwise_conv ops require static shapes");
462 loc, input, weight, resultTy, padAttr.
asArrayRef(),
469 FailureOr<int64_t> maybeIZp = op.getInputZeroPoint();
470 FailureOr<int64_t> maybeWZp = op.getWeightZeroPoint();
472 return rewriter.notifyMatchFailure(
473 op,
"input zero point cannot be statically determined");
475 return rewriter.notifyMatchFailure(
476 op,
"weight zero point cannot be statically determined");
478 const int64_t inputZpVal = *maybeIZp;
479 const int64_t weightZpVal = *maybeWZp;
481 if (op.verifyInputZeroPoint(inputZpVal).failed())
482 return rewriter.notifyMatchFailure(
483 op,
"input zero point must be zero for non-int8 integer types");
485 if (op.verifyWeightZeroPoint(weightZpVal).failed())
486 return rewriter.notifyMatchFailure(
487 op,
"weight zero point must be zero for non-int8 integer types");
489 bool hasNullZps = (inputZpVal == 0) && (weightZpVal == 0);
490 auto weightShape = weightTy.getShape();
491 auto resultShape = resultTy.getShape();
494 TypedAttr zeroAttr = rewriter.getZeroAttr(inputETy);
503 if (inputZpVal < intMin || inputZpVal > intMax)
504 return rewriter.notifyMatchFailure(
505 op,
"tosa.depthwise_conv op quantization has zp outside of input "
508 zeroAttr = rewriter.getIntegerAttr(inputETy, inputZpVal);
511 llvm::SmallVector<int64_t> pad;
513 llvm::append_range(pad, padAttr.
asArrayRef());
514 pad.resize(pad.size() + 2, 0);
516 input =
applyPad(loc, input, pad, zeroAttr, rewriter);
519 ArrayRef<int64_t> stride = strideTosaAttr;
520 ArrayRef<int64_t> dilation = dilationTosaAttr;
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]},
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},
538 Value biasEmptyTensor = tensor::EmptyOp::create(
539 rewriter, loc, resultTy.getShape(), resultETy, filteredDims);
542 SmallVector<AffineMap, 4> indexingMaps;
544 indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank));
545 indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank));
548 Value conv = linalg::DepthwiseConv2DNhwcHwcmOp::create(
549 rewriter, loc, linalgConvTy,
ValueRange{input, weight},
550 ValueRange{zeroTensor}, strideAttr, dilationAttr)
555 if (accETy != resultETy)
556 conv = tosa::CastOp::create(
558 RankedTensorType::get(cast<ShapedType>(conv.
getType()).getShape(),
562 SmallVector<ReassociationExprs, 4> reassociationMap;
564 Value convReshape = tensor::CollapseShapeOp::create(
565 rewriter, loc, resultTy, conv, reassociationMap);
568 linalg::GenericOp::create(
569 rewriter, loc, resultTy,
ValueRange({bias, convReshape}),
571 [&](OpBuilder &nestedBuilder, Location nestedLoc,
574 if (llvm::isa<FloatType>(inputETy))
575 added = arith::AddFOp::create(nestedBuilder, loc, args[0],
578 added = arith::AddIOp::create(nestedBuilder, loc, args[0],
580 linalg::YieldOp::create(nestedBuilder, nestedLoc, added);
583 rewriter.replaceOp(op,
result);
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,
592 ValueRange{zeroTensor}, strideAttr, dilationAttr)
594 SmallVector<ReassociationExprs, 4> reassociationMap;
596 Value convReshape = tensor::CollapseShapeOp::create(
597 rewriter, loc, resultTy, conv, reassociationMap);
599 rewriter, loc, bias, convReshape, biasEmptyTensor, indexingMaps);
600 rewriter.replaceOp(op,
result);
606class MatMulConverter :
public OpConversionPattern<tosa::MatMulOp> {
608 using OpConversionPattern<tosa::MatMulOp>::OpConversionPattern;
610 matchAndRewrite(tosa::MatMulOp op, OpAdaptor adaptor,
611 ConversionPatternRewriter &rewriter)
const final {
612 Location loc = op.getLoc();
614 auto outputTy = cast<ShapedType>(op.getType());
615 auto outputElementTy = outputTy.getElementType();
617 SmallVector<Value> dynDims;
618 dynDims.resize(cast<ShapedType>(op->getResult(0).getType()).getRank());
620 if (!outputTy.hasRank() || outputTy.isDynamicDim(0)) {
621 dynDims[0] = tensor::DimOp::create(rewriter, loc, op->getOperand(0), 0);
624 if (!outputTy.hasRank() || outputTy.isDynamicDim(1)) {
625 dynDims[1] = tensor::DimOp::create(rewriter, loc, op->getOperand(0), 1);
628 if (!outputTy.hasRank() || outputTy.isDynamicDim(2)) {
629 dynDims[2] = tensor::DimOp::create(rewriter, loc, op->getOperand(1), 2);
634 auto zeroAttr = rewriter.getZeroAttr(outputElementTy);
635 Value zero = arith::ConstantOp::create(rewriter, loc, zeroAttr);
637 tensor::EmptyOp::create(rewriter, loc, outputTy.getShape(),
638 outputTy.getElementType(), filteredDims);
639 Value zeroTensor = linalg::FillOp::create(rewriter, loc,
ValueRange{zero},
643 FailureOr<int64_t> maybeAZp = op.getAZeroPoint();
644 FailureOr<int64_t> maybeBZp = op.getBZeroPoint();
646 return rewriter.notifyMatchFailure(
647 op,
"input a zero point cannot be statically determined");
649 return rewriter.notifyMatchFailure(
650 op,
"input b zero point cannot be statically determined");
652 const int64_t aZpVal = *maybeAZp;
653 const int64_t bZpVal = *maybeBZp;
655 if (op.verifyAZeroPoint(aZpVal).failed())
656 return rewriter.notifyMatchFailure(
657 op,
"input a zero point must be zero for non-int8 integer types");
659 if (op.verifyBZeroPoint(bZpVal).failed())
660 return rewriter.notifyMatchFailure(
661 op,
"input b zero point must be zero for non-int8 integer types");
663 if (aZpVal == 0 && bZpVal == 0) {
664 rewriter.replaceOpWithNewOp<linalg::BatchMatmulOp>(
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>(
676 ValueRange{adaptor.getA(), adaptor.getB(), aZp, bZp}, zeroTensor);
682class MaxPool2dConverter :
public OpConversionPattern<tosa::MaxPool2dOp> {
684 using OpConversionPattern::OpConversionPattern;
687 static SmallVector<Value>
688 computeDynamicOutputSizes(tosa::MaxPool2dOp op, OpAdaptor adaptor,
689 ConversionPatternRewriter &rewriter) {
690 TensorType resultTy = op.getType();
691 Location loc = op.getLoc();
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();
698 SmallVector<Value> dynamicDims;
701 if (resultTy.isDynamicDim(0))
702 dynamicDims.push_back(tensor::DimOp::create(rewriter, loc, input, 0));
705 for (int64_t dim : {1, 2}) {
706 if (!resultTy.isDynamicDim(dim))
710 int64_t index = dim - 1;
713 Value ihw = tensor::DimOp::create(rewriter, loc, input, dim);
720 pad[index * 2 + 1], khw, stride[index],
722 dynamicDims.push_back(ohw);
726 if (resultTy.isDynamicDim(3))
727 dynamicDims.push_back(tensor::DimOp::create(rewriter, loc, input, 3));
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());
739 bool isUnsigned = op.getType().getElementType().isUnsignedInteger();
740 ShapedType resultTy =
741 getTypeConverter()->convertType<ShapedType>(op.getType());
743 return rewriter.notifyMatchFailure(op,
"failed to convert type");
744 Type resultETy = inputTy.getElementType();
746 SmallVector<Value> dynamicDims =
747 computeDynamicOutputSizes(op, adaptor, rewriter);
750 TypedAttr initialAttr;
752 initialAttr = rewriter.getFloatAttr(
753 resultETy, APFloat::getLargest(
754 cast<FloatType>(resultETy).getFloatSemantics(),
true));
757 initialAttr = rewriter.getIntegerAttr(
759 else if (isa<IntegerType>(resultETy))
760 initialAttr = rewriter.getIntegerAttr(
765 return rewriter.notifyMatchFailure(
766 op,
"Unsupported initial value for tosa.maxpool_2d op");
769 llvm::SmallVector<int64_t> pad;
771 llvm::append_range(pad, op.getPad());
772 pad.resize(pad.size() + 2, 0);
774 Value paddedInput =
applyPad(loc, input, pad, initialAttr, rewriter);
776 Value initialValue = arith::ConstantOp::create(rewriter, loc, initialAttr);
778 ArrayRef<int64_t> kernel = op.getKernel();
779 ArrayRef<int64_t> stride = op.getStride();
781 Attribute strideAttr = rewriter.getI64VectorAttr(stride);
782 Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
786 tensor::EmptyOp::create(rewriter, loc, resultTy.getShape(),
787 resultTy.getElementType(), dynamicDims);
789 Value filledEmptyTensor =
790 linalg::FillOp::create(rewriter, loc, initialValue, emptyTensor)
793 Value fakeWindowDims =
794 tensor::EmptyOp::create(rewriter, loc, kernel, resultETy);
797 rewriter.replaceOpWithNewOp<linalg::PoolingNhwcMaxUnsignedOp>(
798 op, ArrayRef<Type>{resultTy},
ValueRange{paddedInput, fakeWindowDims},
799 filledEmptyTensor, strideAttr, dilationAttr);
800 return llvm::success();
803 auto resultOp = linalg::PoolingNhwcMaxOp::create(
804 rewriter, op->getLoc(), ArrayRef<Type>{resultTy},
805 ValueRange{paddedInput, fakeWindowDims}, filledEmptyTensor, strideAttr,
808 NanPropagationMode nanMode = op.getNanMode();
809 rewriter.replaceOp(op, resultOp);
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) {
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);
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());
841 rewriter.replaceOp(resultOp, genericOp);
850 using OpRewritePattern<tosa::AvgPool2dOp>::OpRewritePattern;
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();
859 ShapedType resultTy = cast<ShapedType>(op.getType());
860 Type resultETy = cast<ShapedType>(op.getType()).getElementType();
862 Type accETy = op.getAccType();
863 ShapedType accTy = resultTy.clone(accETy);
867 if (!dynamicDimsOr.has_value())
869 SmallVector<Value> dynamicDims = *dynamicDimsOr;
871 FailureOr<int64_t> maybeIZp = op.getInputZeroPoint();
872 FailureOr<int64_t> maybeOZp = op.getOutputZeroPoint();
874 return rewriter.notifyMatchFailure(
875 op,
"input zero point could not be statically determined");
877 return rewriter.notifyMatchFailure(
878 op,
"output zero point could not be statically determined");
880 const int64_t inputZpVal = *maybeIZp;
881 const int64_t outputZpVal = *maybeOZp;
884 llvm::SmallVector<int64_t> pad;
886 llvm::append_range(pad, op.getPad());
887 pad.resize(pad.size() + 2, 0);
888 TypedAttr padAttr = rewriter.getZeroAttr(inElementTy);
892 Value paddedInput =
applyPad(loc, input, pad, padAttr, rewriter);
894 auto initialAttr = rewriter.getZeroAttr(accETy);
895 Value initialValue = arith::ConstantOp::create(rewriter, loc, initialAttr);
897 ArrayRef<int64_t> kernel = op.getKernel();
898 ArrayRef<int64_t> stride = op.getStride();
900 Attribute strideAttr = rewriter.getI64VectorAttr(stride);
901 Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
904 Value poolEmptyTensor = tensor::EmptyOp::create(
905 rewriter, loc, accTy.getShape(), accETy, dynamicDims);
907 Value filledEmptyTensor =
908 linalg::FillOp::create(rewriter, loc,
ValueRange{initialValue},
912 Value fakeWindowDims =
913 tensor::EmptyOp::create(rewriter, loc, kernel, accETy);
916 Value poolingOp = linalg::PoolingNhwcSumOp::create(
917 rewriter, loc, ArrayRef<Type>{accTy},
919 filledEmptyTensor, strideAttr, dilationAttr)
924 Value iH = tensor::DimOp::create(rewriter, loc, poolingOp, 1);
925 Value iW = tensor::DimOp::create(rewriter, loc, poolingOp, 2);
928 iH = arith::SubIOp::create(rewriter, loc, iH, one);
929 iW = arith::SubIOp::create(rewriter, loc, iW, one);
931 Value genericEmptyTensor = tensor::EmptyOp::create(
932 rewriter, loc, resultTy.getShape(), resultETy, dynamicDims);
934 auto affineMap = rewriter.getMultiDimIdentityMap(resultTy.getRank());
935 auto genericOp = linalg::GenericOp::create(
936 rewriter, loc, ArrayRef<Type>({resultTy}),
ValueRange{poolingOp},
938 ArrayRef<AffineMap>({affineMap, affineMap}),
945 auto padFn = [&](Value valid, Value pos, int64_t pad) -> Value {
950 Value dpos = arith::SubIOp::create(rewriter, loc, pos, padVal);
952 Value offset = arith::MinSIOp::create(rewriter, loc, dpos, zero);
953 return arith::AddIOp::create(rewriter, loc, valid, offset)
957 auto coverageFn = [&](int64_t i, Value isize) -> Value {
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);
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);
976 Value kH3 = coverageFn(1, iH);
977 Value kW3 = coverageFn(2, iW);
980 auto count = arith::IndexCastOp::create(
981 rewriter, loc, rewriter.getI32Type(),
982 arith::MulIOp::create(rewriter, loc, kH3, kW3));
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)
995 arith::TruncFOp::create(rewriter, loc, resultETy, poolVal);
1000 if (inputZpVal != 0) {
1001 auto inputZp = arith::ConstantOp::create(
1002 rewriter, loc,
b.getIntegerAttr(accETy, inputZpVal));
1004 arith::MulIOp::create(rewriter, loc, accETy, count, inputZp);
1006 arith::SubIOp::create(rewriter, loc, accETy, poolVal, offset);
1010 Value one32 = arith::ConstantOp::create(
1011 rewriter, loc, rewriter.getI32IntegerAttr(1));
1012 Value thirtyTwo32 = arith::ConstantOp::create(
1013 rewriter, loc, rewriter.getI32IntegerAttr(32));
1016 arith::SubIOp::create(rewriter, loc, count, one32);
1017 Value leadingZeros =
1018 math::CountLeadingZerosOp::create(rewriter, loc, countSubOne);
1020 arith::SubIOp::create(rewriter, loc, thirtyTwo32, leadingZeros);
1024 arith::ExtUIOp::create(rewriter, loc, rewriter.getI64Type(), k);
1025 Value thirtyShiftPlusOne = arith::ConstantOp::create(
1026 rewriter, loc, rewriter.getI64IntegerAttr((1 << 30) + 1));
1028 arith::ShLIOp::create(rewriter, loc, thirtyShiftPlusOne, k64);
1031 Value count64 = arith::ExtUIOp::create(
1032 rewriter, loc, rewriter.getI64Type(), count);
1034 arith::DivUIOp::create(rewriter, loc, numerator, count64);
1035 multiplier = arith::TruncIOp::create(
1036 rewriter, loc, rewriter.getI32Type(), multiplier);
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);
1045 auto roundingAttr = RoundingModeAttr::get(
1046 rewriter.getContext(), RoundingMode::SINGLE_ROUND);
1048 auto scaled = tosa::ApplyScaleOp::create(
1049 rewriter, loc, rewriter.getI32Type(), poolVal,
1050 multiplier, shift, roundingAttr)
1055 if (outputZpVal != 0) {
1056 auto outputZp = arith::ConstantOp::create(
1058 b.getIntegerAttr(scaled.getType(), outputZpVal));
1059 scaled = arith::AddIOp::create(rewriter, loc, scaled, outputZp)
1067 rewriter, loc, accETy,
1068 APInt::getSignedMinValue(outBitwidth).getSExtValue());
1070 rewriter, loc, accETy,
1071 APInt::getSignedMaxValue(outBitwidth).getSExtValue());
1079 arith::TruncIOp::create(rewriter, loc, resultETy, poolVal);
1083 linalg::YieldOp::create(rewriter, loc, poolVal);
1086 rewriter.replaceOp(op, genericOp.getResult(0));
1093 using OpRewritePattern<tosa::TransposeOp>::OpRewritePattern;
1095 LogicalResult matchAndRewrite(tosa::TransposeOp op,
1096 PatternRewriter &rewriter)
const final {
1097 const llvm::ArrayRef<int32_t> constantPerms = op.getPerms();
1099 Location loc = op.getLoc();
1103 SmallVector<OpFoldResult> inputSizes =
1105 auto permutedSizes =
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; }));
1122 const TosaToLinalgNamedOptions &
options) {
1123 if (
options.preferConv2DKernelLayoutHWCF) {
1124 patterns->
add<ConvConverter<tosa::Conv2DOp, linalg::Conv2DNhwcHwcfOp,
1125 linalg::Conv2DNhwcHwcfQOp>>(
1128 patterns->
add<ConvConverter<tosa::Conv2DOp, linalg::Conv2DNhwcFhwcOp,
1129 linalg::Conv2DNhwcFhwcQOp>>(
1134 ConvConverter<tosa::Conv3DOp, linalg::Conv3DNdhwcDhwcfOp, linalg::Conv3DNdhwcDhwcfQOp>,
1135 DepthwiseConvConverter,
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.
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
IntegerAttr getIndexAttr(int64_t value)
AffineMap getMultiDimIdentityMap(unsigned rank)
AffineExpr getAffineConstantExpr(int64_t constant)
AffineExpr getAffineDimExpr(unsigned position)
MLIRContext * getContext() const
ImplicitLocOpBuilder maintains a 'current location', allowing use of the create<> method without spec...
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
This class helps build Operations.
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.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
bool isUnsignedInteger() const
Return true if this is an unsigned integer type (with the specified width).
unsigned getIntOrFloatBitWidth() const
Return the bit width of an integer or a float type, assert failure on other types.
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
static ConstantIndexOp create(OpBuilder &builder, Location location, int64_t value)
static ConstantIntOp create(OpBuilder &builder, Location location, int64_t value, unsigned width)
ArrayRef< T > asArrayRef() const
SmallVector< OpFoldResult > getMixedSizes(OpBuilder &builder, Location loc, Value value)
Return the dimensions of the given tensor value.
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...