[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

196

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

165

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

167

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

169

182

Attributes are value types

[clementval]

Address comments

[clementval]

Rebase

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

string extraSwitchClassDeclaration = [{
     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();
     }
 
     // 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());

flang/lib/Optimizer/CodeGen/CodeGen.cpp

 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/ArrayRef.h"
 
 #define DEBUG_TYPE "flang-codegen"
 
 // fir::LLVMTypeConverter for converting to LLVM IR dialect types.
 #include "TypeConverter.h"
 
+using OperandTy = ArrayRef<mlir::Value>;
+
 namespace {
 /// FIR conversion pattern template
 template <typename FromOp>
 class FIROpConversion : public mlir::ConvertOpToLLVMPattern<FromOp> {
 public:
   explicit FIROpConversion(fir::LLVMTypeConverter &lowering)
       : mlir::ConvertOpToLLVMPattern<FromOp>(lowering) {}
 
 protected:
   mlir::Type convertType(mlir::Type ty) const {
     return lowerTy().convertType(ty);
   }
 
   fir::LLVMTypeConverter &lowerTy() const {
     return *static_cast<fir::LLVMTypeConverter *>(this->getTypeConverter());
   }
 };
 } // namespace
 
+static Block *createBlock(mlir::ConversionPatternRewriter &rewriter,
+                          mlir::Block *insertBefore) {
+  assert(insertBefore && "expected valid insertion block");
+  return rewriter.createBlock(insertBefore->getParent(),
+                              mlir::Region::iterator(insertBefore));
+}
+
 namespace {
 struct AddrOfOpConversion : public FIROpConversion<fir::AddrOfOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::AddrOfOp addr, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     auto ty = convertType(addr.getType());
if (optLinkage.hasValue()) {
         return mlir::LLVM::Linkage::Common;
       if (name == "weak")
         return mlir::LLVM::Linkage::Weak;
     }
     return mlir::LLVM::Linkage::External;
   }
 };
 
+void genCondBrOp(mlir::Location loc, mlir::Value cmp, mlir::Block *dest,
+                 Optional<mlir::ValueRange> destOps,
+                 mlir::ConversionPatternRewriter &rewriter,
+                 mlir::Block *newBlock) {
+  if (destOps.hasValue())
+    rewriter.create<mlir::LLVM::CondBrOp>(loc, cmp, dest, destOps.getValue(),
+                                          newBlock, mlir::ValueRange());
+  else
+    rewriter.create<mlir::LLVM::CondBrOp>(loc, cmp, dest, newBlock);
+}
+
+template <typename A, typename B>
+void genBrOp(A caseOp, mlir::Block *dest, Optional<B> destOps,
+             mlir::ConversionPatternRewriter &rewriter) {
+  if (destOps.hasValue())
+    rewriter.replaceOpWithNewOp<mlir::LLVM::BrOp>(caseOp, destOps.getValue(),
+                                                  dest);
+  else
+    rewriter.replaceOpWithNewOp<mlir::LLVM::BrOp>(caseOp, llvm::None, dest);
+}
+
+void genCaseLadderStep(mlir::Location loc, mlir::Value cmp, mlir::Block *dest,
+                       Optional<mlir::ValueRange> destOps,
+                       mlir::ConversionPatternRewriter &rewriter) {
+  auto *thisBlock = rewriter.getInsertionBlock();
+  auto *newBlock = createBlock(rewriter, dest);
+  rewriter.setInsertionPointToEnd(thisBlock);
+  genCondBrOp(loc, cmp, dest, destOps, rewriter, newBlock);
+  rewriter.setInsertionPointToEnd(newBlock);
+}
+
+template <typename OP>
+void selectMatchAndRewrite(fir::LLVMTypeConverter &lowering, OP select,
+                           typename OP::Adaptor adaptor,
+                           mlir::ConversionPatternRewriter &rewriter) {
+  auto conds = select.getNumConditions();
+  auto caseAttr =
+      select->template getAttrOfType<mlir::ArrayAttr>(OP::getCasesAttr());
+  auto cases = caseAttr.getValue();
+  auto selector = select.getSelector(adaptor.getOperands());
+  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 (decltype(conds) t = 0; t != conds; ++t) {
+    mlir::Block *dest = select.getSuccessor(t);
+    auto destOps = select.getSuccessorOperands(adaptor.getOperands(), t);
+    auto &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:         ]