[flang] Handle special case for SHIFTA intrinsic

[clementval]

Summary

This patch update the lowering of the shifta intrinsic to match
the behvior of gfortran. When the SHIFT value is equal to the
integer bitwidth then we handle it differently.
This is due to the operation used in lowering (mlir::arith::ShRSIOp)
that lowers to ashr.

Before this patch we have the following results:

SHIFTA(  -1, 8) =  0
SHIFTA(  -2, 8) =  0
SHIFTA( -30, 8) =  0
SHIFTA( -31, 8) =  0
SHIFTA( -32, 8) =  0
SHIFTA( -33, 8) =  0
SHIFTA(-126, 8) =  0
SHIFTA(-127, 8) =  0
SHIFTA(-128, 8) =  0

While gfortran is giving this:

SHIFTA(  -1, 8) = -1
SHIFTA(  -2, 8) = -1
SHIFTA( -30, 8) = -1
SHIFTA( -31, 8) = -1
SHIFTA( -32, 8) = -1
SHIFTA( -33, 8) = -1
SHIFTA(-126, 8) = -1
SHIFTA(-127, 8) = -1
SHIFTA(-128, 8) = -1

With this patch flang and gfortran have the same behavior.

[jeanPerier]

flang/lib/Lower/IntrinsicCall.cpp

4075

Do you know if using a fir.if vs a SelectOp impacts the generated code from a performance point of view ? I tend to avoid branches in intrinsic code, but I do not have any data to prove this is a better choice either, and probably easy to later update if this turns out to be sub-optimal.

[clementval]

Use select

[clementval]

flang/lib/Lower/IntrinsicCall.cpp

4075

I think it is better to use SelectOp here. I had a bit more complicated logic at first but with the current test a select is fine. I just updated the patch with this change.

flang/lib/Lower/IntrinsicCall.cpp

