[OMPIRBuilder] - Handle dependencies in `createTarget` #93977

bhandarkar-pranav · 2024-05-31T15:41:55Z

This PR handles dependencies specified by the depend clause on an OpenMP target construct. It does this much the same way clang does it by materializing an OpenMP task that is tagged with the dependencies.

To aid review, here are the relevant implementation details

createTarget - This function itself is largely unchanged except that it now accepts a vector of DependData objects that it simply forwards to emitTargetCall
emitTargetCall - This function has changed now to check if an outer target-task needs to be materialized (Yes, if target construct has nowait or has depend clause). If yes, it calls emitTargetTask to do all the heavy lifting for creating and dispatching the task.
emitTargetTask - Bulk of the PR is here. See the large comment explaining what it does at the beginning of the function.

Things that will/should follow in subsequent PRs

Handle this for target enter data, exit data and update data constructs
(Nice to have) - Outline logic for creating code for a task and use in createTask and emitTargetTask

…ased target codegen path in OMPIRBuilder

…at offloads

…ttests(?) and lit tests - basically more testing before PR

bhandarkar-pranav · 2024-06-04T16:14:50Z

ping! Gentle reminder, I'd appreciate reviews on this.

ergawy

Thanks @pranavb-ca! I did a first round and have a few comments.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

ergawy · 2024-06-05T08:36:14Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+    DepArray = Builder.CreateAlloca(DepArrayTy, nullptr, ".dep.arr.addr");
+
+    unsigned P = 0;
+    for (const OpenMPIRBuilder::DependData &Dep : Dependencies) {


I see there is a similar loop in OpenMPIRBuilder::createTask. Can we outline this to a shared util used in both locations?

Actually, my intention was to switch OpenMPIRBuilder::createTask to use this function but i had a couple of other changes to OpenMPIRBuilder::createTask in mind so I decided to leave that out and roll it into a different PR. If you prefer, I can make the switch in this PR itself though. Just a preference, no strong opinions on this.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

ergawy · 2024-06-05T09:10:10Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+                                     RTLoc, TargetTaskAllocaIP));
+
+  OI.ExitBB = Builder.saveIP().getBlock();
+  OI.PostOutlineCB = [this, ToBeDeleted, Dependencies,


I there are pieces in this lambda that can be outlined into shared utils and called from here and OpenMPIRBuilder::createTask. For example, the logic to setup TaskData (but, other parts as well).

Yes, that is an active Todo which I was planning to add in a different PR because now this PR doesn't touch the handling of the task construct at all. But that's only a preference, I'll be happy to roll those changes into this PR too.

This passes ToBeDeleted and Dependencies by-value to the lambda, i.e. makes a copy of those lists. If called multiple times, items in ToBeDeleted will be erased multiple times. I assume this is not intended?

Suggested change

OI.PostOutlineCB = [this, ToBeDeleted, Dependencies,

OI.PostOutlineCB = [this, &ToBeDeleted, &Dependencies,

This passes ToBeDeleted and Dependencies by-value to the lambda, i.e. makes a copy of those lists. If called multiple times, items in ToBeDeleted will be erased multiple times. I assume this is not intended?

In fact a copy is necessary to "own" the data. PostOutlineCB is called only once so multiple deletions is not a problem. However, and what is more important to note, is that it is called long after ToBeDeleted is destroyed (that is after we have returned from emitTargetTask) making it necessary to keep a copy around.

skatrak

Thank you Pranav, I think this looks good. I just have a few minor comments.

skatrak · 2024-06-05T10:23:49Z

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

-      dds.emplace_back(dd);
-    }
-  }
+  if (!taskOp.getDependVars().empty() && taskOp.getDepends())


Nit: I'd suggest moving this check into buildDependData, since it always needs to be done. It would just return without adding anything to dds if the condition is not met.

It looks like this comment hasn't been addressed. Let me know if you just missed it, if you disagree with it or if it isn't very clear and I should try to explain better.

offload/test/offloading/fortran/target-depend.f90

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Meinersbur

Not a full a review, but the first notes that I started with

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

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Meinersbur

Conceptually, looks quite good. Just some style comments.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Meinersbur · 2024-06-07T17:14:20Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+    AllocaInst *ArgStructAlloca =
+        dyn_cast<AllocaInst>(StaleCI->getArgOperand(1));


Suggested change

AllocaInst *ArgStructAlloca =

dyn_cast<AllocaInst>(StaleCI->getArgOperand(1));

auto *ArgStructAlloca =

dyn_cast<AllocaInst>(StaleCI->getArgOperand(1));

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Meinersbur · 2024-06-07T17:54:30Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+    while (!ToBeDeleted.empty()) {
+      ToBeDeleted.top()->eraseFromParent();
+      ToBeDeleted.pop();
+    }


Suggested change

while (!ToBeDeleted.empty()) {

ToBeDeleted.top()->eraseFromParent();

ToBeDeleted.pop();

}

llvm::for_each(ToBeDeleted, [](Instruction *I) { I->ereaseFromParent(); } );

or a simple foreach-loop. ToBeDeleted goes out of scope anyway.

Meinersbur · 2024-06-07T17:55:44Z

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

+  OI.OuterAllocaBB = AllocaIP.getBlock();
+
+  // Add the thread ID argument.
+  std::stack<Instruction *> ToBeDeleted;


Suggested change

std::stack<Instruction *> ToBeDeleted;

SmallVector<Instruction *> ToBeDeleted;

std::stack uses std::deque (by default), which is quite malloc-heavy. From cppreference:

deques typically have large minimal memory cost; a deque holding just one element has to allocate its full internal array (e.g. 8 times the object size on 64-bit libstdc++; 16 times the object size or 4096 bytes, whichever is larger, on 64-bit libc++).

@Meinersbur -> The only reason, I didnt do it is because createFakeIntVal is used by a few other functions not related to this PR. So changing ToBeDeleted from std::stack to SmallVector would involve changing other functions (such as createTeams) which are not related to this PR. I could go ahead and add that to this PR. Or simply open a separate PR just to make this change which should be quick. Let me know what you prefer.

@Meinersbur - Please take a look. I have gone ahead and made this change. If you agree with my reason above then I'll simply back out that one commit and make it a new PR.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Meinersbur · 2024-06-07T18:02:52Z

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

+static void
+buildDependData(std::optional<ArrayAttr> depends, OperandRange dependVars,
+                LLVM::ModuleTranslation &moduleTranslation,
+                SmallVector<llvm::OpenMPIRBuilder::DependData> &dds) {


Suggested change

SmallVector<llvm::OpenMPIRBuilder::DependData> &dds) {

SmallVectorImpl<llvm::OpenMPIRBuilder::DependData> &dds) {

ergawy

LGTM

… depend clause (sibling tasks)

bhandarkar-pranav · 2024-06-21T05:19:37Z

@Meinersbur, @ergawy, @skatrak - Could you please give this another look. I have updated the testcase target-depend.f90 because the previous version was flaky; I had used depend incorrectly which I have fixed now. (Previously, depend was used on tasks that were not sibling tasks)

skatrak

Thanks for working on my previous comments. I just have a couple more minor ones.

offload/test/offloading/fortran/target-depend.f90

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

skatrak · 2024-06-26T11:05:35Z

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

-      dds.emplace_back(dd);
-    }
-  }
+  if (!taskOp.getDependVars().empty() && taskOp.getDepends())


