Fix nondeterminism, non-trivial nits

Author: krzysz00

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<llvm::MDString>(tag->getOperand(0))->getString(),
+        cast<llvm::MDString>(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<Attribute> 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<llvm::MDNode>(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);

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

@@ -736,13 +736,11 @@ amendOperationImpl(Operation &op, ArrayRef<llvm::Instruction *> instructions,
       tags.emplace_back(oneTag.getPrefix(), oneTag.getSuffix());
     } else if (auto manyTags = dyn_cast<ArrayAttr>(attribute.getValue())) {
       for (Attribute attr : manyTags) {
-        auto tag = dyn_cast<MMRATagAttr>(a);
-        if (tag) {
-          tags.emplace_back(tag.getPrefix(), tag.getSuffix());
-        } else {
+        auto tag = dyn_cast<MMRATagAttr>(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<llvm::Instruction *> 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();
   }

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 = !{}