Page MenuHomePhabricator

[OpenMP][mlir] Lowering for omp.atomic.update
ClosedPublic

Authored by shraiysh on Feb 12 2022, 11:33 PM.

Details

Summary

This patch adds lowering from omp.atomic.update to LLVM IR. Whenever a
special LLVM IR instruction is available for the operation, atomicrmw
instruction is emitted, otherwise a compare-exchange loop based update
is emitted.

Depends on D119522

Diff Detail

Event Timeline

shraiysh created this revision.Feb 12 2022, 11:33 PM
shraiysh requested review of this revision.Feb 12 2022, 11:33 PM
Herald added projects: Restricted Project, Restricted Project. · View Herald Transcript
ftynse requested changes to this revision.Feb 13 2022, 3:21 AM
ftynse added inline comments.
mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
954–955

Please do not commit commented-out code.

968

This comment is not helpful, we can see it is boilerplate. I would rather have most of it factored out into a function like getOpenMPLoc(IRBuilderBase &).

978

There doesn't seem to be a check for the region to only have two operations (binop and terminator), and this translation does not account for it but just ignores the remaining operations.

980

There doesn't seem to be any check that updateOp does have operands at all anywhere.

989

Please use better names than x and X in the same function, these are extremely hard to keep in one's head for more than a couple of lines.

993

Please consider a better name than ao.

998

MLIR uses camelBack names, IRB isn't something that it would use.

mlir/test/Target/LLVMIR/openmp-llvm.mlir
936

Note that MLIR tends to use "lowering" for _conversions_ between dialects rather than for _translations_ between MLIR and other IRs.

This revision now requires changes to proceed.Feb 13 2022, 3:21 AM
shraiysh updated this revision to Diff 408270.Feb 13 2022, 8:38 AM
shraiysh marked 7 inline comments as done.

Thanks for the review @ftynse. Addressed comments.

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

For the function getOpenMPLoc, we would need opInst and moduleTranslation too. Plus, llvm::OpenMPIRBuilder::LocationDescription doesn't have a copy constructor, so that would also have to be passed. Also, the ompBuilder is used later in this function, so that cannot be pushed into the function. At this point, a function like the following doesn't feel like it is worth it. I have removed the comment as suggested. Let me know if you have any suggestions for an alternative.

getOpenMPLoc(Operation &opInst, llvm::IRBuilderBase &builder, LLVM::ModuleTranslation moduleTranslation, llvm::OpenMPIRBuilder::LocationDescription &ompLoc)
ftynse accepted this revision.Feb 14 2022, 1:09 AM
ftynse added inline comments.
mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
968

I don't have a problem with multi-argument signatures. ModuleTranslation needs to be passed by reference as well. Can't we just make LocationDescription copy-constructible (or movable)?

978

Please add a test for user-visible error messages.

This revision is now accepted and ready to land.Feb 14 2022, 1:09 AM
shraiysh updated this revision to Diff 408859.EditedFeb 15 2022, 7:03 AM
shraiysh marked 3 inline comments as done.

Handle multiple operations in update region.
Added tests for same.
Addressed comments.
Rebase with main.

shraiysh requested review of this revision.Feb 15 2022, 7:05 AM

Requesting review again, because added translation for multiple operations in update region.

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

Now the constructor directly accepts builder, so this is not needed. Thanks for the discussion.

shraiysh updated this revision to Diff 410228.Feb 20 2022, 11:36 PM

Rebase with main. Ping!

shraiysh updated this revision to Diff 411746.Feb 28 2022, 12:21 AM

Rebase with main. Ping!

peixin added inline comments.Mar 1 2022, 2:56 AM
mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
961

OpenMP update -> OpenMP atomic update

mlir/test/Target/LLVMIR/openmp-llvm.mlir
1028

I don't think you should process the lowering of the "expr" in MLIR stage. For different programing languages, the expr may have different structure. Can you try to analyze the "expr" when lowering parse-tree to MLIR? That is, the input for the "expr" in atomic update is always one mlir value and the LLVM IR generated for analyzing the "expr" is outside the atomic update region. Another reason for doing this is that "expr" cannot access the storage location designated by x, and the operations inside the atomic update region will be very clean. That is, you will always only have two instructions in atomic update region, one for update "xval", the other one for yield the "newval".

peixin added inline comments.Mar 1 2022, 3:13 AM
mlir/test/Target/LLVMIR/openmp-llvm.mlir
1028
shraiysh updated this revision to Diff 412102.Mar 1 2022, 7:38 AM

Thanks for the review @peixin. Addressed comments.

mlir/test/Target/LLVMIR/openmp-llvm.mlir
1028

The lowering of %expr is not handled in the MLIR region. It is assumed to be a single lowered value, with lowering handled in Flang as you suggested. The idea behind allowing more than one operations was to have the following in MLIR

omp.atomic.update %x {
^bb0(%xval):
  %xnew = f(%xval, %expr)
  omp.yield %xnew
}

