[MLIR][TOSA] Added lowerings for Reduce operations to Linalg

[rsuderman]

Summary

Lowerings for min, max, prod, and sum reduction operations on int and float
values. This includes reduction tests for both cases.

[rsuderman]

Fixed incidental variable names.

[rsuderman]

Fixed clang tidy error.

[bondhugula]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

332

Doc comment please.

367

Likewise - all static methods are to have doc comments as well.

370

You may reduce space and verbosity with a Location loc = op->getLoc() on the top.

652

You don't need getOperation().

[rsuderman]

bondhugula@ comments

[mravishankar]

Looks good. Just one nit.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

377

I think you can drop the getResult() here. It should cast to Value automatically.

[rsuderman]

final nits

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

static mlir::SelectOp clampHelper(Operation *op, ValueRange args,
   auto minOrArg =
       rewriter.create<mlir::SelectOp>(loc, smallerThanMin, min, args[0]);
   auto largerThanMax = rewriter.create<T>(loc, pred, max, args[0]);
   return rewriter.create<mlir::SelectOp>(loc, largerThanMax, max, minOrArg);
 }
 
 static Value
 createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args,
-                                            ArrayRef<Type> resultTypes,
+                                            ArrayRef<Type> resultTys,
                                             PatternRewriter &rewriter) {
   Location loc = op->getLoc();
   auto elementTy =
       op->getOperand(0).getType().cast<ShapedType>().getElementType();
 
   // tosa::AbsOp
   if (isa<tosa::AbsOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::AbsFOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::AbsFOp>(loc, resultTys, args);
 
   // tosa::AddOp
   if (isa<tosa::AddOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::AddFOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::AddFOp>(loc, resultTys, args);
 
   if (isa<tosa::AddOp>(op) && elementTy.isa<IntegerType>())
-    return rewriter.create<mlir::AddIOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::AddIOp>(loc, resultTys, args);
 
   // tosa::SubOp
   if (isa<tosa::SubOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::SubFOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::SubFOp>(loc, resultTys, args);
 
   if (isa<tosa::SubOp>(op) && elementTy.isa<IntegerType>())
-    return rewriter.create<mlir::SubIOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::SubIOp>(loc, resultTys, args);
 
   // tosa::MulOp
   if (isa<tosa::MulOp>(op) && elementTy.isa<FloatType>()) {
     if (dyn_cast<tosa::MulOp>(op).shift() != 0) {
       (void)rewriter.notifyMatchFailure(op,
                                         "Cannot have shift value for float");
       return nullptr;
     }
-    return rewriter.create<mlir::MulFOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::MulFOp>(loc, resultTys, args);
   }
 
   if (isa<tosa::MulOp>(op) && elementTy.isa<IntegerType>()) {
-    auto mul =
-        rewriter.create<mlir::MulIOp>(loc, resultTypes, args[0], args[1]);
+    auto mul = rewriter.create<mlir::MulIOp>(loc, resultTys, args[0], args[1]);
     auto constant =
         rewriter.create<mlir::ConstantOp>(loc, elementTy, op->getAttr("shift"));
-    return rewriter.create<mlir::SignedShiftRightOp>(loc, resultTypes, mul,
+    return rewriter.create<mlir::SignedShiftRightOp>(loc, resultTys, mul,
                                                      constant);
   }
 
   // tosa::NegateOp
   if (isa<tosa::NegateOp>(op) && elementTy.isa<IntegerType>()) {
     auto constant =
         rewriter.create<mlir::ConstantOp>(loc, IntegerAttr::get(elementTy, -1));
-    return rewriter.create<mlir::MulIOp>(loc, resultTypes, args[0], constant);
+    return rewriter.create<mlir::MulIOp>(loc, resultTys, args[0], constant);
   }
 
   if (isa<tosa::NegateOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::NegFOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::NegFOp>(loc, resultTys, args);
 
   // tosa::BitwiseAndOp
   if (isa<tosa::BitwiseAndOp>(op) && elementTy.isa<IntegerType>())
-    return rewriter.create<mlir::AndOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::AndOp>(loc, resultTys, args);
 
   // tosa::BitwiseOrOp
   if (isa<tosa::BitwiseOrOp>(op) && elementTy.isa<IntegerType>())
-    return rewriter.create<mlir::OrOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::OrOp>(loc, resultTys, args);
 
   // tosa::BitwiseXOrOp
   if (isa<tosa::BitwiseXorOp>(op) && elementTy.isa<IntegerType>())
-    return rewriter.create<mlir::XOrOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::XOrOp>(loc, resultTys, args);
 
   // tosa::LogicalLeftShiftOp
   if (isa<tosa::LogicalLeftShiftOp>(op) && elementTy.isa<IntegerType>())
-    return rewriter.create<mlir::ShiftLeftOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::ShiftLeftOp>(loc, resultTys, args);
 
   // tosa::LogicalrightShiftOp
   if (isa<tosa::LogicalRightShiftOp>(op) && elementTy.isa<IntegerType>())
-    return rewriter.create<mlir::UnsignedShiftRightOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::UnsignedShiftRightOp>(loc, resultTys, args);
 
   // tosa::PowOp
   if (isa<tosa::PowOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::math::PowFOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::math::PowFOp>(loc, resultTys, args);
 
   // tosa::RsqrtOp
   if (isa<tosa::RsqrtOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::math::RsqrtOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::math::RsqrtOp>(loc, resultTys, args);
 
   // tosa::LogOp
   if (isa<tosa::LogOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::math::LogOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::math::LogOp>(loc, resultTys, args);
 
   // tosa::ExpOp
   if (isa<tosa::ExpOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::math::ExpOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::math::ExpOp>(loc, resultTys, args);
 
   // tosa::TanhOp
   if (isa<tosa::TanhOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::math::TanhOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::math::TanhOp>(loc, resultTys, args);
 
   // tosa::GreaterOp
   if (isa<tosa::GreaterOp>(op) && elementTy.isa<FloatType>())
     return rewriter.create<mlir::CmpFOp>(loc, CmpFPredicate::OGT, args[0],
                                          args[1]);
 
   if (isa<tosa::GreaterOp>(op) && elementTy.isSignlessInteger())
     return rewriter.create<mlir::CmpIOp>(loc, CmpIPredicate::sgt, args[0],
createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args,
   if (isa<tosa::MinimumOp>(op) && elementTy.isSignlessInteger()) {
     auto predicate = rewriter.create<mlir::CmpIOp>(loc, CmpIPredicate::slt,
                                                    args[0], args[1]);
     return rewriter.create<mlir::SelectOp>(loc, predicate, args[0], args[1]);
   }
 
   // tosa::CeilOp
   if (isa<tosa::CeilOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::CeilFOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::CeilFOp>(loc, resultTys, args);
 
   // tosa::FloorOp
   if (isa<tosa::FloorOp>(op) && elementTy.isa<FloatType>())
-    return rewriter.create<mlir::FloorFOp>(loc, resultTypes, args);
+    return rewriter.create<mlir::FloorFOp>(loc, resultTys, args);
 
   // tosa::ClampOp
   if (isa<tosa::ClampOp>(op) && elementTy.isa<FloatType>()) {
     auto min = rewriter.create<mlir::ConstantOp>(loc, elementTy,
                                                  op->getAttr("min_fp"));
     auto max = rewriter.create<mlir::ConstantOp>(loc, elementTy,
                                                  op->getAttr("max_fp"));
     return clampHelper<mlir::CmpFOp>(op, args, min, max, CmpFPredicate::OLT,
elementwiseMatchAndRewriteHelper(Operation *operation,
   SmallVector<Type> bodyArgTypes;
 
   for (Value in : operation->getOperands())
     bodyArgTypes.emplace_back(getElementTypeOrSelf(in.getType()));
 
   SmallVector<Type> opResultTypes;
   SmallVector<Value> initTensors;
   for (auto result : results) {
-    auto resultType = result.getType().template cast<ShapedType>();
-    if (!resultType.hasStaticShape())
+    auto resultTy = result.getType().template cast<ShapedType>();
+    if (!resultTy.hasStaticShape())
       return rewriter.notifyMatchFailure(
           operation,
           "tosa to linalg conversion expects statically shaped tensors");
 
     initTensors.push_back(rewriter.create<linalg::InitTensorOp>(
-        loc, ArrayRef<Value>({}), resultType.getShape(),
-        resultType.getElementType()));
+        loc, ArrayRef<Value>({}), resultTy.getShape(),
+        resultTy.getElementType()));
     opResultTypes.push_back(result.getType());
   }
 
   auto bodyResultTypes = llvm::to_vector<4>(llvm::map_range(
       initTensors, [](Value v) { return getElementTypeOrSelf(v); }));
 
   unsigned nloops = t0.getRank();
   SmallVector<AffineMap, 2> indexingMaps;
elementwiseMatchAndRewriteHelper(Operation *operation,
 
   if (!didEncounterError)
     return failure();
 
   rewriter.replaceOp(operation, linalgOp->getResults());
   return success();
 }
 
+static Attribute createInitialValueForReduceOp(Operation *op, Type elementTy,
+                                               PatternRewriter &rewriter) {
+  if (isa<tosa::ReduceSumOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.getFloatAttr(elementTy, 0.0);
+
+  if (isa<tosa::ReduceSumOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.getIntegerAttr(elementTy, 0);
+
+  if (isa<tosa::ReduceProdOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.getFloatAttr(elementTy, 1.0);
+
+  if (isa<tosa::ReduceProdOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.getIntegerAttr(elementTy, 1);
+
+  if (isa<tosa::ReduceMinOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.getFloatAttr(
+        elementTy, APFloat::getLargest(
+                       elementTy.cast<FloatType>().getFloatSemantics(), false));
+
+  if (isa<tosa::ReduceMinOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.getIntegerAttr(
+        elementTy, APInt::getSignedMaxValue(elementTy.getIntOrFloatBitWidth()));
+
+  if (isa<tosa::ReduceMaxOp>(op) && elementTy.isa<FloatType>())
+    return rewriter.getFloatAttr(
+        elementTy, APFloat::getLargest(
+                       elementTy.cast<FloatType>().getFloatSemantics(), true));
+
+  if (isa<tosa::ReduceMaxOp>(op) && elementTy.isa<IntegerType>())
+    return rewriter.getIntegerAttr(
+        elementTy, APInt::getSignedMinValue(elementTy.getIntOrFloatBitWidth()));
+
+  return {};
+}
+
+static Value createLinalgBodyCalculationForReduceOp(Operation *op,
+                                                    ValueRange args,
+                                                    Type elementTy,
+                                                    PatternRewriter &rewriter) {
+  if (isa<tosa::ReduceSumOp>(op) && elementTy.isa<FloatType>()) {
+    return rewriter.create<AddFOp>(op->getLoc(), args).getResult();
+  }
+
+  if (isa<tosa::ReduceSumOp>(op) && elementTy.isa<IntegerType>()) {
+    return rewriter.create<AddIOp>(op->getLoc(), args).getResult();
+  }
+
+  if (isa<tosa::ReduceProdOp>(op) && elementTy.isa<FloatType>()) {
+    return rewriter.create<MulFOp>(op->getLoc(), args).getResult();
+  }
+
+  if (isa<tosa::ReduceProdOp>(op) && elementTy.isa<IntegerType>()) {
+    return rewriter.create<MulIOp>(op->getLoc(), args).getResult();
+  }
+
+  if (isa<tosa::ReduceMinOp>(op) && elementTy.isa<FloatType>()) {
+    auto predicate = rewriter.create<mlir::CmpFOp>(
+        op->getLoc(), CmpFPredicate::OLT, args[0], args[1]);
+    return rewriter.create<mlir::SelectOp>(op->getLoc(), predicate, args[0],
+                                           args[1]);
+  }
+
+  if (isa<tosa::ReduceMinOp>(op) && elementTy.isa<IntegerType>()) {
+    auto predicate = rewriter.create<mlir::CmpIOp>(
+        op->getLoc(), CmpIPredicate::slt, args[0], args[1]);
+    return rewriter.create<mlir::SelectOp>(op->getLoc(), predicate, args[0],
+                                           args[1]);
+  }
+
+  if (isa<tosa::ReduceMaxOp>(op) && elementTy.isa<FloatType>()) {
+    auto predicate = rewriter.create<mlir::CmpFOp>(
+        op->getLoc(), CmpFPredicate::OGT, args[0], args[1]);
+    return rewriter.create<mlir::SelectOp>(op->getLoc(), predicate, args[0],
+                                           args[1]);
+  }
+
+  if (isa<tosa::ReduceMaxOp>(op) && elementTy.isa<IntegerType>()) {
+    auto predicate = rewriter.create<mlir::CmpIOp>(
+        op->getLoc(), CmpIPredicate::sgt, args[0], args[1]);
+    return rewriter.create<mlir::SelectOp>(op->getLoc(), predicate, args[0],
+                                           args[1]);
+  }
+
+  return {};
+}
+
+static LogicalResult
+elementwiseMatchAndRewriteHelper(Operation *op, uint64_t axis,
+                                 PatternRewriter &rewriter) {
+  auto loc = op->getLoc();
+  auto inputTy = op->getOperand(0).getType().template cast<ShapedType>();
+  auto resultTy = op->getResult(0).getType().template cast<ShapedType>();
+  auto elementTy = resultTy.getElementType();
+  Value input = op->getOperand(0);
+
+  // First fill the output buffer with the init value.
+  auto initTensor = rewriter
+                        .create<linalg::InitTensorOp>(loc, ArrayRef<Value>({}),
+                                                      resultTy.getShape(),
+                                                      resultTy.getElementType())
+                        .result();
+
+  auto fillValueAttr = createInitialValueForReduceOp(op, elementTy, rewriter);
+  if (!fillValueAttr)
+    return rewriter.notifyMatchFailure(
+        op, "No initial value found for reduction operation");
+
+  auto fillValue = rewriter.create<ConstantOp>(loc, fillValueAttr);
+  auto filledTensor =
+      rewriter.create<linalg::FillOp>(loc, initTensor, fillValue).result();
+
+  SmallVector<AffineExpr, 2> srcExprs;
+  SmallVector<AffineExpr, 2> dstExprs;
+  SmallVector<StringRef, 4> iteratorTypes;
+  for (unsigned int i = 0, rank = inputTy.getRank(); i != rank; ++i) {
+    srcExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext()));
+
+    iteratorTypes.push_back(axis == i ? getReductionIteratorTypeName()
+                                      : getParallelIteratorTypeName());
+    if (axis != i)
+      dstExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext()));
+  }
+
+  bool didEncounterError = false;
+  auto maps = AffineMap::inferFromExprList({srcExprs, dstExprs});
+  auto linalgOp = rewriter.create<linalg::GenericOp>(
+      loc, resultTy, input, filledTensor, maps, iteratorTypes,
+      [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) {
+        auto result = createLinalgBodyCalculationForReduceOp(
+            op, blockArgs, elementTy, rewriter);
+        if (result)
+          didEncounterError = true;
+
+        nestedBuilder.create<linalg::YieldOp>(loc, result);
+      });
+
+  if (!didEncounterError)
+    return failure();
+
+  rewriter.replaceOp(op, linalgOp.getOperation()->getResults());
+  return success();
+}
+
 namespace {
 
 template <typename SrcOp>
 class PointwiseConverter : public OpRewritePattern<SrcOp> {
 public:
   using OpRewritePattern<SrcOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(SrcOp op,
public:
 
   LogicalResult matchAndRewrite(SrcOp op,
                                 PatternRewriter &rewriter) const final {
     rewriter.replaceOp(op, op.getOperation()->getOperands());
     return success();
   }
 };
 
+template <typename SrcOp>
+class ReduceConverter : public OpRewritePattern<SrcOp> {
+public:
+  using OpRewritePattern<SrcOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SrcOp reduce_op,
+                                PatternRewriter &rewriter) const final {
+    return elementwiseMatchAndRewriteHelper(reduce_op.getOperation(),
+                                            reduce_op.axis(), rewriter);
+  }
+};
+
 } // namespace
 
 void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
     MLIRContext *context, OwningRewritePatternList *patterns) {
   patterns->insert<
       PointwiseConverter<tosa::AddOp>, PointwiseConverter<tosa::SubOp>,
       PointwiseConverter<tosa::MulOp>, PointwiseConverter<tosa::NegateOp>,
       PointwiseConverter<tosa::PowOp>, PointwiseConverter<tosa::RsqrtOp>,
       PointwiseConverter<tosa::LogOp>, PointwiseConverter<tosa::ExpOp>,
       PointwiseConverter<tosa::AbsOp>, PointwiseConverter<tosa::TanhOp>,
       PointwiseConverter<tosa::BitwiseAndOp>,
       PointwiseConverter<tosa::BitwiseOrOp>,
       PointwiseConverter<tosa::BitwiseXorOp>,
       PointwiseConverter<tosa::LogicalLeftShiftOp>,
       PointwiseConverter<tosa::LogicalRightShiftOp>,
       PointwiseConverter<tosa::SelectOp>, PointwiseConverter<tosa::GreaterOp>,
       PointwiseConverter<tosa::GreaterEqualOp>,
       PointwiseConverter<tosa::MaximumOp>, PointwiseConverter<tosa::MinimumOp>,
       PointwiseConverter<tosa::CeilOp>, PointwiseConverter<tosa::FloorOp>,
       PointwiseConverter<tosa::ClampOp>, PointwiseConverter<tosa::ReluNOp>,
       IdentityNConverter<tosa::IdentityOp>,
-      IdentityNConverter<tosa::IdentityNOp>,
-      ReshapeOpConverter, TransposeConverter>(context);
+      IdentityNConverter<tosa::IdentityNOp>, ReduceConverter<tosa::ReduceMinOp>,
+      ReduceConverter<tosa::ReduceMaxOp>, ReduceConverter<tosa::ReduceSumOp>,
+      ReduceConverter<tosa::ReduceProdOp>, ReshapeOpConverter,
+      TransposeConverter>(context);
 }

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func @test_transpose(%arg0: tensor<1x2x3xi32>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3, 1]
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel"]} ins([[ARG0]] : tensor<1x2x3xi32>) outs([[OUT:%.+]] : tensor<2x3x1xi32>)
   // CHECK: ^bb0([[ARG1:%.+]]: i32, [[ARG2:%.+]]: i32)
   // CHECK:   linalg.yield [[ARG1]]
   // CHECK: }
   %1 = "tosa.transpose"(%arg0, %0) : (tensor<1x2x3xi32>, tensor<3xi32>) -> (tensor<2x3x1xi32>)
   return
 }
+
+// -----
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
+// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-LABEL: @reduce_float
+// CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xf32>
+func @reduce_float(%arg0: tensor<5x4xf32>) -> () {
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [4]
+  // CHECK: [[CST0:%.+]] = constant 0.0
+  // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST0]])
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<4xf32>)
+  // CHECK: ^bb0(%arg1: f32, %arg2: f32)
+  // CHECK:   [[RES:%.+]] = addf %arg1, %arg2 : f32
+  // CHECK:   linalg.yield [[RES]] : f32
+  %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<4xf32>
+
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [5]
+  // CHECK: [[CST0:%.+]] = constant 0.0
+  // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST0]])
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<5xf32>)
+  // CHECK: ^bb0(%arg1: f32, %arg2: f32)
+  // CHECK:   [[RES:%.+]] = addf %arg1, %arg2 : f32
+  // CHECK:   linalg.yield [[RES]] : f32
+  %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xf32>) -> tensor<5xf32>
+
+  // CHECK: constant 1.0
+  // CHECK: linalg.fill
+  // CHECK: linalg.generic
+  // CHECK: mulf
+  %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<4xf32>
+
+  // CHECK: constant 3.40282347E+38 : f32
+  // CHECK: linalg.fill
+  // CHECK: linalg.generic
+  // CHECK: cmpf olt
+  // CHECK: select
+  %3 = "tosa.reduce_min"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<4xf32>
+
+  // CHECK: constant -3.40282347E+38 : f32
+  // CHECK: linalg.fill
+  // CHECK: linalg.generic
+  // CHECK: cmpf ogt
+  // CHECK: select
+  %4 = "tosa.reduce_max"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<4xf32>
+  return
+}
+
+// -----
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
+// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
+
+// CHECK-LABEL: @reduce_int
+// CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xi32>
+func @reduce_int(%arg0: tensor<5x4xi32>) -> () {
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [4]
+  // CHECK: [[CST0:%.+]] = constant 0
+  // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST0]])
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<4xi32>)
+  // CHECK: ^bb0(%arg1: i32, %arg2: i32)
+  // CHECK:   [[RES:%.+]] = addi %arg1, %arg2 : i32
+  // CHECK:   linalg.yield [[RES]] : i32
+  %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<4xi32>
+
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [5]
+  // CHECK: [[CST0:%.+]] = constant 0
+  // CHECK: [[FILL:%.+]] = linalg.fill([[INIT]], [[CST0]])
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<5xi32>)
+  // CHECK: ^bb0(%arg1: i32, %arg2: i32)
+  // CHECK:   [[RES:%.+]] = addi %arg1, %arg2 : i32
+  // CHECK:   linalg.yield [[RES]] : i32
+  %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5xi32>
+
+  // CHECK: constant 1
+  // CHECK: linalg.fill
+  // CHECK: linalg.generic
+  // CHECK: muli
+  %2 = "tosa.reduce_prod"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<4xi32>
+
+  // CHECK: constant 2147483647 : i32
+  // CHECK: linalg.fill
+  // CHECK: linalg.generic
+  // CHECK: cmpi slt
+  // CHECK: select
+  %3 = "tosa.reduce_min"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<4xi32>
+
+  // CHECK: constant -2147483648 : i32
+  // CHECK: linalg.fill
+  // CHECK: linalg.generic
+  // CHECK: cmpi sgt
+  // CHECK: select
+  %4 = "tosa.reduce_max"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<4xi32>
+  return
+}