Use i64 type for the selector and select values.

Author: vzakhari

flang/lib/Optimizer/CodeGen/CodeGen.cpp

@@ -3573,7 +3573,15 @@ struct SelectOpConversionBase : public fir::FIROpConversion<OP> {
     llvm::SmallVector<mlir::ValueRange> destinationsOperands;
     mlir::Block *defaultDestination;
     mlir::ValueRange defaultOperands;
-    llvm::SmallVector<int32_t> caseValues;
+    // LLVM::SwitchOp selector type and the case values types
+    // must have the same bit width, so cast the selector to i64,
+    // and use i64 for the case values. It is hard to imagine
+    // a computed GO TO with the number of labels in the label-list
+    // bigger than INT_MAX, but let's use i64 to be on the safe side.
+    // Moreover, fir.select operation is more relaxed than
+    // a Fortran computed GO TO, so it may specify such a case value
+    // even if there is just a single label/case.
+    llvm::SmallVector<int64_t> caseValues;
 
     for (unsigned t = 0; t != conds; ++t) {
       mlir::Block *dest = select.getSuccessor(t);
@@ -3600,16 +3608,14 @@ struct SelectOpConversionBase : public fir::FIROpConversion<OP> {
       defaultDestination = *convertedBlock;
     }
 
-    // LLVM::SwitchOp takes a i32 type for the selector.
-    if (select.getSelector().getType() != rewriter.getI32Type())
-      selector =
-          this->integerCast(loc, rewriter, rewriter.getI32Type(), selector);
+    selector =
+        this->integerCast(loc, rewriter, rewriter.getI64Type(), selector);
 
     rewriter.replaceOpWithNewOp<mlir::LLVM::SwitchOp>(
         select, selector,
         /*defaultDestination=*/defaultDestination,
         /*defaultOperands=*/defaultOperands,
-        /*caseValues=*/caseValues,
+        /*caseValues=*/rewriter.getI64VectorAttr(caseValues),
         /*caseDestinations=*/destinations,
         /*caseOperands=*/destinationsOperands,
         /*branchWeights=*/llvm::ArrayRef<std::int32_t>());

flang/test/Fir/convert-to-llvm.fir

@@ -338,8 +338,7 @@ func.func @select(%arg : index, %arg2 : i32) -> i32 {
 // CHECK:         %[[C0:.*]] = llvm.mlir.constant(1 : i32) : i32
 // CHECK:         %[[C1:.*]] = llvm.mlir.constant(2 : i32) : i32
 // CHECK:         %[[C2:.*]] = llvm.mlir.constant(3 : i32) : i32
-// CHECK:         %[[SELECTOR:.*]] = llvm.trunc %[[SELECTVALUE]] : i{{.*}} to i32
-// CHECK:         llvm.switch %[[SELECTOR]] : i32, ^bb5 [
+// CHECK:         llvm.switch %[[SELECTVALUE]] : i64, ^bb5 [
 // CHECK:           1: ^bb1(%[[C0]] : i32),
 // CHECK:           2: ^bb2(%[[C2]], %[[SELECTVALUE]], %[[ARG1]] : i32, [[IDX]], i32),
 // CHECK:           3: ^bb3(%[[ARG1]], %[[C2]] : i32, i32),
@@ -384,7 +383,8 @@ func.func @select_rank(%arg : i32, %arg2 : i32) -> i32 {
 // CHECK:         %[[C0:.*]] = llvm.mlir.constant(1 : i32) : i32
 // CHECK:         %[[C1:.*]] = llvm.mlir.constant(2 : i32) : i32
 // CHECK:         %[[C2:.*]] = llvm.mlir.constant(3 : i32) : i32
-// CHECK:         llvm.switch %[[SELECTVALUE]] : i32, ^bb5 [
+// CHECK:         %[[SELECTOR:.*]] = llvm.sext %[[SELECTVALUE]] : i{{.*}} to i64
+// CHECK:         llvm.switch %[[SELECTOR]] : i64, ^bb5 [
 // CHECK:           1: ^bb1(%[[C0]] : i32),
 // CHECK:           2: ^bb2(%[[C2]], %[[SELECTVALUE]], %[[ARG1]] : i32, i32, i32),
 // CHECK:           3: ^bb3(%[[ARG1]], %[[C2]] : i32, i32),
@@ -2891,32 +2891,36 @@ func.func @select_with_cast(%arg1 : i8, %arg2 : i16, %arg3: i64, %arg4: index) -
   ^bb3:
   fir.select %arg4 : index [ 1, ^bb4, unit, ^bb4 ]
   ^bb4:
+  fir.select %arg3 : i64 [ 4294967296, ^bb5, unit, ^bb5 ]
+  ^bb5:
   return
 }
 // CHECK-LABEL:   llvm.func @select_with_cast(
 // CHECK-SAME:      %[[ARG0:.*]]: i8,
 // CHECK-SAME:      %[[ARG1:.*]]: i16,
 // CHECK-SAME:      %[[ARG2:.*]]: i64,
 // CHECK-SAME:      %[[ARG3:.*]]: i64) {
-// CHECK:           %[[VAL_0:.*]] = llvm.sext %[[ARG0]] : i8 to i32
-// CHECK:           llvm.switch %[[VAL_0]] : i32, ^bb1 [
+// CHECK:           %[[VAL_0:.*]] = llvm.sext %[[ARG0]] : i8 to i64
+// CHECK:           llvm.switch %[[VAL_0]] : i64, ^bb1 [
 // CHECK:             1: ^bb1
 // CHECK:           ]
 // CHECK:         ^bb1:
-// CHECK:           %[[VAL_1:.*]] = llvm.sext %[[ARG1]] : i16 to i32
-// CHECK:           llvm.switch %[[VAL_1]] : i32, ^bb2 [
+// CHECK:           %[[VAL_1:.*]] = llvm.sext %[[ARG1]] : i16 to i64
+// CHECK:           llvm.switch %[[VAL_1]] : i64, ^bb2 [
 // CHECK:             1: ^bb2
 // CHECK:           ]
 // CHECK:         ^bb2:
-// CHECK:           %[[VAL_2:.*]] = llvm.trunc %[[ARG2]] : i64 to i32
-// CHECK:           llvm.switch %[[VAL_2]] : i32, ^bb3 [
+// CHECK:           llvm.switch %[[ARG2]] : i64, ^bb3 [
 // CHECK:             1: ^bb3
 // CHECK:           ]
 // CHECK:         ^bb3:
-// CHECK:           %[[VAL_3:.*]] = llvm.trunc %[[ARG3]] : i64 to i32
-// CHECK:           llvm.switch %[[VAL_3]] : i32, ^bb4 [
+// CHECK:           llvm.switch %[[ARG3]] : i64, ^bb4 [
 // CHECK:             1: ^bb4
 // CHECK:           ]
 // CHECK:         ^bb4:
+// CHECK:           llvm.switch %[[ARG2]] : i64, ^bb5 [
+// CHECK:             4294967296: ^bb5
+// CHECK:           ]
+// CHECK:         ^bb5:
 // CHECK:           llvm.return
 // CHECK:         }

flang/test/Fir/select.fir

@@ -7,8 +7,8 @@
 func.func @f(%a : i32) -> i32 {
    %1 = arith.constant 1 : i32
    %2 = arith.constant 42 : i32
-// CHECK: switch i32 %{{.*}}, label %{{.*}} [
-// CHECK:   i32 1, label %{{.*}}
+// CHECK: switch i64 %{{.*}}, label %{{.*}} [
+// CHECK:   i64 1, label %{{.*}}
 // CHECK: ]
    fir.select %a : i32 [1, ^bb2(%1:i32), unit, ^bb3(%2:i32)]
 ^bb2(%3 : i32) :
@@ -24,9 +24,9 @@ func.func @g(%a : i32) -> i32 {
    %1 = arith.constant 1 : i32
    %2 = arith.constant 42 : i32
 
-// CHECK: switch i32 %{{.*}}, label %{{.*}} [
-// CHECK:  i32 1, label %{{.*}}
-// CHECK:  i32 -1, label %{{.*}}
+// CHECK: switch i64 %{{.*}}, label %{{.*}} [
+// CHECK:  i64 1, label %{{.*}}
+// CHECK:  i64 -1, label %{{.*}}
 // CHECK: ]
    fir.select_rank %a : i32 [1, ^bb2(%1:i32), -1, ^bb4, unit, ^bb3(%2:i32)]
 ^bb2(%3 : i32) :