Reapply "[flang] Lower EOSHIFT into hlfir.eoshift." (#153907)

flang/lib/Lower/HlfirIntrinsics.cpp

@@ -170,6 +170,17 @@ class HlfirCShiftLowering : public HlfirTransformationalIntrinsic {
             mlir::Type stmtResultType) override;
 };
 
+class HlfirEOShiftLowering : public HlfirTransformationalIntrinsic {
+public:
+  using HlfirTransformationalIntrinsic::HlfirTransformationalIntrinsic;
+
+protected:
+  mlir::Value
+  lowerImpl(const Fortran::lower::PreparedActualArguments &loweredActuals,
+            const fir::IntrinsicArgumentLoweringRules *argLowering,
+            mlir::Type stmtResultType) override;
+};
+
 class HlfirReshapeLowering : public HlfirTransformationalIntrinsic {
 public:
   using HlfirTransformationalIntrinsic::HlfirTransformationalIntrinsic;
@@ -430,6 +441,46 @@ mlir::Value HlfirCShiftLowering::lowerImpl(
   return createOp<hlfir::CShiftOp>(resultType, operands);
 }
 
+mlir::Value HlfirEOShiftLowering::lowerImpl(
+    const Fortran::lower::PreparedActualArguments &loweredActuals,
+    const fir::IntrinsicArgumentLoweringRules *argLowering,
+    mlir::Type stmtResultType) {
+  auto operands = getOperandVector(loweredActuals, argLowering);
+  assert(operands.size() == 4);
+  mlir::Value array = operands[0];
+  mlir::Value shift = operands[1];
+  mlir::Value boundary = operands[2];
+  mlir::Value dim = operands[3];
+  // If DIM is present, then dereference it if it is a ref.
+  if (dim)
+    dim = hlfir::loadTrivialScalar(loc, builder, hlfir::Entity{dim});
+
+  mlir::Type resultType = computeResultType(array, stmtResultType);
+
+  if (boundary && fir::isa_trivial(boundary.getType())) {
+    mlir::Type elementType = hlfir::getFortranElementType(resultType);
+    if (auto logicalTy = mlir::dyn_cast<fir::LogicalType>(elementType)) {
+      // Scalar logical constant boundary might be represented using i1, i2, ...
+      // type. We need to cast it to fir.logical type of the ARRAY/result.
+      if (boundary.getType() != logicalTy)
+        boundary = builder.createConvert(loc, logicalTy, boundary);
+    } else {
+      // When the boundary is a constant like '1u', the lowering converts
+      // it into a signless arith.constant value (which is a requirement
+      // of the Arith dialect). If the ARRAY/RESULT is also UNSIGNED,
+      // we have to cast the boundary to the same unsigned type.
+      auto resultIntTy = mlir::dyn_cast<mlir::IntegerType>(elementType);
+      auto boundaryIntTy =
+          mlir::dyn_cast<mlir::IntegerType>(boundary.getType());
+      if (resultIntTy && boundaryIntTy &&
+          resultIntTy.getSignedness() != boundaryIntTy.getSignedness())
+        boundary = builder.createConvert(loc, resultIntTy, boundary);
+    }
+  }
+
+  return createOp<hlfir::EOShiftOp>(resultType, array, shift, boundary, dim);
+}
+
 mlir::Value HlfirReshapeLowering::lowerImpl(
     const Fortran::lower::PreparedActualArguments &loweredActuals,
     const fir::IntrinsicArgumentLoweringRules *argLowering,
@@ -489,6 +540,9 @@ std::optional<hlfir::EntityWithAttributes> Fortran::lower::lowerHlfirIntrinsic(
   if (name == "cshift")
     return HlfirCShiftLowering{builder, loc}.lower(loweredActuals, argLowering,
                                                    stmtResultType);
+  if (name == "eoshift")
+    return HlfirEOShiftLowering{builder, loc}.lower(loweredActuals, argLowering,
+                                                    stmtResultType);
   if (name == "reshape")
     return HlfirReshapeLowering{builder, loc}.lower(loweredActuals, argLowering,
                                                     stmtResultType);

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -147,8 +147,20 @@ mlir::Value fir::FirOpBuilder::createIntegerConstant(mlir::Location loc,
   assert((cst >= 0 || mlir::isa<mlir::IndexType>(ty) ||
           mlir::cast<mlir::IntegerType>(ty).getWidth() <= 64) &&
          "must use APint");
-  return mlir::arith::ConstantOp::create(*this, loc, ty,
-                                         getIntegerAttr(ty, cst));
+
+  mlir::Type cstType = ty;
+  if (auto intType = mlir::dyn_cast<mlir::IntegerType>(ty)) {
+    // Signed and unsigned constants must be encoded as signless
+    // arith.constant followed by fir.convert cast.
+    if (intType.isUnsigned())
+      cstType = mlir::IntegerType::get(getContext(), intType.getWidth());
+    else if (intType.isSigned())
+      TODO(loc, "signed integer constant");
+  }
+
+  mlir::Value cstValue = mlir::arith::ConstantOp::create(
+      *this, loc, cstType, getIntegerAttr(cstType, cst));
+  return createConvert(loc, ty, cstValue);
 }
 
 mlir::Value fir::FirOpBuilder::createAllOnesInteger(mlir::Location loc,

flang/test/HLFIR/simplify-hlfir-intrinsics-eoshift.fir

@@ -697,6 +697,33 @@ func.func @_QPeoshift7(%arg0: !fir.ref<i32> {fir.bindc_name = "n"}, %arg1: !fir.
 // CHECK:           return
 // CHECK:         }
 
+// Test UNSIGNED data type.
+// The default value of the BOUNDARY must be an integer 0
+// converted to ui32 type.
+// subroutine eoshift8(array)
+//   unsigned :: array(:,:)
+//   array = EOSHIFT(array, shift=1, dim=2)
+// end subroutine
+func.func @_QPeoshift8(%arg0: !fir.box<!fir.array<?x?xui32>> {fir.bindc_name = "array"}) {
+  %c2_i32 = arith.constant 2 : i32
+  %c1_i32 = arith.constant 1 : i32
+  %0 = fir.dummy_scope : !fir.dscope
+  %1:2 = hlfir.declare %arg0 dummy_scope %0 {uniq_name = "_QFeoshift8Earray"} : (!fir.box<!fir.array<?x?xui32>>, !fir.dscope) -> (!fir.box<!fir.array<?x?xui32>>, !fir.box<!fir.array<?x?xui32>>)
+  %2 = hlfir.eoshift %1#0 %c1_i32 dim %c2_i32 : (!fir.box<!fir.array<?x?xui32>>, i32, i32) -> !hlfir.expr<?x?xui32>
+  hlfir.assign %2 to %1#0 : !hlfir.expr<?x?xui32>, !fir.box<!fir.array<?x?xui32>>
+  hlfir.destroy %2 : !hlfir.expr<?x?xui32>
+  return
+}
+// CHECK-LABEL:   func.func @_QPeoshift8(
+// CHECK-DAG:           hlfir.elemental %{{.*}} unordered : (!fir.shape<2>) -> !hlfir.expr<?x?xui32> {
+// CHECK-DAG:               %[[VAL_24:.*]] = fir.load %{{.*}} : !fir.ref<ui32>
+// CHECK-DAG:               fir.result %[[VAL_24]] : ui32
+// CHECK-DAG:             } else {
+// CHECK-DAG:               fir.result %[[VAL_12:.*]] : ui32
+// CHECK-DAG:             }
+// CHECK-DAG:           %[[VAL_12]] = fir.convert %[[VAL_1:.*]] : (i32) -> ui32
+// CHECK-DAG:           %[[VAL_1]] = arith.constant 0 : i32
+
 // ! Tests for CHARACTER type (lowered via hlfir.elemental).
 
 // ! Test contiguous 1D array with statically absent boundary.