Skip to content

[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.

Sign up for GitHub

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

Merged
merged 6 commits into from
Jul 11, 2025
Merged

[MLIR][OpenMP] Add canonical loop LLVM-IR lowering #147069

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
4a18faa
Emit named operand indices
Meinersbur Jul 3, 2025
c7ed06d
Add omp.canonical_loop and omp.unroll modelling
Meinersbur Jul 4, 2025
da2613d
omp.canonical_loop and omp.unroll_heuristic lowering
Meinersbur Jul 4, 2025
7a4dff7
Merge branch 'main' into HEAD
Meinersbur Jul 7, 2025
3dda96f
Merge branch 'users/meinersbur/flang_canonical-loop_ops' into HEAD
Meinersbur Jul 7, 2025
7622bf4
Merge branch 'main' into HEAD
Meinersbur Jul 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Copy link
Member

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.

Copy link
Member Author

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.

Copy link
Member Author

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.

Copy link
Member

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) {
      ...
    }
    ...
  }
}

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) {
Expand Down Expand Up @@ -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.
Expand Down
10 changes: 10 additions & 0 deletions mlir/lib/Conversion/OpenMPToLLVM/OpenMPToLLVM.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
78 changes: 78 additions & 0 deletions mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3095,6 +3095,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) {
Expand Down Expand Up @@ -5989,6 +6050,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();
Expand Down
175 changes: 175 additions & 0 deletions mlir/test/Target/LLVMIR/openmp-cli-canonical_loop.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
}
Loading
Loading