[fir] Add IfBuilder and utility functions

[clementval]

Summary

In order to reduct the size of D111337. The IfBuilder and the two
utility functions genIsNotNull and genIsNull have been extracted in
a separate patch with dedicated unittests.

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]>

[Leporacanthicus]

Looks OK to me, but please don't merge until someone else also approves.

[clementval]

In D111796#3064054, @Leporacanthicus wrote:

Looks OK to me, but please don't merge until someone else also approves.

Sure. Thanks for the review.

[awarzynski]

Thank you @clementval , the tests make it so much easier to review!

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

43–47

IMO this is a bit fragile. I appreciate that we won't be able to use llvm::DataLayout for a while, but I fear that we will forget to update this once we can. Maybe we should leave something like static_assert(sizeof(void*) <= 8) here? For the sake of well-being of our future selves.

67

Perhaps I'm missing something here.

flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp

17

Use override in C++11 to make sure you spelled it correctly.

https://github.com/google/googletest/blob/master/docs/primer.md

Not a blocker.

28

Does this have to be a pointer? If yes, then with e.g. unique_ptr you wouldn't have to call delete. Not a blocker for me.

[mehdi_amini]

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

43–47

What will this static assert caught? Building the compiler on a platform that has pointers larger than 8 bytes? The check should be about the target and not the host though.

Now there should be a solution for this in MLIR, I'm wondering why it isn't deployed here (and if this API is the right level for exposing this): https://mlir.llvm.org/docs/DataLayout/

I'd say that at least this API will make it easier to find the places to update in the codebase compared to use directly getI64Type everywhere...

flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp

28

I'd rather see a unique_ptr as well here.

[clementval]

Address most of review comments

[clementval]

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

43–47

The code is probably older than the introduction of the API you mention. I'll check with other whether we can use this in replacement.
As a side note, there is a lot of code to be upstream and probably some of the code is not up to date with the latest MLIR advancement and best practices. We try to update that as we go.

[clementval]

Fix typo

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

 class FirOpBuilder : public mlir::OpBuilder {
 public:
   explicit FirOpBuilder(mlir::Operation *op, const fir::KindMapping &kindMap)
       : OpBuilder{op} {}
   explicit FirOpBuilder(mlir::OpBuilder &builder,
                         const fir::KindMapping &kindMap)
       : OpBuilder{builder} {}
 
+  /// Get the integer type whose bit width corresponds to the width of pointer
+  /// types, or is bigger.
+  mlir::Type getIntPtrType() {
+    // TODO: Delay the need of such type until codegen or find a way to use
+    // llvm::DataLayout::getPointerSizeInBits here.
+    return getI64Type();
+  }
+
   /// Create an integer constant of type \p type and value \p i.
   mlir::Value createIntegerConstant(mlir::Location loc, mlir::Type integerType,
                                     std::int64_t i);
 
   /// Lazy creation of fir.convert op.
   mlir::Value createConvert(mlir::Location loc, mlir::Type toTy,
                             mlir::Value val);
 
   /// Cast the input value to IndexType.
   mlir::Value convertToIndexType(mlir::Location loc, mlir::Value val) {
     return createConvert(loc, getIndexType(), val);
   }
+
+  //===--------------------------------------------------------------------===//
+  // If-Then-Else generation helper
+  //===--------------------------------------------------------------------===//
+
+  /// Helper class to create if-then-else in a structured way:
+  /// Usage: genIfOp().genThen([&](){...}).genElse([&](){...}).end();
+  /// Alternatively, getResults() can be used instead of end() to end the ifOp
+  /// and get the ifOp results.
+  class IfBuilder {
+  public:
+    IfBuilder(fir::IfOp ifOp, FirOpBuilder &builder)
+        : ifOp{ifOp}, builder{builder} {}
+    template <typename CC>
+    IfBuilder &genThen(CC func) {
+      builder.setInsertionPointToStart(&ifOp.thenRegion().front());
+      func();
+      return *this;
+    }
+    template <typename CC>
+    IfBuilder &genElse(CC func) {
+      assert(!ifOp.elseRegion().empty() && "must have else region");
+      builder.setInsertionPointToStart(&ifOp.elseRegion().front());
+      func();
+      return *this;
+    }
+    void end() { builder.setInsertionPointAfter(ifOp); }
+
+    /// End the IfOp and return the results if any.
+    mlir::Operation::result_range getResults() {
+      end();
+      return ifOp.getResults();
+    }
+
+    fir::IfOp &getIfOp() { return ifOp; };
+
+  private:
+    fir::IfOp ifOp;
+    FirOpBuilder &builder;
+  };
+
+  /// Create an IfOp and returns an IfBuilder that can generate the else/then
+  /// bodies.
+  IfBuilder genIfOp(mlir::Location loc, mlir::TypeRange results,
+                    mlir::Value cdt, bool withElseRegion) {
+    auto op = create<fir::IfOp>(loc, results, cdt, withElseRegion);
+    return IfBuilder(op, *this);
+  }
+
+  /// Create an IfOp with no "else" region, and no result values.
+  /// Usage: genIfThen(loc, cdt).genThen(lambda).end();
+  IfBuilder genIfThen(mlir::Location loc, mlir::Value cdt) {
+    auto op = create<fir::IfOp>(loc, llvm::None, cdt, false);
+    return IfBuilder(op, *this);
+  }
+
+  /// Create an IfOp with an "else" region, and no result values.
+  /// Usage: genIfThenElse(loc, cdt).genThen(lambda).genElse(lambda).end();
+  IfBuilder genIfThenElse(mlir::Location loc, mlir::Value cdt) {
+    auto op = create<fir::IfOp>(loc, llvm::None, cdt, true);
+    return IfBuilder(op, *this);
+  }
+
+  /// Generate code testing \p addr is not a null address.
+  mlir::Value genIsNotNull(mlir::Location loc, mlir::Value addr);
+
+  /// Generate code testing \p addr is a null address.
+  mlir::Value genIsNull(mlir::Location loc, mlir::Value addr);
 };
 
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H

