From b6d28d1205981480a2613a0f3217ca006ebd699d Mon Sep 17 00:00:00 2001 From: Krzysztof Drewniak Date: Wed, 10 Sep 2025 00:00:35 +0000 Subject: [PATCH 1/4] [mlir][LLVMIR] Support memory model relaxation annotations (MMRA) This commit adds support for exportind and importing MMRA data in the LLVM dialect. MMRA is a potentilly-discardable piece of metadata that can be placed on any operation that touches memory (fences, loads, stores, atomics, and intrinsics that operate on memory). It includes one (technically zero) ome more prefix:suffix string pairs which indicate ways in which the LLVM memory model can be relaxed for these annotations. At the MLIR level, each tag is represented with a `#llvm.mmra_tag<"prefix":"suffex">` attribute, and the MMRA metadata as a whole is represented as a discardable llvm.mmra attribute. (This discardability both allows us to transparently enable MMRA for wrapper dialects like ROCDL and ensures that MLIR passes which don't know about MMRA combining will, conservatively, discard the annotations, per the LLVM spec). --- .../mlir/Dialect/LLVMIR/LLVMAttrDefs.td | 41 ++++++++++++++++ .../mlir/Dialect/LLVMIR/LLVMDialect.td | 1 + .../LLVMIR/LLVMIRToLLVMTranslation.cpp | 31 +++++++++++- .../LLVMIR/LLVMToLLVMIRTranslation.cpp | 47 +++++++++++++++++++ mlir/test/Dialect/LLVMIR/mmra.mlir | 29 ++++++++++++ .../Target/LLVMIR/Import/metadata-mmra.ll | 22 +++++++++ mlir/test/Target/LLVMIR/mmra.mlir | 35 ++++++++++++++ 7 files changed, 205 insertions(+), 1 deletion(-) create mode 100644 mlir/test/Dialect/LLVMIR/mmra.mlir create mode 100644 mlir/test/Target/LLVMIR/Import/metadata-mmra.ll create mode 100644 mlir/test/Target/LLVMIR/mmra.mlir diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td index ac99b8aba073a..9980e1174d694 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td @@ -1232,6 +1232,47 @@ def LLVM_TBAATagArrayAttr let constBuilderCall = ?; } +//===----------------------------------------------------------------------===// +// Memory Model Relaxation Annations (mmra) Attributes +//===----------------------------------------------------------------------===// + +//===----------------------------------------------------------------------===// +// MMRATagAttr - single MMRA tag +//===----------------------------------------------------------------------===// + +def LLVM_MMRATagAttr : LLVM_Attr<"MMRATag", "mmra_tag"> { + let parameters = (ins + StringRefParameter<>:$prefix, + StringRefParameter<>:$suffix + ); + + let summary = "MLIR wrapper around a prefix:suffix MMRA tag"; + + let description = [{ + Defines a single memory model relaxation annotation (MMRA) entry + with prefix `$prefix` and suffix `$suffix`. This corresponds directly + to a LLVM `!{prefix, suffix}` metadata tuple, which is often written + `prefix:shuffix` as shorthand. + + Example: + ```mlir + #mmra_tag = #llvm.mmmra_tag<"amdgpu-synchronize-as":"local"> + #mmra_tag1 = #llvm.mmra_tag<"foo":"bar"> + ``` + + See the following link for more details: + https://llvm.org/docs/MemoryModelRelaxationAnnotations.html + }]; + + let assemblyFormat = "`<` $prefix `` `:` `` $suffix `>`"; + + let genMnemonicAlias = 1; +} + +def LLVM_MMRATagArrayAttr : TypedArrayAttrBase< + LLVM_MMRATagAttr, + LLVM_MMRATagAttr.summary # " array">; + //===----------------------------------------------------------------------===// // ConstantRangeAttr //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td index ab0462f945a33..d2d71318a6118 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td @@ -36,6 +36,7 @@ def LLVM_Dialect : Dialect { static StringRef getIdentAttrName() { return "llvm.ident"; } static StringRef getModuleFlags() { return "llvm.module.flags"; } static StringRef getCommandlineAttrName() { return "llvm.commandline"; } + static StringRef getMmraAttrName() { return "llvm.mmra"; } /// Names of llvm parameter attributes. static StringRef getAlignAttrName() { return "llvm.align"; } diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp index 8d5d7f9b649f2..9a548cf77e0f5 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp @@ -11,6 +11,7 @@ //===----------------------------------------------------------------------===// #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.h" +#include "mlir/Dialect/LLVMIR/LLVMAttrs.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/LLVMIR/LLVMInterfaces.h" #include "mlir/Support/LLVM.h" @@ -21,6 +22,7 @@ #include "llvm/IR/InlineAsm.h" #include "llvm/IR/Instructions.h" #include "llvm/IR/IntrinsicInst.h" +#include "llvm/IR/MemoryModelRelaxationAnnotations.h" using namespace mlir; using namespace mlir::LLVM; @@ -88,6 +90,7 @@ static ArrayRef getSupportedMetadataImpl(llvm::LLVMContext &context) { llvm::LLVMContext::MD_alias_scope, llvm::LLVMContext::MD_dereferenceable, llvm::LLVMContext::MD_dereferenceable_or_null, + llvm::LLVMContext::MD_mmra, context.getMDKindID(vecTypeHintMDName), context.getMDKindID(workGroupSizeHintMDName), context.getMDKindID(reqdWorkGroupSizeMDName), @@ -212,6 +215,31 @@ static LogicalResult setDereferenceableAttr(const llvm::MDNode *node, return success(); } +/// Convert the given MMRA metadata (either an MMRA tag or an array of rhem) +/// into corresponding MLIR attributes and set them on the given operation as a +/// discardable `llvm.mmra` attribute. +static LogicalResult setMmraAttr(llvm::MDNode *node, Operation *op, + LLVM::ModuleImport &moduleImport) { + llvm::MMRAMetadata wrapper(node); + if (wrapper.empty()) { + return success(); + } + MLIRContext *ctx = op->getContext(); + Attribute mlirMmra; + if (wrapper.size() == 1) { + auto [prefix, suffix] = *wrapper.begin(); + mlirMmra = LLVM::MMRATagAttr::get(ctx, prefix, suffix); + } else { + SmallVector tags; + for (auto [prefix, suffix] : wrapper) { + tags.push_back(LLVM::MMRATagAttr::get(ctx, prefix, suffix)); + } + mlirMmra = ArrayAttr::get(ctx, tags); + } + op->setAttr(LLVMDialect::getMmraAttrName(), mlirMmra); + return success(); +} + /// Converts the given loop metadata node to an MLIR loop annotation attribute /// and attaches it to the imported operation if the translation succeeds. /// Returns failure otherwise. @@ -432,7 +460,8 @@ class LLVMDialectLLVMIRImportInterface : public LLVMImportDialectInterface { return setDereferenceableAttr( node, llvm::LLVMContext::MD_dereferenceable_or_null, op, moduleImport); - + if (kind == llvm::LLVMContext::MD_mmra) + return setMmraAttr(node, op, moduleImport); llvm::LLVMContext &context = node->getContext(); if (kind == context.getMDKindID(vecTypeHintMDName)) return setVecTypeHintAttr(builder, node, op, moduleImport); diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp index fd8463ad1a8e2..ddf9b16b7b552 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp @@ -24,6 +24,8 @@ #include "llvm/IR/Instructions.h" #include "llvm/IR/MDBuilder.h" #include "llvm/IR/MatrixBuilder.h" +#include "llvm/IR/MemoryModelRelaxationAnnotations.h" +#include "llvm/Support/LogicalResult.h" using namespace mlir; using namespace mlir::LLVM; @@ -723,6 +725,43 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder, return failure(); } +static LogicalResult +amendOperationImpl(Operation &op, ArrayRef instructions, + NamedAttribute attribute, + LLVM::ModuleTranslation &moduleTranslation) { + StringRef name = attribute.getName(); + if (name == LLVMDialect::getMmraAttrName()) { + SmallVector tags; + if (auto oneTag = dyn_cast(attribute.getValue())) { + tags.emplace_back(oneTag.getPrefix(), oneTag.getSuffix()); + } else if (auto manyTags = dyn_cast(attribute.getValue())) { + for (auto a : manyTags) { + auto tag = dyn_cast(a); + if (tag) { + tags.emplace_back(tag.getPrefix(), tag.getSuffix()); + } else { + return op.emitOpError( + "MMRA annotations array contains value that isn't an MMRA tag"); + } + } + } else { + return op.emitOpError( + "llvm.mmra is something other than an MMRA tag or an array of them"); + } + llvm::MDTuple *mmraMd = + llvm::MMRAMetadata::getMD(moduleTranslation.getLLVMContext(), tags); + if (!mmraMd) { + // Empty list, canonicalizes to nothing + return success(); + } + for (llvm::Instruction *inst : instructions) { + inst->setMetadata(llvm::LLVMContext::MD_mmra, mmraMd); + } + return success(); + } + return success(); +} + namespace { /// Implementation of the dialect interface that converts operations belonging /// to the LLVM dialect to LLVM IR. @@ -738,6 +777,14 @@ class LLVMDialectLLVMIRTranslationInterface LLVM::ModuleTranslation &moduleTranslation) const final { return convertOperationImpl(*op, builder, moduleTranslation); } + + /// Handle some metadata that is represented as a discardable attribute. + LogicalResult + amendOperation(Operation *op, ArrayRef instructions, + NamedAttribute attribute, + LLVM::ModuleTranslation &moduleTranslation) const final { + return amendOperationImpl(*op, instructions, attribute, moduleTranslation); + } }; } // namespace diff --git a/mlir/test/Dialect/LLVMIR/mmra.mlir b/mlir/test/Dialect/LLVMIR/mmra.mlir new file mode 100644 index 0000000000000..95da9666d053c --- /dev/null +++ b/mlir/test/Dialect/LLVMIR/mmra.mlir @@ -0,0 +1,29 @@ +// RUN: mlir-opt %s -split-input-file --verify-roundtrip --mlir-print-local-scope | FileCheck %s + +// CHECK-LABEL: llvm.func @native +// CHECK: llvm.load +// CHECK-SAME: llvm.mmra = #llvm.mmra_tag<"foo":"bar"> +// CHECK: llvm.fence +// CHECK-SAME: llvm.mmra = [#llvm.mmra_tag<"amdgpu-synchronize-as":"local">, #llvm.mmra_tag<"foo":"bar">] +// CHECK: llvm.store +// CHECK-SAME: llvm.mmra = #llvm.mmra_tag<"foo":"bar"> + +#mmra_tag = #llvm.mmra_tag<"foo":"bar"> + +llvm.func @native(%x: !llvm.ptr, %y: !llvm.ptr) { + %0 = llvm.load %x {llvm.mmra = #mmra_tag} : !llvm.ptr -> i32 + llvm.fence syncscope("workgroup-one-as") release + {llvm.mmra = [#llvm.mmra_tag<"amdgpu-synchronize-as":"local">, #mmra_tag]} + llvm.store %0, %y {llvm.mmra = #llvm.mmra_tag<"foo":"bar">} : i32, !llvm.ptr + llvm.return +} + +// ----- + +// CHECK-LABEL: llvm.func @foreign_op +// CHECK: rocdl.load.to.lds +// CHECK-SAME: llvm.mmra = #llvm.mmra_tag<"fake":"example"> +llvm.func @foreign_op(%g: !llvm.ptr<1>, %l: !llvm.ptr<3>) { + rocdl.load.to.lds %g, %l, 4, 0, 0 {llvm.mmra = #llvm.mmra_tag<"fake":"example">} : !llvm.ptr<1> + llvm.return +} diff --git a/mlir/test/Target/LLVMIR/Import/metadata-mmra.ll b/mlir/test/Target/LLVMIR/Import/metadata-mmra.ll new file mode 100644 index 0000000000000..180d438eca70e --- /dev/null +++ b/mlir/test/Target/LLVMIR/Import/metadata-mmra.ll @@ -0,0 +1,22 @@ +; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s + +; CHECK-DAG: #[[$MMRA0:.+]] = #llvm.mmra_tag<"foo":"bar"> +; CHECK-DAG: #[[$MMRA1:.+]] = #llvm.mmra_tag<"amdgpu-synchronize-as":"local"> + +; CHECK-LABEL: llvm.func @native +define void @native(ptr %x, ptr %y) { + ; CHECK: llvm.load + ; CHECK-SAME: llvm.mmra = #[[$MMRA0]] + %v = load i32, ptr %x, align 4, !mmra !0 + ; CHECK: llvm.fence + ; CHECK-SAME: llvm.mmra = [#[[$MMRA0]], #[[$MMRA1]]] + fence syncscope("workgroup-one-as") release, !mmra !2 + ; CHECK: llvm.store {{.*}}, !llvm.ptr{{$}} + store i32 %v, ptr %y, align 4, !mmra !3 + ret void +} + +!0 = !{!"foo", !"bar"} +!1 = !{!"amdgpu-synchronize-as", !"local"} +!2 = !{!0, !1} +!3 = !{} diff --git a/mlir/test/Target/LLVMIR/mmra.mlir b/mlir/test/Target/LLVMIR/mmra.mlir new file mode 100644 index 0000000000000..5864e0e0759e6 --- /dev/null +++ b/mlir/test/Target/LLVMIR/mmra.mlir @@ -0,0 +1,35 @@ +// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s + +// CHECK-LABEL: define void @native +// CHECK: load +// CHECK-SAME: !mmra ![[MMRA0:[0-9]+]] +// CHECK: fence +// CHECK-SAME: !mmra ![[MMRA1:[0-9]+]] +// CHECK: store {{.*}}, align 4{{$}} + +#mmra_tag = #llvm.mmra_tag<"foo":"bar"> + +llvm.func @native(%x: !llvm.ptr, %y: !llvm.ptr) { + %0 = llvm.load %x {llvm.mmra = #mmra_tag} : !llvm.ptr -> i32 + llvm.fence syncscope("workgroup-one-as") release + {llvm.mmra = [#llvm.mmra_tag<"amdgpu-synchronize-as":"local">, #mmra_tag]} + llvm.store %0, %y {llvm.mmra = []} : i32, !llvm.ptr + llvm.return +} + +// Actual MMRA metadata +// CHECK-DAG: ![[MMRA0]] = !{!"foo", !"bar"} +// CHECK-DAG: ![[MMRA_PART0:[0-9]+]] = !{!"amdgpu-synchronize-as", !"local"} +// CHECK-DAG: ![[MMRA1]] = !{![[MMRA_PART0]], ![[MMRA0]]} + +// ----- + +// CHECK-LABEL: define void @foreign_op +// CHECK: call void @llvm.amdgcn.load.to.lds +// CHECK-SAME: !mmra ![[MMRA0:[0-9]+]] +llvm.func @foreign_op(%g: !llvm.ptr<1>, %l: !llvm.ptr<3>) { + rocdl.load.to.lds %g, %l, 4, 0, 0 {llvm.mmra = #llvm.mmra_tag<"fake":"example">} : !llvm.ptr<1> + llvm.return +} + +// CHECK: ![[MMRA0]] = !{!"fake", !"example"} From ac028712e3f91002edabb0f4aa8dc4dcd424e079 Mon Sep 17 00:00:00 2001 From: Krzysztof Drewniak Date: Wed, 10 Sep 2025 11:05:16 -0500 Subject: [PATCH 2/4] Apply style nits from review Co-authored-by: Tobias Gysi --- .../LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp | 9 ++++----- .../LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp | 5 ++--- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp index 9a548cf77e0f5..736b1fe5b850c 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp @@ -215,15 +215,15 @@ static LogicalResult setDereferenceableAttr(const llvm::MDNode *node, return success(); } -/// Convert the given MMRA metadata (either an MMRA tag or an array of rhem) +/// Convert the given MMRA metadata (either an MMRA tag or an array of them) /// into corresponding MLIR attributes and set them on the given operation as a /// discardable `llvm.mmra` attribute. static LogicalResult setMmraAttr(llvm::MDNode *node, Operation *op, LLVM::ModuleImport &moduleImport) { llvm::MMRAMetadata wrapper(node); - if (wrapper.empty()) { + if (wrapper.empty()) return success(); - } + MLIRContext *ctx = op->getContext(); Attribute mlirMmra; if (wrapper.size() == 1) { @@ -231,9 +231,8 @@ static LogicalResult setMmraAttr(llvm::MDNode *node, Operation *op, mlirMmra = LLVM::MMRATagAttr::get(ctx, prefix, suffix); } else { SmallVector tags; - for (auto [prefix, suffix] : wrapper) { + for (auto [prefix, suffix] : wrapper) tags.push_back(LLVM::MMRATagAttr::get(ctx, prefix, suffix)); - } mlirMmra = ArrayAttr::get(ctx, tags); } op->setAttr(LLVMDialect::getMmraAttrName(), mlirMmra); diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp index ddf9b16b7b552..5093b2a47d10c 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp @@ -735,7 +735,7 @@ amendOperationImpl(Operation &op, ArrayRef instructions, if (auto oneTag = dyn_cast(attribute.getValue())) { tags.emplace_back(oneTag.getPrefix(), oneTag.getSuffix()); } else if (auto manyTags = dyn_cast(attribute.getValue())) { - for (auto a : manyTags) { + for (Attribute attr : manyTags) { auto tag = dyn_cast(a); if (tag) { tags.emplace_back(tag.getPrefix(), tag.getSuffix()); @@ -754,9 +754,8 @@ amendOperationImpl(Operation &op, ArrayRef instructions, // Empty list, canonicalizes to nothing return success(); } - for (llvm::Instruction *inst : instructions) { + for (llvm::Instruction *inst : instructions) inst->setMetadata(llvm::LLVMContext::MD_mmra, mmraMd); - } return success(); } return success(); From 0c749fcace4219918dca7e14a2c8c658a283dfd0 Mon Sep 17 00:00:00 2001 From: Krzysztof Drewniak Date: Wed, 10 Sep 2025 16:31:11 +0000 Subject: [PATCH 3/4] Fix nondeterminism, non-trivial nits --- .../LLVMIR/LLVMIRToLLVMTranslation.cpp | 25 +++++++++++++------ .../LLVMIR/LLVMToLLVMIRTranslation.cpp | 10 +++----- .../Target/LLVMIR/Import/metadata-mmra.ll | 4 +-- 3 files changed, 23 insertions(+), 16 deletions(-) diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp index 736b1fe5b850c..44732d5466f6d 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp @@ -220,19 +220,28 @@ static LogicalResult setDereferenceableAttr(const llvm::MDNode *node, /// discardable `llvm.mmra` attribute. static LogicalResult setMmraAttr(llvm::MDNode *node, Operation *op, LLVM::ModuleImport &moduleImport) { - llvm::MMRAMetadata wrapper(node); - if (wrapper.empty()) + if (!node) return success(); - + + // We don't use the LLVM wrappers here becasue we care about the order + // of the metadata for deterministic roundtripping. MLIRContext *ctx = op->getContext(); + auto toAttribute = [&](llvm::MDNode *tag) -> Attribute { + return LLVM::MMRATagAttr::get( + ctx, cast(tag->getOperand(0))->getString(), + cast(tag->getOperand(1))->getString()); + }; Attribute mlirMmra; - if (wrapper.size() == 1) { - auto [prefix, suffix] = *wrapper.begin(); - mlirMmra = LLVM::MMRATagAttr::get(ctx, prefix, suffix); + if (llvm::MMRAMetadata::isTagMD(node)) { + mlirMmra = toAttribute(node); } else { SmallVector tags; - for (auto [prefix, suffix] : wrapper) - tags.push_back(LLVM::MMRATagAttr::get(ctx, prefix, suffix)); + for (const llvm::MDOperand &operand : node->operands()) { + auto *tagNode = dyn_cast(operand.get()); + if (!tagNode || !llvm::MMRAMetadata::isTagMD(tagNode)) + return failure(); + tags.push_back(toAttribute(tagNode)); + } mlirMmra = ArrayAttr::get(ctx, tags); } op->setAttr(LLVMDialect::getMmraAttrName(), mlirMmra); diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp index 5093b2a47d10c..eaf1d20da63c7 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp @@ -736,13 +736,11 @@ amendOperationImpl(Operation &op, ArrayRef instructions, tags.emplace_back(oneTag.getPrefix(), oneTag.getSuffix()); } else if (auto manyTags = dyn_cast(attribute.getValue())) { for (Attribute attr : manyTags) { - auto tag = dyn_cast(a); - if (tag) { - tags.emplace_back(tag.getPrefix(), tag.getSuffix()); - } else { + auto tag = dyn_cast(attr); + if (!tag) return op.emitOpError( "MMRA annotations array contains value that isn't an MMRA tag"); - } + tags.emplace_back(tag.getPrefix(), tag.getSuffix()); } } else { return op.emitOpError( @@ -754,7 +752,7 @@ amendOperationImpl(Operation &op, ArrayRef instructions, // Empty list, canonicalizes to nothing return success(); } - for (llvm::Instruction *inst : instructions) + for (llvm::Instruction *inst : instructions) inst->setMetadata(llvm::LLVMContext::MD_mmra, mmraMd); return success(); } diff --git a/mlir/test/Target/LLVMIR/Import/metadata-mmra.ll b/mlir/test/Target/LLVMIR/Import/metadata-mmra.ll index 180d438eca70e..5e1ed37d559ef 100644 --- a/mlir/test/Target/LLVMIR/Import/metadata-mmra.ll +++ b/mlir/test/Target/LLVMIR/Import/metadata-mmra.ll @@ -9,7 +9,7 @@ define void @native(ptr %x, ptr %y) { ; CHECK-SAME: llvm.mmra = #[[$MMRA0]] %v = load i32, ptr %x, align 4, !mmra !0 ; CHECK: llvm.fence - ; CHECK-SAME: llvm.mmra = [#[[$MMRA0]], #[[$MMRA1]]] + ; CHECK-SAME: llvm.mmra = [#[[$MMRA1]], #[[$MMRA0]]] fence syncscope("workgroup-one-as") release, !mmra !2 ; CHECK: llvm.store {{.*}}, !llvm.ptr{{$}} store i32 %v, ptr %y, align 4, !mmra !3 @@ -18,5 +18,5 @@ define void @native(ptr %x, ptr %y) { !0 = !{!"foo", !"bar"} !1 = !{!"amdgpu-synchronize-as", !"local"} -!2 = !{!0, !1} +!2 = !{!1, !0} !3 = !{} From de337cb51c72b82bd02f2be31ba56d770013f3c1 Mon Sep 17 00:00:00 2001 From: Krzysztof Drewniak Date: Wed, 10 Sep 2025 17:36:21 +0000 Subject: [PATCH 4/4] Update documentation --- .../mlir/Dialect/LLVMIR/LLVMAttrDefs.td | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td index 9980e1174d694..a8c9ef790cfbd 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td @@ -1233,11 +1233,7 @@ def LLVM_TBAATagArrayAttr } //===----------------------------------------------------------------------===// -// Memory Model Relaxation Annations (mmra) Attributes -//===----------------------------------------------------------------------===// - -//===----------------------------------------------------------------------===// -// MMRATagAttr - single MMRA tag +// MMRATagAttr //===----------------------------------------------------------------------===// def LLVM_MMRATagAttr : LLVM_Attr<"MMRATag", "mmra_tag"> { @@ -1260,6 +1256,14 @@ def LLVM_MMRATagAttr : LLVM_Attr<"MMRATag", "mmra_tag"> { #mmra_tag1 = #llvm.mmra_tag<"foo":"bar"> ``` + Either one MMRA tag or an array of them may be added to any LLVM + operation that operates on memory. + + ```mlir + %v = llvm.load %ptr {llvm.mmra = #mmra_tag} : !llvm.ptr -> i8 + llvm.store %v, %ptr2 {llvm.mmra [#mmra_tag, #mmra_tag1]} : i8, !llvm.ptr + ``` + See the following link for more details: https://llvm.org/docs/MemoryModelRelaxationAnnotations.html }]; @@ -1269,10 +1273,6 @@ def LLVM_MMRATagAttr : LLVM_Attr<"MMRATag", "mmra_tag"> { let genMnemonicAlias = 1; } -def LLVM_MMRATagArrayAttr : TypedArrayAttrBase< - LLVM_MMRATagAttr, - LLVM_MMRATagAttr.summary # " array">; - //===----------------------------------------------------------------------===// // ConstantRangeAttr //===----------------------------------------------------------------------===//