[MLIR] Enable import of non self referential alias scopes #121987

wsmoses · 2025-01-07T19:45:00Z

x/ref #121965

llvmbot · 2025-01-07T19:45:31Z

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: William Moses (wsmoses)

Changes

x/ref #121965

Full diff: https://github.com/llvm/llvm-project/pull/121987.diff

4 Files Affected:

(modified) mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td (+4-2)
(modified) mlir/lib/Target/LLVMIR/ModuleImport.cpp (+29-5)
(modified) mlir/test/Dialect/LLVMIR/roundtrip.mlir (+10)
(modified) mlir/test/Target/LLVMIR/Import/metadata-alias-scopes.ll (+36)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index e8eeafd09a9cba..68df6f64e51ea4 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/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
   }];
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index 95fb673fc72e39..681d4706e1b639 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/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(2));
+        return builder.getAttr<AliasScopeDomainAttr>(
+            builder.getStringAttr(Name->getString()), description);
+    }
   };
 
   // Collect the alias scopes and domains to translate them.
@@ -452,12 +458,21 @@ 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))
+      if (!domain)
+        return emitError(loc) << "unsupported alias scope node: "
+                              << diagMD(scope, llvmModule.get());
+      if (verifySelfRef(scope) && !verifyDescription(scope, 2))
+        return emitError(loc) << "unsupported alias scope node: "
+                              << diagMD(scope, llvmModule.get());
+      if (!verifySelfRef(scope) && (scope->getNumOperands() == 0 || !isa<llvm::MDString>(scope->getOperand(0))))
         return emitError(loc) << "unsupported alias scope node: "
                               << diagMD(scope, llvmModule.get());
-      if (!verifySelfRef(domain) || !verifyDescription(domain, 1))
+      if (verifySelfRef(domain) || !verifyDescription(domain, 1))
         return emitError(loc) << "unsupported alias domain node: "
                               << diagMD(domain, llvmModule.get());
+      if (!verifySelfRef(scope) && (domain->getNumOperands() == 0 || !isa<llvm::MDString>(domain->getOperand(0))))
+        return emitError(loc) << "unsupported alias scope node: "
+                              << diagMD(scope, llvmModule.get());
 
       if (aliasScopeMapping.contains(scope))
         continue;
@@ -473,9 +488,18 @@ ModuleImport::processAliasScopeMetadata(const llvm::MDNode *node) {
       StringAttr description = nullptr;
       if (!aliasScope.getName().empty())
         description = builder.getStringAttr(aliasScope.getName());
-      auto aliasScopeOp = builder.getAttr<AliasScopeAttr>(
+      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(2));
+        aliasScopeOp = builder.getAttr<AliasScopeAttr>(
+          builder.getStringAttr(Name->getString()),
+          cast<AliasScopeDomainAttr>(it->second), description);
+      }
+
       aliasScopeMapping.try_emplace(aliasScope.getNode(), aliasScopeOp);
     }
   }
diff --git a/mlir/test/Dialect/LLVMIR/roundtrip.mlir b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
index aebfd7492093c1..c4023bd63ea2ad 100644
--- a/mlir/test/Dialect/LLVMIR/roundtrip.mlir
+++ b/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 = distinct[0]<>, 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_decl() {
+  // 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
diff --git a/mlir/test/Target/LLVMIR/Import/metadata-alias-scopes.ll b/mlir/test/Target/LLVMIR/Import/metadata-alias-scopes.ll
index f5128ff76bc5ff..e60ac072ea1ca1 100644
--- a/mlir/test/Target/LLVMIR/Import/metadata-alias-scopes.ll
+++ b/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}
+

mlir/lib/Target/LLVMIR/ModuleImport.cpp

gysit

Thanks!

What happens if you export alias scopes with string identifiers. We probably should ensure the string identifiers are persisted. I assume currently they are lost when translating back to LLVM IR.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

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

wsmoses · 2025-01-07T20:10:04Z

Thanks!

What happens if you export alias scopes with string identifiers. We probably should ensure the string identifiers are persisted. I assume currently they are lost when translating back to LLVM IR.

Is there a current export test for non string, I couldn’t find it offhand

gysit · 2025-01-07T20:16:56Z

Is there a current export test for non string, I couldn’t find it offhand

/llvm-project/mlir/test/Target/LLVMIR/attribute-alias-scopes.mlir

contains some tests. Maybe one of them can be modified or a small one can be added there.

wsmoses · 2025-01-07T20:48:07Z

quick question where is the export code here (again a quick skim was fruitlesss, if you happen to know offhand).

Otherwise fixed all of the above (and will push shortly)

gysit · 2025-01-07T20:50:43Z

The export should be here (+-):

llvm-project/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

Line 1715 in ea14bdb

