[flang-commits] [flang] 13dbff3 - [Flang][Openmp] Implement support for partial in unroll construct (#206642)

via flang-commits flang-commits at lists.llvm.org
Mon Jul 13 02:06:52 PDT 2026


Author: ShashwathiNavada
Date: 2026-07-13T14:36:47+05:30
New Revision: 13dbff365eabc45fcc26a36ef077f52b6e168347

URL: https://github.com/llvm/llvm-project/commit/13dbff365eabc45fcc26a36ef077f52b6e168347
DIFF: https://github.com/llvm/llvm-project/commit/13dbff365eabc45fcc26a36ef077f52b6e168347.diff

LOG: [Flang][Openmp] Implement support for partial in unroll construct (#206642)

As described in section 9.2.2 of openmp 5.2 spec, the patch implements
support for partial in unroll construct.

Added: 
    flang/test/Lower/OpenMP/unroll-partial01.f90
    mlir/test/Dialect/OpenMP/cli-unroll-partial.mlir
    mlir/test/Target/LLVMIR/openmp-unroll-partial01.mlir
    mlir/test/Target/LLVMIR/openmp-unroll-partial02.mlir

Modified: 
    flang/lib/Lower/OpenMP/ClauseProcessor.cpp
    flang/lib/Lower/OpenMP/ClauseProcessor.h
    flang/lib/Lower/OpenMP/OpenMP.cpp
    mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
    mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
    mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
    mlir/test/Dialect/OpenMP/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
index 05e7edc8bdc6a..26438ef968008 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.cpp
@@ -818,6 +818,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 1435af3e844e9..ff62d0ead8dc9 100644
--- a/flang/lib/Lower/OpenMP/ClauseProcessor.h
+++ b/flang/lib/Lower/OpenMP/ClauseProcessor.h
@@ -110,6 +110,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 9c6352b73a585..31587c624153c 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -3162,10 +3162,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;
@@ -3177,9 +3184,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 5b4ce7cc214bf..932a4a91cc6ab 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -557,6 +557,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 41d3c547ca618..acf5500205b6f 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -4230,6 +4230,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 6929c75af1e2e..7bb15b74aeac3 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -5188,6 +5188,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,
@@ -9688,6 +9712,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);
           })

diff  --git a/mlir/test/Dialect/OpenMP/cli-unroll-partial.mlir b/mlir/test/Dialect/OpenMP/cli-unroll-partial.mlir
new file mode 100644
index 0000000000000..5ac3a1557ba79
--- /dev/null
+++ b/mlir/test/Dialect/OpenMP/cli-unroll-partial.mlir
@@ -0,0 +1,59 @@
+// RUN: mlir-opt %s            | FileCheck %s --enable-var-scope
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s --enable-var-scope
+
+
+// CHECK-LABEL: @omp_unroll_partial_raw(
+// CHECK-SAME: %[[tc:.+]]: i32) {
+func.func @omp_unroll_partial_raw(%tc : i32) -> () {
+  // CHECK-NEXT: %canonloop = omp.new_cli
+  %canonloop = "omp.new_cli" () : () -> (!omp.cli)
+  // CHECK-NEXT: omp.canonical_loop(%canonloop) %iv : i32 in range(%[[tc]]) {
+  "omp.canonical_loop" (%tc, %canonloop) ({
+    ^bb0(%iv: i32):
+      omp.terminator
+  }) : (i32, !omp.cli) -> ()
+  // CHECK: omp.unroll_partial(%canonloop) {unroll_factor = 4 : i64}
+  "omp.unroll_partial" (%canonloop) {unroll_factor = 4 : i64} : (!omp.cli) -> ()
+  return
+}
+
+
+// CHECK-LABEL: @omp_unroll_partial_pretty(
+// CHECK-SAME: %[[tc:.+]]: i32) {
+func.func @omp_unroll_partial_pretty(%tc : i32) -> () {
+  // CHECK-NEXT: %[[CANONLOOP:.+]] = omp.new_cli
+  %canonloop = omp.new_cli
+  // CHECK-NEXT: omp.canonical_loop(%canonloop) %iv : i32 in range(%[[tc]]) {
+  omp.canonical_loop(%canonloop) %iv : i32 in range(%tc) {
+    omp.terminator
+  }
+  // CHECK: omp.unroll_partial(%canonloop) {unroll_factor = 8 : i64}
+  omp.unroll_partial(%canonloop) {unroll_factor = 8 : i64}
+  return
+}
+
+
+// CHECK-LABEL: @omp_unroll_partial_nested_pretty(
+// CHECK-SAME: %[[tc:.+]]: i32) {
+func.func @omp_unroll_partial_nested_pretty(%tc : i32) -> () {
+  // CHECK-NEXT: %canonloop = omp.new_cli
+  %cli_outer = omp.new_cli
+  // CHECK-NEXT: %canonloop_d1 = omp.new_cli
+  %cli_inner = omp.new_cli
+  // CHECK-NEXT: omp.canonical_loop(%canonloop) %iv : i32 in range(%[[tc]]) {
+  omp.canonical_loop(%cli_outer) %iv_outer : i32 in range(%tc) {
+    // CHECK-NEXT: omp.canonical_loop(%canonloop_d1) %iv_d1 : i32 in range(%[[tc]]) {
+    omp.canonical_loop(%cli_inner) %iv_inner : i32 in range(%tc) {
+      // CHECK: omp.terminator
+      omp.terminator
+    }
+    // CHECK: omp.terminator
+    omp.terminator
+  }
+
+  // CHECK: omp.unroll_partial(%canonloop) {unroll_factor = 4 : i64}
+  omp.unroll_partial(%cli_outer) {unroll_factor = 4 : i64}
+  // CHECK-NEXT: omp.unroll_partial(%canonloop_d1) {unroll_factor = 2 : i64}
+  omp.unroll_partial(%cli_inner) {unroll_factor = 2 : i64}
+  return
+}

