[fir] Add the DoLoopHelper
Add the DoLoopHelper. Some helpers functions
to create fir.do_loop operations.
This code was part of D111337 and was extracted in order to
make the patch easier to review.
This patch is part of the upstreaming effort from fir-dev branch.
Reviewed By: kiranchandramohan
Differential Revision: https://reviews.llvm.org/D111713
Co-authored-by: Valentin Clement <clementval@gmail.com>
Co-authored-by: Jean Perier <jperier@nvidia.com>
GitOrigin-RevId: bde89ac7f1dafb7b7900c52b0d441672ce19c749
diff --git a/include/flang/Optimizer/Builder/DoLoopHelper.h b/include/flang/Optimizer/Builder/DoLoopHelper.h
new file mode 100644
index 0000000..bdd581c
--- /dev/null
+++ b/include/flang/Optimizer/Builder/DoLoopHelper.h
@@ -0,0 +1,50 @@
+//===-- DoLoopHelper.h -- gen fir.do_loop ops -------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_BUILDER_DOLOOPHELPER_H
+#define FORTRAN_OPTIMIZER_BUILDER_DOLOOPHELPER_H
+
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+
+namespace fir::factory {
+
+/// Helper to build fir.do_loop Ops.
+class DoLoopHelper {
+public:
+ explicit DoLoopHelper(fir::FirOpBuilder &builder, mlir::Location loc)
+ : builder(builder), loc(loc) {}
+ DoLoopHelper(const DoLoopHelper &) = delete;
+
+ /// Type of a callback to generate the loop body.
+ using BodyGenerator = std::function<void(fir::FirOpBuilder &, mlir::Value)>;
+
+ /// Build loop [\p lb, \p ub] with step \p step.
+ /// If \p step is an empty value, 1 is used for the step.
+ fir::DoLoopOp createLoop(mlir::Value lb, mlir::Value ub, mlir::Value step,
+ const BodyGenerator &bodyGenerator);
+
+ /// Build loop [\p lb, \p ub] with step 1.
+ fir::DoLoopOp createLoop(mlir::Value lb, mlir::Value ub,
+ const BodyGenerator &bodyGenerator);
+
+ /// Build loop [0, \p count) with step 1.
+ fir::DoLoopOp createLoop(mlir::Value count,
+ const BodyGenerator &bodyGenerator);
+
+private:
+ fir::FirOpBuilder &builder;
+ mlir::Location loc;
+};
+
+} // namespace fir::factory
+
+#endif // FORTRAN_OPTIMIZER_BUILDER_DOLOOPHELPER_H
diff --git a/include/flang/Optimizer/Builder/FIRBuilder.h b/include/flang/Optimizer/Builder/FIRBuilder.h
new file mode 100644
index 0000000..0fc450b
--- /dev/null
+++ b/include/flang/Optimizer/Builder/FIRBuilder.h
@@ -0,0 +1,60 @@
+//===-- FirBuilder.h -- FIR operation builder -------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Builder routines for constructing the FIR dialect of MLIR. As FIR is a
+// dialect of MLIR, it makes extensive use of MLIR interfaces and MLIR's coding
+// style (https://mlir.llvm.org/getting_started/DeveloperGuide/) is used in this
+// module.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H
+#define FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H
+
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Support/KindMapping.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+
+namespace fir {
+
+//===----------------------------------------------------------------------===//
+// FirOpBuilder
+//===----------------------------------------------------------------------===//
+
+/// Extends the MLIR OpBuilder to provide methods for building common FIR
+/// patterns.
+class FirOpBuilder : public mlir::OpBuilder {
+public:
+ explicit FirOpBuilder(mlir::Operation *op, const fir::KindMapping &kindMap)
+ : OpBuilder{op}, kindMap{kindMap} {}
+ explicit FirOpBuilder(mlir::OpBuilder &builder,
+ const fir::KindMapping &kindMap)
+ : OpBuilder{builder}, kindMap{kindMap} {}
+
+ /// 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);
+ }
+
+private:
+ const KindMapping &kindMap;
+};
+
+} // namespace fir
+
+#endif // FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H
diff --git a/lib/Optimizer/Builder/CMakeLists.txt b/lib/Optimizer/Builder/CMakeLists.txt
new file mode 100644
index 0000000..c000c98
--- /dev/null
+++ b/lib/Optimizer/Builder/CMakeLists.txt
@@ -0,0 +1,15 @@
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+add_flang_library(FIRBuilder
+ DoLoopHelper.cpp
+ FIRBuilder.cpp
+
+ DEPENDS
+ FIRDialect
+ FIRSupport
+ ${dialect_libs}
+
+ LINK_LIBS
+ FIRDialect
+ FIRSupport
+ ${dialect_libs}
+)
diff --git a/lib/Optimizer/Builder/DoLoopHelper.cpp b/lib/Optimizer/Builder/DoLoopHelper.cpp
new file mode 100644
index 0000000..54c44f1
--- /dev/null
+++ b/lib/Optimizer/Builder/DoLoopHelper.cpp
@@ -0,0 +1,48 @@
+//===-- DoLoopHelper.cpp --------------------------------------------------===//
+//
+// 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/DoLoopHelper.h"
+
+//===----------------------------------------------------------------------===//
+// DoLoopHelper implementation
+//===----------------------------------------------------------------------===//
+
+fir::DoLoopOp
+fir::factory::DoLoopHelper::createLoop(mlir::Value lb, mlir::Value ub,
+ mlir::Value step,
+ const BodyGenerator &bodyGenerator) {
+ auto lbi = builder.convertToIndexType(loc, lb);
+ auto ubi = builder.convertToIndexType(loc, ub);
+ assert(step && "step must be an actual Value");
+ auto inc = builder.convertToIndexType(loc, step);
+ auto loop = builder.create<fir::DoLoopOp>(loc, lbi, ubi, inc);
+ auto insertPt = builder.saveInsertionPoint();
+ builder.setInsertionPointToStart(loop.getBody());
+ auto index = loop.getInductionVar();
+ bodyGenerator(builder, index);
+ builder.restoreInsertionPoint(insertPt);
+ return loop;
+}
+
+fir::DoLoopOp
+fir::factory::DoLoopHelper::createLoop(mlir::Value lb, mlir::Value ub,
+ const BodyGenerator &bodyGenerator) {
+ return createLoop(
+ lb, ub, builder.createIntegerConstant(loc, builder.getIndexType(), 1),
+ bodyGenerator);
+}
+
+fir::DoLoopOp
+fir::factory::DoLoopHelper::createLoop(mlir::Value count,
+ const BodyGenerator &bodyGenerator) {
+ auto indexType = builder.getIndexType();
+ auto zero = builder.createIntegerConstant(loc, indexType, 0);
+ auto one = builder.createIntegerConstant(loc, count.getType(), 1);
+ auto up = builder.create<arith::SubIOp>(loc, count, one);
+ return createLoop(zero, up, one, bodyGenerator);
+}
diff --git a/lib/Optimizer/Builder/FIRBuilder.cpp b/lib/Optimizer/Builder/FIRBuilder.cpp
new file mode 100644
index 0000000..b800eca
--- /dev/null
+++ b/lib/Optimizer/Builder/FIRBuilder.cpp
@@ -0,0 +1,24 @@
+//===-- FIRBuilder.cpp ----------------------------------------------------===//
+//
+// 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"
+
+mlir::Value fir::FirOpBuilder::createIntegerConstant(mlir::Location loc,
+ mlir::Type ty,
+ std::int64_t cst) {
+ return create<mlir::ConstantOp>(loc, ty, getIntegerAttr(ty, cst));
+}
+
+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;
+}
diff --git a/lib/Optimizer/CMakeLists.txt b/lib/Optimizer/CMakeLists.txt
index 4b37a18..d2da0f9 100644
--- a/lib/Optimizer/CMakeLists.txt
+++ b/lib/Optimizer/CMakeLists.txt
@@ -1,3 +1,4 @@
+add_subdirectory(Builder)
add_subdirectory(CodeGen)
add_subdirectory(Dialect)
add_subdirectory(Support)
diff --git a/unittests/Optimizer/Builder/DoLoopHelperTest.cpp b/unittests/Optimizer/Builder/DoLoopHelperTest.cpp
new file mode 100644
index 0000000..0387d62
--- /dev/null
+++ b/unittests/Optimizer/Builder/DoLoopHelperTest.cpp
@@ -0,0 +1,84 @@
+//===- DoLoopHelper.cpp -- DoLoopHelper 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/DoLoopHelper.h"
+#include "gtest/gtest.h"
+#include "flang/Optimizer/Support/InitFIR.h"
+#include "flang/Optimizer/Support/KindMapping.h"
+#include <string>
+
+struct DoLoopHelperTest : public testing::Test {
+public:
+ void SetUp() {
+ fir::KindMapping kindMap(&context);
+ mlir::OpBuilder builder(&context);
+ firBuilder = new fir::FirOpBuilder(builder, kindMap);
+ fir::support::loadDialects(context);
+ }
+ void TearDown() { delete firBuilder; }
+
+ fir::FirOpBuilder &getBuilder() { return *firBuilder; }
+
+ mlir::MLIRContext context;
+ fir::FirOpBuilder *firBuilder;
+};
+
+void checkConstantValue(const mlir::Value &value, int64_t v) {
+ EXPECT_TRUE(mlir::isa<ConstantOp>(value.getDefiningOp()));
+ auto cstOp = dyn_cast<ConstantOp>(value.getDefiningOp());
+ auto valueAttr = cstOp.getValue().dyn_cast_or_null<IntegerAttr>();
+ EXPECT_EQ(v, valueAttr.getInt());
+}
+
+TEST_F(DoLoopHelperTest, createLoopWithCountTest) {
+ auto firBuilder = getBuilder();
+ fir::factory::DoLoopHelper helper(firBuilder, firBuilder.getUnknownLoc());
+
+ auto c10 = firBuilder.createIntegerConstant(
+ firBuilder.getUnknownLoc(), firBuilder.getIndexType(), 10);
+ auto loop =
+ helper.createLoop(c10, [&](fir::FirOpBuilder &, mlir::Value index) {});
+ checkConstantValue(loop.lowerBound(), 0);
+ EXPECT_TRUE(mlir::isa<arith::SubIOp>(loop.upperBound().getDefiningOp()));
+ auto subOp = dyn_cast<arith::SubIOp>(loop.upperBound().getDefiningOp());
+ EXPECT_EQ(c10, subOp.lhs());
+ checkConstantValue(subOp.rhs(), 1);
+ checkConstantValue(loop.step(), 1);
+}
+
+TEST_F(DoLoopHelperTest, createLoopWithLowerAndUpperBound) {
+ auto firBuilder = getBuilder();
+ fir::factory::DoLoopHelper helper(firBuilder, firBuilder.getUnknownLoc());
+
+ auto lb = firBuilder.createIntegerConstant(
+ firBuilder.getUnknownLoc(), firBuilder.getIndexType(), 1);
+ auto ub = firBuilder.createIntegerConstant(
+ firBuilder.getUnknownLoc(), firBuilder.getIndexType(), 20);
+ auto loop =
+ helper.createLoop(lb, ub, [&](fir::FirOpBuilder &, mlir::Value index) {});
+ checkConstantValue(loop.lowerBound(), 1);
+ checkConstantValue(loop.upperBound(), 20);
+ checkConstantValue(loop.step(), 1);
+}
+
+TEST_F(DoLoopHelperTest, createLoopWithStep) {
+ auto firBuilder = getBuilder();
+ fir::factory::DoLoopHelper helper(firBuilder, firBuilder.getUnknownLoc());
+
+ auto lb = firBuilder.createIntegerConstant(
+ firBuilder.getUnknownLoc(), firBuilder.getIndexType(), 1);
+ auto ub = firBuilder.createIntegerConstant(
+ firBuilder.getUnknownLoc(), firBuilder.getIndexType(), 20);
+ auto step = firBuilder.createIntegerConstant(
+ firBuilder.getUnknownLoc(), firBuilder.getIndexType(), 2);
+ auto loop = helper.createLoop(
+ lb, ub, step, [&](fir::FirOpBuilder &, mlir::Value index) {});
+ checkConstantValue(loop.lowerBound(), 1);
+ checkConstantValue(loop.upperBound(), 20);
+ checkConstantValue(loop.step(), 2);
+}
diff --git a/unittests/Optimizer/CMakeLists.txt b/unittests/Optimizer/CMakeLists.txt
index 8bfb860..aea7482 100644
--- a/unittests/Optimizer/CMakeLists.txt
+++ b/unittests/Optimizer/CMakeLists.txt
@@ -1,6 +1,7 @@
get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
set(LIBS
+ FIRBuilder
FIRCodeGen
FIRDialect
FIRSupport
@@ -8,6 +9,7 @@
)
add_flang_unittest(FlangOptimizerTests
+ Builder/DoLoopHelperTest.cpp
FIRContextTest.cpp
InternalNamesTest.cpp
KindMappingTest.cpp
