Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[MLIR][OpenMP] Support lowering of host_eval to LLVM IR #179

Merged

Conversation

skatrak
Copy link

@skatrak skatrak commented Oct 10, 2024

This patch updates the MLIR to LLVM IR lowering of omp.target to support passing num_teams, num_threads, thread_limit and SPMD loop bounds through the host_eval argument of omp.target.

This replaces the previous implementation where this information was directly attached to the omp.target operation rather than captured to be used by the corresponding nested operation.

The implementation of TargetOp::getInnermostCapturedOmpOp is also improved to address bugs in the detection of target SPMD in MLIR, uncovered by changes to the translation to LLVM IR.

@skatrak skatrak force-pushed the users/skatrak/target-passthrough-02-mlir branch from 44b6230 to 719e50c Compare October 11, 2024 13:10
@skatrak skatrak force-pushed the users/skatrak/target-passthrough-03-mlir-llvm branch from e2ee789 to a9a83c9 Compare October 11, 2024 13:11
@skatrak skatrak force-pushed the users/skatrak/target-passthrough-02-mlir branch from 719e50c to bc6485b Compare October 22, 2024 15:47
@skatrak skatrak force-pushed the users/skatrak/target-passthrough-03-mlir-llvm branch from a9a83c9 to 332c9c5 Compare October 22, 2024 15:47
@skatrak skatrak force-pushed the users/skatrak/target-passthrough-02-mlir branch from bc6485b to ad391a1 Compare October 23, 2024 14:31
@skatrak skatrak force-pushed the users/skatrak/target-passthrough-03-mlir-llvm branch from 332c9c5 to 0443d95 Compare October 23, 2024 14:32
Copy link

@bhandarkar-pranav bhandarkar-pranav left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor comments. I went over the entire PR of course, but I'd prefer if someone more familiar with the SPMD stuff went over it so we are sure it is good.

Copy link
Author

@skatrak skatrak left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you @bhandarkar-pranav for taking a look. I should have addressed all your concerns with the last update. Regarding your comment about finding someone else familiar with the SPMD stuff, I happen to also be the guy who made the previous implementation 😅, and unsurprisingly changes look good to myself! 😄

Copy link

@bhandarkar-pranav bhandarkar-pranav left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@skatrak - Thank you for the changes. LGTM.

Copy link

@Meinersbur Meinersbur left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OpenMPIRBuilder changes LGTM

No test added?

// Consider the following difficulties (assuming 8-bit signed integers):
// * Adding \p Step to the loop counter which passes \p Stop may overflow:
// DO I = 1, 100, 50
/// * A \p Step of INT_MIN cannot not be normalized to a positive direction:
// DO I = 100, 0, -128
updateToLocation(Loc);

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The code below was meant to check basic errors before modifying anything.

It doesn't matter that much where it is, but it should not give the impression that the comment above would apply to updateToLocation.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out. I moved the call to updateToLocation below the asserts, as it was before.

@skatrak skatrak force-pushed the users/skatrak/target-passthrough-02-mlir branch from ad391a1 to 595f1ae Compare October 29, 2024 14:59
@skatrak skatrak force-pushed the users/skatrak/target-passthrough-03-mlir-llvm branch from 36f00b9 to 0e07eda Compare October 29, 2024 15:34
@skatrak
Copy link
Author

skatrak commented Oct 29, 2024

OpenMPIRBuilder changes LGTM

No test added?

Thank you @Meinersbur for taking a look. I added a test that is basically copied from an existing one, since that one used to check all corner cases related to the trip count calculation. I don't know if it makes sense to keep both as they are, but let me know what you think.

@Meinersbur
Copy link

What I had in mind was a test where getHostEvalVars() is non-empty and thus the code-generation changes, or is first possible without crashing.

calculateCanonicalLoopTripCount was already tested indirectly through createCanonicalLoop. These are unit-tests, and since calculateCanonicalLoopTripCount is not available as its own "unit", I think you can remove the equivalent ones for createCanonicalLoop.

@skatrak skatrak force-pushed the users/skatrak/target-passthrough-03-mlir-llvm branch from 0e07eda to a3727a0 Compare October 30, 2024 12:07
@skatrak
Copy link
Author

skatrak commented Oct 30, 2024

What I had in mind was a test where getHostEvalVars() is non-empty and thus the code-generation changes, or is first possible without crashing.

Ah, that makes sense. Sorry for the misunderstanding! I just added an MLIR to LLVM IR translation test for these host evaluated clauses. I thought you were referring to the OMPIRBuilder refactoring because that's what you said you had reviewed 😅.

calculateCanonicalLoopTripCount was already tested indirectly through createCanonicalLoop. These are unit-tests, and since calculateCanonicalLoopTripCount is not available as its own "unit", I think you can remove the equivalent ones for createCanonicalLoop.

What I just did was to remove the CanonicalLoopBounds test and replace it with the CanonicalLoopTripCount test that does the same checks of the trip count, but without creating and looking at the canonical loop structure (which is already done by CanonicalLoopSimple above). Let me know if that's not what you meant, so that I can update it again.

Copy link

@Meinersbur Meinersbur left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks

Comment on lines +1446 to +1447
Value *TripCount = OMPBuilder.calculateCanonicalLoopTripCount(
Loc, StartVal, StopVal, StepVal, IsSigned, InclusiveStop);

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's actually pretty nice

@skatrak skatrak force-pushed the users/skatrak/target-passthrough-02-mlir branch from 595f1ae to f33594a Compare November 5, 2024 13:34
@skatrak skatrak force-pushed the users/skatrak/target-passthrough-03-mlir-llvm branch from a3727a0 to dec7a5e Compare November 5, 2024 13:34
Base automatically changed from users/skatrak/target-passthrough-02-mlir to amd-trunk-dev November 7, 2024 10:51
This patch updates the MLIR to LLVM IR lowering of `omp.target` to support
passing `num_teams`, `num_threads`, `thread_limit` and SPMD loop bounds through
the `host_eval` argument of `omp.target`.

This replaces the previous implementation where this information was directly
attached to the `omp.target` operation rather than captured to be used by the
corresponding nested operation.
@skatrak skatrak force-pushed the users/skatrak/target-passthrough-03-mlir-llvm branch from dec7a5e to 6e921c3 Compare November 7, 2024 10:52
@skatrak skatrak merged commit 5654efd into amd-trunk-dev Nov 7, 2024
2 of 4 checks passed
@skatrak skatrak deleted the users/skatrak/target-passthrough-03-mlir-llvm branch November 7, 2024 10:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants