Skip to content

[mlir] Added Convergent trait that matches LLVM's semantics #152358

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.

Sign up for GitHub

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[mlir] Added Convergent trait that matches LLVM's semantics #152358

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c6779ec
[mlir] Added `Convergent` trait that matches LLVM's semantics
npanchen Aug 6, 2025
058f9ae
Fixed `NVVM_PTXBuilder_Op` and removed bad comment
npanchen Aug 6, 2025
cb8f0a2
Only mark `nvvm.barrier0` as convergent
npanchen Aug 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,13 @@
#define NVVMIR_OPS

include "mlir/IR/EnumAttr.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/InferIntRangeInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
include "mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td"
include "mlir/Interfaces/InferIntRangeInterface.td"
include "mlir/Dialect/LLVMIR/LLVMTypes.td"

def LLVM_PointerGeneric : LLVM_PointerInAddressSpace<0>;
Expand Down Expand Up @@ -561,7 +562,7 @@ def NVVM_MBarrierTestWaitSharedOp : NVVM_Op<"mbarrier.test.wait.shared">,
// NVVM synchronization op definitions
//===----------------------------------------------------------------------===//

def NVVM_Barrier0Op : NVVM_Op<"barrier0"> {
def NVVM_Barrier0Op : NVVM_Op<"barrier0", [Convergent]> {
let assemblyFormat = "attr-dict";
string llvmBuilder = [{
createIntrinsicCall(
Expand Down
6 changes: 5 additions & 1 deletion mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -337,8 +337,12 @@ struct ReturnLike : public TraitBase<ConcreteType, ReturnLike> {
return success();
}
};
} // namespace OpTrait

// The Operation may not be made control-dependent on any additional values.
// See https://llvm.org/docs/ConvergentOperations.html for more details.
template <typename ConcreteType>
struct Convergent : public TraitBase<ConcreteType, Convergent> {};
Copy link
Collaborator

Choose a reason for hiding this comment

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

Is there any impact that MLIR is using structured control-flow in the support for convergence? In particular aren't token inserted during lowering to CFG to preserve loop structures for example?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

There's no direct impact on absence of this in MLIR, but this trait is a prerequisite to properly annotate functions that invoke such operations. OtherwiseLLVM will apply optimizations incorrectly.
Specifically, the use case I saw was:

fn callee() {
  nvvm.barrier0
}

fn caller() {
  if (condition) {
    call callee()
  }
}

callee has to be marked with convergent attribute, otherwise if jump threading is applied, the code will become invalid.

Copy link
Member

Choose a reason for hiding this comment

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

OtherwiseLLVM will apply optimizations incorrectly.

For each Op that MLIR generates, LLVM sets convergency. So LLVM won't do incorrect optimization for MLIR.

Copy link
Collaborator

Choose a reason for hiding this comment

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

For each Op that MLIR generates, LLVM sets convergency.

Are you sure we would do it inter-procedurally and annotate the actual call here?
Or are we relying on some attribute propagation later?

@npanchen : when mapping a structured control-flow to a CFG, just annotating things as "convergent" isn't enough I believe. See this example: https://youtu.be/_Z5DuiVCFAw?t=434 ; basically structured control-flow needs to say "something" about the reconvergence property and the management of anchors (the video is amazing for explaining all this I think!)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

the convergent property is only set on the intrinsic, not on arbitrary function that calls to that intrinsic. Since LLVM does not infer that property automatically, that can result to undesired optimization.
Just to highlight that JT works differently for the case I described above: https://godbolt.org/z/9T4d7MM58. To note, the IR there looks ok, but the problem I observed is similar, but just has more complex CF.

Also, the MLIR trait seems generally useful to prevent, say, loop multiversioning by trip count if loop contains such operation.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@npanchen : when mapping a structured control-flow to a CFG, just annotating things as "convergent" isn't enough I believe. See this example: https://youtu.be/_Z5DuiVCFAw?t=434 ; basically structured control-flow needs to say "something" about the reconvergence property and the management of anchors (the video is amazing for explaining all this I think!)

Thanks, that's quite interesting video. I actually was unable to find when convergence.loop/convergence.entry/convergence.anchor are generated. Are they even still in use ?
If they are, yeah, that might be more interesting from design point of view. As of know I can only think about adding RecursivelyConvergent trait. But not sure that's a good idea for functions.

completely side note: there's also noduplicate property, which seems not to be used a lot, but it's not supported by MLIR and it does directly impact inliner.

Copy link
Collaborator

Choose a reason for hiding this comment

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

They are emitted by clang, search for shouldEmitConvergenceTokens() there, you'll see the kind of things I have in mind for our structured-to-cfg needs (and actually: a semantics definition of convergence guarantee on our structured control-flow)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I see. I was not searching right.

} // namespace OpTrait
} // namespace mlir

//===----------------------------------------------------------------------===//
Expand Down
3 changes: 3 additions & 0 deletions mlir/include/mlir/Interfaces/ControlFlowInterfaces.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -511,4 +511,7 @@ def ReturnLike : TraitList<[
>
]>;

// Op is "convergent".
def Convergent : NativeOpTrait<"Convergent">;

#endif // MLIR_INTERFACES_CONTROLFLOWINTERFACES
5 changes: 5 additions & 0 deletions mlir/test/lib/Dialect/Test/TestOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2113,6 +2113,11 @@ def TestTypeChangerOp : TEST_Op<"type_changer">,
def TestValidOp : TEST_Op<"valid", [Terminator]>,
Arguments<(ins Variadic<AnyType>)>;

def TestConvergentOp : TEST_Op<"convergent", [Convergent]> {
let arguments = (ins AnyType);
let results = (outs AnyType);
}
Copy link
Collaborator

Choose a reason for hiding this comment

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

What is this doing right now? Without an actual test exercising this op, adding an op isn't useful.


def TestMergeBlocksOp : TEST_Op<"merge_blocks"> {
let summary = "merge_blocks operation";
let description = [{
Expand Down