[flang] add nsw to operations in subscripts (#110060)

This patch adds nsw to operations when lowering subscripts.

See also the discussion in the following discourse post.
https://discourse.llvm.org/t/rfc-add-nsw-flags-to-arithmetic-integer-operations-using-the-option-fno-wrapv/77584/9

Author: yus3710-fj

flang/docs/Extensions.md

@@ -780,6 +780,12 @@ character j
 print *, [(j,j=1,10)]
 ```
 
+* The Fortran standard doesn't mention integer overflow explicitly. In many cases,
+  however, integer overflow makes programs non-conforming.
+  F18 follows other widely-used Fortran compilers. Specifically, f18 assumes
+  integer overflow never occurs in address calculations and increment of
+  do-variable unless the option `-fwrapv` is enabled.
+
 ## De Facto Standard Features
 
 * `EXTENDS_TYPE_OF()` returns `.TRUE.` if both of its arguments have the

flang/include/flang/Lower/LoweringOptions.def

@@ -34,8 +34,14 @@ ENUM_LOWERINGOPT(NoPPCNativeVecElemOrder, unsigned, 1, 0)
 /// On by default.
 ENUM_LOWERINGOPT(Underscoring, unsigned, 1, 1)
 
+/// If true, assume the behavior of integer overflow is defined
+/// (i.e. wraps around as two's complement). On by default.
+/// TODO: make the default off
+ENUM_LOWERINGOPT(IntegerWrapAround, unsigned, 1, 1)
+
 /// If true, add nsw flags to loop variable increments.
 /// Off by default.
+/// TODO: integrate this option with the above
 ENUM_LOWERINGOPT(NSWOnLoopVarInc, unsigned, 1, 0)
 
 #undef LOWERINGOPT

flang/include/flang/Optimizer/Builder/FIRBuilder.h

@@ -85,13 +85,16 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
   // The listener self-reference has to be updated in case of copy-construction.
   FirOpBuilder(const FirOpBuilder &other)
       : OpBuilder(other), OpBuilder::Listener(), kindMap{other.kindMap},
-        fastMathFlags{other.fastMathFlags}, symbolTable{other.symbolTable} {
+        fastMathFlags{other.fastMathFlags},
+        integerOverflowFlags{other.integerOverflowFlags},
+        symbolTable{other.symbolTable} {
     setListener(this);
   }
 
   FirOpBuilder(FirOpBuilder &&other)
       : OpBuilder(other), OpBuilder::Listener(),
         kindMap{std::move(other.kindMap)}, fastMathFlags{other.fastMathFlags},
+        integerOverflowFlags{other.integerOverflowFlags},
         symbolTable{other.symbolTable} {
     setListener(this);
   }
@@ -521,6 +524,18 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
     return fmfString;
   }
 
+  /// Set default IntegerOverflowFlags value for all operations
+  /// supporting mlir::arith::IntegerOverflowFlagsAttr that will be created
+  /// by this builder.
+  void setIntegerOverflowFlags(mlir::arith::IntegerOverflowFlags flags) {
+    integerOverflowFlags = flags;
+  }
+
+  /// Get current IntegerOverflowFlags value.
+  mlir::arith::IntegerOverflowFlags getIntegerOverflowFlags() const {
+    return integerOverflowFlags;
+  }
+
   /// Dump the current function. (debug)
   LLVM_DUMP_METHOD void dumpFunc();
 
@@ -547,6 +562,10 @@ class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
   /// mlir::arith::FastMathAttr.
   mlir::arith::FastMathFlags fastMathFlags{};
 
+  /// IntegerOverflowFlags that need to be set for operations that support
+  /// mlir::arith::IntegerOverflowFlagsAttr.
+  mlir::arith::IntegerOverflowFlags integerOverflowFlags{};
+
   /// fir::GlobalOp and func::FuncOp symbol table to speed-up
   /// lookups.
   mlir::SymbolTable *symbolTable = nullptr;

flang/lib/Lower/ConvertCall.cpp

@@ -2570,9 +2570,26 @@ genIntrinsicRef(const Fortran::evaluate::SpecificIntrinsic *intrinsic,
           hlfir::Entity{*var}, /*isPresent=*/std::nullopt});
       continue;
     }
+    // arguments of bitwise comparison functions may not have nsw flag
+    // even if -fno-wrapv is enabled
+    mlir::arith::IntegerOverflowFlags iofBackup{};
+    auto isBitwiseComparison = [](const std::string intrinsicName) -> bool {
+      if (intrinsicName == "bge" || intrinsicName == "bgt" ||
+          intrinsicName == "ble" || intrinsicName == "blt")
+        return true;
+      return false;
+    };
+    if (isBitwiseComparison(callContext.getProcedureName())) {
+      iofBackup = callContext.getBuilder().getIntegerOverflowFlags();
+      callContext.getBuilder().setIntegerOverflowFlags(
+          mlir::arith::IntegerOverflowFlags::none);
+    }
     auto loweredActual = Fortran::lower::convertExprToHLFIR(
         loc, callContext.converter, *expr, callContext.symMap,
         callContext.stmtCtx);
+    if (isBitwiseComparison(callContext.getProcedureName()))
+      callContext.getBuilder().setIntegerOverflowFlags(iofBackup);
+
     std::optional<mlir::Value> isPresent;
     if (argLowering) {
       fir::ArgLoweringRule argRules =

flang/lib/Lower/ConvertExprToHLFIR.cpp

@@ -1584,9 +1584,14 @@ class HlfirBuilder {
       auto rightVal = hlfir::loadTrivialScalar(l, b, rightElement);
       return binaryOp.gen(l, b, op.derived(), leftVal, rightVal);
     };
+    auto iofBackup = builder.getIntegerOverflowFlags();
+    // nsw is never added to operations on vector subscripts
+    // even if -fno-wrapv is enabled.
+    builder.setIntegerOverflowFlags(mlir::arith::IntegerOverflowFlags::none);
     mlir::Value elemental = hlfir::genElementalOp(loc, builder, elementType,
                                                   shape, typeParams, genKernel,
                                                   /*isUnordered=*/true);
+    builder.setIntegerOverflowFlags(iofBackup);
     fir::FirOpBuilder *bldr = &builder;
     getStmtCtx().attachCleanup(
         [=]() { bldr->create<hlfir::DestroyOp>(loc, elemental); });
@@ -1899,10 +1904,17 @@ class HlfirBuilder {
 template <typename T>
 hlfir::Entity
 HlfirDesignatorBuilder::genSubscript(const Fortran::evaluate::Expr<T> &expr) {
+  fir::FirOpBuilder &builder = getBuilder();
+  mlir::arith::IntegerOverflowFlags iofBackup{};
+  if (!getConverter().getLoweringOptions().getIntegerWrapAround()) {
+    iofBackup = builder.getIntegerOverflowFlags();
+    builder.setIntegerOverflowFlags(mlir::arith::IntegerOverflowFlags::nsw);
+  }
   auto loweredExpr =
       HlfirBuilder(getLoc(), getConverter(), getSymMap(), getStmtCtx())
           .gen(expr);
-  fir::FirOpBuilder &builder = getBuilder();
+  if (!getConverter().getLoweringOptions().getIntegerWrapAround())
+    builder.setIntegerOverflowFlags(iofBackup);
   // Skip constant conversions that litters designators and makes generated
   // IR harder to read: directly use index constants for constant subscripts.
   mlir::Type idxTy = builder.getIndexType();

flang/lib/Optimizer/Builder/FIRBuilder.cpp

@@ -768,14 +768,23 @@ mlir::Value fir::FirOpBuilder::genAbsentOp(mlir::Location loc,
 
 void fir::FirOpBuilder::setCommonAttributes(mlir::Operation *op) const {
   auto fmi = mlir::dyn_cast<mlir::arith::ArithFastMathInterface>(*op);
-  if (!fmi)
-    return;
-  // TODO: use fmi.setFastMathFlagsAttr() after D137114 is merged.
-  //       For now set the attribute by the name.
-  llvm::StringRef arithFMFAttrName = fmi.getFastMathAttrName();
-  if (fastMathFlags != mlir::arith::FastMathFlags::none)
-    op->setAttr(arithFMFAttrName, mlir::arith::FastMathFlagsAttr::get(
-                                      op->getContext(), fastMathFlags));
+  if (fmi) {
+    // TODO: use fmi.setFastMathFlagsAttr() after D137114 is merged.
+    //       For now set the attribute by the name.
+    llvm::StringRef arithFMFAttrName = fmi.getFastMathAttrName();
+    if (fastMathFlags != mlir::arith::FastMathFlags::none)
+      op->setAttr(arithFMFAttrName, mlir::arith::FastMathFlagsAttr::get(
+                                        op->getContext(), fastMathFlags));
+  }
+  auto iofi =
+      mlir::dyn_cast<mlir::arith::ArithIntegerOverflowFlagsInterface>(*op);
+  if (iofi) {
+    llvm::StringRef arithIOFAttrName = iofi.getIntegerOverflowAttrName();
+    if (integerOverflowFlags != mlir::arith::IntegerOverflowFlags::none)
+      op->setAttr(arithIOFAttrName,
+                  mlir::arith::IntegerOverflowFlagsAttr::get(
+                      op->getContext(), integerOverflowFlags));
+  }
 }
 
 void fir::FirOpBuilder::setFastMathFlags(

flang/test/Lower/nsw.f90

@@ -0,0 +1,61 @@
+! RUN: bbc -emit-fir %s -o - | FileCheck %s
+! RUN: bbc -emit-fir -fwrapv %s -o - | FileCheck %s --check-prefix=NO-NSW
+
+! NO-NSW-NOT: overflow<nsw>
+
+subroutine subscript(a, i, j, k)
+  integer :: a(:,:,:), i, j, k
+  a(i+1, j-2, k*3) = 5
+end subroutine
+! CHECK-LABEL:   func.func @_QPsubscript(
+! CHECK:           %[[VAL_4:.*]] = arith.constant 3 : i32
+! CHECK:           %[[VAL_5:.*]] = arith.constant 2 : i32
+! CHECK:           %[[VAL_6:.*]] = arith.constant 1 : i32
+! CHECK:           %[[VAL_9:.*]] = fir.declare %{{.*}}a"} : (!fir.box<!fir.array<?x?x?xi32>>, !fir.dscope) -> !fir.box<!fir.array<?x?x?xi32>>
+! CHECK:           %[[VAL_10:.*]] = fir.rebox %[[VAL_9]] : (!fir.box<!fir.array<?x?x?xi32>>) -> !fir.box<!fir.array<?x?x?xi32>>
+! CHECK:           %[[VAL_11:.*]] = fir.declare %{{.*}}i"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
+! CHECK:           %[[VAL_12:.*]] = fir.declare %{{.*}}j"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
+! CHECK:           %[[VAL_13:.*]] = fir.declare %{{.*}}k"} : (!fir.ref<i32>, !fir.dscope) -> !fir.ref<i32>
+! CHECK:           %[[VAL_14:.*]] = fir.load %[[VAL_11]] : !fir.ref<i32>
+! CHECK:           %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_6]] overflow<nsw> : i32
+! CHECK:           %[[VAL_16:.*]] = fir.convert %[[VAL_15]] : (i32) -> i64
+! CHECK:           %[[VAL_17:.*]] = fir.load %[[VAL_12]] : !fir.ref<i32>
+! CHECK:           %[[VAL_18:.*]] = arith.subi %[[VAL_17]], %[[VAL_5]] overflow<nsw> : i32
+! CHECK:           %[[VAL_19:.*]] = fir.convert %[[VAL_18]] : (i32) -> i64
+! CHECK:           %[[VAL_20:.*]] = fir.load %[[VAL_13]] : !fir.ref<i32>
+! CHECK:           %[[VAL_21:.*]] = arith.muli %[[VAL_20]], %[[VAL_4]] overflow<nsw> : i32
+! CHECK:           %[[VAL_22:.*]] = fir.convert %[[VAL_21]] : (i32) -> i64
+! CHECK:           %[[VAL_23:.*]] = fir.array_coor %[[VAL_10]] %[[VAL_16]], %[[VAL_19]], %[[VAL_22]] :
+
+! Test that nsw is never added to arith ops
+! on vector subscripts.
+subroutine vector_subscript_as_value(x, y, z)
+  integer :: x(100)
+  integer(8) :: y(20), z(20)
+  call bar(x(y+z))
+end subroutine
+! CHECK-LABEL:   func.func @_QPvector_subscript_as_value(
+! CHECK-NOT: overflow<nsw>
+! CHECK:           return
+
+subroutine vector_subscript_lhs(x, vector1, vector2)
+  integer(8) :: vector1(10), vector2(10)
+  real :: x(:)
+  x(vector1+vector2) = 42.
+end subroutine
+! CHECK-LABEL:   func.func @_QPvector_subscript_lhs(
+! CHECK-NOT: overflow<nsw>
+! CHECK:           return
+
+! Test that nsw is never added to arith ops
+! on arguments of bitwise comparison intrinsics.
+subroutine bitwise_comparison(a, b)
+  integer :: a, b
+  print *, bge(a+b, a-b)
+  print *, bgt(a+b, a-b)
+  print *, ble(a+b, a-b)
+  print *, blt(a+b, a-b)
+end subroutine
+! CHECK-LABEL:   func.func @_QPbitwise_comparison(
+! CHECK-NOT: overflow<nsw>
+! CHECK:           return

flang/tools/bbc/bbc.cpp

@@ -228,6 +228,12 @@ static llvm::cl::opt<std::string>
                          llvm::cl::desc("Override host target triple"),
                          llvm::cl::init(""));
 
+static llvm::cl::opt<bool> integerWrapAround(
+    "fwrapv",
+    llvm::cl::desc("Treat signed integer overflow as two's complement"),
+    llvm::cl::init(false));
+
+// TODO: integrate this option with the above
 static llvm::cl::opt<bool>
     setNSW("integer-overflow",
            llvm::cl::desc("add nsw flag to internal operations"),
@@ -373,6 +379,7 @@ static llvm::LogicalResult convertFortranSourceToMLIR(
   Fortran::lower::LoweringOptions loweringOptions{};
   loweringOptions.setNoPPCNativeVecElemOrder(enableNoPPCNativeVecElemOrder);
   loweringOptions.setLowerToHighLevelFIR(useHLFIR || emitHLFIR);
+  loweringOptions.setIntegerWrapAround(integerWrapAround);
   loweringOptions.setNSWOnLoopVarInc(setNSW);
   std::vector<Fortran::lower::EnvironmentDefault> envDefaults = {};
   constexpr const char *tuneCPU = "";

flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp

@@ -585,3 +585,62 @@ TEST_F(FIRBuilderTest, genArithFastMath) {
   auto op4_fmf = op4_fmi.getFastMathFlagsAttr().getValue();
   EXPECT_EQ(op4_fmf, FMF1);
 }
+
+TEST_F(FIRBuilderTest, genArithIntegerOverflow) {
+  auto builder = getBuilder();
+  auto ctx = builder.getContext();
+  auto loc = builder.getUnknownLoc();
+
+  auto intTy = IntegerType::get(ctx, 32);
+  auto arg = builder.create<fir::UndefOp>(loc, intTy);
+
+  // Test that IntegerOverflowFlags is 'none' by default.
+  mlir::Operation *op1 = builder.create<mlir::arith::AddIOp>(loc, arg, arg);
+  auto op1_iofi =
+      mlir::dyn_cast_or_null<mlir::arith::ArithIntegerOverflowFlagsInterface>(
+          op1);
+  EXPECT_TRUE(op1_iofi);
+  auto op1_ioff = op1_iofi.getOverflowAttr().getValue();
+  EXPECT_EQ(op1_ioff, arith::IntegerOverflowFlags::none);
+
+  // Test that the builder is copied properly.
+  fir::FirOpBuilder builder_copy(builder);
+
+  arith::IntegerOverflowFlags nsw = arith::IntegerOverflowFlags::nsw;
+  builder.setIntegerOverflowFlags(nsw);
+  arith::IntegerOverflowFlags nuw = arith::IntegerOverflowFlags::nuw;
+  builder_copy.setIntegerOverflowFlags(nuw);
+
+  // Modifying IntegerOverflowFlags for the copy must not affect the original
+  // builder.
+  mlir::Operation *op2 = builder.create<mlir::arith::AddIOp>(loc, arg, arg);
+  auto op2_iofi =
+      mlir::dyn_cast_or_null<mlir::arith::ArithIntegerOverflowFlagsInterface>(
+          op2);
+  EXPECT_TRUE(op2_iofi);
+  auto op2_ioff = op2_iofi.getOverflowAttr().getValue();
+  EXPECT_EQ(op2_ioff, nsw);
+
+  // Modifying IntegerOverflowFlags for the original builder must not affect the
+  // copy.
+  mlir::Operation *op3 =
+      builder_copy.create<mlir::arith::AddIOp>(loc, arg, arg);
+  auto op3_iofi =
+      mlir::dyn_cast_or_null<mlir::arith::ArithIntegerOverflowFlagsInterface>(
+          op3);
+  EXPECT_TRUE(op3_iofi);
+  auto op3_ioff = op3_iofi.getOverflowAttr().getValue();
+  EXPECT_EQ(op3_ioff, nuw);
+
+  // Test that the builder copy inherits IntegerOverflowFlags from the original.
+  fir::FirOpBuilder builder_copy2(builder);
+
+  mlir::Operation *op4 =
+      builder_copy2.create<mlir::arith::AddIOp>(loc, arg, arg);
+  auto op4_iofi =
+      mlir::dyn_cast_or_null<mlir::arith::ArithIntegerOverflowFlagsInterface>(
+          op4);
+  EXPECT_TRUE(op4_iofi);
+  auto op4_ioff = op4_iofi.getOverflowAttr().getValue();
+  EXPECT_EQ(op4_ioff, nsw);
+}