[StableHLO] Port reduce canon patterns (#14045)

Turns out that linalg lowering depends on these.

The code is ported from equivalent folds and rewrite patterns in the
mlir-hlo repo. The only notable difference is that this implementation
fixes bugs with rewrite patterns performing op updates without going
through the pattern rewriter. Also added a new test for the second noop
case.

Fixes: #14042
Issue: #12678

compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/Canonicalization.cpp

@@ -694,6 +694,80 @@ struct DynamicBroadcastInDimAllDimsNonExpanding final
   }
 };
 
+struct NoopReduceOpCanon final : OpRewritePattern<mlir::stablehlo::ReduceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::ReduceOp op,
+                                PatternRewriter &rewriter) const override {
+    // No dimensions to reduce.
+    if (op.getDimensions().empty()) {
+      rewriter.replaceOp(op, op.getInputs());
+      return success();
+    }
+
+    // If all returned values in the ReduceOp region exists outside the
+    // region, replace the ReduceOp with those values.
+    if (auto retOp = dyn_cast<mlir::stablehlo::ReturnOp>(
+            op.getBody().front().getTerminator())) {
+      Region *retRegion = retOp->getParentRegion();
+      if (llvm::any_of(retOp.getResults(), [retRegion](Value result) {
+            return result.getParentRegion() == retRegion;
+          })) {
+        return failure();
+      }
+
+      rewriter.replaceOp(op, retOp.getResults());
+      return success();
+    }
+
+    return failure();
+  }
+};
+
+struct EmptyReduceOpCanon final : OpRewritePattern<mlir::stablehlo::ReduceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::ReduceOp op,
+                                PatternRewriter &rewriter) const override {
+    // We require all reduce shapes to be the same, up to the element types, so
+    // we can just the first operand and the first result as a representative.
+    auto elemTy = dyn_cast<RankedTensorType>(op.getInputs().getType().front());
+    if (!elemTy) {
+      return rewriter.notifyMatchFailure(op.getLoc(),
+                                         "unranked input unsupported");
+    }
+
+    if (!llvm::is_contained(elemTy.getShape(), 0)) return failure();
+
+    Location loc = op.getLoc();
+    DenseIntElementsAttr empty = rewriter.getI64TensorAttr({});
+    if (elemTy.hasStaticShape()) {
+      SmallVector<Value> broadcasts(op.getNumResults());
+      for (auto [bcast, init, outTy] : llvm::zip_equal(
+               broadcasts, op.getInitValues(), op.getResultTypes())) {
+        bcast = rewriter.create<mlir::stablehlo::BroadcastInDimOp>(loc, outTy,
+                                                                   init, empty);
+      }
+      rewriter.replaceOp(op, broadcasts);
+      return success();
+    }
+
+    SmallVector<Value> shapes;
+    if (failed(op.reifyReturnTypeShapes(rewriter, op.getOperands(), shapes))) {
+      return failure();
+    }
+
+    SmallVector<Value> broadcasts(op.getNumResults());
+    for (auto [bcast, init, shape, outTy] : llvm::zip_equal(
+             broadcasts, op.getInitValues(), shapes, op.getResultTypes())) {
+      bcast = rewriter.create<mlir::stablehlo::DynamicBroadcastInDimOp>(
+          loc, outTy, init, shape, empty);
+    }
+    rewriter.replaceOp(op, broadcasts);
+    return success();
+  }
+};
+
 struct DynamicReshapeOpCanon final
     : OpRewritePattern<mlir::stablehlo::DynamicReshapeOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -922,6 +996,8 @@ void populateCanonicalizationPatterns(MLIRContext *context,
       BroadcastInDimOpCanon, DynamicBroadcastInDimOpNotActuallyDynamic,
       ChainedDynamicBroadcastInDimCanonicalization,
       DynamicBroadcastInDimAllDimsNonExpanding,
+      // Reduce op.
+      NoopReduceOpCanon, EmptyReduceOpCanon,
       // Shape manipulation(-ish) ops.
       ConcatenateOpCanon, ConvertOpCanon, DynamicReshapeOpCanon, GatherOpCanon,
       ReshapeOpCanon, TransposeOpCanon>(context, benefit);

compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/test/canonicalization.mlir

@@ -427,7 +427,7 @@ func.func @dynamic_broadcast_in_dim_op_not_actually_dynamic(%arg0: tensor<4xf32>
 
 // CHECK-LABEL: func @dynamic_broadcast_in_dim_op_not_actually_dynamic_constant_shape
 func.func @dynamic_broadcast_in_dim_op_not_actually_dynamic_constant_shape(%arg0: tensor<i32>) -> tensor<4x32xi32> {
-  %0 = mhlo.constant dense<[4, 32]> : tensor<2xi32>
+  %0 = stablehlo.constant dense<[4, 32]> : tensor<2xi32>
   // CHECK: %[[RESULT:.+]] = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<i32>) -> tensor<4x32xi32>
   %1 = stablehlo.dynamic_broadcast_in_dim %arg0, %0, dims = [] : (tensor<i32>, tensor<2xi32>) -> tensor<?x32xi32>
   %2 = stablehlo.dynamic_reshape %1, %0 : (tensor<?x32xi32>, tensor<2xi32>) -> tensor<4x32xi32>
@@ -584,3 +584,48 @@ func.func @gather_to_slice_indices_clamp_lowerbound(%arg0 : tensor<4x2xui32>) ->
   // CHECK-NEXT: %[[V1:.*]] = stablehlo.reshape %[[V0]] : (tensor<1x2xui32>) -> tensor<2xui32>
   // CHECK-NEXT: return %[[V1]] : tensor<2xui32>
 }
+
+// -----
+
+// CHECK-LABEL: func.func @reduce_noop_1
+// CHECK-SAME:   ([[ARG0:%.+]]: tensor<4x8xf32>)
+func.func @reduce_noop_1(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+  %0 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  %1 = stablehlo.reduce(%arg0 init: %0) across dimensions = [] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4x8xf32>
+    reducer(%arg1: tensor<f32>, %arg2: tensor<f32>) {
+    %4 = stablehlo.add %arg1, %arg2 : tensor<f32>
+    stablehlo.return %4 : tensor<f32>
+  }
+  // CHECK: return [[ARG0]] : tensor<4x8xf32>
+  func.return %1 : tensor<4x8xf32>
+}
+
+// CHECK-LABEL: func.func @reduce_noop_2
+// CHECK-SAME:   ([[ARG0:%.+]]: tensor<4x8xi32>, [[ARG1:%.+]]: tensor<i32>)
+func.func @reduce_noop_2(%arg0: tensor<4x8xi32>, %arg1: tensor<i32>) -> tensor<i32> {
+  %0 = stablehlo.constant dense<0> : tensor<i32>
+  %1 = stablehlo.reduce(%arg0 init: %0) across dimensions = [0, 1] : (tensor<4x8xi32>, tensor<i32>) -> tensor<i32>
+    reducer(%b1: tensor<i32>, %b2: tensor<i32>) {
+    stablehlo.return %arg1 : tensor<i32>
+  }
+  // CHECK: return [[ARG1]] : tensor<i32>
+  func.return %1 : tensor<i32>
+}
+
+// CHECK-LABEL: func.func @reduce_zero_ext
+func.func @reduce_zero_ext(%arg0: tensor<0xi1>) -> tensor<i32> {
+  %0 = stablehlo.constant dense<false> : tensor<i1>
+  %1 = stablehlo.constant dense<false> : tensor<0xi1>
+  %2 = stablehlo.compare  NE, %arg0, %1, UNSIGNED : (tensor<0xi1>, tensor<0xi1>) -> tensor<0xi1>
+  %3 = stablehlo.convert %2 : (tensor<0xi1>) -> tensor<0xi32>
+  %4 = stablehlo.constant dense<0> : tensor<i32>
+  %5 = stablehlo.reduce(%3 init: %4) across dimensions = [0] : (tensor<0xi32>, tensor<i32>) -> tensor<i32>
+    reducer(%arg1: tensor<i32>, %arg2: tensor<i32>)  {
+    %6 = stablehlo.add %arg1, %arg2 : tensor<i32>
+    stablehlo.return %6 : tensor<i32>
+  }
+
+  // CHECK: [[CST:%.+]] = stablehlo.constant dense<0> : tensor<i32>
+  // CHECK: return [[CST]] : tensor<i32>
+  return %5 : tensor<i32>
+}