1 //===- TosaOps.cpp - MLIR Dialect for TOSA --------------------------------===//
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 // \file
10 // This file implements the TOSA Specification:
11 // https://developer.mlplatform.org/w/tosa/
12 //
13 //===----------------------------------------------------------------------===//
14 
15 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
16 #include "mlir/Dialect/Quant/QuantOps.h"
17 #include "mlir/Dialect/Tensor/IR/Tensor.h"
18 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
19 #include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
20 #include "mlir/IR/BuiltinTypes.h"
21 #include "mlir/IR/DialectImplementation.h"
22 #include "mlir/IR/Matchers.h"
23 #include "mlir/IR/PatternMatch.h"
24 #include "mlir/Transforms/FoldUtils.h"
25 #include "mlir/Transforms/InliningUtils.h"
26 #include "mlir/Transforms/RegionUtils.h"
27 #include "llvm/ADT/DenseMap.h"
28 #include "llvm/ADT/TypeSwitch.h"
29 
30 using namespace mlir;
31 using namespace mlir::tosa;
32 
33 #include "mlir/Dialect/Tosa/IR/TosaOpsDialect.cpp.inc"
34 
35 //===----------------------------------------------------------------------===//
36 // Tosa dialect structs and interface includes.
37 //===----------------------------------------------------------------------===//
38 
39 #include "mlir/Dialect/Tosa/IR/TosaInterfaces.cpp.inc"
40 
41 namespace {
42 //===----------------------------------------------------------------------===//
43 // Dialect Function Inliner Interface.
44 //===----------------------------------------------------------------------===//
45 struct TosaInlinerInterface : public DialectInlinerInterface {
46   using DialectInlinerInterface::DialectInlinerInterface;
47 
48   //===--------------------------------------------------------------------===//
49   // Analysis Hooks.
50   //===--------------------------------------------------------------------===//
51 
52   /// All operations can be inlined by default.
53   bool isLegalToInline(Operation *op, Region *region, bool wouldBeCloned,
54                        BlockAndValueMapping &map) const final {
55     return true;
56   }
57 
58   /// All regions with If and While parent operators can be inlined.
59   bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
60                        BlockAndValueMapping &map) const final {
61     return (isa<tosa::IfOp>(dest->getParentOp()) ||
62             isa<tosa::WhileOp>(dest->getParentOp()));
63   }
64 };
65 } // namespace
66 
67 //===----------------------------------------------------------------------===//
68 // TOSA control flow support.
69 //===----------------------------------------------------------------------===//
70 
71 /// Returns the while loop body.
72 Region &tosa::WhileOp::getLoopBody() { return body(); }
73 
74 //===----------------------------------------------------------------------===//
75 // Tosa dialect initialization.
76 //===----------------------------------------------------------------------===//
77 
78 void TosaDialect::initialize() {
79   addOperations<
80 #define GET_OP_LIST
81 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"
82       >();
83   addAttributes<
84 #define GET_ATTRDEF_LIST
85 #include "mlir/Dialect/Tosa/IR/TosaAttributes.cpp.inc"
86       >();
87   addInterfaces<TosaInlinerInterface>();
88 }
89 
90 Operation *TosaDialect::materializeConstant(OpBuilder &builder, Attribute value,
91                                             Type type, Location loc) {
92   // Tosa dialect constants only support ElementsAttr unlike standard dialect
93   // constant which supports all attributes.
94   if (value.isa<ElementsAttr>())
95     return builder.create<tosa::ConstOp>(loc, type, value.cast<ElementsAttr>());
96   return nullptr;
97 }
98 
99 //===----------------------------------------------------------------------===//
100 // Operator Canonicalizers.
101 //===----------------------------------------------------------------------===//
102 
103 template <typename... Args>
104 void addOpsCanonicalizations(MLIRContext *ctx, RewritePatternSet &patterns) {
105   (void)std::initializer_list<int>{
106       0, (Args::getCanonicalizationPatterns(patterns, ctx), 0)...};
107 }
108 
109 void mlir::tosa::populateTosaOpsCanonicalizationPatterns(
110     MLIRContext *ctx, RewritePatternSet &patterns) {
111   addOpsCanonicalizations<
112 #define GET_OP_LIST
113 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"
114       >(ctx, patterns);
115 }
116 
117 struct ConcatOptimization : public OpRewritePattern<tosa::ConcatOp> {
118   using OpRewritePattern<tosa::ConcatOp>::OpRewritePattern;
119 
120   LogicalResult matchAndRewrite(tosa::ConcatOp op,
121                                 PatternRewriter &rewriter) const override {
122     if (op.input1().size() != 1)
123       return failure();
124     if (op.input1().front().getType() != op.getType()) {
125       rewriter
126           .replaceOpWithNewOp<tensor::CastOp>(op, op.getType(),
127                                               op.input1().front())
128           .getResult();
129       return success();
130     }
131 
132     rewriter.replaceOp(op, op.input1().front());
133     return success();
134   }
135 };
136 
137 void ConcatOp::getCanonicalizationPatterns(RewritePatternSet &results,
138                                            MLIRContext *context) {
139   results.add<ConcatOptimization>(context);
140 }
141 
142 struct ReshapeReshapeOptimization : public OpRewritePattern<tosa::ReshapeOp> {
143   using OpRewritePattern<tosa::ReshapeOp>::OpRewritePattern;
144 
145   LogicalResult matchAndRewrite(tosa::ReshapeOp op,
146                                 PatternRewriter &rewriter) const override {
147     Value input = op.input1();
148     Operation *definingOp = input.getDefiningOp();
149     if (!definingOp)
150       return failure();
151 
152     if (tosa::ReshapeOp reshapeOp = dyn_cast<tosa::ReshapeOp>(definingOp)) {
153       rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
154           op, op.getType(), reshapeOp.input1(), op.new_shape());
155       return success();
156     }
157 
158     return failure();
159   }
160 };
161 
162 struct ReshapeConstOptimization : public OpRewritePattern<tosa::ReshapeOp> {
163   using OpRewritePattern<tosa::ReshapeOp>::OpRewritePattern;
164 
165   LogicalResult matchAndRewrite(tosa::ReshapeOp op,
166                                 PatternRewriter &rewriter) const override {
167     Value input = op.input1();
168     ArrayAttr newShape = op.new_shape();
169 
170     // Check if input is constant
171     DenseElementsAttr inputAttr;
172     if (!matchPattern(input, m_Constant(&inputAttr)))
173       return failure();
174 
175     // Check if has >1 consumer and is not splat
176     if (!input.hasOneUse() && !inputAttr.isSplat())
177       return failure();
178 
179     // Grab the new shape
180     SmallVector<int64_t> newShapeValues = llvm::to_vector<6>(
181         llvm::map_range(newShape.getValue(), [](const Attribute &val) {
182           return val.cast<IntegerAttr>().getValue().getSExtValue();
183         }));
184 
185     // Build new const op with correct output shape
186     ShapedType inputShape = input.getType().cast<ShapedType>();
187     DenseElementsAttr outputAttr =
188         inputAttr.reshape(inputShape.clone(newShapeValues));
189     rewriter.replaceOpWithNewOp<tosa::ConstOp>(op, outputAttr.getType(),
190                                                outputAttr);
191     return success();
192   }
193 };
194 
195 void ReshapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
196                                             MLIRContext *context) {
197   results.add<ReshapeReshapeOptimization>(context);
198   results.add<ReshapeConstOptimization>(context);
199 }
200 
201 LogicalResult SelectOp::canonicalize(SelectOp op, PatternRewriter &rewriter) {
202   auto notOp = op.pred().getDefiningOp<tosa::LogicalNotOp>();
203   if (!notOp)
204     return failure();
205   rewriter.updateRootInPlace(op, [&]() {
206     op.getOperation()->setOperands(
207         {notOp.input1(), op.on_false(), op.on_true()});
208   });
209   return success();
210 }
211 
212 struct NoOpOptimization : public OpRewritePattern<tosa::TransposeOp> {
213   using OpRewritePattern::OpRewritePattern;
214 
215   LogicalResult matchAndRewrite(tosa::TransposeOp op,
216                                 PatternRewriter &rewriter) const override {
217     auto perm = op.perms();
218 
219     DenseIntElementsAttr permAttr;
220     if (!matchPattern(perm, m_Constant(&permAttr))) {
221       return failure();
222     }
223 
224     SmallVector<int64_t> permValues = llvm::to_vector<6>(
225         llvm::map_range(permAttr.getValues<APInt>(),
226                         [](const APInt &val) { return val.getSExtValue(); }));
227 
228     for (int i = 0, s = permValues.size(); i < s; i++) {
229       if (i != permValues[i]) {
230         return failure();
231       }
232     }
233 
234     rewriter.replaceOp(op, op.input1());
235     return success();
236   }
237 };
238 
239 void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
240                                               MLIRContext *context) {
241   results.add<NoOpOptimization>(context);
242 }
243 
244 struct AddZeroOptimization : public OpRewritePattern<tosa::AddOp> {
245   using OpRewritePattern::OpRewritePattern;
246 
247   LogicalResult matchAndRewrite(tosa::AddOp op,
248                                 PatternRewriter &rewriter) const override {
249     auto input1 = op.input1();
250     auto input2 = op.input2();
251 
252     DenseElementsAttr input1Attr;
253     if (matchPattern(input1, m_Constant(&input1Attr)) && input1Attr.isSplat() &&
254         input2.getType() == op.getType()) {
255       if (input1Attr.getType().getElementType().isa<IntegerType>() &&
256           input1Attr.getSplatValue<APInt>().isZero()) {
257         rewriter.replaceOp(op, op.input2());
258         return success();
259       }
260     }
261 
262     DenseElementsAttr input2Attr;
263     if (matchPattern(input2, m_Constant(&input2Attr)) && input2Attr.isSplat() &&
264         input1.getType() == op.getType()) {
265       if (input2Attr.getType().getElementType().isa<IntegerType>() &&
266           input2Attr.getSplatValue<APInt>().isZero()) {
267         rewriter.replaceOp(op, op.input1());
268         return success();
269       }
270     }
271 
272     return failure();
273   }
274 };
275 
276 void AddOp::getCanonicalizationPatterns(RewritePatternSet &results,
277                                         MLIRContext *context) {
278   results.add<AddZeroOptimization>(context);
279 }
280 
281 struct MulOneOptimization : public OpRewritePattern<tosa::MulOp> {
282   using OpRewritePattern::OpRewritePattern;
283 
284   LogicalResult matchAndRewrite(tosa::MulOp op,
285                                 PatternRewriter &rewriter) const override {
286     auto input1 = op.input1();
287     auto input2 = op.input2();
288 
289     DenseElementsAttr input1Attr;
290     if (matchPattern(input1, m_Constant(&input1Attr)) && input1Attr.isSplat() &&
291         input2.getType() == op.getType()) {
292       if (input1Attr.getType().getElementType().isa<FloatType>() &&
293           input1Attr.getSplatValue<APFloat>().isExactlyValue(1)) {
294         rewriter.replaceOp(op, op.input2());
295         return success();
296       }
297 
298       if (input1Attr.getType().getElementType().isa<IntegerType>() &&
299           matchPattern(input1, m_One())) {
300         rewriter.replaceOp(op, op.input2());
301         return success();
302       }
303     }
304 
305     DenseElementsAttr input2Attr;
306     if (matchPattern(input2, m_Constant(&input2Attr)) && input2Attr.isSplat() &&
307         input1.getType() == op.getType()) {
308       if (input2Attr.getType().getElementType().isa<FloatType>() &&
309           input2Attr.getSplatValue<APFloat>().isExactlyValue(1)) {
310         rewriter.replaceOp(op, op.input1());
311         return success();
312       }
313 
314       if (input2Attr.getType().getElementType().isa<IntegerType>() &&
315           matchPattern(input2, m_One())) {
316         rewriter.replaceOp(op, op.input1());
317         return success();
318       }
319     }
320 
321     return failure();
322   }
323 };
324 
325 void MulOp::getCanonicalizationPatterns(RewritePatternSet &results,
326                                         MLIRContext *context) {
327   results.add<MulOneOptimization>(context);
328 }
329 
330 struct MaterializePadValue : public OpRewritePattern<tosa::PadOp> {
331   using OpRewritePattern::OpRewritePattern;
332 
333   LogicalResult matchAndRewrite(tosa::PadOp op,
334                                 PatternRewriter &rewriter) const override {
335     if (op.pad_const())
336       return failure();
337 
338     auto input = op.input1();
339     auto padding = op.padding();
340 
341     ShapedType inputTy = input.getType().cast<ShapedType>();
342     Type elementTy = inputTy.getElementType();
343 
344     Attribute constantAttr;
345     if (elementTy.isa<FloatType>()) {
346       constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
347     } else if (elementTy.isa<IntegerType>() && !op.quantization_info()) {
348       constantAttr = rewriter.getIntegerAttr(elementTy, 0);
349     } else if (elementTy.isa<IntegerType>() && op.quantization_info()) {
350       auto value = op.quantization_info().getValue().getInput_zp();
351       constantAttr = rewriter.getIntegerAttr(elementTy, value);
352     }
353 
354     if (!constantAttr) {
355       return rewriter.notifyMatchFailure(
356           op,
357           "tosa.pad to linalg lowering encountered an unknown element type");
358     }
359 
360     auto denseAttr = DenseElementsAttr::get(
361         RankedTensorType::get({}, elementTy), constantAttr);
362     auto constantVal = rewriter.create<tosa::ConstOp>(
363         op.getLoc(), denseAttr.getType(), denseAttr);
364 
365     rewriter.replaceOpWithNewOp<tosa::PadOp>(
366         op, op.getType(), ValueRange{input, padding, constantVal},
367         op->getAttrs());
368     return success();
369   }
370 };
371 
372 void PadOp::getCanonicalizationPatterns(RewritePatternSet &results,
373                                         MLIRContext *context) {
374   results.add<MaterializePadValue>(context);
375 }
376 
377 struct MaxPool2dIsNoOp : public OpRewritePattern<tosa::MaxPool2dOp> {
378   using OpRewritePattern::OpRewritePattern;
379 
380   LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
381                                 PatternRewriter &rewriter) const override {
382     Value input = op.input();
383     Value output = op.output();
384     ShapedType inputType = input.getType().cast<ShapedType>();
385     ShapedType outputType = output.getType().cast<ShapedType>();
386 
387     if (!inputType.hasStaticShape() || !outputType.hasStaticShape()) {
388       return failure();
389     }
390 
391     // If the output and input shapes are 1x1, then this is a no op.
392     ArrayRef<int64_t> outputShape = outputType.getShape();
393     if (outputShape[1] != 1 || outputShape[2] != 1) {
394       return failure();
395     }
396 
397     ArrayRef<int64_t> inputShape = inputType.getShape();
398     if (inputShape[1] != 1 || inputShape[2] != 1) {
399       return failure();
400     }
401 
402     rewriter.replaceOp(op, input);
403     return success();
404   }
405 };
406 
407 void MaxPool2dOp::getCanonicalizationPatterns(RewritePatternSet &results,
408                                               MLIRContext *context) {
409   results.add<MaxPool2dIsNoOp>(context);
410 }
411 
412 struct ClampIsNoOp : public OpRewritePattern<tosa::ClampOp> {
413   using OpRewritePattern::OpRewritePattern;
414 
415   LogicalResult matchAndRewrite(tosa::ClampOp op,
416                                 PatternRewriter &rewriter) const override {
417     Value input = op.input();
418     auto inputType = op.input().getType().template dyn_cast<RankedTensorType>();
419     auto inputElementType = inputType.getElementType();
420 
421     if (!inputType.hasStaticShape()) {
422       return failure();
423     }
424 
425     if (inputElementType.isF32()) {
426       auto minClamp = op.min_fp();
427       auto maxClamp = op.max_fp();
428       bool isMin = (minClamp.isLargest() || minClamp.isInfinity()) &&
429                    minClamp.isNegative();
430       bool isMax = (maxClamp.isLargest() || maxClamp.isInfinity()) &&
431                    !maxClamp.isNegative();
432 
433       if (isMin && isMax) {
434         rewriter.replaceOp(op, input);
435         return success();
436       }
437       return failure();
438     }
439 
440     if (inputElementType.isUnsignedInteger()) {
441       int64_t minClamp = op.min_int();
442       int64_t maxClamp = op.max_int();
443 
444       int64_t intMin =
445           APInt::getMinValue(inputElementType.getIntOrFloatBitWidth())
446               .getZExtValue();
447       int64_t intMax =
448           APInt::getMaxValue(inputElementType.getIntOrFloatBitWidth())
449               .getZExtValue();
450 
451       if (minClamp <= intMin && maxClamp >= intMax) {
452         rewriter.replaceOp(op, input);
453         return success();
454       }
455       return failure();
456     }
457 
458     if (inputElementType.isa<IntegerType>()) {
459       int64_t minClamp = op.min_int();
460       int64_t maxClamp = op.max_int();
461 
462       int64_t intMin =
463           APInt::getSignedMinValue(inputElementType.getIntOrFloatBitWidth())
464               .getSExtValue();
465       int64_t intMax =
466           APInt::getSignedMaxValue(inputElementType.getIntOrFloatBitWidth())
467               .getSExtValue();
468 
469       if (minClamp <= intMin && maxClamp >= intMax) {
470         rewriter.replaceOp(op, input);
471         return success();
472       }
473       return failure();
474     }
475 
476     return failure();
477   }
478 };
479 
480 struct ClampClampOptimization : public OpRewritePattern<tosa::ClampOp> {
481   using OpRewritePattern<tosa::ClampOp>::OpRewritePattern;
482 
483   LogicalResult matchAndRewrite(tosa::ClampOp op,
484                                 PatternRewriter &rewriter) const override {
485     Value input = op.input();
486 
487     Operation *definingOp = input.getDefiningOp();
488     if (!definingOp)
489       return failure();
490 
491     if (tosa::ClampOp clampOp = dyn_cast<tosa::ClampOp>(definingOp)) {
492       auto minFp = std::max(op.min_fp(), clampOp.min_fp()).convertToFloat();
493       auto maxFp = std::min(op.max_fp(), clampOp.max_fp()).convertToFloat();
494 
495       auto minInt = std::max(op.min_int(), clampOp.min_int());
496       auto maxInt = std::min(op.max_int(), clampOp.max_int());
497 
498       rewriter.replaceOpWithNewOp<tosa::ClampOp>(
499           op, op.getType(), clampOp.input(), rewriter.getI64IntegerAttr(minInt),
500           rewriter.getI64IntegerAttr(maxInt), rewriter.getF32FloatAttr(minFp),
501           rewriter.getF32FloatAttr(maxFp));
502       return success();
503     }
504 
505     return failure();
506   }
507 };
508 
509 void ClampOp::getCanonicalizationPatterns(RewritePatternSet &results,
510                                           MLIRContext *context) {
511   results.add<ClampIsNoOp>(context);
512   results.add<ClampClampOptimization>(context);
513 }
514 
515 //===----------------------------------------------------------------------===//
516 // Operator Folders.
517 //===----------------------------------------------------------------------===//
518 
519 OpFoldResult CastOp::fold(ArrayRef<Attribute> operands) {
520   if (input().getType() == getType())
521     return input();
522   return {};
523 }
524 
525 OpFoldResult ConstOp::fold(ArrayRef<Attribute> operands) {
526   assert(operands.empty() && "constant has no operands");
527   return valueAttr();
528 }
529 
530 #define REDUCE_FOLDER(OP)                                                      \
531   OpFoldResult OP::fold(ArrayRef<Attribute> operands) {                        \
532     ShapedType inputTy = input().getType().cast<ShapedType>();                 \
533     if (!inputTy.hasRank())                                                    \
534       return {};                                                               \
535     if (inputTy.getDimSize(axis()) == 1)                                       \
536       return input();                                                          \
537     return {};                                                                 \
538   }
539 
540 REDUCE_FOLDER(ReduceAllOp)
541 REDUCE_FOLDER(ReduceAnyOp)
542 REDUCE_FOLDER(ReduceMaxOp)
543 REDUCE_FOLDER(ReduceMinOp)
544 REDUCE_FOLDER(ReduceProdOp)
545 REDUCE_FOLDER(ReduceSumOp)
546 #undef REDUCE_FOLDER
547 
548 OpFoldResult ReshapeOp::fold(ArrayRef<Attribute> operands) {
549   auto inputTy = input1().getType().dyn_cast<RankedTensorType>();
550   auto outputTy = getType().dyn_cast<RankedTensorType>();
551 
552   if (!inputTy || !outputTy || inputTy != outputTy)
553     return {};
554   return input1();
555 }
556 
557 OpFoldResult PadOp::fold(ArrayRef<Attribute> operands) {
558   // If the pad is all zeros we can fold this operation away.
559   if (operands[1]) {
560     auto densePad = operands[1].cast<DenseElementsAttr>();
561     if (densePad.isSplat() && densePad.getSplatValue<APInt>().isZero()) {
562       return input1();
563     }
564   }
565 
566   return {};
567 }
568 
569 OpFoldResult SliceOp::fold(ArrayRef<Attribute> operands) {
570   auto inputTy = input().getType().dyn_cast<RankedTensorType>();
571   auto outputTy = getType().dyn_cast<RankedTensorType>();
572 
573   if (!inputTy || !outputTy || inputTy != outputTy)
574     return {};
575   if (inputTy.hasStaticShape())
576     return input();
577 
578   return {};
579 }
580 
581 OpFoldResult tosa::SelectOp::fold(ArrayRef<Attribute> operands) {
582   if (on_true() == on_false())
583     return on_true();
584 
585   auto predicate = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
586   if (!predicate)
587     return {};
588 
589   if (!predicate.isSplat())
590     return {};
591   return predicate.getSplatValue<APInt>().getBoolValue() ? on_true()
592                                                          : on_false();
593 }
594 
595 OpFoldResult TileOp::fold(ArrayRef<Attribute> operands) {
596   bool allOnes = true;
597   for (Attribute val : multiples().getValue()) {
598     allOnes = allOnes && val.cast<IntegerAttr>().getValue().getSExtValue() == 1;
599   }
600 
601   if (allOnes && input1().getType() == getType())
602     return input1();
603   return {};
604 }
605 
606 OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
607   if (!operands[1])
608     return {};
609 
610   // Transposing splat values just means reshaping.
611   if (auto input = operands[0].dyn_cast_or_null<DenseElementsAttr>()) {
612     if (input.isSplat())
613       return input.reshape(getType().cast<ShapedType>());
614   }
615 
616   auto perms = llvm::to_vector<6>(llvm::map_range(
617       operands[1].cast<DenseIntElementsAttr>().getValues<APInt>(),
618       [](const APInt &val) { return val.getSExtValue(); }));
619 
620   if (llvm::equal(llvm::seq<int64_t>(0, perms.size()), perms) &&
621       input1().getType() == getType())
622     return input1();
623   return {};
624 }
625 
626 //===----------------------------------------------------------------------===//
627 // TOSA Operator Verifiers.
628 //===----------------------------------------------------------------------===//
629 
630 template <typename T>
631 static LogicalResult verifyConvOp(T op) {
632   // All TOSA conv ops have an input() and weight().
633   auto inputType = op.input().getType().template dyn_cast<RankedTensorType>();
634   auto weightType = op.weight().getType().template dyn_cast<RankedTensorType>();
635 
636   // Must be ranked tensor types
637   if (!inputType) {
638     op.emitOpError("expect a ranked tensor for input, got ") << op.input();
639     return failure();
640   }
641   if (!weightType) {
642     op.emitOpError("expect a ranked tensor for weight, got ") << op.weight();
643     return failure();
644   }
645 
646   auto inputEType = inputType.getElementType();
647   auto weightEType = weightType.getElementType();
648 
649   bool inputIsQuant = !inputEType.template isa<FloatType>();
650   bool weightIsQuant = !weightEType.template isa<FloatType>();
651 
652   // Either both must be quantized or both unquantized.
653   if (inputIsQuant != weightIsQuant) {
654     op.emitOpError(
655         "expect both input and weight to be float or not together, got ")
656         << inputEType << " and " << weightEType;
657     return failure();
658   }
659 
660   // Quantized type must have constructed the quantizationattr, and unquantized
661   // types should not have a quantizationattr.
662   if ((inputIsQuant && !op.quantization_info()) ||
663       (!inputIsQuant && op.quantization_info())) {
664     op.emitOpError("quantizationattr is required for quantized type, and not "
665                    "allowed for float type");
666     return failure();
667   }
668 
669   return success();
670 }
671 
672 LogicalResult tosa::AvgPool2dOp::verify() {
673   auto inputETy = input().getType().cast<ShapedType>().getElementType();
674   auto resultETy = getType().cast<ShapedType>().getElementType();
675 
676   if (auto quantType = inputETy.dyn_cast<mlir::quant::UniformQuantizedType>())
677     inputETy = quantType.getStorageType();
678 
679   if (auto quantType = resultETy.dyn_cast<mlir::quant::UniformQuantizedType>())
680     resultETy = quantType.getStorageType();
681 
682   if (inputETy.isF32() && resultETy.isF32())
683     return success();
684   if (inputETy.isInteger(8) && resultETy.isInteger(8))
685     return success();
686   if (inputETy.isInteger(16) && resultETy.isInteger(16))
687     return success();
688 
689   return emitOpError("input/output element types are incompatible.");
690 }
691 
692 //===----------------------------------------------------------------------===//
693 // TOSA Operator Quantization Builders.
694 //===----------------------------------------------------------------------===//
695 
696 /// This builder is called on all convolution operators except TransposeConv,
697 /// which has specialized output shape semantics. The builder also defines the
698 /// bitwidth of the output given the bit width of the input & weight content.
699 static void buildConvOpWithQuantInfo(OpBuilder &builder, OperationState &result,
700                                      Type outputType, Value input, Value weight,
701                                      Value bias, ArrayAttr pad,
702                                      ArrayAttr stride, ArrayAttr dilation) {
703 
704   result.addOperands({input, weight, bias});
705   result.addAttribute("pad", pad);
706   result.addAttribute("stride", stride);
707   result.addAttribute("dilation", dilation);
708 
709   auto quantAttr = buildConvOpQuantizationAttr(builder, input, weight);
710   if (quantAttr) {
711     result.addAttribute("quantization_info", quantAttr);
712     result.addTypes(
713         buildConvOpResultTypeInfo(builder, outputType, input, weight));
714   } else {
715     result.addTypes(outputType);
716   }
717 }
718 
719 /// Handles tosa.transpose_conv2d which has outpad and output shape attributes.
720 static void
721 buildTransConvOpWithQuantInfo(OpBuilder &builder, OperationState &result,
722                               Type outputType, Value input, Value weight,
723                               Value bias, ArrayAttr outpad, ArrayAttr stride,
724                               ArrayAttr dilation, ArrayAttr outputShape) {
725   result.addOperands({input, weight, bias});
726   result.addAttribute("out_pad", outpad);
727   result.addAttribute("stride", stride);
728   result.addAttribute("dilation", dilation);
729   result.addAttribute("out_shape", outputShape);
730   auto quantAttr = ::buildConvOpQuantizationAttr(builder, input, weight);
731 
732   if (quantAttr) {
733     result.addAttribute("quantization_info", quantAttr);
734     result.addTypes(
735         buildConvOpResultTypeInfo(builder, outputType, input, weight));
736   } else {
737     result.addTypes(outputType);
738   }
739 }
740 
741 /// The tosa.fully_connected op has its own builder as it does not have
742 /// strides/dilation/padding.
743 static void buildFCOpWithQuantInfo(OpBuilder &builder, OperationState &result,
744                                    Type outputType, Value input, Value weight,
745                                    Value bias) {
746 
747   result.addOperands({input, weight, bias});
748   auto quantAttr = ::buildConvOpQuantizationAttr(builder, input, weight);
749   if (quantAttr) {
750     result.addAttribute("quantization_info", quantAttr);
751     result.addTypes(
752         buildConvOpResultTypeInfo(builder, outputType, input, weight));
753   } else {
754     result.addTypes(outputType);
755   }
756 }
757 
758 /// The tosa.matmul op is also intended to be generated where a fully_connected
759 /// op must be constructed where the weight is not a constant. In this case,
760 /// the fully_connected op must be expressed using matmul.
761 /// TODO: Add link to the leglization document explaining this.
762 static void buildMatMulOpWithQuantInfo(OpBuilder &builder,
763                                        OperationState &result, Type outputType,
764                                        Value a, Value b) {
765   result.addOperands({a, b});
766   auto quantAttr = ::buildMatMulOpQuantizationAttr(builder, a, b);
767 
768   if (quantAttr) {
769     result.addAttribute("quantization_info", quantAttr);
770 
771     auto inputType = a.getType().dyn_cast<ShapedType>();
772     assert(inputType && "Input must be a shaped tensor type!");
773 
774     auto inputQType = inputType.getElementType()
775                           .dyn_cast<mlir::quant::UniformQuantizedType>();
776     assert(inputQType && "Tensor must have quantized datatype!");
777 
778     unsigned inputBits = inputQType.getStorageTypeIntegralWidth();
779 
780     auto outputShapedType = outputType.dyn_cast<ShapedType>();
781     assert(outputShapedType && "Output must be a shaped type");
782 
783     IntegerType accElementType;
784     if (inputBits == 16)
785       accElementType = builder.getIntegerType(48);
786     else
787       accElementType = builder.getI32Type();
788     auto accType = outputShapedType.clone(accElementType);
789     result.addTypes(accType);
790   } else {
791     result.addTypes(outputType);
792   }
793 }
794 
795 /// Both the tosa.avg_pool2d and unary ops use the same UnaruOpQuantizationAttr
796 /// but avg_pool operator has its own builder as it has additional parameters
797 /// not part of the unary ops.
798 static void buildAvgPool2dOpWithQuantInfo(OpBuilder &builder,
799                                           OperationState &result,
800                                           Type outputType, Value input,
801                                           ArrayAttr kernel, ArrayAttr stride,
802                                           ArrayAttr pad) {
803   result.addOperands(input);
804   result.addAttribute("kernel", kernel);
805   result.addAttribute("stride", stride);
806   result.addAttribute("pad", pad);
807   auto quantAttr = buildUnaryOpQuantizationAttr(builder, input, outputType);
808   if (quantAttr)
809     result.addAttribute("quantization_info", quantAttr);
810   result.types.push_back(outputType);
811 }
812 
813 /// This builder is called on single-parameter unary operators that have scale
814 /// relationship between their input and output, expressed by the
815 /// UnaryOpQuantizationAttr.
816 static void buildUnaryOpWithQuantInfo(OpBuilder &builder,
817                                       OperationState &result, Type outputType,
818                                       Value input) {
819   result.addOperands(input);
820   auto quantAttr = buildUnaryOpQuantizationAttr(builder, input, outputType);
821   if (quantAttr)
822     result.addAttribute("quantization_info", quantAttr);
823   result.types.push_back(outputType);
824 }
825 
826 /// This builder is called on TOSA pad operator that needs to create its own
827 /// OptionalAttr quantization_attr parameter to scale the padding values
828 /// correctly. No pad_const is interpreted as zero-padding.
829 static void buildPadOpWithQuantInfo(OpBuilder &builder, OperationState &result,
830                                     Type outputType, Value input,
831                                     Value paddings) {
832   result.addOperands({input, paddings});
833   auto quantAttr = buildPadOpQuantizationAttr(builder, input);
834   if (quantAttr)
835     result.addAttribute("quantization_info", quantAttr);
836   result.types.push_back(outputType);
837 }
838 
839 /// This builder is called on TOSA pad operator when an explicit pad_const
840 /// value is passed in. It also optionally constructs quantization_attr.
841 static void buildExplicitValuePadOpWithQuantInfo(OpBuilder &builder,
842                                                  OperationState &result,
843                                                  Type outputType, Value input,
844                                                  Value paddings,
845                                                  Value padConst) {
846   result.addOperands({input, paddings, padConst});
847   auto quantAttr = buildPadOpQuantizationAttr(builder, input);
848   if (quantAttr)
849     result.addAttribute("quantization_info", quantAttr);
850   result.types.push_back(outputType);
851 }
852 
853 //===----------------------------------------------------------------------===//
854 // TOSA Operator Return Type Inference.
855 //===----------------------------------------------------------------------===//
856 
857 static void getI64Values(ArrayAttr arrayAttr, SmallVector<int64_t> &values) {
858   for (auto it : arrayAttr) {
859     values.push_back(it.cast<IntegerAttr>().getValue().getSExtValue());
860   }
861 }
862 
863 static void getF64Values(ArrayAttr arrayAttr, SmallVector<double> &values) {
864   for (auto it : arrayAttr) {
865     values.push_back(it.cast<FloatAttr>().getValueAsDouble());
866   }
867 }
868 
869 LogicalResult tosa::ArgMaxOp::inferReturnTypeComponents(
870     MLIRContext *context, ::llvm::Optional<Location> location,
871     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
872     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
873   ShapeAdaptor inputShape = operands.getShape(0);
874   IntegerAttr axis = attributes.get("axis").cast<IntegerAttr>();
875   int32_t axisVal = axis.getValue().getSExtValue();
876 
877   if (!inputShape.hasRank()) {
878     inferredReturnShapes.push_back(ShapedTypeComponents());
879     return success();
880   }
881 
882   SmallVector<int64_t> outShape;
883   outShape.reserve(inputShape.getRank() - 1);
884   for (int i = 0, s = inputShape.getRank(); i < s; i++) {
885     if (i == axisVal)
886       continue;
887     outShape.push_back(inputShape.getDimSize(i));
888   }
889 
890   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
891   return success();
892 }
893 
894 LogicalResult tosa::ConcatOp::inferReturnTypeComponents(
895     MLIRContext *context, ::llvm::Optional<Location> location,
896     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
897     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
898   // Infer all dimension sizes by reducing based on inputs.
899   int32_t axis =
900       attributes.get("axis").cast<IntegerAttr>().getValue().getSExtValue();
901   llvm::SmallVector<int64_t> outputShape;
902   bool hasRankedInput = false;
903   for (auto operand : operands) {
904     ShapeAdaptor operandShape = operands.getShape(operand);
905     if (!operandShape.hasRank())
906       continue;
907 
908     // Copy the Operand's rank.
909     if (!hasRankedInput)
910       outputShape.resize(operandShape.getRank(), ShapedType::kDynamicSize);
911 
912     // Copy shapes until the dim is non-dynamic.
913     for (int i = 0, s = operandShape.getRank(); i < s; i++) {
914       if (i == axis || operandShape.isDynamicDim(i))
915         continue;
916       if (outputShape[i] == ShapedType::kDynamicSize)
917         outputShape[i] = operandShape.getDimSize(i);
918       if (outputShape[i] != operandShape.getDimSize(i))
919         return failure();
920     }
921 
922     hasRankedInput = true;
923   }
924 
925   if (!hasRankedInput) {
926     inferredReturnShapes.push_back(ShapedTypeComponents());
927     return success();
928   }
929 
930   // Determine the dimension size along the concatenation axis.
931   int concatDimSize = 0;
932   for (auto operand : operands) {
933     ShapeAdaptor operandShape = operands.getShape(operand);
934 
935     // We need to know the length of the concatenation axis of all inputs to
936     // determine the dimension size of the output shape.
937     if (!operandShape.hasRank() || operandShape.isDynamicDim(axis)) {
938       concatDimSize = ShapedType::kDynamicSize;
939       break;
940     }
941 
942     concatDimSize += operandShape.getDimSize(axis);
943   }
944 
945   outputShape[axis] = concatDimSize;
946 
947   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
948   return success();
949 }
950 
951 LogicalResult tosa::FullyConnectedOp::inferReturnTypeComponents(
952     MLIRContext *context, ::llvm::Optional<Location> location,
953     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
954     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
955   ShapeAdaptor inputShape = operands.getShape(0);
956   ShapeAdaptor weightShape = operands.getShape(1);
957   ShapeAdaptor biasShape = operands.getShape(2);
958 
959   // All shapes are dynamic.
960   SmallVector<int64_t> outShape;
961   outShape.resize(2, ShapedType::kDynamicSize);
962 
963   if (inputShape.hasRank()) {
964     outShape[0] = inputShape.getDimSize(0);
965   }
966 
967   if (weightShape.hasRank()) {
968     outShape[1] = weightShape.getDimSize(0);
969   }
970 
971   if (biasShape.hasRank()) {
972     outShape[1] = outShape[1] == ShapedType::kDynamicSize
973                       ? biasShape.getDimSize(0)
974                       : outShape[1];
975   }
976 
977   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
978   return success();
979 }
980 
981 LogicalResult FullyConnectedOp::verify() { return verifyConvOp(*this); }
982 
983 LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
984     MLIRContext *context, ::llvm::Optional<Location> location,
985     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
986     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
987   ShapeAdaptor lhsShape = operands.getShape(0);
988   ShapeAdaptor rhsShape = operands.getShape(1);
989 
990   // All shapes are dynamic.
991   SmallVector<int64_t> outShape;
992   outShape.resize(3, ShapedType::kDynamicSize);
993 
994   if (lhsShape.hasRank()) {
995     outShape[0] = lhsShape.getDimSize(0);
996     outShape[1] = lhsShape.getDimSize(1);
997   }
998 
999   if (rhsShape.hasRank()) {
1000     outShape[0] = outShape[0] == ShapedType::kDynamicSize
1001                       ? rhsShape.getDimSize(0)
1002                       : outShape[0];
1003     outShape[2] = rhsShape.getDimSize(2);
1004   }
1005 
1006   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
1007   return success();
1008 }
1009 
1010 LogicalResult tosa::PadOp::inferReturnTypeComponents(
1011     MLIRContext *context, ::llvm::Optional<Location> location,
1012     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1013     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1014   ShapeAdaptor inputShape = operands.getShape(0);
1015   ShapeAdaptor paddingShape = operands.getShape(1);
1016   SmallVector<int64_t> outputShape;
1017 
1018   // If both inputs have unknown shape, we cannot determine the shape of the
1019   // output.
1020   if (!inputShape.hasRank() && !paddingShape.hasRank()) {
1021     inferredReturnShapes.push_back(ShapedTypeComponents());
1022     return success();
1023   }
1024 
1025   // If the input rank is unknown we can info the output rank using the padding
1026   // shape's first dim.
1027   if (!inputShape.hasRank()) {
1028     if (paddingShape.isDynamicDim(0)) {
1029       inferredReturnShapes.push_back(ShapedTypeComponents());
1030       return success();
1031     }
1032 
1033     outputShape.resize(paddingShape.getDimSize(0), ShapedType::kDynamicSize);
1034     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1035     return success();
1036   }
1037 
1038   DenseIntElementsAttr paddings;
1039   // If the paddings value is not a constant, all dimensions must be dynamic.
1040   if (!matchPattern(operands[1], m_Constant(&paddings))) {
1041     outputShape.resize(inputShape.getRank(), ShapedType::kDynamicSize);
1042     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1043     return success();
1044   }
1045 
1046   SmallVector<int64_t> paddingValues;
1047   for (auto val : paddings) {
1048     paddingValues.push_back(val.getSExtValue());
1049   }
1050 
1051   outputShape.reserve(inputShape.getRank());
1052   for (int i = 0, s = inputShape.getRank(); i < s; i++) {
1053     if (inputShape.isDynamicDim(i)) {
1054       outputShape.push_back(ShapedType::kDynamicSize);
1055       continue;
1056     }
1057 
1058     outputShape.push_back(inputShape.getDimSize(i) + paddingValues[i * 2] +
1059                           paddingValues[i * 2 + 1]);
1060   }
1061 
1062   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1063   return success();
1064 }
1065 
1066 LogicalResult tosa::SliceOp::inferReturnTypeComponents(
1067     MLIRContext *context, ::llvm::Optional<Location> location,
1068     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1069     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1070   ArrayAttr sizes = SliceOpAdaptor(operands, attributes).size();
1071   SmallVector<int64_t> outputShape;
1072   outputShape.reserve(sizes.size());
1073   for (auto val : sizes) {
1074     outputShape.push_back(val.cast<IntegerAttr>().getValue().getSExtValue());
1075   }
1076 
1077   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1078   return success();
1079 }
1080 
1081 LogicalResult tosa::TableOp::inferReturnTypeComponents(
1082     MLIRContext *context, ::llvm::Optional<Location> location,
1083     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1084     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1085   ShapeAdaptor inputShape = operands.getShape(0);
1086 
1087   if (!inputShape.hasRank()) {
1088     inferredReturnShapes.push_back(ShapedTypeComponents());
1089     return success();
1090   }
1091 
1092   inferredReturnShapes.resize(1);
1093   inputShape.getDims(inferredReturnShapes[0]);
1094   return success();
1095 }
1096 
1097 LogicalResult tosa::TileOp::inferReturnTypeComponents(
1098     MLIRContext *context, ::llvm::Optional<Location> location,
1099     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1100     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1101   TileOpAdaptor adaptor(operands, attributes);
1102   ArrayAttr multiples = adaptor.multiples();
1103   ShapeAdaptor inputShape = operands.getShape(0);
1104   SmallVector<int64_t> outputShape;
1105   if (!inputShape.hasRank()) {
1106     outputShape.resize(multiples.size(), ShapedType::kDynamicSize);
1107     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1108     return success();
1109   }
1110 
1111   // We need the multiple values to determine the output shape.
1112   SmallVector<int64_t> multipleValues;
1113   multipleValues.reserve(multiples.size());
1114   for (auto val : multiples) {
1115     multipleValues.push_back(val.cast<IntegerAttr>().getValue().getSExtValue());
1116   }
1117 
1118   // Any non dynamic dimension can be multiplied to a known size.
1119   outputShape.reserve(multiples.size());
1120   for (int i = 0, s = inputShape.getRank(); i < s; i++) {
1121     int dim = inputShape.getDimSize(i);
1122     if (dim != ShapedType::kDynamicSize)
1123       dim *= multipleValues[i];
1124     outputShape.push_back(dim);
1125   }
1126 
1127   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1128   return success();
1129 }
1130 
1131 LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
1132     MLIRContext *context, ::llvm::Optional<Location> location,
1133     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1134     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1135   ReshapeOpAdaptor adaptor(operands, attributes);
1136   ShapeAdaptor inputShape = operands.getShape(0);
1137 
1138   ArrayAttr newShape = adaptor.new_shape();
1139   llvm::SmallVector<int64_t> newShapeValue;
1140   getI64Values(newShape, newShapeValue);
1141 
1142   // We cannot infer from the total number of elements so we must take the
1143   // shape attribute as exact.
1144   if (!inputShape.hasRank() || !inputShape.hasStaticShape()) {
1145     inferredReturnShapes.push_back(ShapedTypeComponents(newShapeValue));
1146     return success();
1147   }
1148 
1149   // Determine the number of elements covered by the slice of all static
1150   // dimensions. This allows us to infer the length of the remaining dynamic
1151   // dimension.
1152   int64_t numElements = inputShape.getNumElements();
1153   int64_t staticMul = 1;
1154   for (auto val : newShapeValue) {
1155     if (val != ShapedType::kDynamicSize) {
1156       staticMul *= val;
1157     }
1158   }
1159 
1160   // Determine the length of the dynamic dimension.
1161   for (auto &val : newShapeValue) {
1162     if (val == ShapedType::kDynamicSize)
1163       val = numElements / staticMul;
1164   }
1165 
1166   inferredReturnShapes.push_back(ShapedTypeComponents(newShapeValue));
1167   return success();
1168 }
1169 
1170 LogicalResult tosa::TransposeOp::inferReturnTypeComponents(
1171     MLIRContext *context, ::llvm::Optional<Location> location,
1172     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1173     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1174   ShapeAdaptor inputShape = operands.getShape(0);
1175   ShapeAdaptor permsShape = operands.getShape(1);
1176 
1177   // If input rank and permutation length is unknown, the output rank is
1178   // unknown.
1179   if (!inputShape.hasRank() || !permsShape.hasRank() ||
1180       permsShape.isDynamicDim(0)) {
1181     inferredReturnShapes.push_back(ShapedTypeComponents());
1182     return success();
1183   }
1184 
1185   // This would imply the number of permutations does not match the rank of the
1186   // input which is illegal.
1187   if (permsShape.getDimSize(0) != inputShape.getRank()) {
1188     return failure();
1189   }
1190 
1191   // Without the input dims we cannot determine the output dim sizes but we
1192   // can determine the output rank.
1193   SmallVector<int64_t> outputShape;
1194   if (!inputShape.hasRank()) {
1195     outputShape.resize(permsShape.getDimSize(0), ShapedType::kDynamicSize);
1196     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1197     return success();
1198   }
1199 
1200   // Rank-0 means no permutations matter.
1201   if (inputShape.getRank() == 0) {
1202     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1203     return success();
1204   }
1205 
1206   // Check whether the input dimensions are all the same.
1207   bool allTheSame = true;
1208   for (int i = 1, s = inputShape.getRank(); i < s; i++) {
1209     if (inputShape.getDimSize(0) != inputShape.getDimSize(i)) {
1210       allTheSame = false;
1211       break;
1212     }
1213   }
1214 
1215   // If all of the input dimensions are the same we don't care about the
1216   // permutation.
1217   if (allTheSame) {
1218     outputShape.resize(inputShape.getRank(), inputShape.getDimSize(0));
1219     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1220     return success();
1221   }
1222 
1223   outputShape.resize(inputShape.getRank(), ShapedType::kDynamicSize);
1224   // If the permuations are a constant we can directly determine the output
1225   // shape.
1226   if (ShapeAdaptor permShape = operands.getValueAsShape(1)) {
1227     outputShape.reserve(inputShape.getRank());
1228     for (int i = 0, s = inputShape.getRank(); i < s; i++) {
1229       outputShape[i] = inputShape.getDimSize(permShape.getDimSize(i));
1230     }
1231   }
1232 
1233   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1234   return success();
1235 }
1236 
1237 LogicalResult tosa::GatherOp::inferReturnTypeComponents(
1238     MLIRContext *context, ::llvm::Optional<Location> location,
1239     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1240     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1241   llvm::SmallVector<int64_t> outputShape;
1242   outputShape.resize(3, ShapedType::kDynamicSize);
1243 
1244   ShapeAdaptor valuesShape = operands.getShape(0);
1245   if (valuesShape.hasRank()) {
1246     outputShape[0] = valuesShape.getDimSize(0);
1247     outputShape[2] = valuesShape.getDimSize(2);
1248   }
1249 
1250   ShapeAdaptor indicesShape = operands.getShape(1);
1251   if (indicesShape.hasRank()) {
1252     if (outputShape[0] == ShapedType::kDynamicSize)
1253       outputShape[0] = indicesShape.getDimSize(0);
1254     if (outputShape[1] == ShapedType::kDynamicSize)
1255       outputShape[1] = indicesShape.getDimSize(1);
1256   }
1257 
1258   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1259   return success();
1260 }
1261 
1262 LogicalResult tosa::ResizeOp::inferReturnTypeComponents(
1263     MLIRContext *context, ::llvm::Optional<Location> location,
1264     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1265     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1266   ResizeOpAdaptor adaptor(operands, attributes);
1267   llvm::SmallVector<int64_t, 4> outputShape;
1268   outputShape.resize(4, ShapedType::kDynamicSize);
1269 
1270   int32_t inHeight = ShapedType::kDynamicSize;
1271   int32_t inWidth = ShapedType::kDynamicSize;
1272 
1273   ShapeAdaptor inputShape = operands.getShape(adaptor.input());
1274   if (inputShape.hasRank()) {
1275     outputShape[0] = inputShape.getDimSize(0);
1276     outputShape[3] = inputShape.getDimSize(3);
1277 
1278     inHeight = inputShape.getDimSize(1);
1279     inWidth = inputShape.getDimSize(2);
1280   }
1281 
1282   int32_t shift = adaptor.shift();
1283   llvm::SmallVector<int64_t> newShape;
1284   getI64Values(adaptor.output_size(), newShape);
1285   outputShape[1] = newShape[0];
1286   outputShape[2] = newShape[1];
1287 
1288   llvm::SmallVector<int64_t> strideInt;
1289   llvm::SmallVector<int64_t> offsetInt;
1290   llvm::SmallVector<double> strideFp;
1291   llvm::SmallVector<double> offsetFp;
1292   getI64Values(adaptor.offset(), offsetInt);
1293   getF64Values(adaptor.offset_fp(), offsetFp);
1294   getI64Values(adaptor.stride(), strideInt);
1295   getF64Values(adaptor.stride_fp(), strideFp);
1296 
1297   // If we have a 0 zero in integers we know that the resize indexing needs to
1298   // be performed in floating point. Use the floating point varient to compute
1299   // the resize shape.
1300   bool fpMode = strideInt[0] == 0;
1301 
1302   // We can compute the output shape if attribute specifies unknown dimensions
1303   // based on the offset and stride. If we perfectly line up to the last index
1304   // we need to round up the size to include it.
1305   if (outputShape[1] == ShapedType::kDynamicSize && inHeight >= 0 && fpMode) {
1306     float sizeFp = (inHeight - offsetFp[0] - 1) / strideFp[0];
1307     float round = std::floor(sizeFp) == sizeFp ? 1 : 0;
1308     outputShape[1] = std::ceil(sizeFp) + round;
1309   }
1310 
1311   if (outputShape[2] == ShapedType::kDynamicSize && inWidth >= 0 && fpMode) {
1312     float sizeFp = (inWidth - offsetFp[1] - 1) / strideFp[1];
1313     float round = std::floor(sizeFp) == sizeFp ? 1 : 0;
1314     outputShape[2] = std::ceil(sizeFp) + round;
1315   }
1316 
1317   if (outputShape[1] == ShapedType::kDynamicSize && inHeight >= 0 && !fpMode) {
1318     int64_t size = (inHeight - 1);
1319     size = ((size << shift) - offsetInt[0]) / strideInt[0];
1320     outputShape[1] = size + 1;
1321   }
1322 
1323   if (outputShape[2] == ShapedType::kDynamicSize && inWidth >= 0 && !fpMode) {
1324     int64_t size = (inWidth - 1);
1325     size = ((size << shift) - offsetInt[1]) / strideInt[1];
1326     outputShape[2] = size + 1;
1327   }
1328 
1329   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1330   return success();
1331 }
1332 
1333 LogicalResult tosa::ScatterOp::inferReturnTypeComponents(
1334     MLIRContext *context, ::llvm::Optional<Location> location,
1335     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1336     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1337   llvm::SmallVector<int64_t> outputShape;
1338   outputShape.resize(3, ShapedType::kDynamicSize);
1339 
1340   ShapeAdaptor valuesInShape = operands.getShape(0);
1341   if (valuesInShape.hasRank()) {
1342     outputShape[0] = valuesInShape.getDimSize(0);
1343     outputShape[1] = valuesInShape.getDimSize(1);
1344     outputShape[2] = valuesInShape.getDimSize(2);
1345   }
1346 
1347   ShapeAdaptor indicesShape = operands.getShape(1);
1348   if (indicesShape.hasRank()) {
1349     if (outputShape[0] == ShapedType::kDynamicSize)
1350       outputShape[0] = indicesShape.getDimSize(0);
1351   }
1352 
1353   ShapeAdaptor inputShape = operands.getShape(2);
1354   if (inputShape.hasRank()) {
1355     if (outputShape[0] == ShapedType::kDynamicSize)
1356       outputShape[0] = inputShape.getDimSize(0);
1357     if (outputShape[2] == ShapedType::kDynamicSize)
1358       outputShape[2] = inputShape.getDimSize(2);
1359   }
1360 
1361   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1362   return success();
1363 }
1364 
1365 static LogicalResult ReduceInferReturnTypes(
1366     ShapeAdaptor operandShape, IntegerAttr axis,
1367     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1368   if (!operandShape.hasRank()) {
1369     inferredReturnShapes.push_back(ShapedTypeComponents());
1370     return success();
1371   }
1372 
1373   SmallVector<int64_t> outputShape;
1374   operandShape.getDims(outputShape);
1375   int64_t axisVal = axis.getValue().getSExtValue();
1376   outputShape[axisVal] = 1;
1377   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1378   return success();
1379 }
1380 
1381 #define REDUCE_SHAPE_INFER(OP)                                                 \
1382   LogicalResult OP::inferReturnTypeComponents(                                 \
1383       MLIRContext *context, ::llvm::Optional<Location> location,               \
1384       ValueShapeRange operands, DictionaryAttr attributes,                     \
1385       RegionRange regions,                                                     \
1386       SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {           \
1387     return ReduceInferReturnTypes(operands.getShape(0),                        \
1388                                   attributes.get("axis").cast<IntegerAttr>(),  \
1389                                   inferredReturnShapes);                       \
1390   }
1391 
1392 REDUCE_SHAPE_INFER(tosa::ReduceAllOp)
1393 REDUCE_SHAPE_INFER(tosa::ReduceAnyOp)
1394 REDUCE_SHAPE_INFER(tosa::ReduceMaxOp)
1395 REDUCE_SHAPE_INFER(tosa::ReduceMinOp)
1396 REDUCE_SHAPE_INFER(tosa::ReduceProdOp)
1397 REDUCE_SHAPE_INFER(tosa::ReduceSumOp)
1398 #undef REDUCE_SHAPE_INFER
1399 
1400 static LogicalResult resolveBroadcastShape(const ValueShapeRange &operands,
1401                                            SmallVector<int64_t> &outShape) {
1402   int64_t outRank = 0;
1403   for (int i = 0, e = operands.size(); i != e; ++i) {
1404     auto shape = operands.getShape(i);
1405     if (!shape.hasRank()) {
1406       return failure();
1407     }
1408     outRank = std::max<int64_t>(outRank, shape.getRank());
1409   }
1410 
1411   outShape.resize(outRank, 1);
1412 
1413   for (int i = 0, e = operands.size(); i != e; ++i) {
1414     auto shape = operands.getShape(i);
1415     auto rankDiff = outShape.size() - shape.getRank();
1416 
1417     for (size_t i = 0, e = shape.getRank(); i < e; ++i) {
1418       auto dim1 = outShape[i + rankDiff];
1419       auto dim2 = shape.getDimSize(i);
1420       auto resolvedDim = dim1;
1421 
1422       if (dim1 == 1) {
1423         resolvedDim = dim2;
1424       } else if (dim2 == 1) {
1425         resolvedDim = dim1;
1426       } else if (dim1 != dim2) {
1427         return failure();
1428       }
1429       outShape[i + rankDiff] = resolvedDim;
1430     }
1431   }
1432 
1433   return success();
1434 }
1435 
1436 static LogicalResult NAryInferReturnTypes(
1437     const ValueShapeRange &operands,
1438     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1439   llvm::SmallVector<int64_t> outShape;
1440   if (resolveBroadcastShape(operands, outShape).failed()) {
1441     inferredReturnShapes.push_back(ShapedTypeComponents());
1442   } else {
1443     inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
1444   }
1445   return success();
1446 }
1447 
1448 #define NARY_SHAPE_INFER(OP)                                                   \
1449   LogicalResult OP::inferReturnTypeComponents(                                 \
1450       MLIRContext *context, ::llvm::Optional<Location> location,               \
1451       ValueShapeRange operands, DictionaryAttr attributes,                     \
1452       RegionRange regions,                                                     \
1453       SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {           \
1454     return NAryInferReturnTypes(operands, inferredReturnShapes);               \
1455   }
1456 
1457 NARY_SHAPE_INFER(tosa::AbsOp)
1458 NARY_SHAPE_INFER(tosa::AddOp)
1459 NARY_SHAPE_INFER(tosa::ArithmeticRightShiftOp)
1460 NARY_SHAPE_INFER(tosa::BitwiseAndOp)
1461 NARY_SHAPE_INFER(tosa::BitwiseOrOp)
1462 NARY_SHAPE_INFER(tosa::BitwiseXorOp)
1463 NARY_SHAPE_INFER(tosa::BitwiseNotOp)
1464 NARY_SHAPE_INFER(tosa::CastOp)
1465 NARY_SHAPE_INFER(tosa::CeilOp)
1466 NARY_SHAPE_INFER(tosa::ClampOp)
1467 NARY_SHAPE_INFER(tosa::ClzOp)
1468 NARY_SHAPE_INFER(tosa::DivOp)
1469 NARY_SHAPE_INFER(tosa::EqualOp)
1470 NARY_SHAPE_INFER(tosa::ExpOp)
1471 NARY_SHAPE_INFER(tosa::FloorOp)
1472 NARY_SHAPE_INFER(tosa::GreaterEqualOp)
1473 NARY_SHAPE_INFER(tosa::GreaterOp)
1474 NARY_SHAPE_INFER(tosa::IdentityOp)
1475 NARY_SHAPE_INFER(tosa::LogOp)
1476 NARY_SHAPE_INFER(tosa::LogicalAndOp)
1477 NARY_SHAPE_INFER(tosa::LogicalLeftShiftOp)
1478 NARY_SHAPE_INFER(tosa::LogicalNotOp)
1479 NARY_SHAPE_INFER(tosa::LogicalOrOp)
1480 NARY_SHAPE_INFER(tosa::LogicalRightShiftOp)
1481 NARY_SHAPE_INFER(tosa::LogicalXorOp)
1482 NARY_SHAPE_INFER(tosa::MaximumOp)
1483 NARY_SHAPE_INFER(tosa::MinimumOp)
1484 NARY_SHAPE_INFER(tosa::MulOp)
1485 NARY_SHAPE_INFER(tosa::NegateOp)
1486 NARY_SHAPE_INFER(tosa::PowOp)
1487 NARY_SHAPE_INFER(tosa::ReciprocalOp)
1488 NARY_SHAPE_INFER(tosa::ReluNOp)
1489 NARY_SHAPE_INFER(tosa::RescaleOp)
1490 NARY_SHAPE_INFER(tosa::ReverseOp)
1491 NARY_SHAPE_INFER(tosa::RsqrtOp)
1492 NARY_SHAPE_INFER(tosa::SelectOp)
1493 NARY_SHAPE_INFER(tosa::SubOp)
1494 NARY_SHAPE_INFER(tosa::TanhOp)
1495 NARY_SHAPE_INFER(tosa::SigmoidOp)
1496 #undef PRED_SHAPE_INFER
1497 
1498 static LogicalResult poolingInferReturnTypes(
1499     const ValueShapeRange &operands, DictionaryAttr attributes,
1500     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1501   ShapeAdaptor inputShape = operands.getShape(0);
1502   llvm::SmallVector<int64_t> outputShape;
1503   outputShape.resize(4, -1);
1504 
1505   // We only know the rank if the input type is unranked.
1506   if (!inputShape) {
1507     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1508     return success();
1509   }
1510 
1511   // Batch and number of channels are identical for pooling layer.
1512   outputShape[0] = inputShape.getDimSize(0);
1513   outputShape[3] = inputShape.getDimSize(3);
1514 
1515   int32_t height = inputShape.getDimSize(1);
1516   int32_t width = inputShape.getDimSize(2);
1517 
1518   llvm::SmallVector<int64_t> kernel;
1519   llvm::SmallVector<int64_t> stride;
1520   llvm::SmallVector<int64_t> pad;
1521 
1522   getI64Values(attributes.get("kernel").cast<ArrayAttr>(), kernel);
1523   getI64Values(attributes.get("stride").cast<ArrayAttr>(), stride);
1524   getI64Values(attributes.get("pad").cast<ArrayAttr>(), pad);
1525 
1526   if (height != -1) {
1527     int32_t padded = height + pad[0] + pad[1] - kernel[0];
1528     outputShape[1] = padded / stride[0] + 1;
1529   }
1530 
1531   if (width != -1) {
1532     int32_t padded = width + pad[2] + pad[3] - kernel[1];
1533     outputShape[2] = padded / stride[1] + 1;
1534   }
1535 
1536   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1537   return success();
1538 }
1539 
1540 LogicalResult Conv2DOp::inferReturnTypeComponents(
1541     MLIRContext *context, ::llvm::Optional<Location> location,
1542     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1543     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1544   llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamicSize);
1545   Conv2DOp::Adaptor adaptor(operands.getValues(), attributes);
1546 
1547   int32_t inputWidth = ShapedType::kDynamicSize;
1548   int32_t inputHeight = ShapedType::kDynamicSize;
1549   int32_t weightWidth = ShapedType::kDynamicSize;
1550   int32_t weightHeight = ShapedType::kDynamicSize;
1551 
1552   // Input shape describes input width/height and batch.
1553 
1554   ShapeAdaptor inputShape = operands.getShape(adaptor.input());
1555   if (inputShape.hasRank()) {
1556     outputShape[0] = inputShape.getDimSize(0);
1557     inputHeight = inputShape.getDimSize(1);
1558     inputWidth = inputShape.getDimSize(2);
1559   }
1560 
1561   // Weight shapes describes the filter width/height and the output channels.
1562   ShapeAdaptor weightShape = operands.getShape(adaptor.weight());
1563   if (weightShape.hasRank()) {
1564     outputShape[3] = weightShape.getDimSize(0);
1565     weightHeight = weightShape.getDimSize(1);
1566     weightWidth = weightShape.getDimSize(2);
1567   }
1568 
1569   // Bias shape can describe the output channels.
1570   ShapeAdaptor biasShape = operands.getShape(adaptor.bias());
1571   if (biasShape.hasRank()) {
1572     outputShape[3] = ShapedType::isDynamic(outputShape[3])
1573                          ? biasShape.getDimSize(0)
1574                          : outputShape[3];
1575   }
1576 
1577   llvm::SmallVector<int64_t> dilation;
1578   llvm::SmallVector<int64_t> padding;
1579   llvm::SmallVector<int64_t> stride;
1580 
1581   getI64Values(adaptor.dilation(), dilation);
1582   getI64Values(adaptor.pad(), padding);
1583   getI64Values(adaptor.stride(), stride);
1584 
1585   if (!ShapedType::isDynamic(inputHeight) &&
1586       !ShapedType::isDynamic(weightHeight)) {
1587     int32_t inputSize = inputHeight + padding[0] + padding[1];
1588     int32_t filterSize = (weightHeight - 1) * dilation[0] + 1;
1589     int32_t unstridedResult = inputSize - filterSize + 1;
1590     outputShape[1] = (unstridedResult - 1) / stride[0] + 1;
1591   }
1592 
1593   if (!ShapedType::isDynamic(inputWidth) &&
1594       !ShapedType::isDynamic(weightWidth)) {
1595     int32_t inputSize = inputWidth + padding[2] + padding[3];
1596     int32_t filterSize = (weightWidth - 1) * dilation[1] + 1;
1597     int32_t unstridedResult = inputSize - filterSize + 1;
1598     outputShape[2] = (unstridedResult - 1) / stride[1] + 1;
1599   }
1600 
1601   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1602   return success();
1603 }
1604 
1605 LogicalResult Conv2DOp::verify() { return verifyConvOp(*this); }
1606 
1607 LogicalResult Conv3DOp::inferReturnTypeComponents(
1608     MLIRContext *context, ::llvm::Optional<Location> location,
1609     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1610     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1611   llvm::SmallVector<int64_t> outputShape(5, ShapedType::kDynamicSize);
1612   Conv2DOp::Adaptor adaptor(operands.getValues(), attributes);
1613 
1614   int32_t inputWidth = ShapedType::kDynamicSize;
1615   int32_t inputHeight = ShapedType::kDynamicSize;
1616   int32_t inputDepth = ShapedType::kDynamicSize;
1617 
1618   int32_t weightWidth = ShapedType::kDynamicSize;
1619   int32_t weightHeight = ShapedType::kDynamicSize;
1620   int32_t weightDepth = ShapedType::kDynamicSize;
1621 
1622   // Input shape describes input width/height and batch.
1623   ShapeAdaptor inputShape = operands.getShape(adaptor.input());
1624   if (inputShape.hasRank()) {
1625     outputShape[0] = inputShape.getDimSize(0);
1626     inputHeight = inputShape.getDimSize(1);
1627     inputWidth = inputShape.getDimSize(2);
1628     inputDepth = inputShape.getDimSize(3);
1629   }
1630 
1631   // Weight shapes describes the filter width/height and the output channels.
1632   ShapeAdaptor weightShape = operands.getShape(adaptor.weight());
1633   if (weightShape.hasRank()) {
1634     outputShape[4] = weightShape.getDimSize(0);
1635     weightHeight = weightShape.getDimSize(1);
1636     weightWidth = weightShape.getDimSize(2);
1637     weightDepth = weightShape.getDimSize(3);
1638   }
1639 
1640   // Bias shape can describe the output channels.
1641   ShapeAdaptor biasShape = operands.getShape(adaptor.bias());
1642   if (biasShape.hasRank()) {
1643     outputShape[4] =
1644         (outputShape[4] == -1) ? biasShape.getDimSize(0) : outputShape[4];
1645   }
1646 
1647   llvm::SmallVector<int64_t> dilation;
1648   llvm::SmallVector<int64_t> padding;
1649   llvm::SmallVector<int64_t> stride;
1650 
1651   getI64Values(adaptor.dilation(), dilation);
1652   getI64Values(adaptor.pad(), padding);
1653   getI64Values(adaptor.stride(), stride);
1654 
1655   if (!ShapedType::isDynamic(inputHeight) &&
1656       !ShapedType::isDynamic(weightHeight)) {
1657     int32_t inputSize = inputHeight + padding[0] + padding[1];
1658     int32_t filterSize = (weightHeight - 1) * dilation[0] + 1;
1659     int32_t unstridedResult = inputSize - filterSize + 1;
1660     outputShape[1] = (unstridedResult - 1) / stride[0] + 1;
1661   }
1662 
1663   if (!ShapedType::isDynamic(inputWidth) &&
1664       !ShapedType::isDynamic(weightWidth)) {
1665     int32_t inputSize = inputWidth + padding[2] + padding[3];
1666     int32_t filterSize = (weightWidth - 1) * dilation[1] + 1;
1667     int32_t unstridedResult = inputSize - filterSize + 1;
1668     outputShape[2] = (unstridedResult - 1) / stride[1] + 1;
1669   }
1670 
1671   if (!ShapedType::isDynamic(inputDepth) &&
1672       !ShapedType::isDynamic(weightDepth)) {
1673     int32_t inputSize = inputDepth + padding[4] + padding[5];
1674     int32_t filterSize = (weightDepth - 1) * dilation[2] + 1;
1675     int32_t unstridedResult = inputSize - filterSize + 1;
1676     outputShape[3] = (unstridedResult - 1) / stride[2] + 1;
1677   }
1678 
1679   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1680   return success();
1681 }
1682 
1683 LogicalResult Conv3DOp::verify() { return verifyConvOp(*this); }
1684 
1685 LogicalResult AvgPool2dOp::inferReturnTypeComponents(
1686     MLIRContext *context, ::llvm::Optional<Location> location,
1687     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1688     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1689   return poolingInferReturnTypes(operands, attributes, inferredReturnShapes);
1690 }
1691 
1692 LogicalResult MaxPool2dOp::inferReturnTypeComponents(
1693     MLIRContext *context, ::llvm::Optional<Location> location,
1694     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1695     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1696   return poolingInferReturnTypes(operands, attributes, inferredReturnShapes);
1697 }
1698 
1699 LogicalResult DepthwiseConv2DOp::inferReturnTypeComponents(
1700     MLIRContext *context, ::llvm::Optional<Location> location,
1701     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1702     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1703   llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamicSize);
1704   DepthwiseConv2DOp::Adaptor adaptor(operands.getValues(), attributes);
1705 
1706   int32_t inputWidth = ShapedType::kDynamicSize;
1707   int32_t inputHeight = ShapedType::kDynamicSize;
1708   int32_t inputChannels = ShapedType::kDynamicSize;
1709 
1710   int32_t weightWidth = ShapedType::kDynamicSize;
1711   int32_t weightHeight = ShapedType::kDynamicSize;
1712   int32_t depthChannels = ShapedType::kDynamicSize;
1713 
1714   // Input shape describes input width/height and batch.
1715   ShapeAdaptor inputShape = operands.getShape(adaptor.input());
1716   if (inputShape.hasRank()) {
1717     outputShape[0] = inputShape.getDimSize(0);
1718     inputHeight = inputShape.getDimSize(1);
1719     inputWidth = inputShape.getDimSize(2);
1720     inputChannels = inputShape.getDimSize(3);
1721   }
1722 
1723   // Weight shapes describes the filter width/height and the output channels.
1724   ShapeAdaptor weightShape = operands.getShape(adaptor.weight());
1725   if (weightShape.hasRank()) {
1726     weightHeight = weightShape.getDimSize(0);
1727     weightWidth = weightShape.getDimSize(1);
1728     inputChannels = ShapedType::isDynamic(inputChannels)
1729                         ? weightShape.getDimSize(2)
1730                         : inputChannels;
1731     depthChannels = weightShape.getDimSize(3);
1732   }
1733 
1734   // If both inputChannels and depthChannels are available we can determine
1735   // the output channels.
1736   if (!ShapedType::isDynamic(inputChannels) &&
1737       !ShapedType::isDynamic(depthChannels)) {
1738     outputShape[3] = inputChannels * depthChannels;
1739   }
1740 
1741   // Bias shape can describe the output channels.
1742   ShapeAdaptor biasShape = operands.getShape(adaptor.bias());
1743   if (biasShape.hasRank()) {
1744     outputShape[3] = ShapedType::isDynamic(outputShape[3])
1745                          ? biasShape.getDimSize(0)
1746                          : outputShape[3];
1747   }
1748 
1749   llvm::SmallVector<int64_t> dilation;
1750   llvm::SmallVector<int64_t> padding;
1751   llvm::SmallVector<int64_t> stride;
1752 
1753   getI64Values(adaptor.dilation(), dilation);
1754   getI64Values(adaptor.pad(), padding);
1755   getI64Values(adaptor.stride(), stride);
1756 
1757   if (!ShapedType::isDynamic(inputHeight) &&
1758       !ShapedType::isDynamic(weightHeight)) {
1759     int32_t inputSize = inputHeight + padding[0] + padding[1];
1760     int32_t filterSize = (weightHeight - 1) * dilation[0] + 1;
1761     int32_t unstridedResult = inputSize - filterSize + 1;
1762     outputShape[1] = (unstridedResult - 1) / stride[0] + 1;
1763   }
1764 
1765   if (!ShapedType::isDynamic(inputWidth) &&
1766       !ShapedType::isDynamic(weightWidth)) {
1767     int32_t inputSize = inputWidth + padding[2] + padding[3];
1768     int32_t filterSize = (weightWidth - 1) * dilation[1] + 1;
1769     int32_t unstridedResult = inputSize - filterSize + 1;
1770     outputShape[2] = (unstridedResult - 1) / stride[1] + 1;
1771   }
1772 
1773   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1774   return success();
1775 }
1776 
1777 LogicalResult DepthwiseConv2DOp::verify() { return verifyConvOp(*this); }
1778 
1779 LogicalResult TransposeConv2DOp::inferReturnTypeComponents(
1780     MLIRContext *context, ::llvm::Optional<Location> location,
1781     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1782     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1783   TransposeConv2DOp::Adaptor adaptor(operands.getValues(), attributes);
1784   llvm::SmallVector<int64_t> outputShape;
1785   getI64Values(adaptor.out_shape(), outputShape);
1786 
1787   int32_t inputWidth = ShapedType::kDynamicSize;
1788   int32_t inputHeight = ShapedType::kDynamicSize;
1789   int32_t weightWidth = ShapedType::kDynamicSize;
1790   int32_t weightHeight = ShapedType::kDynamicSize;
1791 
1792   // Input shape describes input width/height and batch.
1793   ShapeAdaptor inputShape = operands.getShape(adaptor.input());
1794   if (inputShape.hasRank()) {
1795     outputShape[0] = ShapedType::isDynamic(outputShape[0])
1796                          ? inputShape.getDimSize(0)
1797                          : outputShape[0];
1798     inputHeight = inputShape.getDimSize(1);
1799     inputWidth = inputShape.getDimSize(2);
1800   }
1801 
1802   // Weight shapes describes the filter width/height and the output channels.
1803   ShapeAdaptor weightShape = operands.getShape(adaptor.filter());
1804   if (weightShape.hasRank()) {
1805     outputShape[3] = ShapedType::isDynamic(outputShape[3])
1806                          ? weightShape.getDimSize(0)
1807                          : outputShape[3];
1808     weightHeight = weightShape.getDimSize(1);
1809     weightWidth = weightShape.getDimSize(2);
1810   }
1811 
1812   // Bias shape can describe the output channels.
1813   ShapeAdaptor biasShape = operands.getShape(adaptor.input());
1814   if (biasShape.hasRank()) {
1815     outputShape[3] = ShapedType::isDynamic(outputShape[3])
1816                          ? biasShape.getDimSize(0)
1817                          : outputShape[3];
1818   }
1819 
1820   llvm::SmallVector<int64_t> dilation;
1821   llvm::SmallVector<int64_t> padding;
1822   llvm::SmallVector<int64_t> stride;
1823 
1824   getI64Values(adaptor.dilation(), dilation);
1825   getI64Values(adaptor.out_pad(), padding);
1826   getI64Values(adaptor.stride(), stride);
1827 
1828   if (!ShapedType::isDynamic(inputHeight) &&
1829       !ShapedType::isDynamic(weightHeight)) {
1830     int32_t dilated = (weightHeight - 1) * dilation[0] + 1;
1831     int32_t calculateSize =
1832         (inputHeight - 1) * stride[0] - padding[0] + dilated;
1833     outputShape[1] = outputShape[1] == -1 ? calculateSize : outputShape[1];
1834   }
1835 
1836   if (!ShapedType::isDynamic(inputWidth) &&
1837       !ShapedType::isDynamic(weightWidth)) {
1838     int32_t dilated = (weightWidth - 1) * dilation[1] + 1;
1839     int32_t calculateSize = (inputWidth - 1) * stride[1] - padding[1] + dilated;
1840     outputShape[2] = outputShape[2] == -1 ? calculateSize : outputShape[2];
1841   }
1842 
1843   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
1844   return success();
1845 }
1846 
1847 LogicalResult IfOp::inferReturnTypeComponents(
1848     MLIRContext *context, ::llvm::Optional<Location> location,
1849     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1850     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1851   llvm::SmallVector<tosa::YieldOp> yieldOps;
1852   for (Region *region : regions) {
1853     for (auto &block : *region)
1854       if (auto returnOp = dyn_cast<tosa::YieldOp>(block.getTerminator()))
1855         yieldOps.push_back(returnOp);
1856   }
1857 
1858   if (yieldOps.empty())
1859     return failure();
1860 
1861   // Get the initial type information for the yield op.
1862   llvm::SmallVector<ValueKnowledge> resultKnowledge;
1863   resultKnowledge.reserve(yieldOps.front().getNumOperands());
1864   for (auto operand : yieldOps.front().getOperands()) {
1865     resultKnowledge.push_back(
1866         ValueKnowledge::getKnowledgeFromType(operand.getType()));
1867   }
1868 
1869   for (auto yieldOp : yieldOps) {
1870     if (resultKnowledge.size() != yieldOp.getNumOperands())
1871       return failure();
1872 
1873     for (const auto &it : llvm::enumerate(yieldOp.getOperands())) {
1874       int32_t index = it.index();
1875       auto meet = ValueKnowledge::meet(
1876           resultKnowledge[index],
1877           ValueKnowledge::getKnowledgeFromType(it.value().getType()));
1878       if (!meet)
1879         continue;
1880       resultKnowledge[index] = meet;
1881     }
1882   }
1883 
1884   for (const ValueKnowledge &result : resultKnowledge) {
1885     inferredReturnShapes.push_back(result.getShapedTypeComponents());
1886   }
1887 
1888   return success();
1889 }
1890 
1891 LogicalResult WhileOp::inferReturnTypeComponents(
1892     MLIRContext *context, ::llvm::Optional<Location> location,
1893     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
1894     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1895   llvm::SmallVector<tosa::YieldOp> yieldOps;
1896   for (auto &block : *regions[1])
1897     if (auto returnOp = dyn_cast<tosa::YieldOp>(block.getTerminator()))
1898       yieldOps.push_back(returnOp);
1899 
1900   // TOSA's while must have a tosa.yield as its terminator. If not found this
1901   // tosa.while is invalid.
1902   if (yieldOps.empty())
1903     return failure();
1904 
1905   // Get the initial type information from the operand types.
1906   llvm::SmallVector<ValueKnowledge> resultKnowledge;
1907   resultKnowledge.reserve(yieldOps.front().getNumOperands());
1908   for (auto operand : yieldOps.front().getOperands()) {
1909     resultKnowledge.push_back(
1910         ValueKnowledge::getKnowledgeFromType(operand.getType()));
1911   }
1912 
1913   for (auto yieldOp : yieldOps) {
1914     if (resultKnowledge.size() != yieldOp.getNumOperands())
1915       return failure();
1916 
1917     for (const auto &it : llvm::enumerate(yieldOp.getOperands())) {
1918       int32_t index = it.index();
1919       if (auto meet = ValueKnowledge::meet(
1920               resultKnowledge[index],
1921               ValueKnowledge::getKnowledgeFromType(it.value().getType()))) {
1922         resultKnowledge[index] = meet;
1923       }
1924     }
1925   }
1926 
1927   for (const ValueKnowledge &result : resultKnowledge) {
1928     inferredReturnShapes.push_back(result.getShapedTypeComponents());
1929   }
1930 
1931   return success();
1932 }
1933 
1934 //===----------------------------------------------------------------------===//
1935 // TOSA Attribute Definitions.
1936 //===----------------------------------------------------------------------===//
1937 
1938 #define GET_ATTRDEF_CLASSES
1939 #include "mlir/Dialect/Tosa/IR/TosaAttributes.cpp.inc"
1940 
1941 //===----------------------------------------------------------------------===//
1942 // TOSA Operator Definitions.
1943 //===----------------------------------------------------------------------===//
1944 
1945 #define GET_OP_CLASSES
1946 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"
1947