flang/lib/Optimizer/Builder/FIRBuilder.cpp

 mlir::Value fir::FirOpBuilder::createConvert(mlir::Location loc,
                                              mlir::Type toTy, mlir::Value val) {
   if (val.getType() != toTy) {
     assert(!fir::isa_derived(toTy));
     return create<fir::ConvertOp>(loc, toTy, val);
   }
   return val;
 }
+
+static mlir::Value genNullPointerComparison(fir::FirOpBuilder &builder,
+                                            mlir::Location loc,
+                                            mlir::Value addr,
+                                            arith::CmpIPredicate condition) {
+  auto intPtrTy = builder.getIntPtrType();
+  auto ptrToInt = builder.createConvert(loc, intPtrTy, addr);
+  auto c0 = builder.createIntegerConstant(loc, intPtrTy, 0);
+  return builder.create<arith::CmpIOp>(loc, condition, ptrToInt, c0);
+}
+
+mlir::Value fir::FirOpBuilder::genIsNotNull(mlir::Location loc,
+                                            mlir::Value addr) {
+  return genNullPointerComparison(*this, loc, addr, arith::CmpIPredicate::ne);
+}
+
+mlir::Value fir::FirOpBuilder::genIsNull(mlir::Location loc, mlir::Value addr) {
+  return genNullPointerComparison(*this, loc, addr, arith::CmpIPredicate::eq);
+}

flang/unittests/Optimizer/Builder/FIRBuilderTest.cpp

