[fir] Add fir.select and fir.select_rank FIR to LLVM IR conversion patterns

[clementval]

Summary

The fir.select and fir.select_rank are lowered to a if-then-else
ladder.

This patch is part of the upstreaming effort from fir-dev branch.

Co-authored-by: Jean Perier <[e-mail]>
Co-authored-by: Eric Schweitz <[e-mail]>

[clementval]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

187

Just realized the LLVM IR dialect has the SwitchOp now. I'll rewrite the conversion to make use of it.

[clementval]

Convert to llvm.switch instead of if-then-else ladder

[clementval]

Remove unused code after switch to llvm.switch

[mehdi_amini]

flang/include/flang/Optimizer/Dialect/FIROps.td

487

If this is just about having a get prepended, there is an option in ODS to do this universally (since https://reviews.llvm.org/D111033 ).
I don't think it is scalable to just alias manually every possible accessor.

You can just add let emitAccessorPrefix = kEmitAccessorPrefix_Both; in the Dialect ODS definition to get both forms generated.

flang/lib/Optimizer/CodeGen/CodeGen.cpp

156

Can we just use int? (or int64_t)
(The code with the decltype seems just really unusual to me)

158

Seems like a missing accessor on fir_SwitchTerminatorOp in the ODS file?

160

173

Attributes are value types

[clementval]

Address comments

[clementval]

Rebase

flang/include/flang/Optimizer/Dialect/FIROps.td

class fir_SwitchTerminatorOp<string mnemonic, list<OpTrait> traits = []> :
 
   string extraSwitchClassDeclaration = [{
     using Conditions = mlir::Value;
 
     static constexpr llvm::StringRef getCasesAttr() { return "case_tags"; }
 
     // The number of destination conditions that may be tested
     unsigned getNumConditions() {
-      return (*this)->getAttrOfType<mlir::ArrayAttr>(getCasesAttr()).size();
+      return getCases().size();
     }
 
     // The selector is the value being tested to determine the destination
     mlir::Value getSelector() { return selector(); }
     mlir::Value getSelector(llvm::ArrayRef<mlir::Value> operands) {
       return operands[0];
     }
+    mlir::Value getSelector(mlir::ValueRange operands) {
+      return operands.front();
+    }
 
     // The number of blocks that may be branched to
     unsigned getNumDest() { return (*this)->getNumSuccessors(); }
 
     llvm::Optional<mlir::OperandRange> getCompareOperands(unsigned cond);
     llvm::Optional<llvm::ArrayRef<mlir::Value>> getCompareOperands(
         llvm::ArrayRef<mlir::Value> operands, unsigned cond);
 
     llvm::Optional<llvm::ArrayRef<mlir::Value>> getSuccessorOperands(
         llvm::ArrayRef<mlir::Value> operands, unsigned cond);
+    llvm::Optional<mlir::ValueRange> getSuccessorOperands(
+        mlir::ValueRange operands, unsigned cond);
     using BranchOpInterfaceTrait::getSuccessorOperands;
 
     // Helper function to deal with Optional operand forms
     void printSuccessorAtIndex(mlir::OpAsmPrinter &p, unsigned i) {
       auto *succ = getSuccessor(i);
       auto ops = getSuccessorOperands(i);
       if (ops.hasValue())
         p.printSuccessorAndUseList(succ, ops.getValue());
       else
         p.printSuccessor(succ);
     }
 
+    mlir::ArrayAttr getCases() {
+      return (*this)->getAttrOfType<mlir::ArrayAttr>(getCasesAttr());
+    }
+
     unsigned targetOffsetSize();
   }];
 }
 
 class fir_IntegralSwitchTerminatorOp<string mnemonic,
     list<OpTrait> traits = []> : fir_SwitchTerminatorOp<mnemonic, traits> {
 
   let skipDefaultBuilders = 1;

flang/lib/Optimizer/CodeGen/CodeGen.cpp

if (optLinkage.hasValue()) {
         return mlir::LLVM::Linkage::Common;
       if (name == "weak")
         return mlir::LLVM::Linkage::Weak;
     }
     return mlir::LLVM::Linkage::External;
   }
 };
 