when we cannot express x = f(x, expr) in MLIR with a single instruction (for example f(x,expr) = (x^expr)^expr)) as we might not have the appropriate single op in the dialect (which can be anything including but not limited to LLVM IR Dialect and FIR Dialect) and that should be okay. In this case, we cannot have f(x, expr) such that the operation can be done in one op. By allowing multiple operations in the region, we provide the liberty of handling this in multiple ops. The generated LLVM IR will work appropriately. Another thing is that the MLIR Operation cannot depend on Flang, because this operation will be used by multiple dialects, so it has to be generic enough to support simple and complicated Dialects. Should I put a comment about this in the code itself?

shraiysh updated this revision to Diff 412103.Mar 1 2022, 7:40 AM

Rebase with main.

shraiysh marked an inline comment as done.Mar 1 2022, 7:41 AM
peixin added inline comments.Mar 1 2022, 7:24 PM
mlir/test/Target/LLVMIR/openmp-llvm.mlir
1028

However, this is for OpenMP dialect. Currently, OpenMP standard only defines C, C++, and fortran, and all the update statements for them can be transformed as one single instruction and one omp.yield operation. f(x,expr) = (x^expr)^expr is obviously not allowed in OpenMP standard. If someone want to use this generic atomic update, they should implement one in other dialects in somewhere else, instead of OpenMP dialect here. What you currently implemented is one generic atomic update, not restricted to OpenMP atomic update.

Herald added a project: Restricted Project. · View Herald TranscriptMar 1 2022, 7:24 PM
shraiysh added inline comments.Mar 1 2022, 8:31 PM
mlir/test/Target/LLVMIR/openmp-llvm.mlir
1028

can be transformed as one single instruction and one omp.yield operation

Yes, in the LLVM IR Dialect. We cannot say the same for all the other dialects with certainty. (or can we? Feel free to correct me if this sounds absurd.) Maybe we can put a check during lowering to LLVM IR for having exactly two operations, but enforcing it in the operation verifier itself is not entirely correct as the internal op can be in any dialect (In some example dialect, a few atomic update statements for some types might need more than a single op in that dialect).

What you currently implemented is one generic atomic update, not restricted to OpenMP atomic update.

That is correct, because only the dialect handling the operation knows how to update it appropriately (including how many ops to take for that update).

peixin added inline comments.Mar 1 2022, 10:34 PM
mlir/test/Target/LLVMIR/openmp-llvm.mlir
1028

I think we can say that. Actually, OpenMP standard says that. Please notice that OpenMPToLLVMIRTranslation.cpp is for OpenMP. In semantics, having more than one single op is not correct. If someone needs it, they should not use OpenMP stomic update since it is not allowed in OpenMP atomic update.

peixin added inline comments.Mar 1 2022, 10:43 PM
mlir/test/Target/LLVMIR/openmp-llvm.mlir
1028

OK. I noticed that you changed the omp.atomic.update dialect. https://github.com/llvm/llvm-project/blob/0e38b295435bf49c21e9ff97354a855ac8c78c7e/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td#L637-L638 Here, you referred to OpenMP standard, but you make it too generic so that this opertion is beyond the OpenMP atomic update. I think we need to discuss about the change in D119522. @kiranchandramohan What do you think?

shraiysh updated this revision to Diff 412633.Mar 3 2022, 1:46 AM

Addressed reviewer comments. Rebase with main.

peixin added a comment.Mar 6 2022, 6:44 PM

The lowering for memory-order-clause and hint clause will be in future work, right?

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

Can you add some comments to describe your design method to tell others why you don't do it similar to atomic read/write? You have expalined this in slack. Can you move those info here? Also, comment that the input is checked in semantic analysis for each programming language.

980

Should this be one if condition or one assert? If I understand correctly, the input not satisfying the if condition is not one standard OpenMP atomic update statement. Right?

983

Will this return BAD_BINOP for int128, fp128, x86_fp80? They seems not to be supported in AtomicRMWInst. Can you add one test case for one of them to represent the corner case?

Also, what about the complex data type?

shraiysh updated this revision to Diff 414034.Mar 9 2022, 1:17 AM
shraiysh marked 8 inline comments as done.

Thanks for the review @peixin. Addressed comments.

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

Sure. Added in the description for omp.atomic.update operation in OpenMPOps.td.

980

Yes, thank you. Added.

983

I have added a TODO in OMPIRBuilder.cpp for this. It goes into error at the moment. This is a bug in OmpIRBuilder.

// TODO: handle the case where XElemTy is not byte-sized or not a power of 2.
peixin accepted this revision.Mar 10 2022, 3:12 AM
This revision is now accepted and ready to land.Mar 10 2022, 3:12 AM
shraiysh updated this revision to Diff 414335.Mar 10 2022, 3:58 AM

Rebase with main

This revision was landed with ongoing or failed builds.Mar 10 2022, 4:59 AM
This revision was automatically updated to reflect the committed changes.