[mlir][LLVMIR] Support memory model relaxation annotations (MMRA) #157770
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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).
|
@llvm/pr-subscribers-mlir-llvm Author: Krzysztof Drewniak (krzysz00) ChangesThis 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 Full diff: https://github.com/llvm/llvm-project/pull/157770.diff 7 Files Affected:
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<unsigned> 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<Attribute> 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<llvm::Instruction *> instructions,
+ NamedAttribute attribute,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ StringRef name = attribute.getName();
+ if (name == LLVMDialect::getMmraAttrName()) {
+ SmallVector<llvm::MMRAMetadata::TagT> tags;
+ if (auto oneTag = dyn_cast<LLVM::MMRATagAttr>(attribute.getValue())) {
+ tags.emplace_back(oneTag.getPrefix(), oneTag.getSuffix());
+ } else if (auto manyTags = dyn_cast<ArrayAttr>(attribute.getValue())) {
+ for (auto a : manyTags) {
+ auto tag = dyn_cast<MMRATagAttr>(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<llvm::Instruction *> 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"}
|
|
@llvm/pr-subscribers-mlir Author: Krzysztof Drewniak (krzysz00) ChangesThis 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 Full diff: https://github.com/llvm/llvm-project/pull/157770.diff 7 Files Affected:
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<unsigned> 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<Attribute> 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<llvm::Instruction *> instructions,
+ NamedAttribute attribute,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ StringRef name = attribute.getName();
+ if (name == LLVMDialect::getMmraAttrName()) {
+ SmallVector<llvm::MMRAMetadata::TagT> tags;
+ if (auto oneTag = dyn_cast<LLVM::MMRATagAttr>(attribute.getValue())) {
+ tags.emplace_back(oneTag.getPrefix(), oneTag.getSuffix());
+ } else if (auto manyTags = dyn_cast<ArrayAttr>(attribute.getValue())) {
+ for (auto a : manyTags) {
+ auto tag = dyn_cast<MMRATagAttr>(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<llvm::Instruction *> 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"}
|
gysit
approved these changes
Sep 10, 2025
mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
Outdated
Show resolved
Hide resolved
mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
Outdated
Show resolved
Hide resolved
mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
Outdated
Show resolved
Hide resolved
mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
Outdated
Show resolved
Hide resolved
mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
Outdated
Show resolved
Hide resolved
| let genMnemonicAlias = 1; | ||
| } | ||
|
|
||
| def LLVM_MMRATagArrayAttr : TypedArrayAttrBase< |
There was a problem hiding this comment.
Are there uses for this array attribute?
There was a problem hiding this comment.
No, good point ( you can't isa<> on this sort of thing because it isn't real)
mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
Outdated
Show resolved
Hide resolved
|
|
||
| //===----------------------------------------------------------------------===// | ||
| // Memory Model Relaxation Annations (mmra) Attributes | ||
| //===----------------------------------------------------------------------===// |
There was a problem hiding this comment.
nit: If we end up with only one attribute one banner comment should be good enough?
|
Re the array attribute ... I'm trying to match LLVM here, which allows either an array or a single tag, so as to not introduce conceptual overheads (since the translator is a decent place to pay for the complexity here) |
Co-authored-by: Tobias Gysi <tobias.gysi@nextsilicon.com>
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
krzysz00
added a commit
to iree-org/llvm-project
that referenced
this pull request
Sep 10, 2025
…vm#157770) 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). --------- Co-authored-by: Tobias Gysi <tobias.gysi@nextsilicon.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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).