+template <typename OP>
+void selectMatchAndRewrite(fir::LLVMTypeConverter &lowering, OP select,
+                           typename OP::Adaptor adaptor,
+                           mlir::ConversionPatternRewriter &rewriter) {
+  unsigned conds = select.getNumConditions();
+  auto cases = select.getCases().getValue();
+  mlir::Value selector = adaptor.selector();
+  auto loc = select.getLoc();
+  assert(conds > 0 && "select must have cases");
+
+  llvm::SmallVector<mlir::Block *> destinations;
+  llvm::SmallVector<mlir::ValueRange> destinationsOperands;
+  mlir::Block *defaultDestination;
+  mlir::ValueRange defaultOperands;
+  llvm::SmallVector<int32_t> caseValues;
+
+  for (unsigned t = 0; t != conds; ++t) {
+    mlir::Block *dest = select.getSuccessor(t);
+    auto destOps = select.getSuccessorOperands(adaptor.getOperands(), t);
+    const mlir::Attribute &attr = cases[t];
+    if (auto intAttr = attr.template dyn_cast<mlir::IntegerAttr>()) {
+      destinations.push_back(dest);
+      destinationsOperands.push_back(destOps.hasValue() ? *destOps
+                                                        : ValueRange());
+      caseValues.push_back(intAttr.getInt());
+      continue;
+    }
+    assert(attr.template dyn_cast_or_null<mlir::UnitAttr>());
+    assert((t + 1 == conds) && "unit must be last");
+    defaultDestination = dest;
+    defaultOperands = destOps.hasValue() ? *destOps : ValueRange();
+  }
+
+  // LLVM::SwitchOp takes a i32 type for the selector.
+  if (select.getSelector().getType() != rewriter.getI32Type())
+    selector =
+        rewriter.create<LLVM::TruncOp>(loc, rewriter.getI32Type(), selector);
+
+  rewriter.replaceOpWithNewOp<mlir::LLVM::SwitchOp>(
+      select, selector,
+      /*defaultDestination=*/defaultDestination,
+      /*defaultOperands=*/defaultOperands,
+      /*caseValues=*/caseValues,
+      /*caseDestinations=*/destinations,
+      /*caseOperands=*/destinationsOperands,
+      /*branchWeights=*/ArrayRef<int32_t>());
+}
+
+/// conversion of fir::SelectOp to an if-then-else ladder
+struct SelectOpConversion : public FIROpConversion<fir::SelectOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::SelectOp op, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    selectMatchAndRewrite<fir::SelectOp>(lowerTy(), op, adaptor, rewriter);
+    return success();
+  }
+};
+
+/// conversion of fir::SelectRankOp to an if-then-else ladder
+struct SelectRankOpConversion : public FIROpConversion<fir::SelectRankOp> {
+  using FIROpConversion::FIROpConversion;
+
+  mlir::LogicalResult
+  matchAndRewrite(fir::SelectRankOp op, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    selectMatchAndRewrite<fir::SelectRankOp>(lowerTy(), op, adaptor, rewriter);
+    return success();
+  }
+};
+
 // convert to LLVM IR dialect `undef`
 struct UndefOpConversion : public FIROpConversion<fir::UndefOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::UndefOp undef, OpAdaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<mlir::LLVM::UndefOp>(

 public:
   mlir::ModuleOp getModule() { return getOperation(); }
 
   void runOnOperation() override final {
     auto *context = getModule().getContext();
     fir::LLVMTypeConverter typeConverter{getModule()};
     auto loc = mlir::UnknownLoc::get(context);
     mlir::OwningRewritePatternList pattern(context);
-    pattern.insert<AddrOfOpConversion, HasValueOpConversion, GlobalOpConversion,
-                   UndefOpConversion>(typeConverter);
+    pattern
+        .insert<AddrOfOpConversion, HasValueOpConversion, GlobalOpConversion,
+                SelectOpConversion, SelectRankOpConversion, UndefOpConversion>(
+            typeConverter);
     mlir::populateStdToLLVMConversionPatterns(typeConverter, pattern);
     mlir::arith::populateArithmeticToLLVMConversionPatterns(typeConverter,
                                                             pattern);
     mlir::ConversionTarget target{*context};
     target.addLegalDialect<mlir::LLVM::LLVMDialect>();
 
     // required NOPs for applying a full conversion
     target.addLegalOp<mlir::ModuleOp>();

flang/lib/Optimizer/Dialect/FIROps.cpp

fir::SelectOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                     unsigned oper) {
   auto a =
       (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
       getOperandSegmentSizeAttr());
   return {getSubOperands(oper, getSubOperands(2, operands, segments), a)};
 }
 
+llvm::Optional<mlir::ValueRange>
+fir::SelectOp::getSuccessorOperands(mlir::ValueRange operands, unsigned oper) {
+  auto a =
+      (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
+  auto segments = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
+      getOperandSegmentSizeAttr());
+  return {getSubOperands(oper, getSubOperands(2, operands, segments), a)};
+}
+
 unsigned fir::SelectOp::targetOffsetSize() {
   return denseElementsSize((*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
       getTargetOffsetAttr()));
 }
 
 //===----------------------------------------------------------------------===//
 // SelectCaseOp
 //===----------------------------------------------------------------------===//
fir::SelectRankOp::getSuccessorOperands(llvm::ArrayRef<mlir::Value> operands,
                                         unsigned oper) {
   auto a =
       (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
   auto segments = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
       getOperandSegmentSizeAttr());
   return {getSubOperands(oper, getSubOperands(2, operands, segments), a)};
 }
 
+llvm::Optional<mlir::ValueRange>
+fir::SelectRankOp::getSuccessorOperands(mlir::ValueRange operands,
+                                        unsigned oper) {
+  auto a =
+      (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(getTargetOffsetAttr());
+  auto segments = (*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
+      getOperandSegmentSizeAttr());
+  return {getSubOperands(oper, getSubOperands(2, operands, segments), a)};
+}
+
 unsigned fir::SelectRankOp::targetOffsetSize() {
   return denseElementsSize((*this)->getAttrOfType<mlir::DenseIntElementsAttr>(
       getTargetOffsetAttr()));
 }
 
 //===----------------------------------------------------------------------===//
 // SelectTypeOp
 //===----------------------------------------------------------------------===//

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

fir.global internal @_QEmultiarray : !fir.array<32x32xi32> {
   %2 = fir.insert_on_range %0, %c0_i32, [0 : index, 31 : index, 0 : index, 31 : index] : (!fir.array<32x32xi32>, i32) -> !fir.array<32x32xi32>
   fir.has_value %2 : !fir.array<32x32xi32>
 }
 
 // CHECK: llvm.mlir.global internal @_QEmultiarray() : !llvm.array<32 x array<32 x i32>> {
 // CHECK:   %[[CST:.*]] = llvm.mlir.constant(dense<1> : vector<32x32xi32>) : !llvm.array<32 x array<32 x i32>>
 // CHECK:   llvm.return %[[CST]] : !llvm.array<32 x array<32 x i32>>
 // CHECK: }
+
+// -----
+
+// Test `fir.select` operation conversion pattern.
+// Check that the if-then-else ladder is correctly constructed and that we
+// branch to the correct block.
+
+func @select(%arg : index, %arg2 : i32) -> i32 {
+  %0 = arith.constant 1 : i32
+  %1 = arith.constant 2 : i32
+  %2 = arith.constant 3 : i32
+  %3 = arith.constant 4 : i32
+  fir.select %arg:index [ 1, ^bb1(%0:i32),
+                          2, ^bb2(%2,%arg,%arg2:i32,index,i32),
+                          3, ^bb3(%arg2,%2:i32,i32),
+                          4, ^bb4(%1:i32),
+                          unit, ^bb5 ]
+  ^bb1(%a : i32) :
+    return %a : i32
+  ^bb2(%b : i32, %b2 : index, %b3:i32) :
+    %castidx = arith.index_cast %b2 : index to i32
+    %4 = arith.addi %b, %castidx : i32
+    %5 = arith.addi %4, %b3 : i32
+    return %5 : i32
+  ^bb3(%c:i32, %c2:i32) :
+    %6 = arith.addi %c, %c2 : i32
+    return %6 : i32
+  ^bb4(%d : i32) :
+    return %d : i32
+  ^bb5 :
+    %zero = arith.constant 0 : i32
+    return %zero : i32
+}
+
+// CHECK-LABEL: func @select(
+// CHECK-SAME:               %[[SELECTVALUE:.*]]: [[IDX:.*]],
+// CHECK-SAME:               %[[ARG1:.*]]: 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]], ^bb5 [
+// CHECK:           1: ^bb1(%[[C0]] : i32),
+// CHECK:           2: ^bb2(%[[C2]], %[[SELECTVALUE]], %[[ARG1]] : i32, [[IDX]], i32),
+// CHECK:           3: ^bb3(%[[ARG1]], %[[C2]] : i32, i32),
+// CHECK:           4: ^bb4(%[[C1]] : i32)
+// CHECK:         ]
+
+// -----
+
+// Test `fir.select_rank` operation conversion pattern.
+// Check that the if-then-else ladder is correctly constructed and that we
+// branch to the correct block.
+
+func @select_rank(%arg : i32, %arg2 : i32) -> i32 {
+  %0 = arith.constant 1 : i32
+  %1 = arith.constant 2 : i32
+  %2 = arith.constant 3 : i32
+  %3 = arith.constant 4 : i32
+  fir.select_rank %arg:i32 [ 1, ^bb1(%0:i32),
+                             2, ^bb2(%2,%arg,%arg2:i32,i32,i32),
+                             3, ^bb3(%arg2,%2:i32,i32),
+                             4, ^bb4(%1:i32),
+                             unit, ^bb5 ]
+  ^bb1(%a : i32) :
+    return %a : i32
+  ^bb2(%b : i32, %b2 : i32, %b3:i32) :
+    %4 = arith.addi %b, %b2 : i32
+    %5 = arith.addi %4, %b3 : i32
+    return %5 : i32
+  ^bb3(%c:i32, %c2:i32) :
+    %6 = arith.addi %c, %c2 : i32
+    return %6 : i32
+  ^bb4(%d : i32) :
+    return %d : i32
+  ^bb5 :
+    %zero = arith.constant 0 : i32
+    return %zero : i32
+}
+
+// CHECK-LABEL: func @select_rank(
+// CHECK-SAME:                    %[[SELECTVALUE:.*]]: i32,
+// CHECK-SAME:                    %[[ARG1:.*]]: 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]], ^bb5 [
+// CHECK:           1: ^bb1(%[[C0]] : i32),
+// CHECK:           2: ^bb2(%[[C2]], %[[SELECTVALUE]], %[[ARG1]] : i32, i32, i32),
+// CHECK:           3: ^bb3(%[[ARG1]], %[[C2]] : i32, i32),
+// CHECK:           4: ^bb4(%[[C1]] : i32)
+// CHECK:         ]