[MLIR][OpenMP] Updates to initial Taskloop Bounds Implementation #3

Stylie777 · 2025-12-19T10:59:17Z

This approach was agreed in the original Taskloop PR here: llvm#166903 (comment). This uses the structArg to transfer the data between outlined functions to better support the bounds in Taskloop.

- Force the first 3 entries to the StructArg to be the bounds info - Ensure it will work when executing the tasks in parallel

tblah

Looks great to me beyond some minor nits. I see O3 can optimize away the extra geps and stores from outlining so I think it is okay to leave them.

tblah · 2025-12-19T11:19:13Z

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

    PostOutlineCBTy PostOutlineCB;
    BasicBlock *EntryBB, *ExitBB, *OuterAllocaBB;
    SmallVector<Value *, 2> ExcludeArgsFromAggregate;
+    SetVector<Value *> Inputs, Outputs;


I know the other members don't have documentation but I think it would be good to add something because I don't think the use of these is immediately obvious.

I don't think we can use Outputs. IIRC there's an assertion somewhere that there are no live out values. Better to remove it.

Added a comment about the use of Inputs

We still need some definition of Outputs somewhere as the CodeExtractor's extractCodeRegion expects there to be a SetVector for both Inputs and Outputs. The API gives 2 options, one where you just pass the CEAC value, and another that includes the inputs and outputs. I am happy to exclude the Outputs from the OutlineInfo struct, but there will need to be a SetVector made before extracting the code region from OpenMPIRBuilder::finalize.

tblah · 2025-12-19T11:24:29Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

                        llvm::SmallVectorImpl<Instruction *> &ToBeDeleted,
                        OpenMPIRBuilder::InsertPointTy InnerAllocaIP,
-                        const Twine &Name = "", bool AsPtr = true) {
+                        const Twine &Name = "", bool AsPtr = true, bool Is64Bit = false) {


Suggested change

const Twine &Name = "", bool AsPtr = true, bool Is64Bit = false) {

const Twine &Name = "", bool AsPtr = true, IntegerType *IntTy) {

Builder.restoreIP(OuterAllocaIP);

IntTy = IntTy ? IntTy : Builder.getInt32Ty();

More flexible.

tblah · 2025-12-19T11:27:36Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

  } else {
    UseFakeVal =
-        cast<BinaryOperator>(Builder.CreateAdd(FakeVal, Builder.getInt32(10)));
+        cast<BinaryOperator>(Builder.CreateAdd(FakeVal, Is64Bit ? Builder.getInt64(10) : Builder.getInt32(10)));


You could do this from IntTy with llvm::ConstantInt::get or IRBuilderBase::getIntN

Done with ConstantInt::get

tblah · 2025-12-19T11:29:17Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

  OI.ExcludeArgsFromAggregate.push_back(createFakeIntVal(
      Builder, AllocaIP, ToBeDeleted, TaskloopAllocaIP, "global.tid", false));
+  Value *FakeLB = createFakeIntVal(Builder, AllocaIP, ToBeDeleted, TaskloopAllocaIP,
+                   "lb", false, true);


Suggested change

"lb", false, true);

"lb", /*AsPtr=*/false, Builder.getInt64Ty());

It isn't obvious what the bool values mean without some extra help

tblah · 2025-12-19T11:30:34Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+                   "ub", false, true);
+  Value *FakeStep = createFakeIntVal(Builder, AllocaIP, ToBeDeleted, TaskloopAllocaIP,
+                   "step", false, true);
+  /* For Taskloop, we want to force the bounds being the first 3 inputs in the aggregate struct*/


nit: llvm style is to use C++ style comments

tblah · 2025-12-19T11:30:42Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

-    // HasShareds is true if any variables are captured in the outlined region,
-    // false otherwise.
-    bool HasShareds = StaleCI->arg_size() > 1;
+    /* Create the casting for the Bounds Values that can be used when outlining to replace the uses of the fakes with real values */


