This is an archive of the discontinued LLVM Phabricator instance.

[MLIR] Add std.atomic_rmw op
ClosedPublic

Authored by flaub on Feb 11 2020, 6:05 AM.

Details

Summary

The RFC for this op is here: https://llvm.discourse.group/t/rfc-add-std-atomic-rmw-op/489

The std.atmomic_rmw op provides a way to support read-modify-write
sequences with data race freedom. It is intended to be used in the lowering
of an upcoming affine.atomic_rmw op which can be used for reductions.

A lowering to LLVM is provided with 2 paths:

  • Simple patterns: llvm.atomicrmw
  • Everything else: llvm.cmpxchg

Diff Detail

Unit TestsFailed

Event Timeline

flaub created this revision.Feb 11 2020, 6:05 AM
Herald added a project: Restricted Project. · View Herald TranscriptFeb 11 2020, 6:05 AM
flaub updated this revision to Diff 243841.Feb 11 2020, 6:09 AM
  • Remove stdx.
jbruestle requested changes to this revision.Feb 11 2020, 10:48 AM
jbruestle added inline comments.
mlir/include/mlir/Dialect/StandardOps/Ops.td
230

I think you just mean 'block argument' not induction variable here (and below in a few places), since it's not iterating over anything.

252

Maybe rename getInitialValue(), or getLoadedValue() or something similar?

This revision now requires changes to proceed.Feb 11 2020, 10:48 AM
flaub updated this revision to Diff 243960.Feb 11 2020, 12:43 PM
  • Review updates
flaub marked 2 inline comments as done.Feb 11 2020, 12:43 PM
flaub updated this revision to Diff 243962.Feb 11 2020, 12:45 PM
  • Remove iv
flaub updated this revision to Diff 243963.Feb 11 2020, 12:48 PM
  • Remove iv
flaub added a comment.Feb 11 2020, 1:01 PM

Review addressed in latest push.

jbruestle accepted this revision.Feb 11 2020, 1:16 PM
This revision is now accepted and ready to land.Feb 11 2020, 1:16 PM
rriddle requested changes to this revision.Feb 11 2020, 1:25 PM
rriddle added inline comments.
mlir/include/mlir/Dialect/StandardOps/Ops.td
227

typo: indicies -> indices

239

Wrap these in a mlir code block

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2518

static functions should be in the global namespace and marked as static. Only classes should be placed within anonymous namespaces.

2551

Drop trivial braces.

2574

Same here.

2630

Top-level comments should be ///

2688

Don't create SmallVectors for things like this. Use ArrayRef or arrays.

2716

Same here.

mlir/lib/Dialect/StandardOps/Ops.cpp
3019

Remove trivial braces.

I would have expected that this could be covered by ODS constraints.

This revision now requires changes to proceed.Feb 11 2020, 1:25 PM
flaub updated this revision to Diff 244020.Feb 11 2020, 3:40 PM
flaub marked 10 inline comments as done.
  • Review updates
flaub marked 3 inline comments as not done.Feb 11 2020, 3:40 PM
flaub added inline comments.
mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2518

I'm OK with this, but I'm unfamiliar with this convention. What's the purpose behind having functions being static instead of being in the anonymous namespace? I was always under the impression that the two were functionally equivalent and that the more 'C++' way was to use anonymous namespaces.

Also, should I close out the namespace here and then add the static functions and then re-open the anonymous namespace? Or would it make sense to move these up above?

mlir/lib/Dialect/StandardOps/Ops.cpp
3019

Did you have a specific trait in mind? I'm comparing the parent's element type to the op's operand type. I didn't see anything in OpBase.td at first glance.

rriddle added inline comments.Feb 11 2020, 3:51 PM
mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2518
mlir/lib/Dialect/StandardOps/Ops.cpp
3019

Hmmm, I thought there was one for this already. I think a lot of the current usages are abusing mlir::getElementTypeOrSelf.

rriddle added inline comments.Feb 11 2020, 3:52 PM
mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2518

Also, should I close out the namespace here and then add the static functions and then re-open the anonymous namespace? Or would it make sense to move these up above?

Whichever one makes sense. You can close the namespace or hoist the functions.

flaub updated this revision to Diff 244026.Feb 11 2020, 4:00 PM
flaub marked an inline comment as not done.
  • Drop trivial braces
  • Trivial braces
flaub marked an inline comment as done.Feb 11 2020, 4:06 PM
flaub added inline comments.
mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2518

Thanks for the link!

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2656

This seems like it'd be better a loop.while + progressive lowering?

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir
869

So this is interesting to me.
Weren't you and/or @jbruestle advocating that we should have the reduction op encoded as an attribute in the case or affine.parallel_for with reduction semantics?
It seems a very similar scenario to me, I'd be interested of where you draw the distinction between "encoded as an attribute" and just use a region?

flaub marked 2 inline comments as done.Feb 11 2020, 5:12 PM
flaub added inline comments.
mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2656

