[Mlir-commits] [mlir] 8949c6d - [MLIR][OpenMP] Add Taskloop Collapse Support (#175924)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Feb 5 00:59:05 PST 2026
Author: Jack Styles
Date: 2026-02-05T08:59:00Z
New Revision: 8949c6d86b65f6a315cd295f955fdf2e4510f7c4
URL: https://github.com/llvm/llvm-project/commit/8949c6d86b65f6a315cd295f955fdf2e4510f7c4
DIFF: https://github.com/llvm/llvm-project/commit/8949c6d86b65f6a315cd295f955fdf2e4510f7c4.diff
LOG: [MLIR][OpenMP] Add Taskloop Collapse Support (#175924)
Following work completed in #174386 and #174623, this patch adds support
for collapse to Taskloop. Collapse allows for the user to compress
multiple loop nests into a single loop, and for this to work with
Taskloop, there needs to be some changes to how we process the loops,
and the tasks that run them.
This patch brings Taskloop equivalent to OpenMP 4.5 support for MLIR and
Flang.
Added:
mlir/test/Target/LLVMIR/openmp-taskloop-collapse.mlir
Modified:
llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
mlir/test/Target/LLVMIR/openmp-todo.mlir
Removed:
################################################################################
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 037fcaa863fe7..8ed4e0ba14502 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -1482,6 +1482,9 @@ class OpenMPIRBuilder {
/// \param Mergeable If the given task is `mergeable`
/// \param Priority `priority-value' specifies the execution order of the
/// tasks that is generated by the construct
+ /// \param NumOfCollapseLoops Defines the number of loops that are being
+ /// collapsed. The default value is 1, as thats the value when collapse is not
+ /// used.
/// \param DupCB The callback to generate the duplication code. See
/// documentation for \ref TaskDupCallbackTy. This can be nullptr.
/// \param TaskContextStructPtrVal If non-null, a pointer to to be placed
@@ -1494,7 +1497,8 @@ class OpenMPIRBuilder {
Value *LBVal, Value *UBVal, Value *StepVal, bool Untied = false,
Value *IfCond = nullptr, Value *GrainSize = nullptr, bool NoGroup = false,
int Sched = 0, Value *Final = nullptr, bool Mergeable = false,
- Value *Priority = nullptr, TaskDupCallbackTy DupCB = nullptr,
+ Value *Priority = nullptr, uint64_t NumOfCollapseLoops = 1,
+ TaskDupCallbackTy DupCB = nullptr,
Value *TaskContextStructPtrVal = nullptr);
/// Generator for `#omp task`
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 99f98d078a482..cefd066d185bb 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -2098,7 +2098,8 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTaskloop(
llvm::function_ref<llvm::Expected<llvm::CanonicalLoopInfo *>()> LoopInfo,
Value *LBVal, Value *UBVal, Value *StepVal, bool Untied, Value *IfCond,
Value *GrainSize, bool NoGroup, int Sched, Value *Final, bool Mergeable,
- Value *Priority, TaskDupCallbackTy DupCB, Value *TaskContextStructPtrVal) {
+ Value *Priority, uint64_t NumOfCollapseLoops, TaskDupCallbackTy DupCB,
+ Value *TaskContextStructPtrVal) {
if (!updateToLocation(Loc))
return InsertPointTy();
@@ -2175,8 +2176,8 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTaskloop(
OI.PostOutlineCB = [this, Ident, LBVal, UBVal, StepVal, Untied,
TaskloopAllocaBB, CLI, Loc, TaskDupFn, ToBeDeleted,
IfCond, GrainSize, NoGroup, Sched, FakeLB, FakeUB,
- FakeStep, Final, Mergeable,
- Priority](Function &OutlinedFn) mutable {
+ FakeStep, Final, Mergeable, Priority,
+ NumOfCollapseLoops](Function &OutlinedFn) mutable {
// Replace the Stale CI by appropriate RTL function call.
assert(OutlinedFn.hasOneUse() &&
"there must be a single user for the outlined function");
@@ -2359,29 +2360,53 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTaskloop(
Builder.SetInsertPoint(CLI->getBody(),
CLI->getBody()->getFirstInsertionPt());
- // The canonical loop is generated with a fixed lower bound. We need to
- // update the index calculation code to use the task's lower bound. The
- // generated code looks like this:
- // %omp_loop.iv = phi ...
- // ...
- // %tmp = mul [type] %omp_loop.iv, step
- // %user_index = add [type] tmp, lb
- // OpenMPIRBuilder constructs canonical loops to have exactly three uses of
- // the normalised induction variable:
- // 1. This one: converting the normalised IV to the user IV
- // 2. The increment (add)
- // 3. The comparison against the trip count (icmp)
- // (1) is the only use that is a mul followed by an add so this cannot match
- // other IR.
- assert(CLI->getIndVar()->getNumUses() == 3 &&
- "Canonical loop should have exactly three uses of the ind var");
- for (User *IVUser : CLI->getIndVar()->users()) {
- if (auto *Mul = dyn_cast<BinaryOperator>(IVUser)) {
- if (Mul->getOpcode() == Instruction::Mul) {
- for (User *MulUser : Mul->users()) {
- if (auto *Add = dyn_cast<BinaryOperator>(MulUser)) {
- if (Add->getOpcode() == Instruction::Add) {
- Add->setOperand(1, CastedTaskLB);
+ if (NumOfCollapseLoops > 1) {
+ llvm::SmallVector<User *> UsersToReplace;
+ // When using the collapse clause, the bounds of the loop have to be
+ // adjusted to properly represent the iterator of the outer loop.
+ Value *IVPlusTaskLB = Builder.CreateAdd(
+ CLI->getIndVar(),
+ Builder.CreateSub(CastedTaskLB, ConstantInt::get(IVTy, 1)));
+ // To ensure every Use is correctly captured, we first want to record
+ // which users to replace the value in, and then replace the value.
+ for (auto IVUse = CLI->getIndVar()->uses().begin();
+ IVUse != CLI->getIndVar()->uses().end(); IVUse++) {
+ User *IVUser = IVUse->getUser();
+ if (auto *Op = dyn_cast<BinaryOperator>(IVUser)) {
+ if (Op->getOpcode() == Instruction::URem ||
+ Op->getOpcode() == Instruction::UDiv) {
+ UsersToReplace.push_back(IVUser);
+ }
+ }
+ }
+ for (User *User : UsersToReplace) {
+ User->replaceUsesOfWith(CLI->getIndVar(), IVPlusTaskLB);
+ }
+ } else {
+ // The canonical loop is generated with a fixed lower bound. We need to
+ // update the index calculation code to use the task's lower bound. The
+ // generated code looks like this:
+ // %omp_loop.iv = phi ...
+ // ...
+ // %tmp = mul [type] %omp_loop.iv, step
+ // %user_index = add [type] tmp, lb
+ // OpenMPIRBuilder constructs canonical loops to have exactly three uses
+ // of the normalised induction variable:
+ // 1. This one: converting the normalised IV to the user IV
+ // 2. The increment (add)
+ // 3. The comparison against the trip count (icmp)
+ // (1) is the only use that is a mul followed by an add so this cannot
+ // match other IR.
+ assert(CLI->getIndVar()->getNumUses() == 3 &&
+ "Canonical loop should have exactly three uses of the ind var");
+ for (User *IVUser : CLI->getIndVar()->users()) {
+ if (auto *Mul = dyn_cast<BinaryOperator>(IVUser)) {
+ if (Mul->getOpcode() == Instruction::Mul) {
+ for (User *MulUser : Mul->users()) {
+ if (auto *Add = dyn_cast<BinaryOperator>(MulUser)) {
+ if (Add->getOpcode() == Instruction::Add) {
+ Add->setOperand(1, CastedTaskLB);
+ }
}
}
}
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index a76419353b1b6..4844cc99b5c38 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -333,10 +333,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
if (op.getBare())
result = todo("ompx_bare");
};
- auto checkCollapse = [&todo](auto op, LogicalResult &result) {
- if (op.getCollapseNumLoops() > 1)
- result = todo("collapse");
- };
auto checkDepend = [&todo](auto op, LogicalResult &result) {
if (!op.getDependVars().empty() || op.getDependKinds())
result = todo("depend");
@@ -400,10 +396,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
checkAllocate(op, result);
checkOrder(op, result);
})
- .Case([&](omp::LoopNestOp op) {
- if (mlir::isa<omp::TaskloopOp>(op.getOperation()->getParentOp()))
- checkCollapse(op, result);
- })
.Case([&](omp::OrderedRegionOp op) { checkParLevelSimd(op, result); })
.Case([&](omp::SectionsOp op) {
checkAllocate(op, result);
@@ -2805,6 +2797,84 @@ convertOmpTaskloopOp(Operation &opInst, llvm::IRBuilderBase &builder,
return loopInfo;
};
+ Operation::operand_range lowerBounds = loopOp.getLoopLowerBounds();
+ Operation::operand_range upperBounds = loopOp.getLoopUpperBounds();
+ Operation::operand_range steps = loopOp.getLoopSteps();
+ llvm::Type *boundType =
+ moduleTranslation.lookupValue(lowerBounds[0])->getType();
+ llvm::Value *lbVal = nullptr;
+ llvm::Value *ubVal = builder.getIntN(boundType->getIntegerBitWidth(), 1);
+ llvm::Value *stepVal = nullptr;
+ if (loopOp.getCollapseNumLoops() > 1) {
+ // In cases where Collapse is used with Taskloop, the upper bound of the
+ // iteration space needs to be recalculated to cater for the collapsed loop.
+ // The Collapsed Loop UpperBound is the product of all collapsed
+ // loop's tripcount.
+ // The LowerBound for collapsed loops is always 1. When the loops are
+ // collapsed, it will reset the bounds and introduce processing to ensure
+ // the index's are presented as expected. As this happens after creating
+ // Taskloop, these bounds need predicting. Example:
+ // !$omp taskloop collapse(2)
+ // do i = 1, 10
+ // do j = 1, 5
+ // ..
+ // end do
+ // end do
+ // This loop above has a total of 50 iterations, so the lb will be 1, and
+ // the ub will be 50. collapseLoops in OMPIRBuilder then handles ensuring
+ // that i and j are properly presented when used in the loop.
+ for (uint64_t i = 0; i < loopOp.getCollapseNumLoops(); i++) {
+ llvm::Value *loopLb = moduleTranslation.lookupValue(lowerBounds[i]);
+ llvm::Value *loopUb = moduleTranslation.lookupValue(upperBounds[i]);
+ llvm::Value *loopStep = moduleTranslation.lookupValue(steps[i]);
+ // In some cases, such as where the ub is less than the lb so the loop
+ // steps down, the calculation for the loopTripCount is swapped. To ensure
+ // the correct value is found, calculate both UB - LB and LB - UB then
+ // select which value to use depending on how the loop has been
+ // configured.
+ llvm::Value *loopLbMinusOne = builder.CreateSub(
+ loopLb, builder.getIntN(boundType->getIntegerBitWidth(), 1));
+ llvm::Value *loopUbMinusOne = builder.CreateSub(
+ loopUb, builder.getIntN(boundType->getIntegerBitWidth(), 1));
+ llvm::Value *boundsCmp = builder.CreateICmpSLT(loopLb, loopUb);
+ llvm::Value *ubMinusLb = builder.CreateSub(loopUb, loopLbMinusOne);
+ llvm::Value *lbMinusUb = builder.CreateSub(loopLb, loopUbMinusOne);
+ llvm::Value *loopTripCount =
+ builder.CreateSelect(boundsCmp, ubMinusLb, lbMinusUb);
+ loopTripCount = builder.CreateBinaryIntrinsic(
+ llvm::Intrinsic::abs, loopTripCount, builder.getFalse());
+ // For loops that have a step value not equal to 1, we need to adjust the
+ // trip count to ensure the correct number of iterations for the loop is
+ // captured.
+ llvm::Value *loopTripCountDivStep =
+ builder.CreateSDiv(loopTripCount, loopStep);
+ loopTripCountDivStep = builder.CreateBinaryIntrinsic(
+ llvm::Intrinsic::abs, loopTripCountDivStep, builder.getFalse());
+ llvm::Value *loopTripCountRem =
+ builder.CreateSRem(loopTripCount, loopStep);
+ loopTripCountRem = builder.CreateBinaryIntrinsic(
+ llvm::Intrinsic::abs, loopTripCountRem, builder.getFalse());
+ llvm::Value *needsRoundUp = builder.CreateICmpNE(
+ loopTripCountRem,
+ builder.getIntN(loopTripCountRem->getType()->getIntegerBitWidth(),
+ 0));
+ loopTripCount =
+ builder.CreateAdd(loopTripCountDivStep,
+ builder.CreateZExtOrTrunc(
+ needsRoundUp, loopTripCountDivStep->getType()));
+ ubVal = builder.CreateMul(ubVal, loopTripCount);
+ }
+ lbVal = builder.getIntN(boundType->getIntegerBitWidth(), 1);
+ stepVal = builder.getIntN(boundType->getIntegerBitWidth(), 1);
+ } else {
+ lbVal = moduleTranslation.lookupValue(lowerBounds[0]);
+ ubVal = moduleTranslation.lookupValue(upperBounds[0]);
+ stepVal = moduleTranslation.lookupValue(steps[0]);
+ }
+ assert(lbVal != nullptr && "Expected value for lbVal");
+ assert(ubVal != nullptr && "Expected value for ubVal");
+ assert(stepVal != nullptr && "Expected value for stepVal");
+
llvm::Value *ifCond = nullptr;
llvm::Value *grainsize = nullptr;
int sched = 0; // default
@@ -2837,15 +2907,13 @@ convertOmpTaskloopOp(Operation &opInst, llvm::IRBuilderBase &builder,
llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
llvm::OpenMPIRBuilder::InsertPointOrErrorTy afterIP =
moduleTranslation.getOpenMPBuilder()->createTaskloop(
- ompLoc, allocaIP, bodyCB, loopInfo,
- moduleTranslation.lookupValue(loopOp.getLoopLowerBounds()[0]),
- moduleTranslation.lookupValue(loopOp.getLoopUpperBounds()[0]),
- moduleTranslation.lookupValue(loopOp.getLoopSteps()[0]),
+ ompLoc, allocaIP, bodyCB, loopInfo, lbVal, ubVal, stepVal,
taskloopOp.getUntied(), ifCond, grainsize, taskloopOp.getNogroup(),
sched, moduleTranslation.lookupValue(taskloopOp.getFinal()),
taskloopOp.getMergeable(),
moduleTranslation.lookupValue(taskloopOp.getPriority()),
- taskDupOrNull, taskStructMgr.getStructPtr());
+ loopOp.getCollapseNumLoops(), taskDupOrNull,
+ taskStructMgr.getStructPtr());
if (failed(handleError(afterIP, opInst)))
return failure();
diff --git a/mlir/test/Target/LLVMIR/openmp-taskloop-collapse.mlir b/mlir/test/Target/LLVMIR/openmp-taskloop-collapse.mlir
new file mode 100644
index 0000000000000..f0abff7e38869
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-taskloop-collapse.mlir
@@ -0,0 +1,331 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+omp.private {type = private} @_QFtestEi_private_i32 : i32
+
+omp.private {type = firstprivate} @_QFtestEa_firstprivate_i32 : i32 copy {
+^bb0(%arg0: !llvm.ptr, %arg1: !llvm.ptr):
+ %0 = llvm.load %arg0 : !llvm.ptr -> i32
+ llvm.store %0, %arg1 : i32, !llvm.ptr
+ omp.yield(%arg1 : !llvm.ptr)
+}
+
+
+llvm.func @_QPtest() {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+ %2 = llvm.alloca %0 x i32 {bindc_name = "j"} : (i64) -> !llvm.ptr
+ %3 = llvm.alloca %0 x i32 {bindc_name = "a"} : (i64) -> !llvm.ptr
+ %6 = llvm.mlir.constant(20 : i32) : i32
+ llvm.store %6, %3 : i32, !llvm.ptr
+ %c1_i32 = llvm.mlir.constant(1 :i32) : i32
+ %c5_i32 = llvm.mlir.constant(5 : i32) : i32
+ %c10_i32 = llvm.mlir.constant(10 : i32) : i32
+ omp.taskloop private(@_QFtestEa_firstprivate_i32 %3 -> %arg0, @_QFtestEi_private_i32 %1 -> %arg1 : !llvm.ptr, !llvm.ptr) {
+ omp.loop_nest (%arg2, %arg3) : i32 = (%c1_i32, %c1_i32) to (%c10_i32, %c5_i32) inclusive step (%c1_i32, %c1_i32) collapse(2) {
+ llvm.store %arg2, %arg1 : i32, !llvm.ptr
+ %10 = llvm.load %arg0 : !llvm.ptr -> i32
+ %11 = llvm.mlir.constant(1 : i32) : i32
+ %12 = llvm.add %10, %11 : i32
+ llvm.store %12, %arg0 : i32, !llvm.ptr
+ omp.yield
+ }
+ }
+ llvm.return
+}
+
+// CHECK: %[[structArg:.*]] = alloca { i64, i64, i64, ptr }, align 8
+// CHECK: %[[ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[structArg]], i32 0, i32 1
+// CHECK: store i64 50, ptr %[[ub]], align 4
+
+// CHECK: %[[VAL_1:.*]] = load ptr, ptr %0, align 8
+// CHECK: %[[gep_task_lb:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 0
+// CHECK: %[[task_lb:.*]] = load i64, ptr %[[gep_task_lb]], align 4
+// CHECK: %[[gep_task_ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 1
+// CHECK: %[[task_ub:.*]] = load i64, ptr %gep_ub.val, align 4
+
+// CHECK: %[[VAL_3:.*]] = sub i64 %[[task_ub]], %[[task_lb]]
+// CHECK: %[[VAL_4:.*]] = sdiv i64 %[[VAL_3]], 1
+// CHECK: %[[trip_cnt:.*]] = add i64 %[[VAL_4]], 1
+// CHECK: %[[VAL_6:.*]] = trunc i64 %[[task_lb]] to i32
+
+// CHECK: %[[VAL_7:.*]] = sub i32 %[[VAL_6]], 1
+// CHECK: %[[VAL_8:.*]] = add i32 %omp_collapsed.iv, %[[VAL_7]]
+// CHECK: %[[VAL_9:.*]] = urem i32 %[[VAL_8]], 5
+// CHECK: %[[VAL_10:.*]] = udiv i32 %[[VAL_8]], 5
+// CHECK: %[[VAL_11:.*]] = mul i32 %[[VAL_10]], 1
+// CHECK: %[[VAL_12:.*]] = add i32 %[[VAL_11]], 1
+// CHECK: %[[VAL_13:.*]] = mul i32 %[[VAL_9]], 1
+// CHECK: %[[VAL_14:.*]] = add i32 %[[VAL_13]], 1
+
+// -----
+
+llvm.func @_QPtest2() {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+ %2 = llvm.alloca %0 x i32 {bindc_name = "j"} : (i64) -> !llvm.ptr
+ %3 = llvm.alloca %0 x i32 {bindc_name = "a"} : (i64) -> !llvm.ptr
+ %6 = llvm.mlir.constant(20 : i32) : i32
+ llvm.store %6, %3 : i32, !llvm.ptr
+ %c1_i32 = llvm.mlir.constant(1 :i32) : i32
+ %c2_i32 = llvm.mlir.constant(2 : i32) : i32
+ %c5_i32 = llvm.mlir.constant(5 : i32) : i32
+ %c10_i32 = llvm.mlir.constant(10 : i32) : i32
+ omp.taskloop private(@_QFtestEa_firstprivate_i32 %3 -> %arg0, @_QFtestEi_private_i32 %1 -> %arg1 : !llvm.ptr, !llvm.ptr) {
+ omp.loop_nest (%arg2, %arg3, %arg4) : i32 = (%c1_i32, %c1_i32, %c2_i32) to (%c10_i32, %c5_i32, %c5_i32) inclusive step (%c1_i32, %c1_i32, %c1_i32) collapse(3) {
+ llvm.store %arg2, %arg1 : i32, !llvm.ptr
+ %10 = llvm.load %arg0 : !llvm.ptr -> i32
+ %11 = llvm.mlir.constant(1 : i32) : i32
+ %12 = llvm.add %10, %11 : i32
+ llvm.store %12, %arg0 : i32, !llvm.ptr
+ omp.yield
+ }
+ }
+ llvm.return
+}
+
+// CHECK: %[[structArg:.*]] = alloca { i64, i64, i64, ptr }, align 8
+// CHECK: %[[ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[structArg]], i32 0, i32 1
+// CHECK: store i64 200, ptr %[[ub]], align 4
+
+// CHECK: %[[VAL_1:.*]] = load ptr, ptr %0, align 8
+// CHECK: %[[gep_task_lb:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 0
+// CHECK: %[[task_lb:.*]] = load i64, ptr %[[gep_task_lb]], align 4
+// CHECK: %[[gep_task_ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 1
+// CHECK: %[[task_ub:.*]] = load i64, ptr %gep_ub.val, align 4
+
+// CHECK: %[[VAL_3:.*]] = sub i64 %[[task_ub]], %[[task_lb]]
+// CHECK: %[[VAL_4:.*]] = sdiv i64 %[[VAL_3]], 1
+// CHECK: %[[trip_cnt:.*]] = add i64 %[[VAL_4]], 1
+// CHECK: %[[VAL_6:.*]] = trunc i64 %[[task_lb]] to i32
+
+// CHECK: %[[VAL_7:.*]] = sub i32 %[[VAL_6]], 1
+// CHECK: %[[VAL_8:.*]] = add i32 %omp_collapsed.iv, %[[VAL_7]]
+// CHECK: %[[VAL_9:.*]] = urem i32 %[[VAL_8]], 4
+// CHECK: %[[VAL_10:.*]] = udiv i32 %[[VAL_8]], 4
+// CHECK: %[[VAL_11:.*]] = urem i32 %[[VAL_10]], 5
+// CHECK: %[[VAL_12:.*]] = udiv i32 %[[VAL_10]], 5
+// CHECK: %[[VAL_13:.*]] = mul i32 %[[VAL_12]], 1
+// CHECK: %[[VAL_14:.*]] = add i32 %[[VAL_13]], 1
+// CHECK: %[[VAL_15:.*]] = mul i32 %[[VAL_11]], 1
+// CHECK: %[[VAL_16:.*]] = add i32 %[[VAL_15]], 1
+// CHECK: %[[VAL_17:.*]] = mul i32 %[[VAL_9]], 1
+// CHECK: %[[VAL_18:.*]] = add i32 %[[VAL_17]], 2
+
+// -----
+
+llvm.func @_QPtest3() {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+ %2 = llvm.alloca %0 x i32 {bindc_name = "j"} : (i64) -> !llvm.ptr
+ %3 = llvm.alloca %0 x i32 {bindc_name = "a"} : (i64) -> !llvm.ptr
+ %6 = llvm.mlir.constant(20 : i32) : i32
+ llvm.store %6, %3 : i32, !llvm.ptr
+ %c1_i32 = llvm.mlir.constant(1 :i32) : i32
+ %c2_i32 = llvm.mlir.constant(2 : i32) : i32
+ %c5_i32 = llvm.mlir.constant(5 : i32) : i32
+ %c10_i32 = llvm.mlir.constant(10 : i32) : i32
+ %c20_i32 = llvm.mlir.constant(20 : i32) : i32
+ omp.taskloop private(@_QFtestEa_firstprivate_i32 %3 -> %arg0, @_QFtestEi_private_i32 %1 -> %arg1 : !llvm.ptr, !llvm.ptr) {
+ omp.loop_nest (%arg2, %arg3) : i32 = (%c10_i32, %c1_i32) to (%c20_i32, %c5_i32) inclusive step (%c1_i32, %c1_i32) collapse(2) {
+ llvm.store %arg2, %arg1 : i32, !llvm.ptr
+ %10 = llvm.load %arg0 : !llvm.ptr -> i32
+ %11 = llvm.mlir.constant(1 : i32) : i32
+ %12 = llvm.add %10, %11 : i32
+ llvm.store %12, %arg0 : i32, !llvm.ptr
+ omp.yield
+ }
+ }
+ llvm.return
+}
+
+// CHECK: %[[structArg:.*]] = alloca { i64, i64, i64, ptr }, align 8
+// CHECK: %[[ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[structArg]], i32 0, i32 1
+// CHECK: store i64 55, ptr %[[ub]], align 4
+
+// CHECK: %[[VAL_1:.*]] = load ptr, ptr %0, align 8
+// CHECK: %[[gep_task_lb:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 0
+// CHECK: %[[task_lb:.*]] = load i64, ptr %[[gep_task_lb]], align 4
+// CHECK: %[[gep_task_ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 1
+// CHECK: %[[task_ub:.*]] = load i64, ptr %[[gep_task_ub]], align 4
+
+// CHECK: %[[VAL_3:.*]] = sub i64 %[[task_ub]], %[[task_lb]]
+// CHECK: %[[VAL_4:.*]] = sdiv i64 %[[VAL_3]], 1
+// CHECK: %[[trip_cnt:.*]] = add i64 %[[VAL_4]], 1
+// CHECK: %[[VAL_5:.*]] = trunc i64 %[[trip_cnt]] to i32
+// CHECK: %6 = trunc i64 %[[task_lb]] to i32
+
+// CHECK: %[[VAL_7:.*]] = sub i32 %[[VAL_6]], 1
+// CHECK: %[[VAL_8:.*]] = add i32 %omp_collapsed.iv, %[[VAL_7]]
+// CHECK: %[[VAL_9:.*]] = urem i32 %[[VAL_8]], 5
+// CHECK: %[[VAL_10:.*]] = udiv i32 %[[VAL_8]], 5
+
+// CHECK: %[[VAL_11:.*]] = mul i32 %[[VAL_10]], 1
+// CHECK: %[[VAL_12:.*]] = add i32 %[[VAL_11]], 10
+
+// CHECK: %[[VAL_13:.*]] = mul i32 %[[VAL_9]], 1
+// CHECK: %[[VAL_14:.*]] = add i32 %[[VAL_13]], 1
+
+// -----
+
+llvm.func @_QPtest4() {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+ %2 = llvm.alloca %0 x i32 {bindc_name = "j"} : (i64) -> !llvm.ptr
+ %3 = llvm.alloca %0 x i32 {bindc_name = "a"} : (i64) -> !llvm.ptr
+ %6 = llvm.mlir.constant(20 : i32) : i32
+ llvm.store %6, %3 : i32, !llvm.ptr
+ %c1_i32 = llvm.mlir.constant(1 :i32) : i32
+ %c2_i32 = llvm.mlir.constant(2 : i32) : i32
+ %c3_i32 = llvm.mlir.constant(3 : i32) : i32
+ %c5_i32 = llvm.mlir.constant(5 : i32) : i32
+ %c10_i32 = llvm.mlir.constant(10 : i32) : i32
+ %c15_i32 = llvm.mlir.constant(15 : i32) : i32
+ omp.taskloop private(@_QFtestEa_firstprivate_i32 %3 -> %arg0, @_QFtestEi_private_i32 %1 -> %arg1 : !llvm.ptr, !llvm.ptr) {
+ omp.loop_nest (%arg2, %arg3) : i32 = (%c2_i32, %c5_i32) to (%c10_i32, %c15_i32) inclusive step (%c2_i32, %c3_i32) collapse(2) {
+ llvm.store %arg2, %arg1 : i32, !llvm.ptr
+ %10 = llvm.load %arg0 : !llvm.ptr -> i32
+ %11 = llvm.mlir.constant(1 : i32) : i32
+ %12 = llvm.add %10, %11 : i32
+ llvm.store %12, %arg0 : i32, !llvm.ptr
+ omp.yield
+ }
+ }
+ llvm.return
+}
+
+// CHECK: %[[structArg:.*]] = alloca { i64, i64, i64, ptr }, align 8
+// CHECK: %[[ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[structArg]], i32 0, i32 1
+// CHECK: store i64 20, ptr %[[ub]], align 4
+
+// CHECK: %[[VAL_1:.*]] = load ptr, ptr %0, align 8
+// CHECK: %[[gep_task_lb:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 0
+// CHECK: %[[task_lb:.*]] = load i64, ptr %[[gep_task_lb]], align 4
+// CHECK: %[[gep_task_ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 1
+// CHECK: %[[task_ub:.*]] = load i64, ptr %[[gep_task_ub]], align 4
+
+// CHECK: %[[VAL_3:.*]] = sub i64 %[[task_ub]], %[[task_lb]]
+// CHECK: %[[VAL_4:.*]] = sdiv i64 %[[VAL_3]], 1
+// CHECK: %[[trip_cnt:.*]] = add i64 %[[VAL_4]], 1
+// CHECK: %[[VAL_5:.*]] = trunc i64 %[[trip_cnt]] to i32
+// CHECK: %6 = trunc i64 %[[task_lb]] to i32
+
+// CHECK: %[[VAL_7:.*]] = sub i32 %[[VAL_6]], 1
+// CHECK: %[[VAL_8:.*]] = add i32 %omp_collapsed.iv, %[[VAL_7]]
+// CHECK: %[[VAL_9:.*]] = urem i32 %[[VAL_8]], 4
+// CHECK: %[[VAL_10:.*]] = udiv i32 %[[VAL_8]], 4
+
+// CHECK: %[[VAL_11:.*]] = mul i32 %[[VAL_10]], 2
+// CHECK: %[[VAL_12:.*]] = add i32 %[[VAL_11]], 2
+
+// CHECK: %[[VAL_13:.*]] = mul i32 %[[VAL_9]], 3
+// CHECK: %[[VAL_14:.*]] = add i32 %[[VAL_13]], 5
+
+
+// -----
+
+llvm.func @_QPtest5() {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+ %2 = llvm.alloca %0 x i32 {bindc_name = "j"} : (i64) -> !llvm.ptr
+ %3 = llvm.alloca %0 x i32 {bindc_name = "a"} : (i64) -> !llvm.ptr
+ %6 = llvm.mlir.constant(20 : i32) : i32
+ llvm.store %6, %3 : i32, !llvm.ptr
+ %cneg2_i32 = llvm.mlir.constant(-2: i32) : i32
+ %c1_i32 = llvm.mlir.constant(1 :i32) : i32
+ %c2_i32 = llvm.mlir.constant(2 : i32) : i32
+ %c3_i32 = llvm.mlir.constant(3 : i32) : i32
+ %c5_i32 = llvm.mlir.constant(5 : i32) : i32
+ %c10_i32 = llvm.mlir.constant(10 : i32) : i32
+ %c15_i32 = llvm.mlir.constant(15 : i32) : i32
+ omp.taskloop private(@_QFtestEa_firstprivate_i32 %3 -> %arg0, @_QFtestEi_private_i32 %1 -> %arg1 : !llvm.ptr, !llvm.ptr) {
+ omp.loop_nest (%arg2, %arg3) : i32 = (%cneg2_i32, %c5_i32) to (%c10_i32, %c15_i32) inclusive step (%c2_i32, %c3_i32) collapse(2) {
+ llvm.store %arg2, %arg1 : i32, !llvm.ptr
+ %10 = llvm.load %arg0 : !llvm.ptr -> i32
+ %11 = llvm.mlir.constant(1 : i32) : i32
+ %12 = llvm.add %10, %11 : i32
+ llvm.store %12, %arg0 : i32, !llvm.ptr
+ omp.yield
+ }
+ }
+ llvm.return
+}
+
+// CHECK: %[[structArg:.*]] = alloca { i64, i64, i64, ptr }, align 8
+// CHECK: %[[ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[structArg]], i32 0, i32 1
+// CHECK: store i64 28, ptr %[[ub]], align 4
+
+// CHECK: %[[VAL_1:.*]] = load ptr, ptr %0, align 8
+// CHECK: %[[gep_task_lb:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 0
+// CHECK: %[[task_lb:.*]] = load i64, ptr %[[gep_task_lb]], align 4
+// CHECK: %[[gep_task_ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 1
+// CHECK: %[[task_ub:.*]] = load i64, ptr %[[gep_task_ub]], align 4
+
+// CHECK: %[[VAL_3:.*]] = sub i64 %[[task_ub]], %[[task_lb]]
+// CHECK: %[[VAL_4:.*]] = sdiv i64 %[[VAL_3]], 1
+// CHECK: %[[trip_cnt:.*]] = add i64 %[[VAL_4]], 1
+// CHECK: %[[VAL_5:.*]] = trunc i64 %[[trip_cnt]] to i32
+// CHECK: %6 = trunc i64 %[[task_lb]] to i32
+
+// CHECK: %[[VAL_7:.*]] = sub i32 %[[VAL_6]], 1
+// CHECK: %[[VAL_8:.*]] = add i32 %omp_collapsed.iv, %[[VAL_7]]
+// CHECK: %[[VAL_9:.*]] = urem i32 %[[VAL_8]], 4
+// CHECK: %[[VAL_10:.*]] = udiv i32 %[[VAL_8]], 4
+
+// CHECK: %[[VAL_11:.*]] = mul i32 %[[VAL_10]], 2
+// CHECK: %[[VAL_12:.*]] = add i32 %[[VAL_11]], -2
+
+// CHECK: %[[VAL_13:.*]] = mul i32 %[[VAL_9]], 3
+// CHECK: %[[VAL_14:.*]] = add i32 %[[VAL_13]], 5
+
+// -----
+
+llvm.func @_QPtest6() {
+ %0 = llvm.mlir.constant(1 : i64) : i64
+ %1 = llvm.alloca %0 x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr
+ %2 = llvm.alloca %0 x i32 {bindc_name = "j"} : (i64) -> !llvm.ptr
+ %3 = llvm.alloca %0 x i32 {bindc_name = "a"} : (i64) -> !llvm.ptr
+ %6 = llvm.mlir.constant(20 : i32) : i32
+ llvm.store %6, %3 : i32, !llvm.ptr
+ %cneg1_i32 = llvm.mlir.constant(-1: i32) : i32
+ %c1_i32 = llvm.mlir.constant(1 :i32) : i32
+ %c5_i32 = llvm.mlir.constant(5 : i32) : i32
+ %c10_i32 = llvm.mlir.constant(10 : i32) : i32
+ omp.taskloop private(@_QFtestEa_firstprivate_i32 %3 -> %arg0, @_QFtestEi_private_i32 %1 -> %arg1 : !llvm.ptr, !llvm.ptr) {
+ omp.loop_nest (%arg2, %arg3) : i32 = (%c10_i32, %c1_i32) to (%c5_i32, %c5_i32) inclusive step (%cneg1_i32, %c1_i32) collapse(2) {
+ llvm.store %arg2, %arg1 : i32, !llvm.ptr
+ %10 = llvm.load %arg0 : !llvm.ptr -> i32
+ %11 = llvm.mlir.constant(1 : i32) : i32
+ %12 = llvm.add %10, %11 : i32
+ llvm.store %12, %arg0 : i32, !llvm.ptr
+ omp.yield
+ }
+ }
+ llvm.return
+}
+
+// CHECK: %[[structArg:.*]] = alloca { i64, i64, i64, ptr }, align 8
+// CHECK: %[[ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[structArg]], i32 0, i32 1
+// CHECK: store i64 30, ptr %[[ub]], align 4
+
+// CHECK: %[[VAL_1:.*]] = load ptr, ptr %0, align 8
+// CHECK: %[[gep_task_lb:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 0
+// CHECK: %[[task_lb:.*]] = load i64, ptr %[[gep_task_lb]], align 4
+// CHECK: %[[gep_task_ub:.*]] = getelementptr { i64, i64, i64, ptr }, ptr %[[VAL_1]], i32 0, i32 1
+// CHECK: %[[task_ub:.*]] = load i64, ptr %[[gep_task_ub]], align 4
+
+// CHECK: %[[VAL_3:.*]] = sub i64 %[[task_ub]], %[[task_lb]]
+// CHECK: %[[VAL_4:.*]] = sdiv i64 %[[VAL_3]], 1
+// CHECK: %[[trip_cnt:.*]] = add i64 %[[VAL_4]], 1
+// CHECK: %[[VAL_5:.*]] = trunc i64 %[[trip_cnt]] to i32
+// CHECK: %6 = trunc i64 %[[task_lb]] to i32
+
+// CHECK: %[[VAL_7:.*]] = sub i32 %[[VAL_6]], 1
+// CHECK: %[[VAL_8:.*]] = add i32 %omp_collapsed.iv, %[[VAL_7]]
+// CHECK: %[[VAL_9:.*]] = urem i32 %[[VAL_8]], 5
+// CHECK: %[[VAL_10:.*]] = udiv i32 %[[VAL_8]], 5
+
+// CHECK: %[[VAL_11:.*]] = mul i32 %[[VAL_10]], -1
+// CHECK: %[[VAL_12:.*]] = add i32 %[[VAL_11]], 10
+
+// CHECK: %[[VAL_13:.*]] = mul i32 %[[VAL_9]], 1
+// CHECK: %[[VAL_14:.*]] = add i32 %[[VAL_13]], 1
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 70e4edf0705f2..9a10ad74baeb6 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -322,20 +322,6 @@ llvm.func @taskloop_allocate(%lb : i32, %ub : i32, %step : i32, %x : !llvm.ptr)
llvm.return
}
-// -----
-
-llvm.func @taskloop_collapse(%lb : i32, %ub : i32, %step : i32, %lb1 : i32, %ub1 : i32, %step1 : i32) {
- // expected-error at below {{LLVM Translation failed for operation: omp.taskloop}}
- omp.taskloop {
- // expected-error at below {{not yet implemented: Unhandled clause collapse in omp.loop_nest operation}}
- // expected-error at below {{LLVM Translation failed for operation: omp.loop_nest}}
- omp.loop_nest (%iv, %iv1) : i32 = (%lb, %lb1) to (%ub, %ub1) inclusive step (%step, %step1) collapse(2) {
- omp.yield
- }
- }
- llvm.return
-}
-
// -----
omp.declare_reduction @add_reduction_i32 : i32 init {
^bb0(%arg0: i32):
More information about the Mlir-commits
mailing list