It looks like this comment hasn't been addressed. Let me know if you just missed it, if you disagree with it or if it isn't very clear and I should try to explain better.

bhandarkar-pranav · 2024-07-02T13:45:10Z

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -3088,10 +3102,13 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
    if (!mapData.IsDeclareTarget[i] && !mapData.IsAMember[i])
      kernelInput.push_back(mapData.OriginalValue[i]);
  }
+  SmallVector<llvm::OpenMPIRBuilder::DependData> dds;
+  buildDependData(targetOp.getOperation(), moduleTranslation, dds);


@skatrak - Apologies for missing your previous comment about moving the empty depends list check into buildDependData. I have done that now in this iteration of the PR.

bhandarkar-pranav · 2024-07-02T13:54:19Z

@skatrak - I have updated the PR with changes based on your feedback. Could you please take a look?

skatrak

Thank you Pranav, LGTM. I have just one more comment but no need to wait for another review by me before merging this. Just make sure @Meinersbur is happy with it as well!

skatrak · 2024-07-03T11:25:18Z

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -682,6 +682,41 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
  return bodyGenStatus;
 }

+static void
+buildDependData(Operation *op, LLVM::ModuleTranslation &moduleTranslation,


I'd suggest taking std::optional<ArrayAttr> depends and OperandRange dependVars as arguments instead of Operation *op here. This way you don't need to try casting op to each supported operation below, but rather make the body of this function what the processDepends lambda inside currently has.

skatrak · 2024-07-03T11:27:06Z

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

-      dds.emplace_back(dd);
-    }
-  }
+  buildDependData(taskOp.getOperation(), moduleTranslation, dds);


Suggestion to how this call could look if following my comment above:

Suggested change

buildDependData(taskOp.getOperation(), moduleTranslation, dds);

buildDependData(moduleTranslation, taskOp.getDepends(), taskOp.getDependVars(), dds);

Meinersbur

LGTM

