-
Notifications
You must be signed in to change notification settings - Fork 14.7k
[MLIR][OpenMP] Add canonical loop LLVM-IR lowering #147069
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[MLIR][OpenMP] Add canonical loop LLVM-IR lowering #147069
Conversation
@llvm/pr-subscribers-flang-openmp @llvm/pr-subscribers-mlir Author: Michael Kruse (Meinersbur) ChangesSupport for translating the operations introduced in #144785 to LLVM-IR. In order to keep the lowering simple, Eventually, all transformations on a loop must be applied directly after emitting Patch is 21.16 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/147069.diff 6 Files Affected:
diff --git a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
index 79e8bb6add0da..5d52cf3f04b6a 100644
--- a/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
+++ b/mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
@@ -15,6 +15,7 @@
#define MLIR_TARGET_LLVMIR_MODULETRANSLATION_H
#include "mlir/Dialect/LLVMIR/LLVMInterfaces.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/Value.h"
@@ -24,6 +25,7 @@
#include "mlir/Target/LLVMIR/TypeToLLVM.h"
#include "llvm/ADT/SetVector.h"
+#include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
#include "llvm/IR/FPEnv.h"
namespace llvm {
@@ -108,6 +110,41 @@ class ModuleTranslation {
return blockMapping.lookup(block);
}
+ /// Find the LLVM-IR loop that represents an MLIR loop.
+ llvm::CanonicalLoopInfo *lookupOMPLoop(omp::NewCliOp mlir) const {
+ llvm::CanonicalLoopInfo *result = loopMapping.lookup(mlir);
+ assert(result && "attempt to get non-existing loop");
+ return result;
+ }
+
+ /// Find the LLVM-IR loop that represents an MLIR loop.
+ llvm::CanonicalLoopInfo *lookupOMPLoop(Value mlir) const {
+ return lookupOMPLoop(mlir.getDefiningOp<omp::NewCliOp>());
+ }
+
+ /// Mark an OpenMP loop as having been consumed.
+ void invalidateOmpLoop(omp::NewCliOp mlir) { loopMapping.erase(mlir); }
+
+ /// Mark an OpenMP loop as having been consumed.
+ void invalidateOmpLoop(Value mlir) {
+ invalidateOmpLoop(mlir.getDefiningOp<omp::NewCliOp>());
+ }
+
+ /// Map an MLIR OpenMP dialect CanonicalLoopInfo to its lowered LLVM-IR
+ /// OpenMPIRBuilder CanonicalLoopInfo
+ void mapOmpLoop(omp::NewCliOp mlir, llvm::CanonicalLoopInfo *llvm) {
+ assert(llvm && "argument must be non-null");
+ llvm::CanonicalLoopInfo *&cur = loopMapping[mlir];
+ assert(cur == nullptr && "attempting to map a loop that is already mapped");
+ cur = llvm;
+ }
+
+ /// Map an MLIR OpenMP dialect CanonicalLoopInfo to its lowered LLVM-IR
+ /// OpenMPIRBuilder CanonicalLoopInfo
+ void mapOmpLoop(Value mlir, llvm::CanonicalLoopInfo *llvm) {
+ mapOmpLoop(mlir.getDefiningOp<omp::NewCliOp>(), llvm);
+ }
+
/// Stores the mapping between an MLIR operation with successors and a
/// corresponding LLVM IR instruction.
void mapBranch(Operation *mlir, llvm::Instruction *llvm) {
@@ -381,6 +418,12 @@ class ModuleTranslation {
DenseMap<Value, llvm::Value *> valueMapping;
DenseMap<Block *, llvm::BasicBlock *> blockMapping;
+ /// List of not yet consumed MLIR loop handles (represented by an omp.new_cli
+ /// operation which creates a value of type CanonicalLoopInfoType) and their
+ /// LLVM-IR representation as CanonicalLoopInfo which is managed by the
+ /// OpenMPIRBuilder.
+ DenseMap<omp::NewCliOp, llvm::CanonicalLoopInfo *> loopMapping;
+
/// A mapping between MLIR LLVM dialect terminators and LLVM IR terminators
/// they are converted to. This allows for connecting PHI nodes to the source
/// values after all operations are converted.
diff --git a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
index 77a2708653576..7ac9687c4eeda 100644
--- a/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
+++ b/mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
@@ -41,6 +41,16 @@ template <typename T>
struct OpenMPOpConversion : public ConvertOpToLLVMPattern<T> {
using ConvertOpToLLVMPattern<T>::ConvertOpToLLVMPattern;
+ OpenMPOpConversion(LLVMTypeConverter &typeConverter,
+ PatternBenefit benefit = 1)
+ : ConvertOpToLLVMPattern<T>(typeConverter, benefit) {
+ // Operations using CanonicalLoopInfoType are lowered only by
+ // mlir::translateModuleToLLVMIR() using the OpenMPIRBuilder. Until then,
+ // the type and operations using it must be preserved.
+ typeConverter.addConversion(
+ [&](::mlir::omp::CanonicalLoopInfoType type) { return type; });
+ }
+
LogicalResult
matchAndRewrite(T op, typename T::Adaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 7a517fb86d1ff..8b9bf7cb2a393 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -3096,6 +3096,67 @@ convertOmpLoopNest(Operation &opInst, llvm::IRBuilderBase &builder,
return success();
}
+/// Convert an omp.canonical_loop to LLVM-IR
+static LogicalResult
+convertOmpCanonicalLoopOp(omp::CanonicalLoopOp op, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+
+ llvm::OpenMPIRBuilder::LocationDescription loopLoc(builder);
+ Value loopIV = op.getInductionVar();
+ Value loopTC = op.getTripCount();
+
+ llvm::Value *llvmTC = moduleTranslation.lookupValue(loopTC);
+
+ llvm::Expected<llvm::CanonicalLoopInfo *> llvmOrError =
+ ompBuilder->createCanonicalLoop(
+ loopLoc,
+ [&](llvm::OpenMPIRBuilder::InsertPointTy ip, llvm::Value *llvmIV) {
+ // Register the mapping of MLIR induction variable to LLVM-IR
+ // induction variable
+ moduleTranslation.mapValue(loopIV, llvmIV);
+
+ builder.restoreIP(ip);
+ llvm::Expected<llvm::BasicBlock *> bodyGenStatus =
+ convertOmpOpRegions(op.getRegion(), "omp.loop.region", builder,
+ moduleTranslation);
+
+ return bodyGenStatus.takeError();
+ },
+ llvmTC, "omp.loop");
+ if (!llvmOrError)
+ return op.emitError(llvm::toString(llvmOrError.takeError()));
+
+ llvm::CanonicalLoopInfo *llvmCLI = *llvmOrError;
+ llvm::IRBuilderBase::InsertPoint afterIP = llvmCLI->getAfterIP();
+ builder.restoreIP(afterIP);
+
+ // Register the mapping of MLIR loop to LLVM-IR OpenMPIRBuilder loop
+ if (Value cli = op.getCli())
+ moduleTranslation.mapOmpLoop(cli, llvmCLI);
+
+ return success();
+}
+
+/// Apply a `#pragma omp unroll` / "!$omp unroll" transformation using the
+/// OpenMPIRBuilder.
+static LogicalResult
+applyUnrollHeuristic(omp::UnrollHeuristicOp op, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ llvm::OpenMPIRBuilder *ompBuilder = moduleTranslation.getOpenMPBuilder();
+
+ Value applyee = op.getApplyee();
+ assert(applyee && "Loop to apply unrolling on required");
+
+ llvm::CanonicalLoopInfo *consBuilderCLI =
+ moduleTranslation.lookupOMPLoop(applyee);
+ llvm::OpenMPIRBuilder::LocationDescription loc(builder);
+ ompBuilder->unrollLoopHeuristic(loc.DL, consBuilderCLI);
+
+ moduleTranslation.invalidateOmpLoop(applyee);
+ return success();
+}
+
/// Convert an Atomic Ordering attribute to llvm::AtomicOrdering.
static llvm::AtomicOrdering
convertAtomicOrdering(std::optional<omp::ClauseMemoryOrderKind> ao) {
@@ -5988,6 +6049,23 @@ convertHostOrTargetOperation(Operation *op, llvm::IRBuilderBase &builder,
// etc. and then discarded
return success();
})
+ .Case([&](omp::NewCliOp op) {
+ // Meta-operation: Doesn't do anything by itself, but used to
+ // identify a loop.
+ return success();
+ })
+ .Case([&](omp::CanonicalLoopOp op) {
+ return convertOmpCanonicalLoopOp(op, builder, moduleTranslation);
+ })
+ .Case([&](omp::UnrollHeuristicOp op) {
+ // FIXME: Handling omp.unroll_heuristic as an executable requires
+ // that the generator (e.g. omp.canonical_loop) has been seen first.
+ // For construct that require all codegen to occur inside a callback
+ // (e.g. OpenMPIRBilder::createParallel), all codegen of that
+ // contained region including their transformations must occur at
+ // the omp.canonical_loop.
+ return applyUnrollHeuristic(op, builder, moduleTranslation);
+ })
.Default([&](Operation *inst) {
return inst->emitError()
<< "not yet implemented: " << inst->getName();
diff --git a/mlir/test/Target/LLVMIR/openmp-cli-canonical_loop.mlir b/mlir/test/Target/LLVMIR/openmp-cli-canonical_loop.mlir
new file mode 100644
index 0000000000000..9abef003d6183
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-cli-canonical_loop.mlir
@@ -0,0 +1,175 @@
+// Test lowering of standalone omp.canonical_loop
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: define void @anon_loop(
+// CHECK-SAME: ptr %[[ptr:.+]],
+// CHECK-SAME: i32 %[[tc:.+]]) {
+// CHECK-NEXT: br label %omp_omp.loop.preheader
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.preheader:
+// CHECK-NEXT: br label %omp_omp.loop.header
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.header:
+// CHECK-NEXT: %omp_omp.loop.iv = phi i32 [ 0, %omp_omp.loop.preheader ], [ %omp_omp.loop.next, %omp_omp.loop.inc ]
+// CHECK-NEXT: br label %omp_omp.loop.cond
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.cond:
+// CHECK-NEXT: %omp_omp.loop.cmp = icmp ult i32 %omp_omp.loop.iv, %[[tc]]
+// CHECK-NEXT: br i1 %omp_omp.loop.cmp, label %omp_omp.loop.body, label %omp_omp.loop.exit
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.body:
+// CHECK-NEXT: br label %omp.loop.region
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.loop.region:
+// CHECK-NEXT: store float 4.200000e+01, ptr %[[ptr]], align 4
+// CHECK-NEXT: br label %omp.region.cont
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.region.cont:
+// CHECK-NEXT: br label %omp_omp.loop.inc
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.inc:
+// CHECK-NEXT: %omp_omp.loop.next = add nuw i32 %omp_omp.loop.iv, 1
+// CHECK-NEXT: br label %omp_omp.loop.header
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.exit:
+// CHECK-NEXT: br label %omp_omp.loop.after
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.after:
+// CHECK-NEXT: ret void
+// CHECK-NEXT: }
+llvm.func @anon_loop(%ptr: !llvm.ptr, %tc : i32) -> () {
+ omp.canonical_loop %iv : i32 in range(%tc) {
+ %val = llvm.mlir.constant(42.0 : f32) : f32
+ llvm.store %val, %ptr : f32, !llvm.ptr
+ omp.terminator
+ }
+ llvm.return
+}
+
+
+
+// CHECK-LABEL: define void @trivial_loop(
+// CHECK-SAME: ptr %[[ptr:.+]],
+// CHECK-SAME: i32 %[[tc:.+]]) {
+// CHECK-NEXT: br label %omp_omp.loop.preheader
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.preheader:
+// CHECK-NEXT: br label %omp_omp.loop.header
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.header:
+// CHECK-NEXT: %omp_omp.loop.iv = phi i32 [ 0, %omp_omp.loop.preheader ], [ %omp_omp.loop.next, %omp_omp.loop.inc ]
+// CHECK-NEXT: br label %omp_omp.loop.cond
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.cond:
+// CHECK-NEXT: %omp_omp.loop.cmp = icmp ult i32 %omp_omp.loop.iv, %[[tc]]
+// CHECK-NEXT: br i1 %omp_omp.loop.cmp, label %omp_omp.loop.body, label %omp_omp.loop.exit
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.body:
+// CHECK-NEXT: br label %omp.loop.region
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.loop.region:
+// CHECK-NEXT: store float 4.200000e+01, ptr %[[ptr]], align 4
+// CHECK-NEXT: br label %omp.region.cont
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.region.cont:
+// CHECK-NEXT: br label %omp_omp.loop.inc
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.inc:
+// CHECK-NEXT: %omp_omp.loop.next = add nuw i32 %omp_omp.loop.iv, 1
+// CHECK-NEXT: br label %omp_omp.loop.header
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.exit:
+// CHECK-NEXT: br label %omp_omp.loop.after
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.after:
+// CHECK-NEXT: ret void
+// CHECK-NEXT: }
+llvm.func @trivial_loop(%ptr: !llvm.ptr, %tc : i32) -> () {
+ %cli = omp.new_cli
+ omp.canonical_loop(%cli) %iv : i32 in range(%tc) {
+ %val = llvm.mlir.constant(42.0 : f32) : f32
+ llvm.store %val, %ptr : f32, !llvm.ptr
+ omp.terminator
+ }
+ llvm.return
+}
+
+
+// CHECK-LABEL: define void @nested_loop(
+// CHECK-SAME: ptr %[[ptr:.+]], i32 %[[outer_tc:.+]], i32 %[[inner_tc:.+]]) {
+// CHECK-NEXT: br label %omp_omp.loop.preheader
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.preheader:
+// CHECK-NEXT: br label %omp_omp.loop.header
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.header:
+// CHECK-NEXT: %omp_omp.loop.iv = phi i32 [ 0, %omp_omp.loop.preheader ], [ %omp_omp.loop.next, %omp_omp.loop.inc ]
+// CHECK-NEXT: br label %omp_omp.loop.cond
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.cond:
+// CHECK-NEXT: %omp_omp.loop.cmp = icmp ult i32 %omp_omp.loop.iv, %[[outer_tc]]
+// CHECK-NEXT: br i1 %omp_omp.loop.cmp, label %omp_omp.loop.body, label %omp_omp.loop.exit
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.body:
+// CHECK-NEXT: br label %omp.loop.region
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.loop.region:
+// CHECK-NEXT: br label %omp_omp.loop.preheader1
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.preheader1:
+// CHECK-NEXT: br label %omp_omp.loop.header2
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.header2:
+// CHECK-NEXT: %omp_omp.loop.iv8 = phi i32 [ 0, %omp_omp.loop.preheader1 ], [ %omp_omp.loop.next10, %omp_omp.loop.inc5 ]
+// CHECK-NEXT: br label %omp_omp.loop.cond3
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.cond3:
+// CHECK-NEXT: %omp_omp.loop.cmp9 = icmp ult i32 %omp_omp.loop.iv8, %[[inner_tc]]
+// CHECK-NEXT: br i1 %omp_omp.loop.cmp9, label %omp_omp.loop.body4, label %omp_omp.loop.exit6
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.body4:
+// CHECK-NEXT: br label %omp.loop.region12
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.loop.region12:
+// CHECK-NEXT: store float 4.200000e+01, ptr %[[ptr]], align 4
+// CHECK-NEXT: br label %omp.region.cont11
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.region.cont11:
+// CHECK-NEXT: br label %omp_omp.loop.inc5
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.inc5:
+// CHECK-NEXT: %omp_omp.loop.next10 = add nuw i32 %omp_omp.loop.iv8, 1
+// CHECK-NEXT: br label %omp_omp.loop.header2
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.exit6:
+// CHECK-NEXT: br label %omp_omp.loop.after7
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.after7:
+// CHECK-NEXT: br label %omp.region.cont
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.region.cont:
+// CHECK-NEXT: br label %omp_omp.loop.inc
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.inc:
+// CHECK-NEXT: %omp_omp.loop.next = add nuw i32 %omp_omp.loop.iv, 1
+// CHECK-NEXT: br label %omp_omp.loop.header
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.exit:
+// CHECK-NEXT: br label %omp_omp.loop.after
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.after:
+// CHECK-NEXT: ret void
+// CHECK-NEXT: }
+llvm.func @nested_loop(%ptr: !llvm.ptr, %outer_tc : i32, %inner_tc : i32) -> () {
+ %outer_cli = omp.new_cli
+ %inner_cli = omp.new_cli
+ omp.canonical_loop(%outer_cli) %outer_iv : i32 in range(%outer_tc) {
+ omp.canonical_loop(%inner_cli) %inner_iv : i32 in range(%inner_tc) {
+ %val = llvm.mlir.constant(42.0 : f32) : f32
+ llvm.store %val, %ptr : f32, !llvm.ptr
+ omp.terminator
+ }
+ omp.terminator
+ }
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/openmp-cli-unroll-heuristic01.mlir b/mlir/test/Target/LLVMIR/openmp-cli-unroll-heuristic01.mlir
new file mode 100644
index 0000000000000..0f0448e15f983
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-cli-unroll-heuristic01.mlir
@@ -0,0 +1,56 @@
+// Test lowering of the omp.unroll_heuristic
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+
+// CHECK-LABEL: define void @unroll_heuristic_trivial_loop(
+// CHECK-SAME: ptr %[[ptr:.+]], i32 %[[tc:.+]]) {
+// CHECK-NEXT: br label %omp_omp.loop.preheader
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.preheader:
+// CHECK-NEXT: br label %omp_omp.loop.header
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.header:
+// CHECK-NEXT: %omp_omp.loop.iv = phi i32 [ 0, %omp_omp.loop.preheader ], [ %omp_omp.loop.next, %omp_omp.loop.inc ]
+// CHECK-NEXT: br label %omp_omp.loop.cond
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.cond:
+// CHECK-NEXT: %omp_omp.loop.cmp = icmp ult i32 %omp_omp.loop.iv, %[[tc]]
+// CHECK-NEXT: br i1 %omp_omp.loop.cmp, label %omp_omp.loop.body, label %omp_omp.loop.exit
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.body:
+// CHECK-NEXT: br label %omp.loop.region
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.loop.region:
+// CHECK-NEXT: store float 4.200000e+01, ptr %[[ptr]], align 4
+// CHECK-NEXT: br label %omp.region.cont
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.region.cont:
+// CHECK-NEXT: br label %omp_omp.loop.inc
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.inc:
+// CHECK-NEXT: %omp_omp.loop.next = add nuw i32 %omp_omp.loop.iv, 1
+// CHECK-NEXT: br label %omp_omp.loop.header, !llvm.loop ![[$MD1:[0-9]+]]
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.exit:
+// CHECK-NEXT: br label %omp_omp.loop.after
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.after:
+// CHECK-NEXT: ret void
+// CHECK-NEXT: }
+llvm.func @unroll_heuristic_trivial_loop(%ptr: !llvm.ptr, %tc: i32) -> () {
+ %literal_cli = omp.new_cli
+ omp.canonical_loop(%literal_cli) %iv : i32 in range(%tc) {
+ %val = llvm.mlir.constant(42.0 : f32) : f32
+ llvm.store %val, %ptr : f32, !llvm.ptr
+ omp.terminator
+ }
+ omp.unroll_heuristic(%literal_cli)
+ llvm.return
+}
+
+
+// Start of metadata
+// CHECK-LABEL: !llvm.module.flags
+
+// CHECK: ![[$MD1]] = distinct !{![[$MD1]], ![[$MD2:[0-9]+]]}
+// CHECK: ![[$MD2]] = !{!"llvm.loop.unroll.enable"}
diff --git a/mlir/test/Target/LLVMIR/openmp-cli-unroll-heuristic02.mlir b/mlir/test/Target/LLVMIR/openmp-cli-unroll-heuristic02.mlir
new file mode 100644
index 0000000000000..f82b4990e378e
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-cli-unroll-heuristic02.mlir
@@ -0,0 +1,93 @@
+// Test lowering of the omp.unroll_heuristic
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+
+// CHECK-LABEL: define void @unroll_heuristic_nested_loop(
+// CHECK-SAME: ptr %[[ptr:.+]], i32 %[[outer_tc:.+]], i32 %[[inner_tc:.+]]) {
+// CHECK-NEXT: br label %omp_omp.loop.preheader
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.preheader:
+// CHECK-NEXT: br label %omp_omp.loop.header
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.header:
+// CHECK-NEXT: %omp_omp.loop.iv = phi i32 [ 0, %omp_omp.loop.preheader ], [ %omp_omp.loop.next, %omp_omp.loop.inc ]
+// CHECK-NEXT: br label %omp_omp.loop.cond
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.cond:
+// CHECK-NEXT: %omp_omp.loop.cmp = icmp ult i32 %omp_omp.loop.iv, %[[outer_tc]]
+// CHECK-NEXT: br i1 %omp_omp.loop.cmp, label %omp_omp.loop.body, label %omp_omp.loop.exit
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.body:
+// CHECK-NEXT: br label %omp.loop.region
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.loop.region:
+// CHECK-NEXT: br label %omp_omp.loop.preheader1
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.preheader1:
+// CHECK-NEXT: br label %omp_omp.loop.header2
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.header2:
+// CHECK-NEXT: %omp_omp.loop.iv8 = phi i32 [ 0, %omp_omp.loop.preheader1 ], [ %omp_omp.loop.next10, %omp_omp.loop.inc5 ]
+// CHECK-NEXT: br label %omp_omp.loop.cond3
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.cond3:
+// CHECK-NEXT: %omp_omp.loop.cmp9 = icmp ult i32 %omp_omp.loop.iv8, %[[inner_tc]]
+// CHECK-NEXT: br i1 %omp_omp.loop.cmp9, label %omp_omp.loop.body4, label %omp_omp.loop.exit6
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.body4:
+// CHECK-NEXT: br label %omp.loop.region12
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.loop.region12:
+// CHECK-NEXT: store float 4.200000e+01, ptr %[[ptr]], align 4
+// CHECK-NEXT: br label %omp.region.cont11
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.region.cont11:
+// CHECK-NEXT: br label %omp_omp.loop.inc5
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.inc5:
+// CHECK-NEXT: %omp_omp.loop.next10 = add nuw i32 %omp_omp.loop.iv8, 1
+// CHECK-NEXT: br label %omp_omp.loop.header2, !llvm.loop ![[$MD1:[0-9]+]]
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.exit6:
+// CHECK-NEXT: br label %omp_omp.loop.after7
+// CHECK-EMPTY:
+// CHECK-NEXT: omp_omp.loop.after7:
+// CHECK-NEXT: br label %omp.region.cont
+// CHECK-EMPTY:
+// CHECK-NEXT: omp.region.cont:
+// CHECK-NEXT: br label %...
[truncated]
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM. Thanks for splitting the PR. Please wait for another approval before merging.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM, thanks!
Hi all, in our downstream project we build the libMLIRLLVMIRTransforms target separately (e.g. with
We configure the build with the following CMake command:
using GCC v11.4 and CMake v3.31 on Ubuntu 24.04. This issue shows when doing a clean build in Release mode. For reference, this quick-and-dirty fix seems to allow the target to be built:
although we probably shouldn't do this as MLIRLLVMIRTransforms doesn't depend directly on OpenMP headers, if I understand correctly. @Meinersbur is it expected that mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h now includes OpenMP specific code, given that the remainder of the header seems to be mostly on LLVM IR? If so, could there be a missing dependency in some CMakeLists file? Thank you in advance :) |
The only thing that changed is that previously I will add the additional dependency. |
Turns out it is sufficient to forward-declare Applied here: #151302 |
`#include "llvm/Frontend/OpenMP/OMPIRBuilder.h"` can be replaced with forward-declarations of `OpenMPIRBuilder` and `CanonicalLoopInfo`. This also avoids a dependency to `omp_gen` of the LLVMFrontendOpenMP component which is included indirectly in `OMPIRBuilder.h`. Since its inclusion in #147069, additional indirect dependencies on headers included by `OMPIRBuilder.h` were introduced as well. These are now included directly. Reported-by: fabrizio-indirli See #147069 (comment)
`#include "llvm/Frontend/OpenMP/OMPIRBuilder.h"` can be replaced with forward-declarations of `OpenMPIRBuilder` and `CanonicalLoopInfo`. This also avoids a dependency to `omp_gen` of the LLVMFrontendOpenMP component which is included indirectly in `OMPIRBuilder.h`. Since its inclusion in #147069, additional indirect dependencies on headers included by `OMPIRBuilder.h` were introduced as well. These are now included directly. Reported-by: fabrizio-indirli See llvm/llvm-project#147069 (comment)
Indeed #151302 fixed it, thank you! |
@@ -108,6 +110,41 @@ class ModuleTranslation { | |||
return blockMapping.lookup(block); | |||
} | |||
|
|||
/// Find the LLVM-IR loop that represents an MLIR loop. | |||
llvm::CanonicalLoopInfo *lookupOMPLoop(omp::NewCliOp mlir) const { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sorry @Meinersbur for the late comment, was just now able to take a brief look at this. I don't love the idea of introducing OpenMP dialect-specific handling into this generic MLIR translation class, did you consider using stack frames (OpenMPLoopInfoStackFrame
) similarly to what is done to omp.loop_nest
instead?
The idea there was to create a stack frame holding a single null llvm::CanonicalLoopInfo *
when finding the top-level loop wrapper (which we know is associated to a single omp.loop_nest
), then populating it when that operation is translated to LLVM IR and finally updating it with the processing of each loop wrapper. It would probably work somewhat differently in this case, but if we could use a similar method or even reuse/extend what we have, that might be helpful when supporting e.g. worksharing of transformed loops. Plus, I think we should really try to avoid dialect-specific handling here.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Post-commit reviews are very welcome!
Since we need a namespace for omp.new_cli
mappings, a potential StateStackFrameBase
would need to be the top of the ModuleTranslation/function's stack frame. That requires additional code during function creation, i.e. there is still OpenMP-specific handling in otherwise LLVM dialect-centric translation. As being mandatory, it would not be much different than adding the namespace-lookup object (be it a class derived from StateStackFrameBase or the DenseMap directly) as a member of the ModuleTranslation class.
There is an LLVMTranslationInterface iface
field in ModuleTranslation which looked promising at first, but also only reacts when encountering a dialect's operation or attribute.
Any other ideas? I could hide the namespace lookup object from ModuleTranslation.h
by using the pImpl idiom.
The abstraction seems to have been broken anyway. ModuleTranslation already has a getOpenMPBuilder()
member and has several lookupXYZ
members. MLIR could have translated to the LLVM Dialect in an iterative loweing step, but this idea was not followed in a compromise in order to use the OpenMPIRBuilder. Hence, ModuleTranslation will always require some OpenMP-specific code.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Considering that the LLVMTranslationDialectInterface was introduced for OpenMP (https://reviews.llvm.org/D96504) I could try to extend it with a callback on function creation. It could also hide getOpenMPBuilder()
completely from ModuleTranslation. But that should probably be discussed with the MLIR community.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
One idea that may work for this would be to introduce yet another operation to the dialect to serve as a context where loop transformations are performed. Something like this:
omp.wsloop {
omp.loop_context {
%outer_cli = omp.new_cli
%inner_cli = omp.new_cli
omp.canonical_loop(%outer_cli) %outer_iv : i32 in range(%outer_tc) {
omp.canonical_loop(%inner_cli) %inner_iv : i32 in range(%inner_tc) {
%val = llvm.mlir.constant(42.0 : f32) : f32
llvm.store %val, %ptr : f32, !llvm.ptr
omp.terminator
}
omp.terminator
}
omp.unroll_heuristic(%outer_cli)
omp.unroll_heuristic(%inner_cli)
omp.terminator
}
}
Here, the omp.loop_context
op would take the place of the current omp.loop_nest
, so it integrates seamlessly with all loop wrappers, and it would have a single block region containing only omp.new_cli
, omp.canonical_loop
and loop transformation operations. That would give us a place to create a stack frame holding all information needed to handle loop transformations, and perhaps even let us remove the omp.new_cli
and instead make CLIs entry block arguments:
omp.wsloop {
omp.loop_context(%outer_cli, %inner_cli) {
omp.canonical_loop(%outer_cli) %outer_iv : i32 in range(%outer_tc) {
...
}
...
}
}
Support for translating the operations introduced in #144785 to LLVM-IR.
In order to keep the lowering simple,
OpenMPIRBuider::unrollLoopHeuristic
is applied when encountering theomp.unroll_heuristic
op. As a result, the operation that unrolling is applied to (omp.canonical_loop
) must have been emitted before even though logically there is no such requirement.Eventually, all transformations on a loop must be applied directly after emitting
omp.canonical_loop
, i.e. future transformations must be looked-up when encounteringomp.canonical_loop
itself. This is because many OpenMPIRBuilder methods (e.g.createParallel
) expect all the region code to be emitted withing a callback. In the case ofcreateParallel
, the region code is getting outlined into a new function. Therefore, making the operation order a formal requirement would not make the implementation any easier.