I agree that it'd be a lot nicer to write the loop at a higher level, but it wasn't clear to me where this would go. Also, the determination of whether to use a cmpxchg or atomic_rmw is really specific to LLVM lowering. So we wouldn't want to use a loop in the case of simple bodies that can map to a single intrinsic/op at the lower level. I'm also thinking about how we will want to have a lowering from std to SPIR-V or OpenMP, in which case a loop may or may not make sense for those lowerings.

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir
869

I think we initially thought having a closed attribute would be good, but it seemed that providing the ability to lower arbitrary reductions into cmpxchg wasn't too hard to do and it was easy enough to identify these simple cases that do lower to a single intrinsic. Our current plan is to still use an enum at the top level (the tile dialect), but then use a region for affine and below. The upcoming affine.atomic_rmw should basically mirror the standard one in regards to region vs attribute.

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2656

Not for this revision but in the future there are some nice tricks we could play here.
Your Conversion could very well decide to introduce a higher level loop construct, rewrite the region into that and let the rewrite infrastructure stitch the pieces together.
In other words, I wasn't advocating that the atomic behavior should leak into the other targets, but that you could decide during lowering to introduce a higher level construct that is implemented and tested independently.

Great work! I only have some nits.

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2521

I'd appreciate some documentation on this function and below.

2618

Type conversion may fail and return null, please check the result.

2621

Nit: rewriter.replaceOpWithNewOp<LLVM::AtomicRMWOp>(op, resultType, ...);

2656

This is similar to the discussion @nicolasvasilache and I had about memory copies, and is worse discussing in general. One practical thing I'd like to point out: we need to make sure we don't introduce cyclic library dependencies, and it may be tricky.

2713

Nit: normally, single-result ops should be convertible to Value so you shouldn't be needing the .res() part

flaub updated this revision to Diff 244317.Feb 12 2020, 6:32 PM
flaub marked 4 inline comments as done.
  • Remove stdx.
  • Review updates
  • Remove iv
  • Review updates
  • Drop trivial braces
  • Trivial braces
  • Address feedback
flaub added a comment.Feb 12 2020, 6:33 PM

@rriddle Are there anymore blockers that need to be addressed?

ftynse accepted this revision.Feb 14 2020, 1:55 AM

LGTM when extra documentation is added as discussed on Discourse.

mlir/include/mlir/Dialect/StandardOps/Ops.td
235

Could you please describe the restrictions on the body contents of the atomic region as discussed in the RFC?

rriddle accepted this revision.Feb 14 2020, 9:28 PM

LGTM after comments are resolved.

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2512

Can you document this struct please.

2519

Use /// for all top-level comments. Here and below.

2738

For single element things you should be able to pass the values directly. Is there a problem you are running into doing that?

This revision is now accepted and ready to land.Feb 14 2020, 9:28 PM
flaub updated this revision to Diff 245584.EditedFeb 19 2020, 11:08 PM

Simplified design based on RFC feedback

flaub updated this revision to Diff 245585.Feb 19 2020, 11:12 PM
  • Comments
flaub updated this revision to Diff 245586.Feb 19 2020, 11:13 PM
  • Fix example
rriddle requested changes to this revision.Feb 19 2020, 11:51 PM

Thanks for the update Frank! Added a few comments.

mlir/include/mlir/Dialect/StandardOps/Ops.td
262

Should this be a more constrained type, like 'IntegerOrFloatLike'?

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2792

Can we keep this ordered?

2890

This is invalid, *all* IR mutations need to go through the rewriter. Some pattern drivers, like DialectConversion which is being used here, will undo transformations if something goes wrong. If something is done outside of the rewriter, this leads to invalid code/crashes. Seems like you want to do rewriter.replaceOp here.

mlir/lib/Dialect/StandardOps/Ops.cpp
2959

Could you switch to the declarative format? It will format the enum as a string instead of a keyword for now, but that is worth remove all of this parsing code.

This revision now requires changes to proceed.Feb 19 2020, 11:51 PM
flaub marked 4 inline comments as done.Feb 20 2020, 12:13 AM
flaub added inline comments.
mlir/include/mlir/Dialect/StandardOps/Ops.td
262

OK, I suppose that will work here for now. I could see this being expanded later if the lowering supported more types (which in theory it could by lowering to cmpxchg or others).

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp
2792

Will do.

2890

OK, thanks for clarifying. I'll try to use rewriter.replaceOp instead.

mlir/lib/Dialect/StandardOps/Ops.cpp
2959

Cool, I will try that, thanks (sounds like a worthwhile tradeoff).

flaub updated this revision to Diff 245598.Feb 20 2020, 12:57 AM
  • Review feedback
flaub updated this revision to Diff 245599.Feb 20 2020, 1:03 AM
  • Fix example
flaub updated this revision to Diff 246342.Feb 24 2020, 4:32 PM
  • Merge branch 'master' into arcpatch-D74401
  • Merge branch 'master' into arcpatch-D74401
  • Update with master
flaub added a comment.Feb 24 2020, 4:33 PM

@rriddle Thanks for the review, does this look better now?

rriddle accepted this revision.Feb 24 2020, 4:39 PM

Thanks!

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
273 ↗(On Diff #246342)

nit: Can you keep these two lines aligned?

This revision is now accepted and ready to land.Feb 24 2020, 4:39 PM
This revision was automatically updated to reflect the committed changes.