[MLIR] Enable import of non self referential alias scopes

Author: wsmoses

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

@@ -825,7 +825,7 @@ def LLVM_MemoryEffectsAttr : LLVM_Attr<"MemoryEffects", "memory_effects"> {
 def LLVM_AliasScopeDomainAttr : LLVM_Attr<"AliasScopeDomain",
                                           "alias_scope_domain"> {
   let parameters = (ins
-    "DistinctAttr":$id,
+    "Attribute":$id,
     OptionalParameter<"StringAttr">:$description
   );
 
@@ -853,7 +853,7 @@ def LLVM_AliasScopeDomainAttr : LLVM_Attr<"AliasScopeDomain",
 
 def LLVM_AliasScopeAttr : LLVM_Attr<"AliasScope", "alias_scope"> {
   let parameters = (ins
-    "DistinctAttr":$id,
+    "Attribute":$id,
     "AliasScopeDomainAttr":$domain,
     OptionalParameter<"StringAttr">:$description
   );
@@ -891,6 +891,8 @@ def LLVM_AliasScopeAttr : LLVM_Attr<"AliasScope", "alias_scope"> {
     }
     ```
 
+    The first attribute can either be a DistinctAttribute or a StringAttribute.
+
     See the following link for more details:
     https://llvm.org/docs/LangRef.html#noalias-and-alias-scope-metadata
   }];

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -438,8 +438,14 @@ ModuleImport::processAliasScopeMetadata(const llvm::MDNode *node) {
     if (aliasDomain->getNumOperands() >= 2)
       if (auto *operand = dyn_cast<llvm::MDString>(aliasDomain->getOperand(1)))
         description = builder.getStringAttr(operand->getString());
-    return builder.getAttr<AliasScopeDomainAttr>(
-        DistinctAttr::create(builder.getUnitAttr()), description);
+    if (verifySelfRef(aliasDomain))
+      return builder.getAttr<AliasScopeDomainAttr>(
+          DistinctAttr::create(builder.getUnitAttr()), description);
+    else {
+      auto Name = cast<llvm::MDString>(aliasDomain->getOperand(0));
+      return builder.getAttr<AliasScopeDomainAttr>(
+          builder.getStringAttr(Name->getString()), description);
+    }
   };
 
   // Collect the alias scopes and domains to translate them.
@@ -452,12 +458,28 @@ ModuleImport::processAliasScopeMetadata(const llvm::MDNode *node) {
       // verifying its domain. Perform the verification before looking it up in
       // the alias scope mapping since it could have been inserted as a domain
       // node before.
-      if (!verifySelfRef(scope) || !domain || !verifyDescription(scope, 2))
-        return emitError(loc) << "unsupported alias scope node: "
+      if (!domain)
+        return emitError(loc) << "unsupported alias scope node (no domain): "
                               << diagMD(scope, llvmModule.get());
-      if (!verifySelfRef(domain) || !verifyDescription(domain, 1))
-        return emitError(loc) << "unsupported alias domain node: "
-                              << diagMD(domain, llvmModule.get());
+      if (verifySelfRef(scope) && !verifyDescription(scope, 2))
+        return emitError(loc)
+               << "unsupported alias scope node (bad description): "
+               << diagMD(scope, llvmModule.get());
+      if (!verifySelfRef(scope) && (scope->getNumOperands() == 0 ||
+                                    !isa<llvm::MDString>(scope->getOperand(0))))
+        return emitError(loc)
+               << "unsupported alias scope node (not self-ref and no string): "
+               << diagMD(scope, llvmModule.get());
+      if (!verifyDescription(domain, 1))
+        return emitError(loc)
+               << "unsupported alias domain node (bad description): "
+               << diagMD(domain, llvmModule.get());
+      if (!verifySelfRef(domain) &&
+          (domain->getNumOperands() == 0 ||
+           !isa<llvm::MDString>(domain->getOperand(0))))
+        return emitError(loc)
+               << "unsupported alias domain node (not self ref and no string): "
+               << diagMD(scope, llvmModule.get());
 
       if (aliasScopeMapping.contains(scope))
         continue;
@@ -473,9 +495,18 @@ ModuleImport::processAliasScopeMetadata(const llvm::MDNode *node) {
       StringAttr description = nullptr;
       if (!aliasScope.getName().empty())
         description = builder.getStringAttr(aliasScope.getName());
-      auto aliasScopeOp = builder.getAttr<AliasScopeAttr>(
-          DistinctAttr::create(builder.getUnitAttr()),
-          cast<AliasScopeDomainAttr>(it->second), description);
+      AliasScopeAttr aliasScopeOp;
+      if (verifySelfRef(scope))
+        aliasScopeOp = builder.getAttr<AliasScopeAttr>(
+            DistinctAttr::create(builder.getUnitAttr()),
+            cast<AliasScopeDomainAttr>(it->second), description);
+      else {
+        auto Name = cast<llvm::MDString>(scope->getOperand(0));
+        aliasScopeOp = builder.getAttr<AliasScopeAttr>(
+            builder.getStringAttr(Name->getString()),
+            cast<AliasScopeDomainAttr>(it->second), description);
+      }
+
       aliasScopeMapping.try_emplace(aliasScope.getNode(), aliasScopeOp);
     }
   }

mlir/test/Dialect/LLVMIR/roundtrip.mlir

@@ -750,6 +750,16 @@ llvm.func @experimental_noalias_scope_decl() {
   llvm.return
 }
 
+#alias_scope_domain2 = #llvm.alias_scope_domain<id = "domainid", description = "The domain">
+#alias_scope2 = #llvm.alias_scope<id = "stringid", domain = #alias_scope_domain2, description = "The domain">
+
+// CHECK-LABEL: @experimental_noalias_scope_decl
+llvm.func @experimental_noalias_scope_decl2() {
+  // CHECK: llvm.intr.experimental.noalias.scope.decl #{{.*}}
+  llvm.intr.experimental.noalias.scope.decl #alias_scope2
+  llvm.return
+}
+
 // CHECK-LABEL: @experimental_constrained_fptrunc
 llvm.func @experimental_constrained_fptrunc(%in: f64) {
   // CHECK: llvm.intr.experimental.constrained.fptrunc %{{.*}} towardzero ignore : f64 to f32

mlir/test/Target/LLVMIR/Import/metadata-alias-scopes.ll

@@ -92,3 +92,39 @@ declare void @foo(ptr %arg1)
 !0 = distinct !{!0, !"The domain"}
 !1 = !{!1, !0}
 !2 = !{!1}
+
+; // -----
+
+; CHECK: #[[DOMAIN:.*]] = #llvm.alias_scope_domain<id = "domain1">
+; CHECK: #[[$SCOPE0:.*]] = #llvm.alias_scope<id = "scopeid1", domain = #[[DOMAIN]], description = "The first scope">
+; CHECK: #[[$SCOPE1:.*]] = #llvm.alias_scope<id = "scopeid2", domain = #[[DOMAIN]]>
+; CHECK: #[[$SCOPE2:.*]] = #llvm.alias_scope<id = "scopeid3", domain = #[[DOMAIN]]>
+
+; CHECK-LABEL: llvm.func @alias_scope
+define void @alias_scope(ptr %arg1) {
+  ; CHECK: llvm.load
+  ; CHECK-SAME:  alias_scopes = [#[[$SCOPE0]]]
+  ; CHECK-SAME:  noalias_scopes = [#[[$SCOPE1]], #[[$SCOPE2]]]
+  %1 = load i32, ptr %arg1, !alias.scope !4, !noalias !7
+  ; CHECK: llvm.load
+  ; CHECK-SAME:  alias_scopes = [#[[$SCOPE1]]]
+  ; CHECK-SAME:  noalias_scopes = [#[[$SCOPE0]], #[[$SCOPE2]]]
+  %2 = load i32, ptr %arg1, !alias.scope !5, !noalias !8
+  ; CHECK: llvm.load
+  ; CHECK-SAME:  alias_scopes = [#[[$SCOPE2]]]
+  ; CHECK-SAME:  noalias_scopes = [#[[$SCOPE0]], #[[$SCOPE1]]]
+  %3 = load i32, ptr %arg1, !alias.scope !6, !noalias !9
+  ret void
+}
+
+!0 = !{!"domain1"}
+!1 = !{!"scopeid1", !0, !"The first scope"}
+!2 = !{!"scopeid2", !0}
+!3 = !{!"scopeid3", !0}
+!4 = !{!1}
+!5 = !{!2}
+!6 = !{!3}
+!7 = !{!2, !3}
+!8 = !{!1, !3}
+!9 = !{!1, !2}
+