[flang-commits] [flang] [mlir] [Flang][Openmp] Implement support for partial in unroll construct (PR #206642)
via flang-commits
flang-commits at lists.llvm.org
Tue Jun 30 02:57:21 PDT 2026
https://github.com/ShashwathiNavada updated https://github.com/llvm/llvm-project/pull/206642
>From d6cdc9ec1ffc9ad1b70b2c8cbfb0430548c3b724 Mon Sep 17 00:00:00 2001
From: ShashwathiNavada <shashwathinavada at gmail.com>
Date: Thu, 25 Jun 2026 01:40:11 -0500
Subject: [PATCH] Implement support for partial in unroll construct
---
flang/lib/Lower/OpenMP/ClauseProcessor.cpp | 9 ++++
flang/lib/Lower/OpenMP/ClauseProcessor.h | 1 +
flang/lib/Lower/OpenMP/OpenMP.cpp | 24 +++++++--
flang/test/Lower/OpenMP/unroll-partial01.f90 | 54 +++++++++++++++++++
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 35 ++++++++++++
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 50 +++++++++++++++++
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 27 ++++++++++
7 files changed, 195 insertions(+), 5 deletions(-)
create mode 100644 flang/test/Lower/OpenMP/unroll-partial01.f90
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index e2a2c8575ffc8..9031f6ddf0744 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -795,6 +795,15 @@ bool ClauseProcessor::processOrdered(
return false;
}
+bool ClauseProcessor::processPartial(std::optional<int64_t> &result) const {
+ if (auto *clause = findUniqueClause<omp::clause::Partial>()) {
+ if (clause->v.has_value())
+ result = evaluate::ToInt64(*clause->v);
+ return true;
+ }
+ return false;
+}
+
bool ClauseProcessor::processPriority(
lower::StatementContext &stmtCtx,
mlir::omp::PriorityClauseOps &result) const {
diff --git a/flang/lib/Lower/OpenMP/ClauseProcessor.h b/flang/lib/Lower/OpenMP/ClauseProcessor.h
index 1fc221b721ebf..e92db70ca707b 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.h
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.h
@@ -111,6 +111,7 @@ class ClauseProcessor {
mlir::omp::NumThreadsClauseOps &result) const;
bool processOrder(mlir::omp::OrderClauseOps &result) const;
bool processOrdered(mlir::omp::OrderedClauseOps &result) const;
+ bool processPartial(std::optional<int64_t> &result) const;
bool processPriority(lower::StatementContext &stmtCtx,
mlir::omp::PriorityClauseOps &result) const;
bool processProcBind(mlir::omp::ProcBindClauseOps &result) const;
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 266b06f353675..c81e4a59de197 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2539,10 +2539,17 @@ static void genUnrollOp(Fortran::lower::AbstractConverter &converter,
ConstructQueue::const_iterator item) {
fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
- // Clauses for unrolling not yet implemnted
ClauseProcessor cp(converter, semaCtx, item->clauses);
- cp.processTODO<clause::Partial, clause::Full>(
- loc, llvm::omp::Directive::OMPD_unroll);
+
+ // The `full` clause is not yet implemented.
+ cp.processTODO<clause::Full>(loc, llvm::omp::Directive::OMPD_unroll);
+
+ // Process the `partial` clause. If present, it may carry a constant unroll
+ // factor.
+ std::optional<int64_t> partialFactor;
+ bool hasPartial = cp.processPartial(partialFactor);
+ if (hasPartial && !partialFactor.has_value())
+ TODO(loc, "PARTIAL clause on UNROLL without a constant factor");
// Emit the associated loop
llvm::SmallVector<mlir::omp::CanonicalLoopOp, 1> canonLoops;
@@ -2554,9 +2561,16 @@ static void genUnrollOp(Fortran::lower::AbstractConverter &converter,
for (auto &&canonLoop : canonLoops)
applyees.push_back(canonLoop.getCli());
- // Apply unrolling to it
auto cli = llvm::getSingleElement(canonLoops).getCli();
- mlir::omp::UnrollHeuristicOp::create(firOpBuilder, loc, cli);
+
+ if (partialFactor.has_value()) {
+ // Partially unroll the loop by the given constant factor.
+ mlir::omp::UnrollPartialOp::create(firOpBuilder, loc, cli,
+ static_cast<uint64_t>(*partialFactor));
+ } else {
+ // Apply heuristic unrolling to it.
+ mlir::omp::UnrollHeuristicOp::create(firOpBuilder, loc, cli);
+ }
}
static mlir::omp::MaskedOp
diff --git a/flang/test/Lower/OpenMP/unroll-partial01.f90 b/flang/test/Lower/OpenMP/unroll-partial01.f90
new file mode 100644
index 0000000000000..fe76d8f077753
--- /dev/null
+++ b/flang/test/Lower/OpenMP/unroll-partial01.f90
@@ -0,0 +1,54 @@
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=51 -o - %s 2>&1 | FileCheck %s
+
+
+subroutine omp_unroll_partial01(lb, ub, inc)
+ integer res, i, lb, ub, inc
+
+ !$omp unroll partial(4)
+ do i = lb, ub, inc
+ res = i
+ end do
+ !$omp end unroll
+
+end subroutine omp_unroll_partial01
+
+
+! CHECK-LABEL: func.func @_QPomp_unroll_partial01(
+! CHECK-SAME: %[[ARG0:.*]]: !fir.ref<i32> {fir.bindc_name = "lb"},
+! CHECK-SAME: %[[ARG1:.*]]: !fir.ref<i32> {fir.bindc_name = "ub"},
+! CHECK-SAME: %[[ARG2:.*]]: !fir.ref<i32> {fir.bindc_name = "inc"}) {
+! CHECK: %[[VAL_0:.*]] = fir.dummy_scope : !fir.dscope
+! CHECK: %[[VAL_1:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFomp_unroll_partial01Ei"}
+! CHECK: %[[VAL_2:.*]]:2 = hlfir.declare %[[VAL_1]] {uniq_name = "_QFomp_unroll_partial01Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: %[[VAL_3:.*]]:2 = hlfir.declare %[[ARG2]] dummy_scope %[[VAL_0]] arg {{[0-9]+}} {uniq_name = "_QFomp_unroll_partial01Einc"} : (!fir.ref<i32>, !fir.dscope) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: %[[VAL_4:.*]]:2 = hlfir.declare %[[ARG0]] dummy_scope %[[VAL_0]] arg {{[0-9]+}} {uniq_name = "_QFomp_unroll_partial01Elb"} : (!fir.ref<i32>, !fir.dscope) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: %[[VAL_5:.*]] = fir.alloca i32 {bindc_name = "res", uniq_name = "_QFomp_unroll_partial01Eres"}
+! CHECK: %[[VAL_6:.*]]:2 = hlfir.declare %[[VAL_5]] {uniq_name = "_QFomp_unroll_partial01Eres"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: %[[VAL_7:.*]]:2 = hlfir.declare %[[ARG1]] dummy_scope %[[VAL_0]] arg {{[0-9]+}} {uniq_name = "_QFomp_unroll_partial01Eub"} : (!fir.ref<i32>, !fir.dscope) -> (!fir.ref<i32>, !fir.ref<i32>)
+! CHECK: %[[VAL_8:.*]] = fir.load %[[VAL_4]]#0 : !fir.ref<i32>
+! CHECK: %[[VAL_9:.*]] = fir.load %[[VAL_7]]#0 : !fir.ref<i32>
+! CHECK: %[[VAL_10:.*]] = fir.load %[[VAL_3]]#0 : !fir.ref<i32>
+! CHECK: %[[VAL_11:.*]] = arith.constant 0 : i32
+! CHECK: %[[VAL_12:.*]] = arith.constant 1 : i32
+! CHECK: %[[VAL_13:.*]] = arith.cmpi slt, %[[VAL_10]], %[[VAL_11]] : i32
+! CHECK: %[[VAL_14:.*]] = arith.subi %[[VAL_11]], %[[VAL_10]] : i32
+! CHECK: %[[VAL_15:.*]] = arith.select %[[VAL_13]], %[[VAL_14]], %[[VAL_10]] : i32
+! CHECK: %[[VAL_16:.*]] = arith.select %[[VAL_13]], %[[VAL_9]], %[[VAL_8]] : i32
+! CHECK: %[[VAL_17:.*]] = arith.select %[[VAL_13]], %[[VAL_8]], %[[VAL_9]] : i32
+! CHECK: %[[VAL_18:.*]] = arith.subi %[[VAL_17]], %[[VAL_16]] overflow<nuw> : i32
+! CHECK: %[[VAL_19:.*]] = arith.divui %[[VAL_18]], %[[VAL_15]] : i32
+! CHECK: %[[VAL_20:.*]] = arith.addi %[[VAL_19]], %[[VAL_12]] overflow<nuw> : i32
+! CHECK: %[[VAL_21:.*]] = arith.cmpi slt, %[[VAL_17]], %[[VAL_16]] : i32
+! CHECK: %[[VAL_22:.*]] = arith.select %[[VAL_21]], %[[VAL_11]], %[[VAL_20]] : i32
+! CHECK: %[[VAL_23:.*]] = omp.new_cli
+! CHECK: omp.canonical_loop(%[[VAL_23]]) %[[VAL_24:.*]] : i32 in range(%[[VAL_22]]) {
+! CHECK: %[[VAL_25:.*]] = arith.muli %[[VAL_24]], %[[VAL_10]] : i32
+! CHECK: %[[VAL_26:.*]] = arith.addi %[[VAL_8]], %[[VAL_25]] : i32
+! CHECK: hlfir.assign %[[VAL_26]] to %[[VAL_2]]#0 : i32, !fir.ref<i32>
+! CHECK: %[[VAL_27:.*]] = fir.load %[[VAL_2]]#0 : !fir.ref<i32>
+! CHECK: hlfir.assign %[[VAL_27]] to %[[VAL_6]]#0 : i32, !fir.ref<i32>
+! CHECK: omp.terminator
+! CHECK: }
+! CHECK: omp.unroll_partial(%[[VAL_23]]) {unroll_factor = 4 : i64}
+! CHECK: return
+! CHECK: }
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 0962b330e2f23..b5e4cefa1ce72 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -556,6 +556,41 @@ def UnrollHeuristicOp : OpenMPTransform_Op<"unroll_heuristic", []> {
let hasCustomAssemblyFormat = 1;
}
+//===----------------------------------------------------------------------===//
+// OpenMP unroll_partial operation
+//===----------------------------------------------------------------------===//
+
+def UnrollPartialOp : OpenMPTransform_Op<"unroll_partial", []> {
+ let summary = "OpenMP partial unroll operation";
+ let description = [{
+ Represents a `#pragma omp unroll partial` construct introduced in OpenMP
+ 5.1.
+
+ The operation has one applyee and no generatees. The applyee is partially
+ unrolled by the factor given in the `unroll_factor` attribute, which must be
+ a positive constant. That is, the operation replaces `unroll_factor`
+ consecutive logical iterations of the applyee by a single iteration that
+ executes the original body `unroll_factor` times.
+
+ The resulting unrolled loop is currently not accessible as a generatee, so
+ no further loop-associated transformations can be chained onto it.
+
+ The lowering is implemented using `OpenMPIRBuilder::unrollLoopPartial`.
+
+ Assembly formats:
+ omp.unroll_partial(%cli) {unroll_factor = 4 : i64}
+ }];
+
+ let arguments = (ins CanonicalLoopInfoType:$applyee,
+ ConfinedAttr<I64Attr, [IntPositive]>:$unroll_factor);
+
+ let builders = [
+ OpBuilder<(ins "::mlir::Value":$cli, "uint64_t":$unrollFactor)>,
+ ];
+
+ let hasCustomAssemblyFormat = 1;
+}
+
//===----------------------------------------------------------------------===//
// OpenMP tile operation
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index db5fd8f2e3230..27296190d6ac3 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -4193,6 +4193,56 @@ UnrollHeuristicOp::getGenerateesODSOperandIndexAndLength() {
return {0, 0};
}
+//===----------------------------------------------------------------------===//
+// UnrollPartialOp
+//===----------------------------------------------------------------------===//
+
+void UnrollPartialOp::build(::mlir::OpBuilder &odsBuilder,
+ ::mlir::OperationState &odsState, ::mlir::Value cli,
+ uint64_t unrollFactor) {
+ odsState.addOperands(cli);
+ Properties &props = odsState.getOrAddProperties<Properties>();
+ props.unroll_factor = odsBuilder.getI64IntegerAttr(unrollFactor);
+}
+
+void UnrollPartialOp::print(OpAsmPrinter &p) {
+ p << '(' << getApplyee() << ')';
+
+ p.printOptionalAttrDict((*this)->getAttrs());
+}
+
+mlir::ParseResult UnrollPartialOp::parse(::mlir::OpAsmParser &parser,
+ ::mlir::OperationState &result) {
+ auto cliType = CanonicalLoopInfoType::get(parser.getContext());
+
+ if (parser.parseLParen())
+ return failure();
+
+ OpAsmParser::UnresolvedOperand applyee;
+ if (parser.parseOperand(applyee) ||
+ parser.resolveOperand(applyee, cliType, result.operands))
+ return failure();
+
+ if (parser.parseRParen())
+ return failure();
+
+ // The unroll factor is carried by the `unroll_factor` attribute.
+ if (parser.parseOptionalAttrDict(result.attributes))
+ return failure();
+
+ return mlir::success();
+}
+
+std::pair<unsigned, unsigned>
+UnrollPartialOp::getApplyeesODSOperandIndexAndLength() {
+ return getODSOperandIndexAndLength(odsIndex_applyee);
+}
+
+std::pair<unsigned, unsigned>
+UnrollPartialOp::getGenerateesODSOperandIndexAndLength() {
+ return {0, 0};
+}
+
//===----------------------------------------------------------------------===//
// TileOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index f0511bb4be7dd..d15c7f01c3674 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -4490,6 +4490,30 @@ applyUnrollHeuristic(omp::UnrollHeuristicOp op, llvm::IRBuilderBase &builder,
return success();
}
+/// Apply a `#pragma omp unroll partial` / `!$omp unroll partial`
+/// transformation using the OpenMPIRBuilder.
+static LogicalResult
+applyUnrollPartial(omp::UnrollPartialOp 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);
+
+ // No generatee is supported yet, so the unrolled loop's CanonicalLoopInfo is
+ // not requested and unrolling is deferred to LLVM's LoopUnroll pass.
+ int32_t factor = static_cast<int32_t>(op.getUnrollFactor());
+ ompBuilder->unrollLoopPartial(loc.DL, consBuilderCLI, factor,
+ /*UnrolledCLI=*/nullptr);
+
+ moduleTranslation.invalidateOmpLoop(applyee);
+ return success();
+}
+
/// Apply a `#pragma omp tile` / `!$omp tile` transformation using the
/// OpenMPIRBuilder.
static LogicalResult applyTile(omp::TileOp op, llvm::IRBuilderBase &builder,
@@ -8889,6 +8913,9 @@ LogicalResult OpenMPDialectLLVMIRTranslationInterface::convertOperation(
// the omp.canonical_loop.
return applyUnrollHeuristic(op, builder, moduleTranslation);
})
+ .Case([&](omp::UnrollPartialOp op) {
+ return applyUnrollPartial(op, builder, moduleTranslation);
+ })
.Case([&](omp::TileOp op) {
return applyTile(op, builder, moduleTranslation);
})
More information about the flang-commits
mailing list