nit: comment style

tblah · 2025-12-19T11:34:30Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

                      /*sizeof_task=*/TaskSize, /*sizeof_shared=*/SharedsSize,
                      /*task_func=*/&OutlinedFn});

+    Value *Shareds = StaleCI->getArgOperand(1);


This could be moved above line 2035 to make that section clearer

I've moved this above the declaration for ArgStructAlloca so this can use the Shareds variable rather than calling the getArgOperand function

tblah · 2025-12-19T11:37:53Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+        if (ConstantInt *CI = dyn_cast<ConstantInt>(Gep.getOperand(2))) {
+          switch (CI->getZExtValue()) {
+          case 0:
+            TaskLB = &I;


It would also be good to check that the value being indexed is the right one, not just the numeric value of the index.

Added a check to make sure the GEP Instruction being checked is using the Shared's as its first operand.

kaviya2510 · 2025-12-22T09:34:44Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

    Value *Flags = Builder.getInt32(Tied);

    Value *TaskSize = Builder.getInt64(
        divideCeil(M.getDataLayout().getTypeSizeInBits(Taskloop), 8));


Suggested change

divideCeil(M.getDataLayout().getTypeSizeInBits(Taskloop), 8));

divideCeil(M.getDataLayout().getTypeSizeInBits(Task), 8));

kaviya2510

Thank you for the work @Stylie777.

There are some changes that need to be updated for passing the arguments to runtime call __kmpc_taskloop(..). Kindly address those changes.

Rest of the work looks good to me.

kaviya2510 · 2025-12-22T09:35:38Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

    Value *Flags = Builder.getInt32(Tied);

    Value *TaskSize = Builder.getInt64(
        divideCeil(M.getDataLayout().getTypeSizeInBits(Taskloop), 8));


Suggested change

divideCeil(M.getDataLayout().getTypeSizeInBits(Taskloop), 8));

divideCeil(M.getDataLayout().getTypeSizeInBits(Task), 8));

As we are utilizing structArg to store the the loopbound values, the struct __OMP_STRUCT_TYPE(Taskloop, kmp_task_info, false, VoidPtr, VoidPtr, Int32, VoidPtr, VoidPtr, Int64, Int64, Int64) is no longer needed.

The required size for storing loop bounds can be reserved in kmp_task_t by strutArg itself.

kaviya2510 · 2025-12-22T10:00:10Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

-    Builder.CreateStore(Step_ext, step);
-    llvm::Value *loadstep = Builder.CreateLoad(Builder.getInt64Ty(), step);
+    llvm::Value *Lb = Builder.CreateStructGEP(ArgStructType, TaskShareds, 0);
+    Builder.CreateStore(CastedLBVal, Lb);


Suggested change

Builder.CreateStore(CastedLBVal, Lb);

auto *Idx0 = Builder.getInt32(0);

llvm::Value *Lb = Builder.CreateGEP(ArgStructType, TaskShareds, {Idx0, Builder.getInt32(0)});

The values of lb,ub and step are already populated in StructArg. You can directly access it and pass the pointer to the runtime call __kmpc_taskloop(...)

kaviya2510 · 2025-12-22T10:00:14Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

-                           SharedsSize);
-    }
+    llvm::Value *Ub = Builder.CreateStructGEP(ArgStructType, TaskShareds, 1);
+    Builder.CreateStore(CastedUBVal, Ub);


Same here. Remove the store instruction.

kaviya2510 · 2025-12-22T10:00:35Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

-      Builder.CreateMemCpy(TaskShareds, Alignment, Shareds, Alignment,
-                           SharedsSize);
-    }
+    llvm::Value *Ub = Builder.CreateStructGEP(ArgStructType, TaskShareds, 1);


Suggested change

llvm::Value *Ub = Builder.CreateStructGEP(ArgStructType, TaskShareds, 1);

