diff --git a/flang/lib/Optimizer/HLFIR/Transforms/SimplifyHLFIRIntrinsics.cpp b/flang/lib/Optimizer/HLFIR/Transforms/SimplifyHLFIRIntrinsics.cpp
index a90c5a9ef5a4..28325bc8e548 100644
--- a/flang/lib/Optimizer/HLFIR/Transforms/SimplifyHLFIRIntrinsics.cpp
+++ b/flang/lib/Optimizer/HLFIR/Transforms/SimplifyHLFIRIntrinsics.cpp
@@ -19,11 +19,9 @@
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/HLFIR/Passes.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 namespace hlfir {
 #define GEN_PASS_DEF_SIMPLIFYHLFIRINTRINSICS
@@ -45,9 +43,15 @@ public:
   llvm::LogicalResult
   matchAndRewrite(hlfir::TransposeOp transpose,
                   mlir::PatternRewriter &rewriter) const override {
+    hlfir::ExprType expr = transpose.getType();
+    // TODO: hlfir.elemental supports polymorphic data types now,
+    // so this can be supported.
+    if (expr.isPolymorphic())
+      return rewriter.notifyMatchFailure(transpose,
+                                         "TRANSPOSE of polymorphic type");
+
     mlir::Location loc = transpose.getLoc();
     fir::FirOpBuilder builder{rewriter, transpose.getOperation()};
-    hlfir::ExprType expr = transpose.getType();
     mlir::Type elementType = expr.getElementType();
     hlfir::Entity array = hlfir::Entity{transpose.getArray()};
     mlir::Value resultShape = genResultShape(loc, builder, array);
@@ -105,15 +109,32 @@ public:
   llvm::LogicalResult
   matchAndRewrite(hlfir::SumOp sum,
                   mlir::PatternRewriter &rewriter) const override {
+    if (!simplifySum)
+      return rewriter.notifyMatchFailure(sum, "SUM simplification is disabled");
+
+    hlfir::Entity array = hlfir::Entity{sum.getArray()};
+    bool isTotalReduction = hlfir::Entity{sum}.getRank() == 0;
+    mlir::Value dim = sum.getDim();
+    int64_t dimVal = 0;
+    if (!isTotalReduction) {
+      // In case of partial reduction we should ignore the operations
+      // with invalid DIM values. They may appear in dead code
+      // after constant propagation.
+      auto constDim = fir::getIntIfConstant(dim);
+      if (!constDim)
+        return rewriter.notifyMatchFailure(sum, "Nonconstant DIM for SUM");
+      dimVal = *constDim;
+
+      if ((dimVal <= 0 || dimVal > array.getRank()))
+        return rewriter.notifyMatchFailure(
+            sum, "Invalid DIM for partial SUM reduction");
+    }
+
     mlir::Location loc = sum.getLoc();
     fir::FirOpBuilder builder{rewriter, sum.getOperation()};
     mlir::Type elementType = hlfir::getFortranElementType(sum.getType());
-    hlfir::Entity array = hlfir::Entity{sum.getArray()};
     mlir::Value mask = sum.getMask();
-    mlir::Value dim = sum.getDim();
-    bool isTotalReduction = hlfir::Entity{sum}.getRank() == 0;
-    int64_t dimVal =
-        isTotalReduction ? 0 : fir::getIntIfConstant(dim).value_or(0);
+
     mlir::Value resultShape, dimExtent;
     llvm::SmallVector<mlir::Value> arrayExtents;
     if (isTotalReduction)
@@ -360,27 +381,38 @@ class CShiftAsElementalConversion
 public:
   using mlir::OpRewritePattern<hlfir::CShiftOp>::OpRewritePattern;
 
-  explicit CShiftAsElementalConversion(mlir::MLIRContext *ctx)
-      : OpRewritePattern(ctx) {
-    setHasBoundedRewriteRecursion();
-  }
-
   llvm::LogicalResult
   matchAndRewrite(hlfir::CShiftOp cshift,
                   mlir::PatternRewriter &rewriter) const override {
     using Fortran::common::maxRank;
 
-    mlir::Location loc = cshift.getLoc();
-    fir::FirOpBuilder builder{rewriter, cshift.getOperation()};
     hlfir::ExprType expr = mlir::dyn_cast<hlfir::ExprType>(cshift.getType());
     assert(expr &&
            "expected an expression type for the result of hlfir.cshift");
+    unsigned arrayRank = expr.getRank();
+    // When it is a 1D CSHIFT, we may assume that the DIM argument
+    // (whether it is present or absent) is equal to 1, otherwise,
+    // the program is illegal.
+    int64_t dimVal = 1;
+    if (arrayRank != 1)
+      if (mlir::Value dim = cshift.getDim()) {
+        auto constDim = fir::getIntIfConstant(dim);
+        if (!constDim)
+          return rewriter.notifyMatchFailure(cshift,
+                                             "Nonconstant DIM for CSHIFT");
+        dimVal = *constDim;
+      }
+
+    if (dimVal <= 0 || dimVal > arrayRank)
+      return rewriter.notifyMatchFailure(cshift, "Invalid DIM for CSHIFT");
+
+    mlir::Location loc = cshift.getLoc();
+    fir::FirOpBuilder builder{rewriter, cshift.getOperation()};
     mlir::Type elementType = expr.getElementType();
     hlfir::Entity array = hlfir::Entity{cshift.getArray()};
     mlir::Value arrayShape = hlfir::genShape(loc, builder, array);
     llvm::SmallVector<mlir::Value> arrayExtents =
         hlfir::getExplicitExtentsFromShape(arrayShape, builder);
-    unsigned arrayRank = expr.getRank();
     llvm::SmallVector<mlir::Value, 1> typeParams;
     hlfir::genLengthParameters(loc, builder, array, typeParams);
     hlfir::Entity shift = hlfir::Entity{cshift.getShift()};
@@ -395,20 +427,6 @@ public:
       shiftVal = builder.createConvert(loc, calcType, shiftVal);
     }
 
-    int64_t dimVal = 1;
-    if (arrayRank == 1) {
-      // When it is a 1D CSHIFT, we may assume that the DIM argument
-      // (whether it is present or absent) is equal to 1, otherwise,
-      // the program is illegal.
-      assert(shiftVal && "SHIFT must be scalar");
-    } else {
-      if (mlir::Value dim = cshift.getDim())
-        dimVal = fir::getIntIfConstant(dim).value_or(0);
-      assert(dimVal > 0 && dimVal <= arrayRank &&
-             "DIM must be present and a positive constant not exceeding "
-             "the array's rank");
-    }
-
     auto genKernel = [&](mlir::Location loc, fir::FirOpBuilder &builder,
                          mlir::ValueRange inputIndices) -> hlfir::Entity {
       llvm::SmallVector<mlir::Value, maxRank> indices{inputIndices};
@@ -462,68 +480,19 @@ class SimplifyHLFIRIntrinsics
 public:
   void runOnOperation() override {
     mlir::MLIRContext *context = &getContext();
+
+    mlir::GreedyRewriteConfig config;
+    // Prevent the pattern driver from merging blocks
+    config.enableRegionSimplification =
+        mlir::GreedySimplifyRegionLevel::Disabled;
+
     mlir::RewritePatternSet patterns(context);
     patterns.insert<TransposeAsElementalConversion>(context);
     patterns.insert<SumAsElementalConversion>(context);
     patterns.insert<CShiftAsElementalConversion>(context);
-    mlir::ConversionTarget target(*context);
-    // don't transform transpose of polymorphic arrays (not currently supported
-    // by hlfir.elemental)
-    target.addDynamicallyLegalOp<hlfir::TransposeOp>(
-        [](hlfir::TransposeOp transpose) {
-          return mlir::cast<hlfir::ExprType>(transpose.getType())
-              .isPolymorphic();
-        });
-    // Handle only SUM(DIM=CONSTANT) case for now.
-    // It may be beneficial to expand the non-DIM case as well.
-    // E.g. when the input array is an elemental array expression,
-    // expanding the SUM into a total reduction loop nest
-    // would avoid creating a temporary for the elemental array expression.
-    target.addDynamicallyLegalOp<hlfir::SumOp>([](hlfir::SumOp sum) {
-      if (!simplifySum)
-        return true;
 
-      // Always inline total reductions.
-      if (hlfir::Entity{sum}.getRank() == 0)
-        return false;
-      mlir::Value dim = sum.getDim();
-      if (!dim)
-        return false;
-
-      if (auto dimVal = fir::getIntIfConstant(dim)) {
-        fir::SequenceType arrayTy = mlir::cast<fir::SequenceType>(
-            hlfir::getFortranElementOrSequenceType(sum.getArray().getType()));
-        if (*dimVal > 0 && *dimVal <= arrayTy.getDimension()) {
-          // Ignore SUMs with illegal DIM values.
-          // They may appear in dead code,
-          // and they do not have to be converted.
-          return false;
-        }
-      }
-      return true;
-    });
-    target.addDynamicallyLegalOp<hlfir::CShiftOp>([](hlfir::CShiftOp cshift) {
-      unsigned resultRank = hlfir::Entity{cshift}.getRank();
-      if (resultRank == 1)
-        return false;
-
-      mlir::Value dim = cshift.getDim();
-      if (!dim)
-        return false;
-
-      // If DIM is present, then it must be constant to please
-      // the conversion. In addition, ignore cases with
-      // illegal DIM values.
-      if (auto dimVal = fir::getIntIfConstant(dim))
-        if (*dimVal > 0 && *dimVal <= resultRank)
-          return false;
-
-      return true;
-    });
-    target.markUnknownOpDynamicallyLegal(
-        [](mlir::Operation *) { return true; });
-    if (mlir::failed(mlir::applyFullConversion(getOperation(), target,
-                                               std::move(patterns)))) {
+    if (mlir::failed(mlir::applyPatternsAndFoldGreedily(
+            getOperation(), std::move(patterns), config))) {
       mlir::emitError(getOperation()->getLoc(),
                       "failure in HLFIR intrinsic simplification");
       signalPassFailure();
diff --git a/flang/test/HLFIR/simplify-hlfir-intrinsics-cshift.fir b/flang/test/HLFIR/simplify-hlfir-intrinsics-cshift.fir
index acb89c0719aa..d21d7755062b 100644
--- a/flang/test/HLFIR/simplify-hlfir-intrinsics-cshift.fir
+++ b/flang/test/HLFIR/simplify-hlfir-intrinsics-cshift.fir
@@ -1,17 +1,19 @@
 // Test hlfir.cshift simplification to hlfir.elemental:
 // RUN: fir-opt --simplify-hlfir-intrinsics %s | FileCheck %s
 
-func.func @cshift_vector(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: !fir.ref<i32>) {
+func.func @cshift_vector(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: !fir.ref<i32>) -> !hlfir.expr<?xi32>{
   %res = hlfir.cshift %arg0 %arg1 : (!fir.box<!fir.array<?xi32>>, !fir.ref<i32>) -> !hlfir.expr<?xi32>
-  return
+  return %res : !hlfir.expr<?xi32>
 }
 // CHECK-LABEL:   func.func @cshift_vector(
 // CHECK-SAME:                             %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>,
-// CHECK-SAME:                             %[[VAL_1:.*]]: !fir.ref<i32>) {
+// CHECK-SAME:                             %[[VAL_1:.*]]: !fir.ref<i32>) -> !hlfir.expr<?xi32> {
+// CHECK:           %[[VAL_26:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_16:.*]] = arith.constant 0 : i64
+// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : i64
 // CHECK:           %[[VAL_2:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_3:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
 // CHECK:           %[[VAL_4:.*]] = fir.shape %[[VAL_3]]#1 : (index) -> !fir.shape<1>
-// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : i64
 // CHECK:           %[[VAL_6:.*]] = fir.load %[[VAL_1]] : !fir.ref<i32>
 // CHECK:           %[[VAL_7:.*]] = fir.convert %[[VAL_6]] : (i32) -> i64
 // CHECK:           %[[VAL_8:.*]] = hlfir.elemental %[[VAL_4]] unordered : (!fir.shape<1>) -> !hlfir.expr<?xi32> {
@@ -22,7 +24,6 @@ func.func @cshift_vector(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: !fir.ref<i32
 // CHECK:             %[[VAL_13:.*]] = fir.convert %[[VAL_3]]#1 : (index) -> i64
 // CHECK:             %[[VAL_14:.*]] = arith.remsi %[[VAL_12]], %[[VAL_13]] : i64
 // CHECK:             %[[VAL_15:.*]] = arith.xori %[[VAL_12]], %[[VAL_13]] : i64
-// CHECK:             %[[VAL_16:.*]] = arith.constant 0 : i64
 // CHECK:             %[[VAL_17:.*]] = arith.cmpi slt, %[[VAL_15]], %[[VAL_16]] : i64
 // CHECK:             %[[VAL_18:.*]] = arith.cmpi ne, %[[VAL_14]], %[[VAL_16]] : i64
 // CHECK:             %[[VAL_19:.*]] = arith.andi %[[VAL_18]], %[[VAL_17]] : i1
@@ -30,9 +31,7 @@ func.func @cshift_vector(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: !fir.ref<i32
 // CHECK:             %[[VAL_21:.*]] = arith.select %[[VAL_19]], %[[VAL_20]], %[[VAL_14]] : i64
 // CHECK:             %[[VAL_22:.*]] = arith.addi %[[VAL_21]], %[[VAL_5]] : i64
 // CHECK:             %[[VAL_23:.*]] = fir.convert %[[VAL_22]] : (i64) -> index
-// CHECK:             %[[VAL_24:.*]] = arith.constant 0 : index
-// CHECK:             %[[VAL_25:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_24]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_26:.*]] = arith.constant 1 : index
+// CHECK:             %[[VAL_25:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
 // CHECK:             %[[VAL_27:.*]] = arith.subi %[[VAL_25]]#0, %[[VAL_26]] : index
 // CHECK:             %[[VAL_28:.*]] = arith.addi %[[VAL_23]], %[[VAL_27]] : index
 // CHECK:             %[[VAL_29:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_28]])  : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
@@ -42,21 +41,21 @@ func.func @cshift_vector(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: !fir.ref<i32
 // CHECK:           return
 // CHECK:         }
 
-func.func @cshift_2d_by_scalar(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir.ref<i32>) {
+func.func @cshift_2d_by_scalar(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir.ref<i32>) -> !hlfir.expr<?x?xi32> {
   %dim = arith.constant 2 : i32
   %res = hlfir.cshift %arg0 %arg1 dim %dim : (!fir.box<!fir.array<?x?xi32>>, !fir.ref<i32>, i32) -> !hlfir.expr<?x?xi32>
-  return
+  return %res : !hlfir.expr<?x?xi32>
 }
 // CHECK-LABEL:   func.func @cshift_2d_by_scalar(
 // CHECK-SAME:                                   %[[VAL_0:.*]]: !fir.box<!fir.array<?x?xi32>>,
-// CHECK-SAME:                                   %[[VAL_1:.*]]: !fir.ref<i32>) {
-// CHECK:           %[[VAL_2:.*]] = arith.constant 2 : i32
+// CHECK-SAME:                                   %[[VAL_1:.*]]: !fir.ref<i32>) -> !hlfir.expr<?x?xi32> {
+// CHECK:           %[[VAL_20:.*]] = arith.constant 0 : i64
+// CHECK:           %[[VAL_8:.*]] = arith.constant 1 : i64
+// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_4:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_3]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
-// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_6:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_5]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
 // CHECK:           %[[VAL_7:.*]] = fir.shape %[[VAL_4]]#1, %[[VAL_6]]#1 : (index, index) -> !fir.shape<2>
-// CHECK:           %[[VAL_8:.*]] = arith.constant 1 : i64
 // CHECK:           %[[VAL_9:.*]] = fir.load %[[VAL_1]] : !fir.ref<i32>
 // CHECK:           %[[VAL_10:.*]] = fir.convert %[[VAL_9]] : (i32) -> i64
 // CHECK:           %[[VAL_11:.*]] = hlfir.elemental %[[VAL_7]] unordered : (!fir.shape<2>) -> !hlfir.expr<?x?xi32> {
@@ -67,7 +66,6 @@ func.func @cshift_2d_by_scalar(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir
 // CHECK:             %[[VAL_17:.*]] = fir.convert %[[VAL_6]]#1 : (index) -> i64
 // CHECK:             %[[VAL_18:.*]] = arith.remsi %[[VAL_16]], %[[VAL_17]] : i64
 // CHECK:             %[[VAL_19:.*]] = arith.xori %[[VAL_16]], %[[VAL_17]] : i64
-// CHECK:             %[[VAL_20:.*]] = arith.constant 0 : i64
 // CHECK:             %[[VAL_21:.*]] = arith.cmpi slt, %[[VAL_19]], %[[VAL_20]] : i64
 // CHECK:             %[[VAL_22:.*]] = arith.cmpi ne, %[[VAL_18]], %[[VAL_20]] : i64
 // CHECK:             %[[VAL_23:.*]] = arith.andi %[[VAL_22]], %[[VAL_21]] : i1
@@ -75,14 +73,11 @@ func.func @cshift_2d_by_scalar(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir
 // CHECK:             %[[VAL_25:.*]] = arith.select %[[VAL_23]], %[[VAL_24]], %[[VAL_18]] : i64
 // CHECK:             %[[VAL_26:.*]] = arith.addi %[[VAL_25]], %[[VAL_8]] : i64
 // CHECK:             %[[VAL_27:.*]] = fir.convert %[[VAL_26]] : (i64) -> index
-// CHECK:             %[[VAL_28:.*]] = arith.constant 0 : index
-// CHECK:             %[[VAL_29:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_28]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_30:.*]] = arith.constant 1 : index
-// CHECK:             %[[VAL_31:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_30]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_32:.*]] = arith.constant 1 : index
-// CHECK:             %[[VAL_33:.*]] = arith.subi %[[VAL_29]]#0, %[[VAL_32]] : index
+// CHECK:             %[[VAL_29:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_3]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
+// CHECK:             %[[VAL_31:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_5]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
+// CHECK:             %[[VAL_33:.*]] = arith.subi %[[VAL_29]]#0, %[[VAL_5]] : index
 // CHECK:             %[[VAL_34:.*]] = arith.addi %[[VAL_12]], %[[VAL_33]] : index
-// CHECK:             %[[VAL_35:.*]] = arith.subi %[[VAL_31]]#0, %[[VAL_32]] : index
+// CHECK:             %[[VAL_35:.*]] = arith.subi %[[VAL_31]]#0, %[[VAL_5]] : index
 // CHECK:             %[[VAL_36:.*]] = arith.addi %[[VAL_27]], %[[VAL_35]] : index
 // CHECK:             %[[VAL_37:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_34]], %[[VAL_36]])  : (!fir.box<!fir.array<?x?xi32>>, index, index) -> !fir.ref<i32>
 // CHECK:             %[[VAL_38:.*]] = fir.load %[[VAL_37]] : !fir.ref<i32>
@@ -91,27 +86,25 @@ func.func @cshift_2d_by_scalar(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir
 // CHECK:           return
 // CHECK:         }
 
-func.func @cshift_2d_by_vector(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir.box<!fir.array<?xi32>>) {
+func.func @cshift_2d_by_vector(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir.box<!fir.array<?xi32>>) -> !hlfir.expr<?x?xi32> {
   %dim = arith.constant 2 : i32
   %res = hlfir.cshift %arg0 %arg1 dim %dim : (!fir.box<!fir.array<?x?xi32>>, !fir.box<!fir.array<?xi32>>, i32) -> !hlfir.expr<?x?xi32>
-  return
+  return %res : !hlfir.expr<?x?xi32>
 }
 // CHECK-LABEL:   func.func @cshift_2d_by_vector(
 // CHECK-SAME:                                   %[[VAL_0:.*]]: !fir.box<!fir.array<?x?xi32>>,
-// CHECK-SAME:                                   %[[VAL_1:.*]]: !fir.box<!fir.array<?xi32>>) {
-// CHECK:           %[[VAL_2:.*]] = arith.constant 2 : i32
+// CHECK-SAME:                                   %[[VAL_1:.*]]: !fir.box<!fir.array<?xi32>>) -> !hlfir.expr<?x?xi32> {
+// CHECK:           %[[VAL_26:.*]] = arith.constant 0 : i64
+// CHECK:           %[[VAL_8:.*]] = arith.constant 1 : i64
+// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_4:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_3]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
-// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_6:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_5]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
 // CHECK:           %[[VAL_7:.*]] = fir.shape %[[VAL_4]]#1, %[[VAL_6]]#1 : (index, index) -> !fir.shape<2>
-// CHECK:           %[[VAL_8:.*]] = arith.constant 1 : i64
 // CHECK:           %[[VAL_9:.*]] = hlfir.elemental %[[VAL_7]] unordered : (!fir.shape<2>) -> !hlfir.expr<?x?xi32> {
 // CHECK:           ^bb0(%[[VAL_10:.*]]: index, %[[VAL_11:.*]]: index):
-// CHECK:             %[[VAL_12:.*]] = arith.constant 0 : index
-// CHECK:             %[[VAL_13:.*]]:3 = fir.box_dims %[[VAL_1]], %[[VAL_12]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_14:.*]] = arith.constant 1 : index
-// CHECK:             %[[VAL_15:.*]] = arith.subi %[[VAL_13]]#0, %[[VAL_14]] : index
+// CHECK:             %[[VAL_13:.*]]:3 = fir.box_dims %[[VAL_1]], %[[VAL_3]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
+// CHECK:             %[[VAL_15:.*]] = arith.subi %[[VAL_13]]#0, %[[VAL_5]] : index
 // CHECK:             %[[VAL_16:.*]] = arith.addi %[[VAL_10]], %[[VAL_15]] : index
 // CHECK:             %[[VAL_17:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_16]])  : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
 // CHECK:             %[[VAL_18:.*]] = fir.load %[[VAL_17]] : !fir.ref<i32>
@@ -122,7 +115,6 @@ func.func @cshift_2d_by_vector(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir
 // CHECK:             %[[VAL_23:.*]] = fir.convert %[[VAL_6]]#1 : (index) -> i64
 // CHECK:             %[[VAL_24:.*]] = arith.remsi %[[VAL_22]], %[[VAL_23]] : i64
 // CHECK:             %[[VAL_25:.*]] = arith.xori %[[VAL_22]], %[[VAL_23]] : i64
-// CHECK:             %[[VAL_26:.*]] = arith.constant 0 : i64
 // CHECK:             %[[VAL_27:.*]] = arith.cmpi slt, %[[VAL_25]], %[[VAL_26]] : i64
 // CHECK:             %[[VAL_28:.*]] = arith.cmpi ne, %[[VAL_24]], %[[VAL_26]] : i64
 // CHECK:             %[[VAL_29:.*]] = arith.andi %[[VAL_28]], %[[VAL_27]] : i1
@@ -130,14 +122,11 @@ func.func @cshift_2d_by_vector(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir
 // CHECK:             %[[VAL_31:.*]] = arith.select %[[VAL_29]], %[[VAL_30]], %[[VAL_24]] : i64
 // CHECK:             %[[VAL_32:.*]] = arith.addi %[[VAL_31]], %[[VAL_8]] : i64
 // CHECK:             %[[VAL_33:.*]] = fir.convert %[[VAL_32]] : (i64) -> index
-// CHECK:             %[[VAL_34:.*]] = arith.constant 0 : index
-// CHECK:             %[[VAL_35:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_34]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_36:.*]] = arith.constant 1 : index
-// CHECK:             %[[VAL_37:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_36]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_38:.*]] = arith.constant 1 : index
-// CHECK:             %[[VAL_39:.*]] = arith.subi %[[VAL_35]]#0, %[[VAL_38]] : index
+// CHECK:             %[[VAL_35:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_3]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
+// CHECK:             %[[VAL_37:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_5]] : (!fir.box<!fir.array<?x?xi32>>, index) -> (index, index, index)
+// CHECK:             %[[VAL_39:.*]] = arith.subi %[[VAL_35]]#0, %[[VAL_5]] : index
 // CHECK:             %[[VAL_40:.*]] = arith.addi %[[VAL_10]], %[[VAL_39]] : index
-// CHECK:             %[[VAL_41:.*]] = arith.subi %[[VAL_37]]#0, %[[VAL_38]] : index
+// CHECK:             %[[VAL_41:.*]] = arith.subi %[[VAL_37]]#0, %[[VAL_5]] : index
 // CHECK:             %[[VAL_42:.*]] = arith.addi %[[VAL_33]], %[[VAL_41]] : index
 // CHECK:             %[[VAL_43:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_40]], %[[VAL_42]])  : (!fir.box<!fir.array<?x?xi32>>, index, index) -> !fir.ref<i32>
 // CHECK:             %[[VAL_44:.*]] = fir.load %[[VAL_43]] : !fir.ref<i32>
@@ -146,20 +135,22 @@ func.func @cshift_2d_by_vector(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: !fir
 // CHECK:           return
 // CHECK:         }
 
-func.func @cshift_vector_char(%arg0: !fir.box<!fir.array<?x!fir.char<2,?>>>, %arg1: !fir.ref<i32>) {
+func.func @cshift_vector_char(%arg0: !fir.box<!fir.array<?x!fir.char<2,?>>>, %arg1: !fir.ref<i32>) -> !hlfir.expr<?x!fir.char<2,?>> {
   %res = hlfir.cshift %arg0 %arg1 : (!fir.box<!fir.array<?x!fir.char<2,?>>>, !fir.ref<i32>) -> !hlfir.expr<?x!fir.char<2,?>>
-  return
+  return %res : !hlfir.expr<?x!fir.char<2,?>>
 }
 // CHECK-LABEL:   func.func @cshift_vector_char(
 // CHECK-SAME:                                  %[[VAL_0:.*]]: !fir.box<!fir.array<?x!fir.char<2,?>>>,
-// CHECK-SAME:                                  %[[VAL_1:.*]]: !fir.ref<i32>) {
+// CHECK-SAME:                                  %[[VAL_1:.*]]: !fir.ref<i32>) -> !hlfir.expr<?x!fir.char<2,?>> {
+// CHECK:           %[[VAL_32:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_19:.*]] = arith.constant 0 : i64
+// CHECK:           %[[VAL_8:.*]] = arith.constant 1 : i64
+// CHECK:           %[[VAL_6:.*]] = arith.constant 2 : index
 // CHECK:           %[[VAL_2:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_3:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x!fir.char<2,?>>>, index) -> (index, index, index)
 // CHECK:           %[[VAL_4:.*]] = fir.shape %[[VAL_3]]#1 : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_5:.*]] = fir.box_elesize %[[VAL_0]] : (!fir.box<!fir.array<?x!fir.char<2,?>>>) -> index
-// CHECK:           %[[VAL_6:.*]] = arith.constant 2 : index
 // CHECK:           %[[VAL_7:.*]] = arith.divsi %[[VAL_5]], %[[VAL_6]] : index
-// CHECK:           %[[VAL_8:.*]] = arith.constant 1 : i64
 // CHECK:           %[[VAL_9:.*]] = fir.load %[[VAL_1]] : !fir.ref<i32>
 // CHECK:           %[[VAL_10:.*]] = fir.convert %[[VAL_9]] : (i32) -> i64
 // CHECK:           %[[VAL_11:.*]] = hlfir.elemental %[[VAL_4]] typeparams %[[VAL_7]] unordered : (!fir.shape<1>, index) -> !hlfir.expr<?x!fir.char<2,?>> {
@@ -170,7 +161,6 @@ func.func @cshift_vector_char(%arg0: !fir.box<!fir.array<?x!fir.char<2,?>>>, %ar
 // CHECK:             %[[VAL_16:.*]] = fir.convert %[[VAL_3]]#1 : (index) -> i64
 // CHECK:             %[[VAL_17:.*]] = arith.remsi %[[VAL_15]], %[[VAL_16]] : i64
 // CHECK:             %[[VAL_18:.*]] = arith.xori %[[VAL_15]], %[[VAL_16]] : i64
-// CHECK:             %[[VAL_19:.*]] = arith.constant 0 : i64
 // CHECK:             %[[VAL_20:.*]] = arith.cmpi slt, %[[VAL_18]], %[[VAL_19]] : i64
 // CHECK:             %[[VAL_21:.*]] = arith.cmpi ne, %[[VAL_17]], %[[VAL_19]] : i64
 // CHECK:             %[[VAL_22:.*]] = arith.andi %[[VAL_21]], %[[VAL_20]] : i1
@@ -179,11 +169,8 @@ func.func @cshift_vector_char(%arg0: !fir.box<!fir.array<?x!fir.char<2,?>>>, %ar
 // CHECK:             %[[VAL_25:.*]] = arith.addi %[[VAL_24]], %[[VAL_8]] : i64
 // CHECK:             %[[VAL_26:.*]] = fir.convert %[[VAL_25]] : (i64) -> index
 // CHECK:             %[[VAL_27:.*]] = fir.box_elesize %[[VAL_0]] : (!fir.box<!fir.array<?x!fir.char<2,?>>>) -> index
-// CHECK:             %[[VAL_28:.*]] = arith.constant 2 : index
-// CHECK:             %[[VAL_29:.*]] = arith.divsi %[[VAL_27]], %[[VAL_28]] : index
-// CHECK:             %[[VAL_30:.*]] = arith.constant 0 : index
-// CHECK:             %[[VAL_31:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_30]] : (!fir.box<!fir.array<?x!fir.char<2,?>>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_32:.*]] = arith.constant 1 : index
+// CHECK:             %[[VAL_29:.*]] = arith.divsi %[[VAL_27]], %[[VAL_6]] : index
+// CHECK:             %[[VAL_31:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x!fir.char<2,?>>>, index) -> (index, index, index)
 // CHECK:             %[[VAL_33:.*]] = arith.subi %[[VAL_31]]#0, %[[VAL_32]] : index
 // CHECK:             %[[VAL_34:.*]] = arith.addi %[[VAL_26]], %[[VAL_33]] : index
 // CHECK:             %[[VAL_35:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_34]])  typeparams %[[VAL_29]] : (!fir.box<!fir.array<?x!fir.char<2,?>>>, index, index) -> !fir.boxchar<2>
@@ -192,17 +179,19 @@ func.func @cshift_vector_char(%arg0: !fir.box<!fir.array<?x!fir.char<2,?>>>, %ar
 // CHECK:           return
 // CHECK:         }
 
-func.func @cshift_vector_poly(%arg0: !fir.class<!fir.array<?x!fir.type<_QFFtestTt>>>, %arg1: i32) {
+func.func @cshift_vector_poly(%arg0: !fir.class<!fir.array<?x!fir.type<_QFFtestTt>>>, %arg1: i32) -> !hlfir.expr<?x!fir.type<_QFFtestTt>?> {
   %res = hlfir.cshift %arg0 %arg1 : (!fir.class<!fir.array<?x!fir.type<_QFFtestTt>>>, i32) -> !hlfir.expr<?x!fir.type<_QFFtestTt>?>
-  return
+  return %res : !hlfir.expr<?x!fir.type<_QFFtestTt>?>
 }
 // CHECK-LABEL:   func.func @cshift_vector_poly(
 // CHECK-SAME:                                  %[[VAL_0:.*]]: !fir.class<!fir.array<?x!fir.type<_QFFtestTt>>>,
-// CHECK-SAME:                                  %[[VAL_1:.*]]: i32) {
+// CHECK-SAME:                                  %[[VAL_1:.*]]: i32) -> !hlfir.expr<?x!fir.type<_QFFtestTt>?> {
+// CHECK:           %[[VAL_25:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_15:.*]] = arith.constant 0 : i64
+// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : i64
 // CHECK:           %[[VAL_2:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_3:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.class<!fir.array<?x!fir.type<_QFFtestTt>>>, index) -> (index, index, index)
 // CHECK:           %[[VAL_4:.*]] = fir.shape %[[VAL_3]]#1 : (index) -> !fir.shape<1>
-// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : i64
 // CHECK:           %[[VAL_6:.*]] = fir.convert %[[VAL_1]] : (i32) -> i64
 // CHECK:           %[[VAL_7:.*]] = hlfir.elemental %[[VAL_4]] mold %[[VAL_0]] unordered : (!fir.shape<1>, !fir.class<!fir.array<?x!fir.type<_QFFtestTt>>>) -> !hlfir.expr<?x!fir.type<_QFFtestTt>?> {
 // CHECK:           ^bb0(%[[VAL_8:.*]]: index):
@@ -212,7 +201,6 @@ func.func @cshift_vector_poly(%arg0: !fir.class<!fir.array<?x!fir.type<_QFFtestT
 // CHECK:             %[[VAL_12:.*]] = fir.convert %[[VAL_3]]#1 : (index) -> i64
 // CHECK:             %[[VAL_13:.*]] = arith.remsi %[[VAL_11]], %[[VAL_12]] : i64
 // CHECK:             %[[VAL_14:.*]] = arith.xori %[[VAL_11]], %[[VAL_12]] : i64
-// CHECK:             %[[VAL_15:.*]] = arith.constant 0 : i64
 // CHECK:             %[[VAL_16:.*]] = arith.cmpi slt, %[[VAL_14]], %[[VAL_15]] : i64
 // CHECK:             %[[VAL_17:.*]] = arith.cmpi ne, %[[VAL_13]], %[[VAL_15]] : i64
 // CHECK:             %[[VAL_18:.*]] = arith.andi %[[VAL_17]], %[[VAL_16]] : i1
@@ -220,9 +208,7 @@ func.func @cshift_vector_poly(%arg0: !fir.class<!fir.array<?x!fir.type<_QFFtestT
 // CHECK:             %[[VAL_20:.*]] = arith.select %[[VAL_18]], %[[VAL_19]], %[[VAL_13]] : i64
 // CHECK:             %[[VAL_21:.*]] = arith.addi %[[VAL_20]], %[[VAL_5]] : i64
 // CHECK:             %[[VAL_22:.*]] = fir.convert %[[VAL_21]] : (i64) -> index
-// CHECK:             %[[VAL_23:.*]] = arith.constant 0 : index
-// CHECK:             %[[VAL_24:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_23]] : (!fir.class<!fir.array<?x!fir.type<_QFFtestTt>>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_25:.*]] = arith.constant 1 : index
+// CHECK:             %[[VAL_24:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.class<!fir.array<?x!fir.type<_QFFtestTt>>>, index) -> (index, index, index)
 // CHECK:             %[[VAL_26:.*]] = arith.subi %[[VAL_24]]#0, %[[VAL_25]] : index
 // CHECK:             %[[VAL_27:.*]] = arith.addi %[[VAL_22]], %[[VAL_26]] : index
 // CHECK:             %[[VAL_28:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_27]])  : (!fir.class<!fir.array<?x!fir.type<_QFFtestTt>>>, index) -> !fir.class<!fir.type<_QFFtestTt>>
@@ -232,48 +218,39 @@ func.func @cshift_vector_poly(%arg0: !fir.class<!fir.array<?x!fir.type<_QFFtestT
 // CHECK:         }
 
 // negative: non-constant dim argument
-func.func @cshift_nonconst_dim(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: i32, %dim : i32) {
+func.func @cshift_nonconst_dim(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: i32, %dim : i32) -> !hlfir.expr<?x?xi32> {
   %res = hlfir.cshift %arg0 %arg1 dim %dim : (!fir.box<!fir.array<?x?xi32>>, i32, i32) -> !hlfir.expr<?x?xi32>
-  return
+  return %res : !hlfir.expr<?x?xi32>
 }
 // CHECK-LABEL:   func.func @cshift_nonconst_dim(
-// CHECK-SAME:                                   %[[VAL_0:.*]]: !fir.box<!fir.array<?x?xi32>>,
-// CHECK-SAME:                                   %[[VAL_1:.*]]: i32,
-// CHECK-SAME:                                   %[[VAL_2:.*]]: i32) {
-// CHECK:           %[[VAL_3:.*]] = hlfir.cshift %[[VAL_0]] %[[VAL_1]] dim %[[VAL_2]] : (!fir.box<!fir.array<?x?xi32>>, i32, i32) -> !hlfir.expr<?x?xi32>
-// CHECK:           return
-// CHECK:         }
+// CHECK:           hlfir.cshift {{.*}} : (!fir.box<!fir.array<?x?xi32>>, i32, i32) -> !hlfir.expr<?x?xi32>
 
 // negative: invalid constant dim argument
-func.func @cshift_invalid_dim(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: i32) {
+func.func @cshift_invalid_dim(%arg0: !fir.box<!fir.array<?x?xi32>>, %arg1: i32) -> !hlfir.expr<?x?xi32> {
   %dim = arith.constant 3 : i32
   %res = hlfir.cshift %arg0 %arg1 dim %dim : (!fir.box<!fir.array<?x?xi32>>, i32, i32) -> !hlfir.expr<?x?xi32>
-  return
+  return %res : !hlfir.expr<?x?xi32>
 }
 // CHECK-LABEL:   func.func @cshift_invalid_dim(
-// CHECK-SAME:                                  %[[VAL_0:.*]]: !fir.box<!fir.array<?x?xi32>>,
-// CHECK-SAME:                                  %[[VAL_1:.*]]: i32) {
-// CHECK:           %[[VAL_2:.*]] = arith.constant 3 : i32
-// CHECK:           %[[VAL_3:.*]] = hlfir.cshift %[[VAL_0]] %[[VAL_1]] dim %[[VAL_2]] : (!fir.box<!fir.array<?x?xi32>>, i32, i32) -> !hlfir.expr<?x?xi32>
-// CHECK:           return
-// CHECK:         }
+// CHECK:           hlfir.cshift {{.*}} : (!fir.box<!fir.array<?x?xi32>>, i32, i32) -> !hlfir.expr<?x?xi32>
 
 // When the input array is 1D, we may assume that DIM==1,
 // otherwise the program is illegal, and we can do anything
 // about it.
-func.func @cshift_vector_assumed_dim_1(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: i32) {
+func.func @cshift_vector_assumed_dim_1(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: i32) -> !hlfir.expr<?xi32> {
   %dim = arith.constant 3 : i32
   %res = hlfir.cshift %arg0 %arg1 dim %dim : (!fir.box<!fir.array<?xi32>>, i32, i32) -> !hlfir.expr<?xi32>
-  return
+  return %res : !hlfir.expr<?xi32>
 }
 // CHECK-LABEL:   func.func @cshift_vector_assumed_dim_1(
 // CHECK-SAME:                                           %[[VAL_0:.*]]: !fir.box<!fir.array<?xi32>>,
-// CHECK-SAME:                                           %[[VAL_1:.*]]: i32) {
-// CHECK:           %[[VAL_2:.*]] = arith.constant 3 : i32
+// CHECK-SAME:                                           %[[VAL_1:.*]]: i32) -> !hlfir.expr<?xi32> {
+// CHECK:           %[[VAL_26:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_16:.*]] = arith.constant 0 : i64
+// CHECK:           %[[VAL_6:.*]] = arith.constant 1 : i64
 // CHECK:           %[[VAL_3:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_4:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_3]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
 // CHECK:           %[[VAL_5:.*]] = fir.shape %[[VAL_4]]#1 : (index) -> !fir.shape<1>
-// CHECK:           %[[VAL_6:.*]] = arith.constant 1 : i64
 // CHECK:           %[[VAL_7:.*]] = fir.convert %[[VAL_1]] : (i32) -> i64
 // CHECK:           %[[VAL_8:.*]] = hlfir.elemental %[[VAL_5]] unordered : (!fir.shape<1>) -> !hlfir.expr<?xi32> {
 // CHECK:           ^bb0(%[[VAL_9:.*]]: index):
@@ -283,7 +260,6 @@ func.func @cshift_vector_assumed_dim_1(%arg0: !fir.box<!fir.array<?xi32>>, %arg1
 // CHECK:             %[[VAL_13:.*]] = fir.convert %[[VAL_4]]#1 : (index) -> i64
 // CHECK:             %[[VAL_14:.*]] = arith.remsi %[[VAL_12]], %[[VAL_13]] : i64
 // CHECK:             %[[VAL_15:.*]] = arith.xori %[[VAL_12]], %[[VAL_13]] : i64
-// CHECK:             %[[VAL_16:.*]] = arith.constant 0 : i64
 // CHECK:             %[[VAL_17:.*]] = arith.cmpi slt, %[[VAL_15]], %[[VAL_16]] : i64
 // CHECK:             %[[VAL_18:.*]] = arith.cmpi ne, %[[VAL_14]], %[[VAL_16]] : i64
 // CHECK:             %[[VAL_19:.*]] = arith.andi %[[VAL_18]], %[[VAL_17]] : i1
@@ -291,9 +267,7 @@ func.func @cshift_vector_assumed_dim_1(%arg0: !fir.box<!fir.array<?xi32>>, %arg1
 // CHECK:             %[[VAL_21:.*]] = arith.select %[[VAL_19]], %[[VAL_20]], %[[VAL_14]] : i64
 // CHECK:             %[[VAL_22:.*]] = arith.addi %[[VAL_21]], %[[VAL_6]] : i64
 // CHECK:             %[[VAL_23:.*]] = fir.convert %[[VAL_22]] : (i64) -> index
-// CHECK:             %[[VAL_24:.*]] = arith.constant 0 : index
-// CHECK:             %[[VAL_25:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_24]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_26:.*]] = arith.constant 1 : index
+// CHECK:             %[[VAL_25:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_3]] : (!fir.box<!fir.array<?xi32>>, index) -> (index, index, index)
 // CHECK:             %[[VAL_27:.*]] = arith.subi %[[VAL_25]]#0, %[[VAL_26]] : index
 // CHECK:             %[[VAL_28:.*]] = arith.addi %[[VAL_23]], %[[VAL_27]] : index
 // CHECK:             %[[VAL_29:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_28]])  : (!fir.box<!fir.array<?xi32>>, index) -> !fir.ref<i32>
diff --git a/flang/test/HLFIR/simplify-hlfir-intrinsics-sum.fir b/flang/test/HLFIR/simplify-hlfir-intrinsics-sum.fir
index bb406a6ad359..58a2144947b1 100644
--- a/flang/test/HLFIR/simplify-hlfir-intrinsics-sum.fir
+++ b/flang/test/HLFIR/simplify-hlfir-intrinsics-sum.fir
@@ -1,30 +1,27 @@
 // RUN: fir-opt --simplify-hlfir-intrinsics -flang-simplify-hlfir-sum %s | FileCheck %s
 
 // box with known extents
-func.func @sum_box_known_extents(%arg0: !fir.box<!fir.array<2x3xi32>>) {
+func.func @sum_box_known_extents(%arg0: !fir.box<!fir.array<2x3xi32>>) -> !hlfir.expr<2xi32> {
   %cst = arith.constant 2 : i32
   %res = hlfir.sum %arg0 dim %cst : (!fir.box<!fir.array<2x3xi32>>, i32) -> !hlfir.expr<2xi32>
-  return
+  return %res : !hlfir.expr<2xi32>
 }
 // CHECK-LABEL:   func.func @sum_box_known_extents(
-// CHECK-SAME:                                     %[[VAL_0:.*]]: !fir.box<!fir.array<2x3xi32>>) {
-// CHECK:           %[[VAL_1:.*]] = arith.constant 2 : i32
+// CHECK-SAME:                                     %[[VAL_0:.*]]: !fir.box<!fir.array<2x3xi32>>) -> !hlfir.expr<2xi32> {
+// CHECK:           %[[VAL_12:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_8:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_7:.*]] = arith.constant 0 : i32
 // CHECK:           %[[VAL_2:.*]] = arith.constant 2 : index
 // CHECK:           %[[VAL_3:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_4:.*]] = fir.shape %[[VAL_2]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_5:.*]] = hlfir.elemental %[[VAL_4]] unordered : (!fir.shape<1>) -> !hlfir.expr<2xi32> {
 // CHECK:           ^bb0(%[[VAL_6:.*]]: index):
-// CHECK:             %[[VAL_7:.*]] = arith.constant 0 : i32
-// CHECK:             %[[VAL_8:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_9:.*]] = fir.do_loop %[[VAL_10:.*]] = %[[VAL_8]] to %[[VAL_3]] step %[[VAL_8]] unordered iter_args(%[[VAL_11:.*]] = %[[VAL_7]]) -> (i32) {
-// CHECK:               %[[VAL_12:.*]] = arith.constant 0 : index
 // CHECK:               %[[VAL_13:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_12]] : (!fir.box<!fir.array<2x3xi32>>, index) -> (index, index, index)
-// CHECK:               %[[VAL_14:.*]] = arith.constant 1 : index
-// CHECK:               %[[VAL_15:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_14]] : (!fir.box<!fir.array<2x3xi32>>, index) -> (index, index, index)
-// CHECK:               %[[VAL_16:.*]] = arith.constant 1 : index
-// CHECK:               %[[VAL_17:.*]] = arith.subi %[[VAL_13]]#0, %[[VAL_16]] : index
+// CHECK:               %[[VAL_15:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_8]] : (!fir.box<!fir.array<2x3xi32>>, index) -> (index, index, index)
+// CHECK:               %[[VAL_17:.*]] = arith.subi %[[VAL_13]]#0, %[[VAL_8]] : index
 // CHECK:               %[[VAL_18:.*]] = arith.addi %[[VAL_6]], %[[VAL_17]] : index
-// CHECK:               %[[VAL_19:.*]] = arith.subi %[[VAL_15]]#0, %[[VAL_16]] : index
+// CHECK:               %[[VAL_19:.*]] = arith.subi %[[VAL_15]]#0, %[[VAL_8]] : index
 // CHECK:               %[[VAL_20:.*]] = arith.addi %[[VAL_10]], %[[VAL_19]] : index
 // CHECK:               %[[VAL_21:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_18]], %[[VAL_20]])  : (!fir.box<!fir.array<2x3xi32>>, index, index) -> !fir.ref<i32>
 // CHECK:               %[[VAL_22:.*]] = fir.load %[[VAL_21]] : !fir.ref<i32>
@@ -37,21 +34,20 @@ func.func @sum_box_known_extents(%arg0: !fir.box<!fir.array<2x3xi32>>) {
 // CHECK:         }
 
 // expr with known extents
-func.func @sum_expr_known_extents(%arg0: !hlfir.expr<2x3xi32>) {
+func.func @sum_expr_known_extents(%arg0: !hlfir.expr<2x3xi32>) -> !hlfir.expr<3xi32> {
   %cst = arith.constant 1 : i32
   %res = hlfir.sum %arg0 dim %cst : (!hlfir.expr<2x3xi32>, i32) -> !hlfir.expr<3xi32>
-  return
+  return %res : !hlfir.expr<3xi32>
 }
 // CHECK-LABEL:   func.func @sum_expr_known_extents(
-// CHECK-SAME:                                      %[[VAL_0:.*]]: !hlfir.expr<2x3xi32>) {
-// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i32
+// CHECK-SAME:                                      %[[VAL_0:.*]]: !hlfir.expr<2x3xi32>) -> !hlfir.expr<3xi32> {
+// CHECK:           %[[VAL_8:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_7:.*]] = arith.constant 0 : i32
 // CHECK:           %[[VAL_2:.*]] = arith.constant 2 : index
 // CHECK:           %[[VAL_3:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_4:.*]] = fir.shape %[[VAL_3]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_5:.*]] = hlfir.elemental %[[VAL_4]] unordered : (!fir.shape<1>) -> !hlfir.expr<3xi32> {
 // CHECK:           ^bb0(%[[VAL_6:.*]]: index):
-// CHECK:             %[[VAL_7:.*]] = arith.constant 0 : i32
-// CHECK:             %[[VAL_8:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_9:.*]] = fir.do_loop %[[VAL_10:.*]] = %[[VAL_8]] to %[[VAL_2]] step %[[VAL_8]] unordered iter_args(%[[VAL_11:.*]] = %[[VAL_7]]) -> (i32) {
 // CHECK:               %[[VAL_12:.*]] = hlfir.apply %[[VAL_0]], %[[VAL_10]], %[[VAL_6]] : (!hlfir.expr<2x3xi32>, index, index) -> i32
 // CHECK:               %[[VAL_13:.*]] = arith.addi %[[VAL_11]], %[[VAL_12]] : i32
@@ -63,34 +59,30 @@ func.func @sum_expr_known_extents(%arg0: !hlfir.expr<2x3xi32>) {
 // CHECK:         }
 
 // box with unknown extent
-func.func @sum_box_unknown_extent1(%arg0: !fir.box<!fir.array<?x3xcomplex<f64>>>) {
+func.func @sum_box_unknown_extent1(%arg0: !fir.box<!fir.array<?x3xcomplex<f64>>>) -> !hlfir.expr<3xcomplex<f64>> {
   %cst = arith.constant 1 : i32
   %res = hlfir.sum %arg0 dim %cst : (!fir.box<!fir.array<?x3xcomplex<f64>>>, i32) -> !hlfir.expr<3xcomplex<f64>>
-  return
+  return %res : !hlfir.expr<3xcomplex<f64>>
 }
 // CHECK-LABEL:   func.func @sum_box_unknown_extent1(
-// CHECK-SAME:                                       %[[VAL_0:.*]]: !fir.box<!fir.array<?x3xcomplex<f64>>>) {
-// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i32
+// CHECK-SAME:                                       %[[VAL_0:.*]]: !fir.box<!fir.array<?x3xcomplex<f64>>>) -> !hlfir.expr<3xcomplex<f64>> {
+// CHECK:           %[[VAL_12:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_8:.*]] = arith.constant 0.000000e+00 : f64
+// CHECK:           %[[VAL_4:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_2:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_3:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
-// CHECK:           %[[VAL_4:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_5:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_6:.*]] = hlfir.elemental %[[VAL_5]] unordered : (!fir.shape<1>) -> !hlfir.expr<3xcomplex<f64>> {
 // CHECK:           ^bb0(%[[VAL_7:.*]]: index):
-// CHECK:             %[[VAL_8:.*]] = arith.constant 0.000000e+00 : f64
 // CHECK:             %[[VAL_9:.*]] = fir.undefined complex<f64>
 // CHECK:             %[[VAL_10:.*]] = fir.insert_value %[[VAL_9]], %[[VAL_8]], [0 : index] : (complex<f64>, f64) -> complex<f64>
 // CHECK:             %[[VAL_11:.*]] = fir.insert_value %[[VAL_10]], %[[VAL_8]], [1 : index] : (complex<f64>, f64) -> complex<f64>
-// CHECK:             %[[VAL_12:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_13:.*]] = fir.do_loop %[[VAL_14:.*]] = %[[VAL_12]] to %[[VAL_3]]#1 step %[[VAL_12]] iter_args(%[[VAL_15:.*]] = %[[VAL_11]]) -> (complex<f64>) {
-// CHECK:               %[[VAL_16:.*]] = arith.constant 0 : index
-// CHECK:               %[[VAL_17:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_16]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
-// CHECK:               %[[VAL_18:.*]] = arith.constant 1 : index
-// CHECK:               %[[VAL_19:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_18]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
-// CHECK:               %[[VAL_20:.*]] = arith.constant 1 : index
-// CHECK:               %[[VAL_21:.*]] = arith.subi %[[VAL_17]]#0, %[[VAL_20]] : index
+// CHECK:               %[[VAL_17:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK:               %[[VAL_19:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_12]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK:               %[[VAL_21:.*]] = arith.subi %[[VAL_17]]#0, %[[VAL_12]] : index
 // CHECK:               %[[VAL_22:.*]] = arith.addi %[[VAL_14]], %[[VAL_21]] : index
-// CHECK:               %[[VAL_23:.*]] = arith.subi %[[VAL_19]]#0, %[[VAL_20]] : index
+// CHECK:               %[[VAL_23:.*]] = arith.subi %[[VAL_19]]#0, %[[VAL_12]] : index
 // CHECK:               %[[VAL_24:.*]] = arith.addi %[[VAL_7]], %[[VAL_23]] : index
 // CHECK:               %[[VAL_25:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_22]], %[[VAL_24]])  : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index, index) -> !fir.ref<complex<f64>>
 // CHECK:               %[[VAL_26:.*]] = fir.load %[[VAL_25]] : !fir.ref<complex<f64>>
@@ -102,34 +94,30 @@ func.func @sum_box_unknown_extent1(%arg0: !fir.box<!fir.array<?x3xcomplex<f64>>>
 // CHECK:           return
 // CHECK:         }
 
-func.func @sum_box_unknown_extent2(%arg0: !fir.box<!fir.array<?x3xcomplex<f64>>>) {
+func.func @sum_box_unknown_extent2(%arg0: !fir.box<!fir.array<?x3xcomplex<f64>>>) -> !hlfir.expr<?xcomplex<f64>> {
   %cst = arith.constant 2 : i32
   %res = hlfir.sum %arg0 dim %cst : (!fir.box<!fir.array<?x3xcomplex<f64>>>, i32) -> !hlfir.expr<?xcomplex<f64>>
-  return
+  return %res : !hlfir.expr<?xcomplex<f64>>
 }
 // CHECK-LABEL:   func.func @sum_box_unknown_extent2(
-// CHECK-SAME:                                       %[[VAL_0:.*]]: !fir.box<!fir.array<?x3xcomplex<f64>>>) {
-// CHECK:           %[[VAL_1:.*]] = arith.constant 2 : i32
+// CHECK-SAME:                                       %[[VAL_0:.*]]: !fir.box<!fir.array<?x3xcomplex<f64>>>) -> !hlfir.expr<?xcomplex<f64>> {
+// CHECK:           %[[VAL_12:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_8:.*]] = arith.constant 0.000000e+00 : f64
+// CHECK:           %[[VAL_4:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_2:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_3:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
-// CHECK:           %[[VAL_4:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_5:.*]] = fir.shape %[[VAL_3]]#1 : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_6:.*]] = hlfir.elemental %[[VAL_5]] unordered : (!fir.shape<1>) -> !hlfir.expr<?xcomplex<f64>> {
 // CHECK:           ^bb0(%[[VAL_7:.*]]: index):
-// CHECK:             %[[VAL_8:.*]] = arith.constant 0.000000e+00 : f64
 // CHECK:             %[[VAL_9:.*]] = fir.undefined complex<f64>
 // CHECK:             %[[VAL_10:.*]] = fir.insert_value %[[VAL_9]], %[[VAL_8]], [0 : index] : (complex<f64>, f64) -> complex<f64>
 // CHECK:             %[[VAL_11:.*]] = fir.insert_value %[[VAL_10]], %[[VAL_8]], [1 : index] : (complex<f64>, f64) -> complex<f64>
-// CHECK:             %[[VAL_12:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_13:.*]] = fir.do_loop %[[VAL_14:.*]] = %[[VAL_12]] to %[[VAL_4]] step %[[VAL_12]] iter_args(%[[VAL_15:.*]] = %[[VAL_11]]) -> (complex<f64>) {
-// CHECK:               %[[VAL_16:.*]] = arith.constant 0 : index
-// CHECK:               %[[VAL_17:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_16]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
-// CHECK:               %[[VAL_18:.*]] = arith.constant 1 : index
-// CHECK:               %[[VAL_19:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_18]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
-// CHECK:               %[[VAL_20:.*]] = arith.constant 1 : index
-// CHECK:               %[[VAL_21:.*]] = arith.subi %[[VAL_17]]#0, %[[VAL_20]] : index
+// CHECK:               %[[VAL_17:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_2]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK:               %[[VAL_19:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_12]] : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index) -> (index, index, index)
+// CHECK:               %[[VAL_21:.*]] = arith.subi %[[VAL_17]]#0, %[[VAL_12]] : index
 // CHECK:               %[[VAL_22:.*]] = arith.addi %[[VAL_7]], %[[VAL_21]] : index
-// CHECK:               %[[VAL_23:.*]] = arith.subi %[[VAL_19]]#0, %[[VAL_20]] : index
+// CHECK:               %[[VAL_23:.*]] = arith.subi %[[VAL_19]]#0, %[[VAL_12]] : index
 // CHECK:               %[[VAL_24:.*]] = arith.addi %[[VAL_14]], %[[VAL_23]] : index
 // CHECK:               %[[VAL_25:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_22]], %[[VAL_24]])  : (!fir.box<!fir.array<?x3xcomplex<f64>>>, index, index) -> !fir.ref<complex<f64>>
 // CHECK:               %[[VAL_26:.*]] = fir.load %[[VAL_25]] : !fir.ref<complex<f64>>
@@ -142,22 +130,21 @@ func.func @sum_box_unknown_extent2(%arg0: !fir.box<!fir.array<?x3xcomplex<f64>>>
 // CHECK:         }
 
 // expr with unknown extent
-func.func @sum_expr_unknown_extent1(%arg0: !hlfir.expr<?x3xf32>) {
+func.func @sum_expr_unknown_extent1(%arg0: !hlfir.expr<?x3xf32>) -> !hlfir.expr<3xf32> {
   %cst = arith.constant 1 : i32
   %res = hlfir.sum %arg0 dim %cst : (!hlfir.expr<?x3xf32>, i32) -> !hlfir.expr<3xf32>
-  return
+  return %res : !hlfir.expr<3xf32>
 }
 // CHECK-LABEL:   func.func @sum_expr_unknown_extent1(
-// CHECK-SAME:                                        %[[VAL_0:.*]]: !hlfir.expr<?x3xf32>) {
-// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i32
+// CHECK-SAME:                                        %[[VAL_0:.*]]: !hlfir.expr<?x3xf32>) -> !hlfir.expr<3xf32> {
+// CHECK:           %[[VAL_9:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_8:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_4:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_2:.*]] = hlfir.shape_of %[[VAL_0]] : (!hlfir.expr<?x3xf32>) -> !fir.shape<2>
 // CHECK:           %[[VAL_3:.*]] = hlfir.get_extent %[[VAL_2]] {dim = 0 : index} : (!fir.shape<2>) -> index
-// CHECK:           %[[VAL_4:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_5:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_6:.*]] = hlfir.elemental %[[VAL_5]] unordered : (!fir.shape<1>) -> !hlfir.expr<3xf32> {
 // CHECK:           ^bb0(%[[VAL_7:.*]]: index):
-// CHECK:             %[[VAL_8:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK:             %[[VAL_9:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_10:.*]] = fir.do_loop %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_3]] step %[[VAL_9]] iter_args(%[[VAL_12:.*]] = %[[VAL_8]]) -> (f32) {
 // CHECK:               %[[VAL_13:.*]] = hlfir.apply %[[VAL_0]], %[[VAL_11]], %[[VAL_7]] : (!hlfir.expr<?x3xf32>, index, index) -> f32
 // CHECK:               %[[VAL_14:.*]] = arith.addf %[[VAL_12]], %[[VAL_13]] : f32
@@ -168,22 +155,21 @@ func.func @sum_expr_unknown_extent1(%arg0: !hlfir.expr<?x3xf32>) {
 // CHECK:           return
 // CHECK:         }
 
-func.func @sum_expr_unknown_extent2(%arg0: !hlfir.expr<?x3xf32>) {
+func.func @sum_expr_unknown_extent2(%arg0: !hlfir.expr<?x3xf32>) -> !hlfir.expr<?xf32> {
   %cst = arith.constant 2 : i32
   %res = hlfir.sum %arg0 dim %cst : (!hlfir.expr<?x3xf32>, i32) -> !hlfir.expr<?xf32>
-  return
+  return %res : !hlfir.expr<?xf32>
 }
 // CHECK-LABEL:   func.func @sum_expr_unknown_extent2(
-// CHECK-SAME:                                        %[[VAL_0:.*]]: !hlfir.expr<?x3xf32>) {
-// CHECK:           %[[VAL_1:.*]] = arith.constant 2 : i32
+// CHECK-SAME:                                        %[[VAL_0:.*]]: !hlfir.expr<?x3xf32>) -> !hlfir.expr<?xf32> {
+// CHECK:           %[[VAL_9:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_8:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_4:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_2:.*]] = hlfir.shape_of %[[VAL_0]] : (!hlfir.expr<?x3xf32>) -> !fir.shape<2>
 // CHECK:           %[[VAL_3:.*]] = hlfir.get_extent %[[VAL_2]] {dim = 0 : index} : (!fir.shape<2>) -> index
-// CHECK:           %[[VAL_4:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_5:.*]] = fir.shape %[[VAL_3]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_6:.*]] = hlfir.elemental %[[VAL_5]] unordered : (!fir.shape<1>) -> !hlfir.expr<?xf32> {
 // CHECK:           ^bb0(%[[VAL_7:.*]]: index):
-// CHECK:             %[[VAL_8:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK:             %[[VAL_9:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_10:.*]] = fir.do_loop %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_4]] step %[[VAL_9]] iter_args(%[[VAL_12:.*]] = %[[VAL_8]]) -> (f32) {
 // CHECK:               %[[VAL_13:.*]] = hlfir.apply %[[VAL_0]], %[[VAL_7]], %[[VAL_11]] : (!hlfir.expr<?x3xf32>, index, index) -> f32
 // CHECK:               %[[VAL_14:.*]] = arith.addf %[[VAL_12]], %[[VAL_13]] : f32
@@ -195,24 +181,23 @@ func.func @sum_expr_unknown_extent2(%arg0: !hlfir.expr<?x3xf32>) {
 // CHECK:         }
 
 // scalar mask
-func.func @sum_scalar_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.ref<!fir.logical<1>>) {
+func.func @sum_scalar_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.ref<!fir.logical<1>>) -> !hlfir.expr<3xf32> {
   %cst = arith.constant 1 : i32
   %res = hlfir.sum %arg0 dim %cst mask %mask : (!hlfir.expr<?x3xf32>, i32, !fir.ref<!fir.logical<1>>) -> !hlfir.expr<3xf32>
-  return
+  return %res : !hlfir.expr<3xf32>
 }
 // CHECK-LABEL:   func.func @sum_scalar_mask(
 // CHECK-SAME:                               %[[VAL_0:.*]]: !hlfir.expr<?x3xf32>,
-// CHECK-SAME:                               %[[VAL_1:.*]]: !fir.ref<!fir.logical<1>>) {
-// CHECK:           %[[VAL_2:.*]] = arith.constant 1 : i32
+// CHECK-SAME:                               %[[VAL_1:.*]]: !fir.ref<!fir.logical<1>>) -> !hlfir.expr<3xf32> {
+// CHECK:             %[[VAL_11:.*]] = arith.constant 1 : index
+// CHECK:             %[[VAL_10:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_5:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_3:.*]] = hlfir.shape_of %[[VAL_0]] : (!hlfir.expr<?x3xf32>) -> !fir.shape<2>
 // CHECK:           %[[VAL_4:.*]] = hlfir.get_extent %[[VAL_3]] {dim = 0 : index} : (!fir.shape<2>) -> index
-// CHECK:           %[[VAL_5:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_6:.*]] = fir.shape %[[VAL_5]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_7:.*]] = fir.load %[[VAL_1]] : !fir.ref<!fir.logical<1>>
 // CHECK:           %[[VAL_8:.*]] = hlfir.elemental %[[VAL_6]] unordered : (!fir.shape<1>) -> !hlfir.expr<3xf32> {
 // CHECK:           ^bb0(%[[VAL_9:.*]]: index):
-// CHECK:             %[[VAL_10:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK:             %[[VAL_11:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_12:.*]] = fir.do_loop %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_4]] step %[[VAL_11]] iter_args(%[[VAL_14:.*]] = %[[VAL_10]]) -> (f32) {
 // CHECK:               %[[VAL_15:.*]] = fir.convert %[[VAL_7]] : (!fir.logical<1>) -> i1
 // CHECK:               %[[VAL_16:.*]] = fir.if %[[VAL_15]] -> (f32) {
@@ -230,18 +215,20 @@ func.func @sum_scalar_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.ref<!fir.log
 // CHECK:         }
 
 // scalar boxed mask
-func.func @sum_scalar_boxed_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.box<!fir.logical<1>>) {
+func.func @sum_scalar_boxed_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.box<!fir.logical<1>>) -> !hlfir.expr<3xf32> {
   %cst = arith.constant 1 : i32
   %res = hlfir.sum %arg0 dim %cst mask %mask : (!hlfir.expr<?x3xf32>, i32, !fir.box<!fir.logical<1>>) -> !hlfir.expr<3xf32>
-  return
+  return %res : !hlfir.expr<3xf32>
 }
 // CHECK-LABEL:   func.func @sum_scalar_boxed_mask(
 // CHECK-SAME:                                     %[[VAL_0:.*]]: !hlfir.expr<?x3xf32>,
-// CHECK-SAME:                                     %[[VAL_1:.*]]: !fir.box<!fir.logical<1>>) {
-// CHECK:           %[[VAL_2:.*]] = arith.constant 1 : i32
+// CHECK-SAME:                                     %[[VAL_1:.*]]: !fir.box<!fir.logical<1>>) -> !hlfir.expr<3xf32> {
+// CHECK:           %[[VAL_16:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_15:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_11:.*]] = arith.constant true
+// CHECK:           %[[VAL_5:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_3:.*]] = hlfir.shape_of %[[VAL_0]] : (!hlfir.expr<?x3xf32>) -> !fir.shape<2>
 // CHECK:           %[[VAL_4:.*]] = hlfir.get_extent %[[VAL_3]] {dim = 0 : index} : (!fir.shape<2>) -> index
-// CHECK:           %[[VAL_5:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_6:.*]] = fir.shape %[[VAL_5]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_7:.*]] = fir.is_present %[[VAL_1]] : (!fir.box<!fir.logical<1>>) -> i1
 // CHECK:           %[[VAL_8:.*]] = fir.if %[[VAL_7]] -> (!fir.logical<1>) {
@@ -249,14 +236,11 @@ func.func @sum_scalar_boxed_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.box<!f
 // CHECK:             %[[VAL_10:.*]] = fir.load %[[VAL_9]] : !fir.ref<!fir.logical<1>>
 // CHECK:             fir.result %[[VAL_10]] : !fir.logical<1>
 // CHECK:           } else {
-// CHECK:             %[[VAL_11:.*]] = arith.constant true
 // CHECK:             %[[VAL_12:.*]] = fir.convert %[[VAL_11]] : (i1) -> !fir.logical<1>
 // CHECK:             fir.result %[[VAL_12]] : !fir.logical<1>
 // CHECK:           }
 // CHECK:           %[[VAL_13:.*]] = hlfir.elemental %[[VAL_6]] unordered : (!fir.shape<1>) -> !hlfir.expr<3xf32> {
 // CHECK:           ^bb0(%[[VAL_14:.*]]: index):
-// CHECK:             %[[VAL_15:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK:             %[[VAL_16:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_17:.*]] = fir.do_loop %[[VAL_18:.*]] = %[[VAL_16]] to %[[VAL_4]] step %[[VAL_16]] iter_args(%[[VAL_19:.*]] = %[[VAL_15]]) -> (f32) {
 // CHECK:               %[[VAL_20:.*]] = fir.convert %[[VAL_8]] : (!fir.logical<1>) -> i1
 // CHECK:               %[[VAL_21:.*]] = fir.if %[[VAL_20]] -> (f32) {
@@ -274,40 +258,37 @@ func.func @sum_scalar_boxed_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.box<!f
 // CHECK:         }
 
 // array mask
-func.func @sum_array_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.box<!fir.array<?x3x!fir.logical<1>>>) {
+func.func @sum_array_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.box<!fir.array<?x3x!fir.logical<1>>>) -> !hlfir.expr<?xf32> {
   %cst = arith.constant 2 : i32
   %res = hlfir.sum %arg0 dim %cst mask %mask : (!hlfir.expr<?x3xf32>, i32, !fir.box<!fir.array<?x3x!fir.logical<1>>>) -> !hlfir.expr<?xf32>
-  return
+  return %res : !hlfir.expr<?xf32>
 }
 // CHECK-LABEL:   func.func @sum_array_mask(
 // CHECK-SAME:                              %[[VAL_0:.*]]: !hlfir.expr<?x3xf32>,
-// CHECK-SAME:                              %[[VAL_1:.*]]: !fir.box<!fir.array<?x3x!fir.logical<1>>>) {
-// CHECK:           %[[VAL_2:.*]] = arith.constant 2 : i32
+// CHECK-SAME:                              %[[VAL_1:.*]]: !fir.box<!fir.array<?x3x!fir.logical<1>>>) -> !hlfir.expr<?xf32> {
+// CHECK:           %[[VAL_27:.*]] = arith.constant true
+// CHECK:           %[[VAL_16:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_11:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_10:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_5:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_3:.*]] = hlfir.shape_of %[[VAL_0]] : (!hlfir.expr<?x3xf32>) -> !fir.shape<2>
 // CHECK:           %[[VAL_4:.*]] = hlfir.get_extent %[[VAL_3]] {dim = 0 : index} : (!fir.shape<2>) -> index
-// CHECK:           %[[VAL_5:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_6:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_7:.*]] = fir.is_present %[[VAL_1]] : (!fir.box<!fir.array<?x3x!fir.logical<1>>>) -> i1
 // CHECK:           %[[VAL_8:.*]] = hlfir.elemental %[[VAL_6]] unordered : (!fir.shape<1>) -> !hlfir.expr<?xf32> {
 // CHECK:           ^bb0(%[[VAL_9:.*]]: index):
-// CHECK:             %[[VAL_10:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK:             %[[VAL_11:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_12:.*]] = fir.do_loop %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_5]] step %[[VAL_11]] iter_args(%[[VAL_14:.*]] = %[[VAL_10]]) -> (f32) {
 // CHECK:               %[[VAL_15:.*]] = fir.if %[[VAL_7]] -> (!fir.logical<1>) {
-// CHECK:                 %[[VAL_16:.*]] = arith.constant 0 : index
 // CHECK:                 %[[VAL_17:.*]]:3 = fir.box_dims %[[VAL_1]], %[[VAL_16]] : (!fir.box<!fir.array<?x3x!fir.logical<1>>>, index) -> (index, index, index)
-// CHECK:                 %[[VAL_18:.*]] = arith.constant 1 : index
-// CHECK:                 %[[VAL_19:.*]]:3 = fir.box_dims %[[VAL_1]], %[[VAL_18]] : (!fir.box<!fir.array<?x3x!fir.logical<1>>>, index) -> (index, index, index)
-// CHECK:                 %[[VAL_20:.*]] = arith.constant 1 : index
-// CHECK:                 %[[VAL_21:.*]] = arith.subi %[[VAL_17]]#0, %[[VAL_20]] : index
+// CHECK:                 %[[VAL_19:.*]]:3 = fir.box_dims %[[VAL_1]], %[[VAL_11]] : (!fir.box<!fir.array<?x3x!fir.logical<1>>>, index) -> (index, index, index)
+// CHECK:                 %[[VAL_21:.*]] = arith.subi %[[VAL_17]]#0, %[[VAL_11]] : index
 // CHECK:                 %[[VAL_22:.*]] = arith.addi %[[VAL_9]], %[[VAL_21]] : index
-// CHECK:                 %[[VAL_23:.*]] = arith.subi %[[VAL_19]]#0, %[[VAL_20]] : index
+// CHECK:                 %[[VAL_23:.*]] = arith.subi %[[VAL_19]]#0, %[[VAL_11]] : index
 // CHECK:                 %[[VAL_24:.*]] = arith.addi %[[VAL_13]], %[[VAL_23]] : index
 // CHECK:                 %[[VAL_25:.*]] = hlfir.designate %[[VAL_1]] (%[[VAL_22]], %[[VAL_24]])  : (!fir.box<!fir.array<?x3x!fir.logical<1>>>, index, index) -> !fir.ref<!fir.logical<1>>
 // CHECK:                 %[[VAL_26:.*]] = fir.load %[[VAL_25]] : !fir.ref<!fir.logical<1>>
 // CHECK:                 fir.result %[[VAL_26]] : !fir.logical<1>
 // CHECK:               } else {
-// CHECK:                 %[[VAL_27:.*]] = arith.constant true
 // CHECK:                 %[[VAL_28:.*]] = fir.convert %[[VAL_27]] : (i1) -> !fir.logical<1>
 // CHECK:                 fir.result %[[VAL_28]] : !fir.logical<1>
 // CHECK:               }
@@ -327,23 +308,22 @@ func.func @sum_array_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !fir.box<!fir.arra
 // CHECK:         }
 
 // array expr mask
-func.func @sum_array_expr_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !hlfir.expr<?x3x!fir.logical<1>>) {
+func.func @sum_array_expr_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !hlfir.expr<?x3x!fir.logical<1>>) -> !hlfir.expr<?xf32> {
   %cst = arith.constant 2 : i32
   %res = hlfir.sum %arg0 dim %cst mask %mask : (!hlfir.expr<?x3xf32>, i32, !hlfir.expr<?x3x!fir.logical<1>>) -> !hlfir.expr<?xf32>
-  return
+  return %res : !hlfir.expr<?xf32>
 }
 // CHECK-LABEL:   func.func @sum_array_expr_mask(
 // CHECK-SAME:                                   %[[VAL_0:.*]]: !hlfir.expr<?x3xf32>,
-// CHECK-SAME:                                   %[[VAL_1:.*]]: !hlfir.expr<?x3x!fir.logical<1>>) {
-// CHECK:           %[[VAL_2:.*]] = arith.constant 2 : i32
+// CHECK-SAME:                                   %[[VAL_1:.*]]: !hlfir.expr<?x3x!fir.logical<1>>) -> !hlfir.expr<?xf32> {
+// CHECK:           %[[VAL_10:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_9:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK:           %[[VAL_5:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_3:.*]] = hlfir.shape_of %[[VAL_0]] : (!hlfir.expr<?x3xf32>) -> !fir.shape<2>
 // CHECK:           %[[VAL_4:.*]] = hlfir.get_extent %[[VAL_3]] {dim = 0 : index} : (!fir.shape<2>) -> index
-// CHECK:           %[[VAL_5:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_6:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_7:.*]] = hlfir.elemental %[[VAL_6]] unordered : (!fir.shape<1>) -> !hlfir.expr<?xf32> {
 // CHECK:           ^bb0(%[[VAL_8:.*]]: index):
-// CHECK:             %[[VAL_9:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK:             %[[VAL_10:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_11:.*]] = fir.do_loop %[[VAL_12:.*]] = %[[VAL_10]] to %[[VAL_5]] step %[[VAL_10]] iter_args(%[[VAL_13:.*]] = %[[VAL_9]]) -> (f32) {
 // CHECK:               %[[VAL_14:.*]] = hlfir.apply %[[VAL_1]], %[[VAL_8]], %[[VAL_12]] : (!hlfir.expr<?x3x!fir.logical<1>>, index, index) -> !fir.logical<1>
 // CHECK:               %[[VAL_15:.*]] = fir.convert %[[VAL_14]] : (!fir.logical<1>) -> i1
@@ -362,21 +342,20 @@ func.func @sum_array_expr_mask(%arg0: !hlfir.expr<?x3xf32>, %mask: !hlfir.expr<?
 // CHECK:         }
 
 // unordered floating point reduction
-func.func @sum_unordered_reduction(%arg0: !hlfir.expr<2x3xf32>) {
+func.func @sum_unordered_reduction(%arg0: !hlfir.expr<2x3xf32>) -> !hlfir.expr<3xf32> {
   %cst = arith.constant 1 : i32
   %res = hlfir.sum %arg0 dim %cst {fastmath = #arith.fastmath<reassoc>} : (!hlfir.expr<2x3xf32>, i32) -> !hlfir.expr<3xf32>
-  return
+  return %res : !hlfir.expr<3xf32>
 }
 // CHECK-LABEL:   func.func @sum_unordered_reduction(
-// CHECK-SAME:                                       %[[VAL_0:.*]]: !hlfir.expr<2x3xf32>) {
-// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i32
+// CHECK-SAME:                                       %[[VAL_0:.*]]: !hlfir.expr<2x3xf32>) -> !hlfir.expr<3xf32> {
+// CHECK:           %[[VAL_8:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_7:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK:           %[[VAL_2:.*]] = arith.constant 2 : index
 // CHECK:           %[[VAL_3:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_4:.*]] = fir.shape %[[VAL_3]] : (index) -> !fir.shape<1>
 // CHECK:           %[[VAL_5:.*]] = hlfir.elemental %[[VAL_4]] unordered : (!fir.shape<1>) -> !hlfir.expr<3xf32> {
 // CHECK:           ^bb0(%[[VAL_6:.*]]: index):
-// CHECK:             %[[VAL_7:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK:             %[[VAL_8:.*]] = arith.constant 1 : index
 // CHECK:             %[[VAL_9:.*]] = fir.do_loop %[[VAL_10:.*]] = %[[VAL_8]] to %[[VAL_2]] step %[[VAL_8]] unordered iter_args(%[[VAL_11:.*]] = %[[VAL_7]]) -> (f32) {
 // CHECK:               %[[VAL_12:.*]] = hlfir.apply %[[VAL_0]], %[[VAL_10]], %[[VAL_6]] : (!hlfir.expr<2x3xf32>, index, index) -> f32
 // CHECK:               %[[VAL_13:.*]] = arith.addf %[[VAL_11]], %[[VAL_12]] fastmath<reassoc> : f32
@@ -388,22 +367,20 @@ func.func @sum_unordered_reduction(%arg0: !hlfir.expr<2x3xf32>) {
 // CHECK:         }
 
 // total 1d reduction
-func.func @sum_total_1d_reduction(%arg0: !fir.box<!fir.array<3xi32>>) {
+func.func @sum_total_1d_reduction(%arg0: !fir.box<!fir.array<3xi32>>) -> i32 {
   %cst = arith.constant 1 : i32
   %res = hlfir.sum %arg0 dim %cst : (!fir.box<!fir.array<3xi32>>, i32) -> i32
-  return
+  return %res : i32
 }
 // CHECK-LABEL:   func.func @sum_total_1d_reduction(
-// CHECK-SAME:                                      %[[VAL_0:.*]]: !fir.box<!fir.array<3xi32>>) {
-// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i32
+// CHECK-SAME:                                      %[[VAL_0:.*]]: !fir.box<!fir.array<3xi32>>) -> i32 {
+// CHECK:           %[[VAL_8:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_2:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_3:.*]] = arith.constant 0 : i32
 // CHECK:           %[[VAL_4:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_5:.*]] = fir.do_loop %[[VAL_6:.*]] = %[[VAL_4]] to %[[VAL_2]] step %[[VAL_4]] unordered iter_args(%[[VAL_7:.*]] = %[[VAL_3]]) -> (i32) {
-// CHECK:             %[[VAL_8:.*]] = arith.constant 0 : index
 // CHECK:             %[[VAL_9:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_8]] : (!fir.box<!fir.array<3xi32>>, index) -> (index, index, index)
-// CHECK:             %[[VAL_10:.*]] = arith.constant 1 : index
-// CHECK:             %[[VAL_11:.*]] = arith.subi %[[VAL_9]]#0, %[[VAL_10]] : index
+// CHECK:             %[[VAL_11:.*]] = arith.subi %[[VAL_9]]#0, %[[VAL_4]] : index
 // CHECK:             %[[VAL_12:.*]] = arith.addi %[[VAL_6]], %[[VAL_11]] : index
 // CHECK:             %[[VAL_13:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_12]])  : (!fir.box<!fir.array<3xi32>>, index) -> !fir.ref<i32>
 // CHECK:             %[[VAL_14:.*]] = fir.load %[[VAL_13]] : !fir.ref<i32>
@@ -414,27 +391,24 @@ func.func @sum_total_1d_reduction(%arg0: !fir.box<!fir.array<3xi32>>) {
 // CHECK:         }
 
 // total 2d reduction
-func.func @sum_total_2d_reduction(%arg0: !fir.box<!fir.array<?x3xi32>>) {
+func.func @sum_total_2d_reduction(%arg0: !fir.box<!fir.array<?x3xi32>>) -> i32 {
   %res = hlfir.sum %arg0 : (!fir.box<!fir.array<?x3xi32>>) -> i32
-  return
+  return %res : i32
 }
 // CHECK-LABEL:   func.func @sum_total_2d_reduction(
-// CHECK-SAME:                                      %[[VAL_0:.*]]: !fir.box<!fir.array<?x3xi32>>) {
+// CHECK-SAME:                                      %[[VAL_0:.*]]: !fir.box<!fir.array<?x3xi32>>) -> i32 {
+// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_4:.*]] = arith.constant 0 : i32
+// CHECK:           %[[VAL_3:.*]] = arith.constant 3 : index
 // CHECK:           %[[VAL_1:.*]] = arith.constant 0 : index
 // CHECK:           %[[VAL_2:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_1]] : (!fir.box<!fir.array<?x3xi32>>, index) -> (index, index, index)
-// CHECK:           %[[VAL_3:.*]] = arith.constant 3 : index
-// CHECK:           %[[VAL_4:.*]] = arith.constant 0 : i32
-// CHECK:           %[[VAL_5:.*]] = arith.constant 1 : index
 // CHECK:           %[[VAL_6:.*]] = fir.do_loop %[[VAL_7:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_5]] unordered iter_args(%[[VAL_8:.*]] = %[[VAL_4]]) -> (i32) {
 // CHECK:             %[[VAL_9:.*]] = fir.do_loop %[[VAL_10:.*]] = %[[VAL_5]] to %[[VAL_2]]#1 step %[[VAL_5]] unordered iter_args(%[[VAL_11:.*]] = %[[VAL_8]]) -> (i32) {
-// CHECK:               %[[VAL_12:.*]] = arith.constant 0 : index
-// CHECK:               %[[VAL_13:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_12]] : (!fir.box<!fir.array<?x3xi32>>, index) -> (index, index, index)
-// CHECK:               %[[VAL_14:.*]] = arith.constant 1 : index
-// CHECK:               %[[VAL_15:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_14]] : (!fir.box<!fir.array<?x3xi32>>, index) -> (index, index, index)
-// CHECK:               %[[VAL_16:.*]] = arith.constant 1 : index
-// CHECK:               %[[VAL_17:.*]] = arith.subi %[[VAL_13]]#0, %[[VAL_16]] : index
+// CHECK:               %[[VAL_13:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_1]] : (!fir.box<!fir.array<?x3xi32>>, index) -> (index, index, index)
+// CHECK:               %[[VAL_15:.*]]:3 = fir.box_dims %[[VAL_0]], %[[VAL_5]] : (!fir.box<!fir.array<?x3xi32>>, index) -> (index, index, index)
+// CHECK:               %[[VAL_17:.*]] = arith.subi %[[VAL_13]]#0, %[[VAL_5]] : index
 // CHECK:               %[[VAL_18:.*]] = arith.addi %[[VAL_10]], %[[VAL_17]] : index
-// CHECK:               %[[VAL_19:.*]] = arith.subi %[[VAL_15]]#0, %[[VAL_16]] : index
+// CHECK:               %[[VAL_19:.*]] = arith.subi %[[VAL_15]]#0, %[[VAL_5]] : index
 // CHECK:               %[[VAL_20:.*]] = arith.addi %[[VAL_7]], %[[VAL_19]] : index
 // CHECK:               %[[VAL_21:.*]] = hlfir.designate %[[VAL_0]] (%[[VAL_18]], %[[VAL_20]])  : (!fir.box<!fir.array<?x3xi32>>, index, index) -> !fir.ref<i32>
 // CHECK:               %[[VAL_22:.*]] = fir.load %[[VAL_21]] : !fir.ref<i32>
@@ -447,19 +421,19 @@ func.func @sum_total_2d_reduction(%arg0: !fir.box<!fir.array<?x3xi32>>) {
 // CHECK:         }
 
 // negative: invalid dim==0
-func.func @sum_invalid_dim0(%arg0: !hlfir.expr<2x3xi32>) {
+func.func @sum_invalid_dim0(%arg0: !hlfir.expr<2x3xi32>) -> !hlfir.expr<3xi32> {
   %cst = arith.constant 0 : i32
   %res = hlfir.sum %arg0 dim %cst : (!hlfir.expr<2x3xi32>, i32) -> !hlfir.expr<3xi32>
-  return
+  return %res : !hlfir.expr<3xi32>
 }
 // CHECK-LABEL:   func.func @sum_invalid_dim0(
 // CHECK:           hlfir.sum %{{.*}} dim %{{.*}} : (!hlfir.expr<2x3xi32>, i32) -> !hlfir.expr<3xi32>
 
 // negative: invalid dim>rank
-func.func @sum_invalid_dim_big(%arg0: !hlfir.expr<2x3xi32>) {
+func.func @sum_invalid_dim_big(%arg0: !hlfir.expr<2x3xi32>) -> !hlfir.expr<3xi32> {
   %cst = arith.constant 3 : i32
   %res = hlfir.sum %arg0 dim %cst : (!hlfir.expr<2x3xi32>, i32) -> !hlfir.expr<3xi32>
-  return
+  return %res : !hlfir.expr<3xi32>
 }
 // CHECK-LABEL:   func.func @sum_invalid_dim_big(
 // CHECK:           hlfir.sum %{{.*}} dim %{{.*}} : (!hlfir.expr<2x3xi32>, i32) -> !hlfir.expr<3xi32>
diff --git a/flang/test/HLFIR/simplify-hlfir-intrinsics.fir b/flang/test/HLFIR/simplify-hlfir-intrinsics.fir
index aeea8bfc9732..5fd21e98a8d8 100644
--- a/flang/test/HLFIR/simplify-hlfir-intrinsics.fir
+++ b/flang/test/HLFIR/simplify-hlfir-intrinsics.fir
@@ -1,25 +1,23 @@
 // RUN: fir-opt --simplify-hlfir-intrinsics %s | FileCheck %s
 
 // box with known extents
-func.func @transpose0(%arg0: !fir.box<!fir.array<1x2xi32>>) {
+func.func @transpose0(%arg0: !fir.box<!fir.array<1x2xi32>>) -> !hlfir.expr<2x1xi32> {
   %res = hlfir.transpose %arg0 : (!fir.box<!fir.array<1x2xi32>>) -> !hlfir.expr<2x1xi32>
-  return
+  return %res : !hlfir.expr<2x1xi32>
 }
 // CHECK-LABEL:   func.func @transpose0(
-// CHECK-SAME:                          %[[ARG0:.*]]: !fir.box<!fir.array<1x2xi32>>) {
+// CHECK-SAME:                          %[[ARG0:.*]]: !fir.box<!fir.array<1x2xi32>>) -> !hlfir.expr<2x1xi32> {
+// CHECK:           %[[C0:.*]] = arith.constant 0 : index
 // CHECK:           %[[C1:.*]] = arith.constant 1 : index
 // CHECK:           %[[C2:.*]] = arith.constant 2 : index
 // CHECK:           %[[SHAPE:.*]] = fir.shape %[[C2]], %[[C1]] : (index, index) -> !fir.shape<2>
 // CHECK:           %[[EXPR:.*]] = hlfir.elemental %[[SHAPE]] unordered : (!fir.shape<2>) -> !hlfir.expr<2x1xi32> {
 // CHECK:           ^bb0(%[[I:.*]]: index, %[[J:.*]]: index):
-// CHECK:             %[[C0:.*]] = arith.constant 0 : index
 // CHECK:             %[[DIMS0:.*]]:3 = fir.box_dims %[[ARG0]], %[[C0]] : (!fir.box<!fir.array<1x2xi32>>, index) -> (index, index, index)
-// CHECK:             %[[C1_1:.*]] = arith.constant 1 : index
-// CHECK:             %[[DIMS1:.*]]:3 = fir.box_dims %[[ARG0]], %[[C1_1]] : (!fir.box<!fir.array<1x2xi32>>, index) -> (index, index, index)
-// CHECK:             %[[C1_2:.*]] = arith.constant 1 : index
-// CHECK:             %[[LOWER_BOUND0:.*]] = arith.subi %[[DIMS0]]#0, %[[C1_2]] : index
+// CHECK:             %[[DIMS1:.*]]:3 = fir.box_dims %[[ARG0]], %[[C1]] : (!fir.box<!fir.array<1x2xi32>>, index) -> (index, index, index)
+// CHECK:             %[[LOWER_BOUND0:.*]] = arith.subi %[[DIMS0]]#0, %[[C1]] : index
 // CHECK:             %[[J_OFFSET:.*]] = arith.addi %[[J]], %[[LOWER_BOUND0]] : index
-// CHECK:             %[[LOWER_BOUND1:.*]] = arith.subi %[[DIMS1]]#0, %[[C1_2]] : index
+// CHECK:             %[[LOWER_BOUND1:.*]] = arith.subi %[[DIMS1]]#0, %[[C1]] : index
 // CHECK:             %[[I_OFFSET:.*]] = arith.addi %[[I]], %[[LOWER_BOUND1]] : index
 // CHECK:             %[[ELEMENT_REF:.*]] = hlfir.designate %[[ARG0]] (%[[J_OFFSET]], %[[I_OFFSET]])  : (!fir.box<!fir.array<1x2xi32>>, index, index) -> !fir.ref<i32>
 // CHECK:             %[[ELEMENT:.*]] = fir.load %[[ELEMENT_REF]] : !fir.ref<i32>
@@ -29,12 +27,12 @@ func.func @transpose0(%arg0: !fir.box<!fir.array<1x2xi32>>) {
 // CHECK:         }
 
 // expr with known extents
-func.func @transpose1(%arg0: !hlfir.expr<1x2xi32>) {
+func.func @transpose1(%arg0: !hlfir.expr<1x2xi32>) -> !hlfir.expr<2x1xi32> {
   %res = hlfir.transpose %arg0 : (!hlfir.expr<1x2xi32>) -> !hlfir.expr<2x1xi32>
-  return
+  return %res : !hlfir.expr<2x1xi32>
 }
 // CHECK-LABEL:   func.func @transpose1(
-// CHECK-SAME:                          %[[ARG0:.*]]: !hlfir.expr<1x2xi32>) {
+// CHECK-SAME:                          %[[ARG0:.*]]: !hlfir.expr<1x2xi32>) -> !hlfir.expr<2x1xi32> {
 // CHECK:           %[[C1:.*]] = arith.constant 1 : index
 // CHECK:           %[[C2:.*]] = arith.constant 2 : index
 // CHECK:           %[[SHAPE:.*]] = fir.shape %[[C2]], %[[C1]] : (index, index) -> !fir.shape<2>
@@ -47,26 +45,24 @@ func.func @transpose1(%arg0: !hlfir.expr<1x2xi32>) {
 // CHECK:         }
 
 // box with unknown extent
-func.func @transpose2(%arg0: !fir.box<!fir.array<?x2xi32>>) {
+func.func @transpose2(%arg0: !fir.box<!fir.array<?x2xi32>>) -> !hlfir.expr<2x?xi32> {
   %res = hlfir.transpose %arg0 : (!fir.box<!fir.array<?x2xi32>>) -> !hlfir.expr<2x?xi32>
-  return
+  return %res : !hlfir.expr<2x?xi32>
 }
 // CHECK-LABEL:   func.func @transpose2(
-// CHECK-SAME:                          %[[ARG0:.*]]: !fir.box<!fir.array<?x2xi32>>) {
+// CHECK-SAME:                          %[[ARG0:.*]]: !fir.box<!fir.array<?x2xi32>>) -> !hlfir.expr<2x?xi32> {
+// CHECK:           %[[C1:.*]] = arith.constant 1 : index
+// CHECK:           %[[C2:.*]] = arith.constant 2 : index
 // CHECK:           %[[C0:.*]] = arith.constant 0 : index
 // CHECK:           %[[DIMS0:.*]]:3 = fir.box_dims %[[ARG0]], %[[C0]] : (!fir.box<!fir.array<?x2xi32>>, index) -> (index, index, index)
-// CHECK:           %[[C2:.*]] = arith.constant 2 : index
 // CHECK:           %[[SHAPE:.*]] = fir.shape %[[C2]], %[[DIMS0]]#1 : (index, index) -> !fir.shape<2>
 // CHECK:           %[[EXPR:.*]] = hlfir.elemental %[[SHAPE]] unordered : (!fir.shape<2>) -> !hlfir.expr<2x?xi32> {
 // CHECK:           ^bb0(%[[I:.*]]: index, %[[J:.*]]: index):
-// CHECK:             %[[C0_1:.*]] = arith.constant 0 : index
-// CHECK:             %[[DIMS0:.*]]:3 = fir.box_dims %[[ARG0]], %[[C0_1]] : (!fir.box<!fir.array<?x2xi32>>, index) -> (index, index, index)
-// CHECK:             %[[C1_1:.*]] = arith.constant 1 : index
-// CHECK:             %[[DIMS1_1:.*]]:3 = fir.box_dims %[[ARG0]], %[[C1_1]] : (!fir.box<!fir.array<?x2xi32>>, index) -> (index, index, index)
-// CHECK:             %[[C1_2:.*]] = arith.constant 1 : index
-// CHECK:             %[[LOWER_BOUND0:.*]] = arith.subi %[[DIMS0]]#0, %[[C1_2]] : index
+// CHECK:             %[[DIMS0:.*]]:3 = fir.box_dims %[[ARG0]], %[[C0]] : (!fir.box<!fir.array<?x2xi32>>, index) -> (index, index, index)
+// CHECK:             %[[DIMS1_1:.*]]:3 = fir.box_dims %[[ARG0]], %[[C1]] : (!fir.box<!fir.array<?x2xi32>>, index) -> (index, index, index)
+// CHECK:             %[[LOWER_BOUND0:.*]] = arith.subi %[[DIMS0]]#0, %[[C1]] : index
 // CHECK:             %[[J_OFFSET:.*]] = arith.addi %[[J]], %[[LOWER_BOUND0]] : index
-// CHECK:             %[[LOWER_BOUND1:.*]] = arith.subi %[[DIMS1_1]]#0, %[[C1_2]] : index
+// CHECK:             %[[LOWER_BOUND1:.*]] = arith.subi %[[DIMS1_1]]#0, %[[C1]] : index
 // CHECK:             %[[I_OFFSET:.*]] = arith.addi %[[I]], %[[LOWER_BOUND1]] : index
 // CHECK:             %[[ELE_REF:.*]] = hlfir.designate %[[ARG0]] (%[[J_OFFSET]], %[[I_OFFSET]])  : (!fir.box<!fir.array<?x2xi32>>, index, index) -> !fir.ref<i32>
 // CHECK:             %[[ELEMENT:.*]] = fir.load %[[ELE_REF]] : !fir.ref<i32>
@@ -76,15 +72,15 @@ func.func @transpose2(%arg0: !fir.box<!fir.array<?x2xi32>>) {
 // CHECK:         }
 
 // expr with unknown extent
-func.func @transpose3(%arg0: !hlfir.expr<?x2xi32>) {
+func.func @transpose3(%arg0: !hlfir.expr<?x2xi32>) -> !hlfir.expr<2x?xi32> {
   %res = hlfir.transpose %arg0 : (!hlfir.expr<?x2xi32>) -> !hlfir.expr<2x?xi32>
-  return
+  return %res : !hlfir.expr<2x?xi32>
 }
 // CHECK-LABEL:   func.func @transpose3(
-// CHECK-SAME:                          %[[ARG0:.*]]: !hlfir.expr<?x2xi32>) {
+// CHECK-SAME:                          %[[ARG0:.*]]: !hlfir.expr<?x2xi32>) -> !hlfir.expr<2x?xi32> {
+// CHECK:           %[[C2:.*]] = arith.constant 2 : index
 // CHECK:           %[[IN_SHAPE:.*]] = hlfir.shape_of %[[ARG0]] : (!hlfir.expr<?x2xi32>) -> !fir.shape<2>
 // CHECK:           %[[EXTENT0:.*]] = hlfir.get_extent %[[IN_SHAPE]] {dim = 0 : index} : (!fir.shape<2>) -> index
-// CHECK:           %[[C2:.*]] = arith.constant 2 : index
 // CHECK:           %[[OUT_SHAPE:.*]] = fir.shape %[[C2]], %[[EXTENT0]] : (index, index) -> !fir.shape<2>
 // CHECK:           %[[EXPR:.*]] = hlfir.elemental %[[OUT_SHAPE]] unordered : (!fir.shape<2>) -> !hlfir.expr<2x?xi32> {
 // CHECK:           ^bb0(%[[I:.*]]: index, %[[J:.*]]: index):
@@ -118,8 +114,7 @@ func.func @transpose4(%arg0: !hlfir.expr<2x2xf32>, %arg1: !fir.ref<!fir.box<!fir
 // CHECK:           %[[ELE:.*]] = hlfir.apply %[[ARG0]], %[[J]], %[[I]] : (!hlfir.expr<2x2xf32>, index, index) -> f32
 // CHECK:           hlfir.yield_element %[[ELE]] : f32
 // CHECK:         }
-// CHECK:         %[[SHAPE1:.*]] = hlfir.shape_of %[[TRANSPOSE]] : (!hlfir.expr<2x2xf32>) -> !fir.shape<2>
-// CHECK:         %[[COS:.*]] = hlfir.elemental %[[SHAPE1]] : (!fir.shape<2>) -> !hlfir.expr<2x2xf32> {
+// CHECK:         %[[COS:.*]] = hlfir.elemental %[[SHAPE0]] : (!fir.shape<2>) -> !hlfir.expr<2x2xf32> {
 // CHECK:         ^bb0(%[[I:.*]]: index, %[[J:.*]]: index):
 // CHECK:           %[[ELE:.*]] = hlfir.apply %[[TRANSPOSE]], %[[I]], %[[J]] : (!hlfir.expr<2x2xf32>, index, index) -> f32
 // CHECK:           %[[COS_ELE:.*]] = math.cos %[[ELE]] fastmath<contract> : f32
@@ -170,14 +165,13 @@ func.func @transpose5(%arg0: !fir.ref<tuple<!fir.box<!fir.array<2x2xf64>>, !fir.
 }
 // CHECK-LABEL: func.func @transpose5(
 // ...
-// CHECK:         %[[TRANSPOSE:.*]] = hlfir.elemental %[[SHAPE0:.*]]
+// CHECK:         %[[TRANSPOSE:.*]] = hlfir.elemental %[[SHAPE0:[A-Za-z._]*]]
 // CHECK:         ^bb0(%[[I:.*]]: index, %[[J:.*]]: index):
 // CHECK:           %[[ELE:.*]] = hlfir.designate %[[ARRAY:.*]] (%[[J]], %[[I]])
 // CHECK:           %[[LOAD:.*]] = fir.load %[[ELE]]
 // CHECK:           hlfir.yield_element %[[LOAD]]
 // CHECK:         }
-// CHECK:         %[[SHAPE1:.*]] = hlfir.shape_of %[[TRANSPOSE]]
-// CHECK:         %[[COS:.*]] = hlfir.elemental %[[SHAPE1]]
+// CHECK:         %[[COS:.*]] = hlfir.elemental %[[SHAPE0]]
 // ...
 // CHECK:         hlfir.assign %[[COS]] to %{{.*}} realloc
 // CHECK:         hlfir.destroy %[[COS]]