+//===- FIRBuilderTest.cpp -- FIRBuilder unit tests ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "gtest/gtest.h"
+#include "flang/Optimizer/Support/InitFIR.h"
+#include "flang/Optimizer/Support/KindMapping.h"
+
+struct FIRBuilderTest : public testing::Test {
+public:
+  void SetUp() override {
+    fir::KindMapping kindMap(&context);
+    mlir::OpBuilder builder(&context);
+    firBuilder = std::make_unique<fir::FirOpBuilder>(builder, kindMap);
+    fir::support::loadDialects(context);
+  }
+
+  fir::FirOpBuilder &getBuilder() { return *firBuilder; }
+
+  mlir::MLIRContext context;
+  std::unique_ptr<fir::FirOpBuilder> firBuilder;
+};
+
+static arith::CmpIOp createCondition(fir::FirOpBuilder &builder) {
+  auto loc = builder.getUnknownLoc();
+  auto zero1 = builder.createIntegerConstant(loc, builder.getIndexType(), 0);
+  auto zero2 = builder.createIntegerConstant(loc, builder.getIndexType(), 0);
+  return builder.create<arith::CmpIOp>(
+      loc, arith::CmpIPredicate::eq, zero1, zero2);
+}
+
+//===----------------------------------------------------------------------===//
+// IfBuilder tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(FIRBuilderTest, genIfThen) {
+  auto builder = getBuilder();
+  auto loc = builder.getUnknownLoc();
+  auto cdt = createCondition(builder);
+  auto ifBuilder = builder.genIfThen(loc, cdt);
+  EXPECT_FALSE(ifBuilder.getIfOp().thenRegion().empty());
+  EXPECT_TRUE(ifBuilder.getIfOp().elseRegion().empty());
+}
+
+TEST_F(FIRBuilderTest, genIfThenElse) {
+  auto builder = getBuilder();
+  auto loc = builder.getUnknownLoc();
+  auto cdt = createCondition(builder);
+  auto ifBuilder = builder.genIfThenElse(loc, cdt);
+  EXPECT_FALSE(ifBuilder.getIfOp().thenRegion().empty());
+  EXPECT_FALSE(ifBuilder.getIfOp().elseRegion().empty());
+}
+
+TEST_F(FIRBuilderTest, genIfWithThen) {
+  auto builder = getBuilder();
+  auto loc = builder.getUnknownLoc();
+  auto cdt = createCondition(builder);
+  auto ifBuilder = builder.genIfOp(loc, {}, cdt, false);
+  EXPECT_FALSE(ifBuilder.getIfOp().thenRegion().empty());
+  EXPECT_TRUE(ifBuilder.getIfOp().elseRegion().empty());
+}
+
+TEST_F(FIRBuilderTest, genIfWithThenAndElse) {
+  auto builder = getBuilder();
+  auto loc = builder.getUnknownLoc();
+  auto cdt = createCondition(builder);
+  auto ifBuilder = builder.genIfOp(loc, {}, cdt, true);
+  EXPECT_FALSE(ifBuilder.getIfOp().thenRegion().empty());
+  EXPECT_FALSE(ifBuilder.getIfOp().elseRegion().empty());
+}
+
+//===----------------------------------------------------------------------===//
+// Helper functions tests
+//===----------------------------------------------------------------------===//
+
+TEST_F(FIRBuilderTest, genIsNotNull) {
+  auto builder = getBuilder();
+  auto loc = builder.getUnknownLoc();
+  auto dummyValue =
+      builder.createIntegerConstant(loc, builder.getIndexType(), 0);
+  auto res = builder.genIsNotNull(loc, dummyValue);
+  EXPECT_TRUE(mlir::isa<arith::CmpIOp>(res.getDefiningOp()));
+  auto cmpOp = dyn_cast<arith::CmpIOp>(res.getDefiningOp());
+  EXPECT_EQ(arith::CmpIPredicate::ne, cmpOp.predicate());
+}
+
+TEST_F(FIRBuilderTest, genIsNull) {
+  auto builder = getBuilder();
+  auto loc = builder.getUnknownLoc();
+  auto dummyValue =
+      builder.createIntegerConstant(loc, builder.getIndexType(), 0);
+  auto res = builder.genIsNull(loc, dummyValue);
+  EXPECT_TRUE(mlir::isa<arith::CmpIOp>(res.getDefiningOp()));
+  auto cmpOp = dyn_cast<arith::CmpIOp>(res.getDefiningOp());
+  EXPECT_EQ(arith::CmpIPredicate::eq, cmpOp.predicate());
+}

flang/unittests/Optimizer/CMakeLists.txt

 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
 set(LIBS
   FIRBuilder
   FIRCodeGen
   FIRDialect
   FIRSupport
   ${dialect_libs}
 )
 
 add_flang_unittest(FlangOptimizerTests
   Builder/DoLoopHelperTest.cpp
+  Builder/FIRBuilderTest.cpp
   FIRContextTest.cpp
   InternalNamesTest.cpp
   KindMappingTest.cpp
   RTBuilder.cpp
 )
 target_link_libraries(FlangOptimizerTests
   PRIVATE
   ${LIBS})