Summary: This patch handles dependencies specified by the `depend` clause on an OpenMP target construct. It does this much the same way clang does it by materializing an OpenMP `task` that is tagged with the dependencies. The following functions are relevant to this patch - 1) `createTarget` - This function itself is largely unchanged except that it now accepts a vector of `DependData` objects that it simply forwards to `emitTargetCall` 2) `emitTargetCall` - This function has changed now to check if an outer target-task needs to be materialized (i.e if `target` construct has `nowait` or has `depend` clause). If yes, it calls `emitTargetTask` to do all the heavy lifting for creating and dispatching the task. 3) `emitTargetTask` - Bulk of the change is here. See the large comment explaining what it does at the beginning of this function Test Plan: Reviewers: Subscribers: Tasks: Tags: Differential Revision: https://phabricator.intern.facebook.com/D60251117

bhandarkar-pranav added 19 commits May 28, 2024 16:06

Add a flag to choose new codegen

8060e0b

clang prints for debugging

83f09a5

add an option -new-ompirbuilder-target-codegen to enable dependency-b…

e2aa768

…ased target codegen path in OMPIRBuilder

checkpoint commit -> able to create an inlined version of the task th…

e15bb96

…at offloads

checkpoint commit

53759cc

checkpoint commit

83cea83

Simple test working. checkpoint commit. next steps clean up code, uni…

d5c2449

…ttests(?) and lit tests - basically more testing before PR

clean up, clean up, everybody clean up

e323356

Add an MLIR lit test

264dfa6

add an end-to-end offloading test for target depend

25021eb

Clean up, clean up, everybody clean up (some more)

e484ee2

Add back some ws that was removed

febda4f

More cleanup and comments

9cc0234

clang-format fixes

141dccc

remove commented out createFakeIntVal

6697f1e

more cleanup

e75c853

Fix dependencies when nowait is used on target construct

799751f

Add comments for emitTargetTask and emitDepArray

77e5753

Fix comment in target-depend.f90

073d194

bhandarkar-pranav requested review from ergawy, Meinersbur, skatrak, jsjodin, kiranchandramohan, abidh, mjklemm, jdoerfert, raghavendhra, agozillon and DominikAdamski May 31, 2024 15:41

bhandarkar-pranav requested review from dpalermo, TIFitis, kiranktp and kparzysz May 31, 2024 15:41

ergawy reviewed Jun 5, 2024

View reviewed changes

skatrak reviewed Jun 5, 2024

View reviewed changes

Meinersbur reviewed Jun 6, 2024

View reviewed changes

Meinersbur reviewed Jun 7, 2024

View reviewed changes

ergawy approved these changes Jun 11, 2024

View reviewed changes

bhandarkar-pranav added 2 commits June 11, 2024 11:53

Incorporate changes for review comments

29566b1

Update testcase to make it more reliable by correcting the use of the…

a12ffb4

… depend clause (sibling tasks)

skatrak reviewed Jun 26, 2024

View reviewed changes

Address more review comments

ac34fd1

bhandarkar-pranav commented Jul 2, 2024

View reviewed changes

Incorporate one more review comment

c0e2ceb

skatrak approved these changes Jul 3, 2024

View reviewed changes

bhandarkar-pranav added 4 commits July 9, 2024 01:09

Incorporate changes from review comments

9b49c09

Simplify test to make it faster yet effective

cda5bc6

Use SmallVector instead of std::stack in CreateFakeIntVal

a314a84

Use SmallVectorImpl in emitTaskDependencies

b060e36

Meinersbur approved these changes Jul 22, 2024

View reviewed changes

bhandarkar-pranav merged commit d7e185c into llvm:main Jul 22, 2024
7 checks passed

	OI.PostOutlineCB = [this, ToBeDeleted, Dependencies,
	OI.PostOutlineCB = [this, &ToBeDeleted, &Dependencies,

		AllocaInst *ArgStructAlloca =
		dyn_cast<AllocaInst>(StaleCI->getArgOperand(1));

	std::stack<Instruction *> ToBeDeleted;
	SmallVector<Instruction *> ToBeDeleted;

	SmallVector<llvm::OpenMPIRBuilder::DependData> &dds) {
	SmallVectorImpl<llvm::OpenMPIRBuilder::DependData> &dds) {

	buildDependData(taskOp.getOperation(), moduleTranslation, dds);
	buildDependData(moduleTranslation, taskOp.getDepends(), taskOp.getDependVars(), dds);

[OMPIRBuilder] - Handle dependencies in createTarget #93977

[OMPIRBuilder] - Handle dependencies in createTarget #93977

Uh oh!

Conversation

bhandarkar-pranav commented May 31, 2024

Uh oh!

bhandarkar-pranav commented Jun 4, 2024

Uh oh!

ergawy left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

skatrak 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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Meinersbur left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Meinersbur left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Meinersbur Jun 7, 2024 • 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!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ergawy left a comment

Choose a reason for hiding this comment

Uh oh!

bhandarkar-pranav commented Jun 21, 2024

[OMPIRBuilder] - Handle dependencies in `createTarget` #93977

[OMPIRBuilder] - Handle dependencies in `createTarget` #93977

Meinersbur Jun 7, 2024 •

edited

Loading