struct IntrinsicLibrary {
   mlir::Value genScale(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genScan(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genSelectedIntKind(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genSelectedRealKind(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genSetExponent(mlir::Type resultType,
                              llvm::ArrayRef<mlir::Value> args);
   template <typename Shift>
   mlir::Value genShift(mlir::Type resultType, llvm::ArrayRef<mlir::Value>);
+  mlir::Value genShiftA(mlir::Type resultType, llvm::ArrayRef<mlir::Value>);
   mlir::Value genSign(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genSize(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genSpacing(mlir::Type resultType,
                          llvm::ArrayRef<mlir::Value> args);
   fir::ExtendedValue genSpread(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genSum(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   void genSystemClock(llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genTrailz(mlir::Type, llvm::ArrayRef<mlir::Value>);
static constexpr IntrinsicHandler handlers[]{
      /*isElemental=*/false},
     {"selected_real_kind",
      &I::genSelectedRealKind,
      {{{"precision", asAddr, handleDynamicOptional},
        {"range", asAddr, handleDynamicOptional},
        {"radix", asAddr, handleDynamicOptional}}},
      /*isElemental=*/false},
     {"set_exponent", &I::genSetExponent},
-    {"shifta", &I::genShift<mlir::arith::ShRSIOp>},
+    {"shifta", &I::genShiftA},
     {"shiftl", &I::genShift<mlir::arith::ShLIOp>},
     {"shiftr", &I::genShift<mlir::arith::ShRUIOp>},
     {"sign", &I::genSign},
     {"size",
      &I::genSize,
      {{{"array", asBox},
        {"dim", asAddr, handleDynamicOptional},
        {"kind", asValue}}},
mlir::Value IntrinsicLibrary::genSetExponent(mlir::Type resultType,
   assert(args.size() == 2);
 
   return builder.createConvert(
       loc, resultType,
       fir::runtime::genSetExponent(builder, loc, fir::getBase(args[0]),
                                    fir::getBase(args[1])));
 }
 
-// SHIFTA, SHIFTL, SHIFTR
+// SHIFTL, SHIFTR
 template <typename Shift>
 mlir::Value IntrinsicLibrary::genShift(mlir::Type resultType,
                                        llvm::ArrayRef<mlir::Value> args) {
   assert(args.size() == 2);
 
   // If SHIFT < 0 or SHIFT >= BIT_SIZE(I), return 0. This is not required by
   // the standard. However, several other compilers behave this way, so try and
   // maintain compatibility with them to an extent.
mlir::Value tooLarge = builder.create<mlir::arith::CmpIOp>(
       loc, mlir::arith::CmpIPredicate::sge, shift, bitSize);
   mlir::Value outOfBounds =
       builder.create<mlir::arith::OrIOp>(loc, tooSmall, tooLarge);
 
   mlir::Value shifted = builder.create<Shift>(loc, args[0], shift);
   return builder.create<mlir::arith::SelectOp>(loc, outOfBounds, zero, shifted);
 }
 
+// SHIFTA
+mlir::Value IntrinsicLibrary::genShiftA(mlir::Type resultType,
+                                        llvm::ArrayRef<mlir::Value> args) {
+  unsigned bits = resultType.getIntOrFloatBitWidth();
+  mlir::Value bitSize = builder.createIntegerConstant(loc, resultType, bits);
+  mlir::Value shift = builder.createConvert(loc, resultType, args[1]);
+  mlir::Value shiftEqBitSize = builder.create<mlir::arith::CmpIOp>(
+      loc, mlir::arith::CmpIPredicate::eq, shift, bitSize);
+
+  // Lowering of mlir::arith::ShRSIOp is using `ashr`. `ashr` is undefined when
+  // the shift amount is equal to the element size.
+  // So if SHIFT is equal to the bit width then it is handled as a special case.
+  return builder
+      .genIfOp(loc, {resultType}, shiftEqBitSize,
+               /*withElseRegion=*/true)
+      .genThen([&]() {
+        mlir::Value zero = builder.createIntegerConstant(loc, resultType, 0);
+        mlir::Value minusOne =
+            builder.createIntegerConstant(loc, resultType, -1);
+        mlir::Value valueIsNeg = builder.create<mlir::arith::CmpIOp>(
+            loc, mlir::arith::CmpIPredicate::slt, args[0], zero);
+        mlir::Value res = builder.create<mlir::arith::SelectOp>(loc, valueIsNeg,
+                                                                minusOne, zero);
+        builder.create<fir::ResultOp>(loc, res);
+      })
+      .genElse([&] {
+        mlir::Value shifted =
+            builder.create<mlir::arith::ShRSIOp>(loc, args[0], shift);
+        builder.create<fir::ResultOp>(loc, shifted);
+      })
+      .getResults()[0];
+}
+
 // SIGN
 mlir::Value IntrinsicLibrary::genSign(mlir::Type resultType,
                                       llvm::ArrayRef<mlir::Value> args) {
   assert(args.size() == 2);
   if (resultType.isa<mlir::IntegerType>()) {
     mlir::Value abs = genAbs(resultType, {args[0]});
     mlir::Value zero = builder.createIntegerConstant(loc, resultType, 0);
     auto neg = builder.create<mlir::arith::SubIOp>(loc, zero, abs);

flang/test/Lower/Intrinsics/shifta.f90

 ! RUN: bbc -emit-fir %s -o - | FileCheck %s
 ! RUN: %flang_fc1 -emit-fir %s -o - | FileCheck %s
 
 ! CHECK-LABEL: shifta1_test
 ! CHECK-SAME: %[[A:.*]]: !fir.ref<i8>{{.*}}, %[[B:.*]]: !fir.ref<i32>{{.*}}, %[[C:.*]]: !fir.ref<i8>{{.*}}
 subroutine shifta1_test(a, b, c)
   integer(kind=1) :: a
   integer :: b
   integer(kind=1) :: c
 
   ! CHECK: %[[A_VAL:.*]] = fir.load %[[A]] : !fir.ref<i8>
   ! CHECK: %[[B_VAL:.*]] = fir.load %[[B]] : !fir.ref<i32>
   c = shifta(a, b)
   ! CHECK: %[[C_BITS:.*]] = arith.constant 8 : i8
-  ! CHECK: %[[C_0:.*]] = arith.constant 0 : i8
   ! CHECK: %[[B_CONV:.*]] = fir.convert %[[B_VAL]] : (i32) -> i8
-  ! CHECK: %[[UNDER:.*]] = arith.cmpi slt, %[[B_CONV]], %[[C_0]] : i8
-  ! CHECK: %[[OVER:.*]] = arith.cmpi sge, %[[B_CONV]], %[[C_BITS]] : i8
-  ! CHECK: %[[INVALID:.*]] = arith.ori %[[UNDER]], %[[OVER]] : i1
-  ! CHECK: %[[SHIFT:.*]] = arith.shrsi %[[A_VAL]], %[[B_CONV]] : i8
-  ! CHECK: %[[RES:.*]] = arith.select %[[INVALID]], %[[C_0]], %[[SHIFT]] : i8
+  ! CHECK: %[[SHIFT_IS_BITWIDTH:.*]] = arith.cmpi eq, %[[B_CONV]], %[[C_BITS]] : i8
+  ! CHECK: %[[RESULT:.*]] = fir.if %[[SHIFT_IS_BITWIDTH]] -> (i8) {
+  ! CHECK: %[[C0:.*]] = arith.constant 0 : i8
+  ! CHECK: %[[CM1:.*]] = arith.constant -1 : i8
+  ! CHECK: %[[IS_NEG:.*]] = arith.cmpi slt, %[[A_VAL]], %[[C0]] : i8
+  ! CHECK:   %[[RES:.*]] = arith.select %[[IS_NEG]], %[[CM1]], %[[C0]] : i8
+  ! CHECK:   fir.result %[[RES]] : i8
+  ! CHECK: } else {
+  ! CHECK:   %[[SHIFTED:.*]] = arith.shrsi %[[A_VAL]], %[[B_CONV]] : i8
+  ! CHECK:   fir.result %[[SHIFTED]] : i8
+  ! CHECK: }
 end subroutine shifta1_test
 
 ! CHECK-LABEL: shifta2_test
 ! CHECK-SAME: %[[A:.*]]: !fir.ref<i16>{{.*}}, %[[B:.*]]: !fir.ref<i32>{{.*}}, %[[C:.*]]: !fir.ref<i16>{{.*}}
 subroutine shifta2_test(a, b, c)
   integer(kind=2) :: a
   integer :: b
   integer(kind=2) :: c
 
   ! CHECK: %[[A_VAL:.*]] = fir.load %[[A]] : !fir.ref<i16>
   ! CHECK: %[[B_VAL:.*]] = fir.load %[[B]] : !fir.ref<i32>
   c = shifta(a, b)
   ! CHECK: %[[C_BITS:.*]] = arith.constant 16 : i16
-  ! CHECK: %[[C_0:.*]] = arith.constant 0 : i16
   ! CHECK: %[[B_CONV:.*]] = fir.convert %[[B_VAL]] : (i32) -> i16
-  ! CHECK: %[[UNDER:.*]] = arith.cmpi slt, %[[B_CONV]], %[[C_0]] : i16
-  ! CHECK: %[[OVER:.*]] = arith.cmpi sge, %[[B_CONV]], %[[C_BITS]] : i16
-  ! CHECK: %[[INVALID:.*]] = arith.ori %[[UNDER]], %[[OVER]] : i1
-  ! CHECK: %[[SHIFT:.*]] = arith.shrsi %[[A_VAL]], %[[B_CONV]] : i16
-  ! CHECK: %[[RES:.*]] = arith.select %[[INVALID]], %[[C_0]], %[[SHIFT]] : i16
+  ! CHECK: %[[SHIFT_IS_BITWIDTH:.*]] = arith.cmpi eq, %[[B_CONV]], %[[C_BITS]] : i16
+  ! CHECK: %[[RESULT:.*]] = fir.if %[[SHIFT_IS_BITWIDTH]] -> (i16) {
+  ! CHECK: %[[C0:.*]] = arith.constant 0 : i16
+  ! CHECK: %[[CM1:.*]] = arith.constant -1 : i16
+  ! CHECK: %[[IS_NEG:.*]] = arith.cmpi slt, %[[A_VAL]], %[[C0]] : i16
+  ! CHECK:   %[[RES:.*]] = arith.select %[[IS_NEG]], %[[CM1]], %[[C0]] : i16
+  ! CHECK:   fir.result %[[RES]] : i16
+  ! CHECK: } else {
+  ! CHECK:   %[[SHIFTED:.*]] = arith.shrsi %[[A_VAL]], %[[B_CONV]] : i16
+  ! CHECK:   fir.result %[[SHIFTED]] : i16
+  ! CHECK: }
 end subroutine shifta2_test
 
 ! CHECK-LABEL: shifta4_test
 ! CHECK-SAME: %[[A:.*]]: !fir.ref<i32>{{.*}}, %[[B:.*]]: !fir.ref<i32>{{.*}}, %[[C:.*]]: !fir.ref<i32>{{.*}}
 subroutine shifta4_test(a, b, c)
   integer(kind=4) :: a
   integer :: b
   integer(kind=4) :: c
 
   ! CHECK: %[[A_VAL:.*]] = fir.load %[[A]] : !fir.ref<i32>
   ! CHECK: %[[B_VAL:.*]] = fir.load %[[B]] : !fir.ref<i32>
   c = shifta(a, b)
   ! CHECK: %[[C_BITS:.*]] = arith.constant 32 : i32
-  ! CHECK: %[[C_0:.*]] = arith.constant 0 : i32
-  ! CHECK: %[[UNDER:.*]] = arith.cmpi slt, %[[B_VAL]], %[[C_0]] : i32
-  ! CHECK: %[[OVER:.*]] = arith.cmpi sge, %[[B_VAL]], %[[C_BITS]] : i32
-  ! CHECK: %[[INVALID:.*]] = arith.ori %[[UNDER]], %[[OVER]] : i1
-  ! CHECK: %[[SHIFT:.*]] = arith.shrsi %[[A_VAL]], %[[B_VAL]] : i32
-  ! CHECK: %[[RES:.*]] = arith.select %[[INVALID]], %[[C_0]], %[[SHIFT]] : i32
+  ! CHECK: %[[SHIFT_IS_BITWIDTH:.*]] = arith.cmpi eq, %[[B_VAL]], %[[C_BITS]] : i32
+  ! CHECK: %[[RESULT:.*]] = fir.if %[[SHIFT_IS_BITWIDTH]] -> (i32) {
+  ! CHECK: %[[C0:.*]] = arith.constant 0 : i32
+  ! CHECK: %[[CM1:.*]] = arith.constant -1 : i32
+  ! CHECK: %[[IS_NEG:.*]] = arith.cmpi slt, %[[A_VAL]], %[[C0]] : i32
+  ! CHECK:   %[[RES:.*]] = arith.select %[[IS_NEG]], %[[CM1]], %[[C0]] : i32
+  ! CHECK:   fir.result %[[RES]] : i32
+  ! CHECK: } else {
+  ! CHECK:   %[[SHIFTED:.*]] = arith.shrsi %[[A_VAL]], %[[B_VAL]] : i32
+  ! CHECK:   fir.result %[[SHIFTED]] : i32
+  ! CHECK: }
 end subroutine shifta4_test
 
 ! CHECK-LABEL: shifta8_test
 ! CHECK-SAME: %[[A:.*]]: !fir.ref<i64>{{.*}}, %[[B:.*]]: !fir.ref<i32>{{.*}}, %[[C:.*]]: !fir.ref<i64>{{.*}}
 subroutine shifta8_test(a, b, c)
   integer(kind=8) :: a
   integer :: b
   integer(kind=8) :: c
 
   ! CHECK: %[[A_VAL:.*]] = fir.load %[[A]] : !fir.ref<i64>
   ! CHECK: %[[B_VAL:.*]] = fir.load %[[B]] : !fir.ref<i32>
   c = shifta(a, b)
   ! CHECK: %[[C_BITS:.*]] = arith.constant 64 : i64
-  ! CHECK: %[[C_0:.*]] = arith.constant 0 : i64
   ! CHECK: %[[B_CONV:.*]] = fir.convert %[[B_VAL]] : (i32) -> i64
-  ! CHECK: %[[UNDER:.*]] = arith.cmpi slt, %[[B_CONV]], %[[C_0]] : i64
-  ! CHECK: %[[OVER:.*]] = arith.cmpi sge, %[[B_CONV]], %[[C_BITS]] : i64
-  ! CHECK: %[[INVALID:.*]] = arith.ori %[[UNDER]], %[[OVER]] : i1
-  ! CHECK: %[[SHIFT:.*]] = arith.shrsi %[[A_VAL]], %[[B_CONV]] : i64
-  ! CHECK: %[[RES:.*]] = arith.select %[[INVALID]], %[[C_0]], %[[SHIFT]] : i64
+  ! CHECK: %[[SHIFT_IS_BITWIDTH:.*]] = arith.cmpi eq, %[[B_CONV]], %[[C_BITS]] : i64
+  ! CHECK: %[[RESULT:.*]] = fir.if %[[SHIFT_IS_BITWIDTH]] -> (i64) {
+  ! CHECK: %[[C0:.*]] = arith.constant 0 : i64
+  ! CHECK: %[[CM1:.*]] = arith.constant -1 : i64
+  ! CHECK: %[[IS_NEG:.*]] = arith.cmpi slt, %[[A_VAL]], %[[C0]] : i64
+  ! CHECK:   %[[RES:.*]] = arith.select %[[IS_NEG]], %[[CM1]], %[[C0]] : i64
+  ! CHECK:   fir.result %[[RES]] : i64
+  ! CHECK: } else {
+  ! CHECK:   %[[SHIFTED:.*]] = arith.shrsi %[[A_VAL]], %[[B_CONV]] : i64
+  ! CHECK:   fir.result %[[SHIFTED]] : i64
+  ! CHECK: }
 end subroutine shifta8_test
 
 ! CHECK-LABEL: shifta16_test
 ! CHECK-SAME: %[[A:.*]]: !fir.ref<i128>{{.*}}, %[[B:.*]]: !fir.ref<i32>{{.*}}, %[[C:.*]]: !fir.ref<i128>{{.*}}
 subroutine shifta16_test(a, b, c)
   integer(kind=16) :: a
   integer :: b
   integer(kind=16) :: c
 
   ! CHECK: %[[A_VAL:.*]] = fir.load %[[A]] : !fir.ref<i128>
   ! CHECK: %[[B_VAL:.*]] = fir.load %[[B]] : !fir.ref<i32>
   c = shifta(a, b)
   ! CHECK: %[[C_BITS:.*]] = arith.constant 128 : i128
-  ! CHECK: %[[C_0:.*]] = arith.constant 0 : i128
   ! CHECK: %[[B_CONV:.*]] = fir.convert %[[B_VAL]] : (i32) -> i128
-  ! CHECK: %[[UNDER:.*]] = arith.cmpi slt, %[[B_CONV]], %[[C_0]] : i128
-  ! CHECK: %[[OVER:.*]] = arith.cmpi sge, %[[B_CONV]], %[[C_BITS]] : i128
-  ! CHECK: %[[INVALID:.*]] = arith.ori %[[UNDER]], %[[OVER]] : i1
-  ! CHECK: %[[SHIFT:.*]] = arith.shrsi %[[A_VAL]], %[[B_CONV]] : i128
-  ! CHECK: %[[RES:.*]] = arith.select %[[INVALID]], %[[C_0]], %[[SHIFT]] : i128
+  ! CHECK: %[[SHIFT_IS_BITWIDTH:.*]] = arith.cmpi eq, %[[B_CONV]], %[[C_BITS]] : i128
+  ! CHECK: %[[RESULT:.*]] = fir.if %[[SHIFT_IS_BITWIDTH]] -> (i128) {
+  ! CHECK: %[[C0:.*]] = arith.constant 0 : i128
+  ! CHECK: %[[CM1:.*]] = arith.constant {{.*}} : i128
+  ! CHECK: %[[IS_NEG:.*]] = arith.cmpi slt, %[[A_VAL]], %[[C0]] : i128
+  ! CHECK:   %[[RES:.*]] = arith.select %[[IS_NEG]], %[[CM1]], %[[C0]] : i128
+  ! CHECK:   fir.result %[[RES]] : i128
+  ! CHECK: } else {
+  ! CHECK:   %[[SHIFTED:.*]] = arith.shrsi %[[A_VAL]], %[[B_CONV]] : i128
+  ! CHECK:   fir.result %[[SHIFTED]] : i128
+  ! CHECK: }
 end subroutine shifta16_test