ModuleTranslation::getOrCreateAliasScope(AliasScopeAttr aliasScopeAttr) {

github-actions · 2025-01-07T20:53:50Z

✅ With the latest revision this PR passed the C/C++ code formatter.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

ftynse

LGTM with comments addressed. Please give Tobias a chance to re-review if needed.

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

mlir/lib/Target/LLVMIR/ModuleImport.cpp

wsmoses · 2025-01-07T23:49:11Z

Okay @gysit added and fixed the export test

Co-authored-by: Oleksandr "Alex" Zinenko <[email protected]>

Dinistro

Thanks for the extension. I've added mainly nit comments and some verification question.

Dinistro · 2025-01-08T06:58:27Z

mlir/test/Dialect/LLVMIR/roundtrip.mlir

Not needed in this PR: We should consider to introduce splits in these files, as debugging issues in this file is a menace.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

mlir/lib/Target/LLVMIR/ModuleImport.cpp

Dinistro · 2025-01-08T07:02:07Z

mlir/lib/Target/LLVMIR/ModuleImport.cpp

+    if (verifySelfRef(aliasDomain)) {
+      idAttr = DistinctAttr::create(builder.getUnitAttr());
+    } else {
+      auto name = cast<llvm::MDString>(aliasDomain->getOperand(0));
+      idAttr = builder.getStringAttr(name->getString());
+    }


Maybe consider to also move this to a lambda, as this is exact snippet is duplicated.

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

gysit · 2025-01-08T07:18:38Z

mlir/test/Dialect/LLVMIR/roundtrip.mlir

+#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


Suggested change

// CHECK-LABEL: @experimental_noalias_scope_decl

// CHECK-LABEL: @experimental_noalias_scope_decl2

nit: I would probably rename to @experimental_noalias_scope_decl_with_string_id rather than postfixing with a number. In any case the check line should match the function name.

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

gysit

LGTM once the non-optional remaining comments have been addressed.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

Co-authored-by: Christian Ulmann <[email protected]>

wsmoses requested review from ftynse and gysit January 7, 2025 19:45

llvmbot added mlir:llvm mlir labels Jan 7, 2025

wsmoses force-pushed the stringop branch 3 times, most recently from 1f984eb to 74beb3e Compare January 7, 2025 19:53

gysit reviewed Jan 7, 2025

View reviewed changes

mlir/lib/Target/LLVMIR/ModuleImport.cpp Outdated Show resolved Hide resolved

gysit reviewed Jan 7, 2025

View reviewed changes

mlir/lib/Target/LLVMIR/ModuleImport.cpp Outdated Show resolved Hide resolved

mlir/lib/Target/LLVMIR/ModuleImport.cpp Outdated Show resolved Hide resolved

mlir/test/Target/LLVMIR/Import/metadata-alias-scopes.ll Outdated Show resolved Hide resolved

wsmoses force-pushed the stringop branch from 74beb3e to 24f8c87 Compare January 7, 2025 20:46

wsmoses force-pushed the stringop branch 2 times, most recently from c4a752f to cf01a1b Compare January 7, 2025 20:50

gysit reviewed Jan 7, 2025

View reviewed changes

ftynse approved these changes Jan 7, 2025

View reviewed changes

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td Outdated Show resolved Hide resolved

mlir/lib/Target/LLVMIR/ModuleImport.cpp Show resolved Hide resolved

wsmoses force-pushed the stringop branch from cf01a1b to d4eb0db Compare January 7, 2025 23:12

[MLIR] Enable import of non self referential alias scopes

023dbba

wsmoses force-pushed the stringop branch from d4eb0db to 023dbba Compare January 7, 2025 23:42

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

a2f845a

Co-authored-by: Oleksandr "Alex" Zinenko <[email protected]>

wsmoses mentioned this pull request Jan 8, 2025

Initial Oceananigans Kernel EnzymeAD/Reactant.jl#499

Open

Dinistro reviewed Jan 8, 2025

View reviewed changes

gysit reviewed Jan 8, 2025

View reviewed changes

gysit approved these changes Jan 8, 2025

View reviewed changes

ftynse reviewed Jan 8, 2025

View reviewed changes

mlir/lib/Target/LLVMIR/ModuleImport.cpp Outdated Show resolved Hide resolved

ftynse and others added 3 commits January 8, 2025 10:21

Apply suggestions from code review

efe67be

Co-authored-by: Christian Ulmann <[email protected]>

extract lambda

bd0bb3c

add verifier

b91ba34

ftynse merged commit 1c067a5 into llvm:main Jan 8, 2025
8 checks passed

	// CHECK-LABEL: @experimental_noalias_scope_decl
	// CHECK-LABEL: @experimental_noalias_scope_decl2

[MLIR] Enable import of non self referential alias scopes #121987

[MLIR] Enable import of non self referential alias scopes #121987

Uh oh!

Conversation

wsmoses commented Jan 7, 2025

Uh oh!

llvmbot commented Jan 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

gysit left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wsmoses commented Jan 7, 2025

Uh oh!

gysit commented Jan 7, 2025

Uh oh!

wsmoses commented Jan 7, 2025

Uh oh!

gysit commented Jan 7, 2025

Uh oh!

github-actions bot commented Jan 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ftynse left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

wsmoses commented Jan 7, 2025

Uh oh!

Dinistro left a comment

Choose a reason for hiding this comment

Uh oh!

Dinistro Jan 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Dinistro Jan 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gysit Jan 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

gysit left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

llvmbot commented Jan 7, 2025 •

edited

Loading

github-actions bot commented Jan 7, 2025 •

edited

Loading