diff  --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 2db7f953fa3bb..668c9fb6dc2b3 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -4812,3 +4812,19 @@ func.func @free_shared_mem_invalid_alignment2(%n: i32, %ptr : !llvm.ptr) -> () {
   omp.free_shared_mem [%n x i64 : (i32) align(3)] %ptr : !llvm.ptr
   return
 }
+
+// -----
+// unroll_factor = 0 is not a strictly positive integer.
+func.func @omp_unroll_partial_factor_zero(%cli : !omp.cli) -> () {
+  // expected-error @below {{op attribute 'unroll_factor' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive}}
+  omp.unroll_partial(%cli) {unroll_factor = 0 : i64}
+  return
+}
+
+// -----
+// unroll_factor = -1 is not a strictly positive integer.
+func.func @omp_unroll_partial_factor_negative(%cli : !omp.cli) -> () {
+  // expected-error @below {{op attribute 'unroll_factor' failed to satisfy constraint: 64-bit signless integer attribute whose value is positive}}
+  omp.unroll_partial(%cli) {unroll_factor = -1 : i64}
+  return
+}

diff  --git a/mlir/test/Target/LLVMIR/openmp-unroll-partial01.mlir b/mlir/test/Target/LLVMIR/openmp-unroll-partial01.mlir
new file mode 100644
index 0000000000000..b9938afea0c30
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-unroll-partial01.mlir
@@ -0,0 +1,57 @@
+// Test lowering of the omp.unroll_partial (single loop)
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+
+// CHECK-LABEL: define void @unroll_partial_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_partial_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_partial(%literal_cli) {unroll_factor = 4 : i64}
+  llvm.return
+}
+
+
+// Start of metadata
+// CHECK-LABEL: !llvm.module.flags
+
+// CHECK: ![[$MD1]] = distinct !{![[$MD1]], ![[$MD2:[0-9]+]], ![[$MD3:[0-9]+]]}
+// CHECK: ![[$MD2]] = !{!"llvm.loop.unroll.enable"}
+// CHECK: ![[$MD3]] = !{!"llvm.loop.unroll.count", i32 4}

diff  --git a/mlir/test/Target/LLVMIR/openmp-unroll-partial02.mlir b/mlir/test/Target/LLVMIR/openmp-unroll-partial02.mlir
new file mode 100644
index 0000000000000..a62bc499255db
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-unroll-partial02.mlir
@@ -0,0 +1,97 @@
+// Test lowering of the omp.unroll_partial (nested loops)
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+
+// CHECK-LABEL: define void @unroll_partial_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 %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 ![[$MD4:[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_partial_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
+  }
+  omp.unroll_partial(%outer_cli) {unroll_factor = 2 : i64}
+  omp.unroll_partial(%inner_cli) {unroll_factor = 4 : i64}
+  llvm.return
+}
+
+
+// Start of metadata
+// CHECK-LABEL: !llvm.module.flags
+
+// The inner loop back-edge (header2) is encountered first in the IR dump, so
+// its metadata node gets a lower ID than the outer loop's metadata node.
+// Both loops share the same llvm.loop.unroll.enable node (uniqued by content).
+// CHECK: ![[$MD1]] = distinct !{![[$MD1]], ![[$MD2:[0-9]+]], ![[$MD3:[0-9]+]]}
+// CHECK: ![[$MD2]] = !{!"llvm.loop.unroll.enable"}
+// CHECK: ![[$MD3]] = !{!"llvm.loop.unroll.count", i32 4}
+// CHECK: ![[$MD4]] = distinct !{![[$MD4]], ![[$MD2]], ![[$MD5:[0-9]+]]}
+// CHECK: ![[$MD5]] = !{!"llvm.loop.unroll.count", i32 2}


        


More information about the flang-commits mailing list