llvm::Value *Ub =Builder.CreateGEP(ArgStructType, TaskShareds, {Idx0, Builder.getInt32(1)});

GEP to StructArg and get the upper bound value.

kaviya2510 · 2025-12-22T10:01:49Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+    llvm::Value *Ub = Builder.CreateStructGEP(ArgStructType, TaskShareds, 1);
+    Builder.CreateStore(CastedUBVal, Ub);
+
+    llvm::Value *Step = Builder.CreateStructGEP(ArgStructType, TaskShareds, 2);


Suggested change

llvm::Value *Step = Builder.CreateStructGEP(ArgStructType, TaskShareds, 2);

llvm::Value *Step =Builder.CreateGEP(ArgStructType, TaskShareds, {Idx0, Builder.getInt32(2)});

kaviya2510 · 2025-12-22T10:02:02Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+    Builder.CreateStore(CastedUBVal, Ub);
+
+    llvm::Value *Step = Builder.CreateStructGEP(ArgStructType, TaskShareds, 2);
+    Builder.CreateStore(CastedStepVal, Step);


Remove the store.

kaviya2510 · 2025-12-22T10:18:33Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

    llvm::BasicBlock *Body = CLI->getBody();
    for (llvm::Instruction &I : *Body) {
      if (auto *Add = llvm::dyn_cast<llvm::BinaryOperator>(&I)) {
        if (Add->getOpcode() == llvm::Instruction::Add) {


Tom raised a concern that this add instruction pattern might also match other unrelated add instructions, and we discussed this in my PR: llvm#166903 (comment)

He suggested looking at the wsloop and distribute implementations for guidance on how this is handled there. I have not had a chance to dig into that yet. Could you please check this once?

Stylie777 added 2 commits December 19, 2025 10:52

[Flang][OpenMP] Rework bounds in Taskloop lowering

de7fff6

Updates to bounds rework

b30007a

- Force the first 3 entries to the StructArg to be the bounds info - Ensure it will work when executing the tasks in parallel

tblah reviewed Dec 19, 2025

View reviewed changes

Respond to review comments

ba5f1eb

kaviya2510 reviewed Dec 22, 2025

View reviewed changes

-                        const Twine &Name = "", bool AsPtr = true, bool Is64Bit = false) {
+                        const Twine &Name = "", bool AsPtr = true, IntegerType *IntTy) {
+Builder.restoreIP(OuterAllocaIP);
+IntTy = IntTy ? IntTy : Builder.getInt32Ty();

	"lb", false, true);
	"lb", /AsPtr=/false, Builder.getInt64Ty());

	divideCeil(M.getDataLayout().getTypeSizeInBits(Taskloop), 8));
	divideCeil(M.getDataLayout().getTypeSizeInBits(Task), 8));

	Builder.CreateStore(CastedLBVal, Lb);
	auto *Idx0 = Builder.getInt32(0);
	llvm::Value *Lb = Builder.CreateGEP(ArgStructType, TaskShareds, {Idx0, Builder.getInt32(0)});

	llvm::Value *Ub = Builder.CreateStructGEP(ArgStructType, TaskShareds, 1);
	llvm::Value *Ub =Builder.CreateGEP(ArgStructType, TaskShareds, {Idx0, Builder.getInt32(1)});

	llvm::Value *Step = Builder.CreateStructGEP(ArgStructType, TaskShareds, 2);
	llvm::Value *Step =Builder.CreateGEP(ArgStructType, TaskShareds, {Idx0, Builder.getInt32(2)});

[MLIR][OpenMP] Updates to initial Taskloop Bounds Implementation #3

Are you sure you want to change the base?

[MLIR][OpenMP] Updates to initial Taskloop Bounds Implementation #3

Uh oh!

Conversation

Stylie777 commented Dec 19, 2025

Uh oh!

tblah left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kaviya2510 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kaviya2510 Dec 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

kaviya2510 Dec 22, 2025 •

edited

Loading