[mlir][LLVM] Convert alias metadata to using attributes instead of ops

Using MLIR attributes instead of metadata has many advantages:
* No indirection: Attributes can simply refer to each other seemlessly without having to use the indirection of `SymbolRefAttr`. This also gives us correctness by construction in a lot of places as well
* Multithreading save: The Attribute infrastructure gives us thread-safety for free. Creating operations and inserting them into a block is not thread-safe. This is a major use case for e.g. the inliner in MLIR which runs in parallel
* Easier to create: There is no need for a builder or a metadata region

This patch therefore does exactly that. It leverages the new distinct attributes to create distinct alias domains and scopes in a deterministic and threadsafe manner.

Differential Revision: https://reviews.llvm.org/D155159

Author: zero9178

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

@@ -562,4 +562,92 @@ def LLVM_MemoryEffectsAttr : LLVM_Attr<"MemoryEffects", "memory_effects"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
+//===----------------------------------------------------------------------===//
+// AliasScopeDomainAttr
+//===----------------------------------------------------------------------===//
+
+def LLVM_AliasScopeDomainAttr : LLVM_Attr<"AliasScopeDomain",
+                                          "alias_scope_domain"> {
+  let parameters = (ins
+    "DistinctAttr":$id,
+    OptionalParameter<"StringAttr">:$description
+  );
+
+  let builders = [
+    AttrBuilder<(ins CArg<"StringAttr", "{}">:$description), [{
+      return $_get($_ctxt, DistinctAttr::create(UnitAttr::get($_ctxt)), description);
+    }]>
+  ];
+
+  let summary = "LLVM dialect alias scope domain metadata";
+
+  let description = [{
+    Defines a domain that may be associated with an alias scope.
+
+    See the following link for more details:
+    https://llvm.org/docs/LangRef.html#noalias-and-alias-scope-metadata
+  }];
+
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
+//===----------------------------------------------------------------------===//
+// AliasScopeAttr
+//===----------------------------------------------------------------------===//
+
+def LLVM_AliasScopeAttr : LLVM_Attr<"AliasScope", "alias_scope"> {
+  let parameters = (ins
+    "DistinctAttr":$id,
+    "AliasScopeDomainAttr":$domain,
+    OptionalParameter<"StringAttr">:$description
+  );
+
+  let builders = [
+    AttrBuilderWithInferredContext<(ins
+      "AliasScopeDomainAttr":$domain,
+      CArg<"StringAttr", "{}">:$description
+    ), [{
+      MLIRContext *ctx = domain.getContext();
+      return $_get(ctx, DistinctAttr::create(UnitAttr::get(ctx)), domain, description);
+    }]>
+  ];
+
+  let description = [{
+    Defines an alias scope that can be attached to a memory-accessing operation.
+    Such scopes can be used in combination with `noalias` metadata to indicate
+    that sets of memory-affecting operations in one scope do not alias with
+    memory-affecting operations in another scope.
+
+    Example:
+    ```mlir
+    #domain = #llvm.alias_scope_domain<id = distinct[1]<>, description = "Optional domain description">
+    #scope1 = #llvm.alias_scope<id = distinct[2]<>, domain = #domain>
+    #scope2 = #llvm.alias_scope<id = distinct[3]<>, domain = #domain, description = "Optional scope description">
+    llvm.func @foo(%ptr1 : !llvm.ptr<i32>) {
+        %c0 = llvm.mlir.constant(0 : i32) : i32
+        %c4 = llvm.mlir.constant(4 : i32) : i32
+        %1 = llvm.ptrtoint %ptr1 : !llvm.ptr<i32> to i32
+        %2 = llvm.add %1, %c1 : i32
+        %ptr2 = llvm.inttoptr %2 : i32 to !llvm.ptr<i32>
+        llvm.store %c0, %ptr1 { alias_scopes = [#scope1], llvm.noalias = [#scope2] } : !llvm.ptr<i32>
+        llvm.store %c4, %ptr2 { alias_scopes = [#scope2], llvm.noalias = [#scope1] } : !llvm.ptr<i32>
+        llvm.return
+    }
+    ```
+
+    See the following link for more details:
+    https://llvm.org/docs/LangRef.html#noalias-and-alias-scope-metadata
+  }];
+
+  let summary = "LLVM dialect alias scope";
+
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
+def LLVM_AliasScopeArrayAttr
+    : TypedArrayAttrBase<LLVM_AliasScopeAttr,
+                         LLVM_AliasScopeAttr.summary # " array"> {
+  let constBuilderCall = ?;
+}
+
 #endif // LLVMIR_ATTRDEFS

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

@@ -325,15 +325,15 @@ def LLVM_MemsetOp : LLVM_ZeroResultIntrOp<"memset", [0, 2],
 
 def LLVM_NoAliasScopeDeclOp
     : LLVM_ZeroResultIntrOp<"experimental.noalias.scope.decl"> {
-  let arguments = (ins SymbolRefAttr:$scope);
+  let arguments = (ins LLVM_AliasScopeAttr:$scope);
   string llvmBuilder = [{
     // Wrap the scope argument into a list since the LLVM IR intrinsic takes
     // a list containing exactly one scope rather than a scope itself.
-    llvm::MDNode* node = moduleTranslation.getAliasScopes(op, {$scope});
+    llvm::MDNode* node = moduleTranslation.getAliasScopes({$scope});
     builder.CreateNoAliasScopeDeclaration(node);
   }];
   string mlirBuilder = [{
-    FailureOr<SmallVector<SymbolRefAttr>> scopeAttrs =
+    FailureOr<SmallVector<LLVM::AliasScopeAttr>> scopeAttrs =
       moduleImport.matchAliasScopeAttrs(llvmOperands[0]);
     // Drop the intrinsic if the alias scope translation fails since the scope
     // is not used by an aliasing operation, such as a load or store, that is

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

@@ -279,8 +279,8 @@ class LLVM_MemAccessOpBase<string mnemonic, list<Trait> traits = []> :
       DeclareOpInterfaceMethods<AliasAnalysisOpInterface>], traits)>,
     LLVM_MemOpPatterns {
   dag aliasAttrs = (ins OptionalAttr<SymbolRefArrayAttr>:$access_groups,
-                    OptionalAttr<SymbolRefArrayAttr>:$alias_scopes,
-                    OptionalAttr<SymbolRefArrayAttr>:$noalias_scopes,
+                    OptionalAttr<LLVM_AliasScopeArrayAttr>:$alias_scopes,
+                    OptionalAttr<LLVM_AliasScopeArrayAttr>:$noalias_scopes,
                     OptionalAttr<SymbolRefArrayAttr>:$tbaa);
 }
 
@@ -324,8 +324,8 @@ class LLVM_IntrOpBase<Dialect dialect, string opName, string enumName,
             (ins OptionalAttr<SymbolRefArrayAttr>:$access_groups),
             (ins )),
         !if(!gt(requiresAliasAnalysis, 0),
-            (ins OptionalAttr<SymbolRefArrayAttr>:$alias_scopes,
-                 OptionalAttr<SymbolRefArrayAttr>:$noalias_scopes,
+            (ins OptionalAttr<LLVM_AliasScopeArrayAttr>:$alias_scopes,
+                 OptionalAttr<LLVM_AliasScopeArrayAttr>:$noalias_scopes,
                  OptionalAttr<SymbolRefArrayAttr>:$tbaa),
             (ins )));
   string resultPattern = !if(!gt(numResults, 1),

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -1101,68 +1101,6 @@ def LLVM_MetadataOp : LLVM_Op<"metadata", [
   let hasRegionVerifier = 1;
 }
 
-def LLVM_AliasScopeDomainMetadataOp : LLVM_Op<"alias_scope_domain", [
-  HasParent<"MetadataOp">, Symbol
-]> {
-  let arguments = (ins
-    SymbolNameAttr:$sym_name,
-    OptionalAttr<StrAttr>:$description
-  );
-  let summary = "LLVM dialect alias.scope domain metadata.";
-  let description = [{
-    Defines a domain that may be associated with an alias scope.
-
-    See the following link for more details:
-    https://llvm.org/docs/LangRef.html#noalias-and-alias-scope-metadata
-  }];
-  let assemblyFormat = "$sym_name attr-dict";
-}
-
-def LLVM_AliasScopeMetadataOp : LLVM_Op<"alias_scope", [
-  HasParent<"MetadataOp">, Symbol
-]> {
-  let arguments = (ins
-    SymbolNameAttr:$sym_name,
-    FlatSymbolRefAttr:$domain,
-    OptionalAttr<StrAttr>:$description
-  );
-  let summary = "LLVM dialect alias.scope metadata.";
-  let description = [{
-    Defines an alias scope that can be attached to a memory-accessing operation.
-    Such scopes can be used in combination with `noalias` metadata to indicate
-    that sets of memory-affecting operations in one scope do not alias with
-    memory-affecting operations in another scope.
-
-    Example:
-    ```mlir
-    module {
-      llvm.func @foo(%ptr1 : !llvm.ptr<i32>) {
-          %c0 = llvm.mlir.constant(0 : i32) : i32
-          %c4 = llvm.mlir.constant(4 : i32) : i32
-          %1 = llvm.ptrtoint %ptr1 : !llvm.ptr<i32> to i32
-          %2 = llvm.add %1, %c1 : i32
-          %ptr2 = llvm.inttoptr %2 : i32 to !llvm.ptr<i32>
-          llvm.store %c0, %ptr1 { alias_scopes = [@metadata::@scope1], llvm.noalias = [@metadata::@scope2] } : !llvm.ptr<i32>
-          llvm.store %c4, %ptr2 { alias_scopes = [@metadata::@scope2], llvm.noalias = [@metadata::@scope1] } : !llvm.ptr<i32>
-          llvm.return
-      }
-
-      llvm.metadata @metadata {
-        llvm.alias_scope_domain @unused_domain
-        llvm.alias_scope_domain @domain { description = "Optional domain description"}
-        llvm.alias_scope @scope1 { domain = @domain }
-        llvm.alias_scope @scope2 { domain = @domain, description = "Optional scope description" }
-      }
-    }
-    ```
-
-    See the following link for more details:
-    https://llvm.org/docs/LangRef.html#noalias-and-alias-scope-metadata
-  }];
-  let assemblyFormat = "$sym_name attr-dict";
-  let hasVerifier = 1;
-}
-
 def LLVM_AccessGroupMetadataOp : LLVM_Op<"access_group", [
   HasParent<"MetadataOp">, Symbol
 ]> {

mlir/include/mlir/Target/LLVMIR/ModuleImport.h

@@ -151,9 +151,9 @@ class ModuleImport {
   /// Converts `value` to a label attribute. Asserts if the matching fails.
   DILabelAttr matchLabelAttr(llvm::Value *value);
 
-  /// Converts `value` to an array of symbol references pointing to alias scope
-  /// operations, or returns failure if the conversion fails.
-  FailureOr<SmallVector<SymbolRefAttr>>
+  /// Converts `value` to an array of alias scopes or returns failure if the
+  /// conversion fails.
+  FailureOr<SmallVector<AliasScopeAttr>>
   matchAliasScopeAttrs(llvm::Value *value);
 
   /// Translates the debug location.
@@ -202,10 +202,10 @@ class ModuleImport {
   LoopAnnotationAttr translateLoopAnnotationAttr(const llvm::MDNode *node,
                                                  Location loc) const;
 
-  /// Returns the symbol references pointing to the alias scope operations that
-  /// map to the alias scope nodes starting from the metadata `node`. Returns
-  /// failure, if any of the symbol references cannot be found.
-  FailureOr<SmallVector<SymbolRefAttr>>
+  /// Returns the alias scope attributes that map to the alias scope nodes
+  /// starting from the metadata `node`. Returns failure, if any of the
+  /// attributes cannot be found.
+  FailureOr<SmallVector<AliasScopeAttr>>
   lookupAliasScopeAttrs(const llvm::MDNode *node) const;
 
 private:
@@ -349,9 +349,9 @@ class ModuleImport {
   /// operations for all operations that return no result. All operations that
   /// return a result have a valueMapping entry instead.
   DenseMap<llvm::Instruction *, Operation *> noResultOpMapping;
-  /// Mapping between LLVM alias scope and domain metadata nodes and symbol
-  /// references to the LLVM dialect operations corresponding to these nodes.
-  DenseMap<const llvm::MDNode *, SymbolRefAttr> aliasScopeMapping;
+  /// Mapping between LLVM alias scope and domain metadata nodes and
+  /// attributes in the LLVM dialect corresponding to these nodes.
+  DenseMap<const llvm::MDNode *, Attribute> aliasScopeMapping;
   /// Mapping between LLVM TBAA metadata nodes and symbol references to the LLVM
   /// dialect TBAA operations corresponding to these nodes.
   DenseMap<const llvm::MDNode *, SymbolRefAttr> tbaaMapping;

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

@@ -122,14 +122,13 @@ class ModuleTranslation {
   /// in these blocks.
   void forgetMapping(Region &region);
 
-  /// Returns the LLVM metadata corresponding to a symbol reference to an mlir
-  /// LLVM dialect alias scope operation.
-  llvm::MDNode *getAliasScope(Operation *op, SymbolRefAttr aliasScopeRef) const;
+  /// Returns the LLVM metadata corresponding to a mlir LLVM dialect alias scope
+  /// attribute.
+  llvm::MDNode *getAliasScope(AliasScopeAttr aliasScopeAttr) const;
 
-  /// Returns the LLVM metadata corresponding to an array of symbol references
-  /// to mlir LLVM dialect alias scope operations.
-  llvm::MDNode *getAliasScopes(Operation *op,
-                               ArrayRef<SymbolRefAttr> aliasScopeRefs) const;
+  /// Returns the LLVM metadata corresponding to an array of mlir LLVM dialect
+  /// alias scope attributes.
+  llvm::MDNode *getAliasScopes(ArrayRef<AliasScopeAttr> aliasScopeAttrs) const;
 
   // Sets LLVM metadata for memory operations that are in a parallel loop.
   void setAccessGroupsMetadata(AccessGroupOpInterface op,
@@ -335,7 +334,7 @@ class ModuleTranslation {
 
   /// Mapping from an alias scope metadata operation to its LLVM metadata.
   /// This map is populated on module entry.
-  DenseMap<Operation *, llvm::MDNode *> aliasScopeMetadataMapping;
+  DenseMap<Attribute, llvm::MDNode *> aliasScopeMetadataMapping;
 
   /// Mapping from a tbaa metadata operation to its LLVM metadata.
   /// This map is populated on module entry.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

@@ -3084,21 +3084,6 @@ LogicalResult TBAATypeDescriptorOp::verify() {
   return success();
 }
 
-//===----------------------------------------------------------------------===//
-// AliasScopeMetadataOp
-//===----------------------------------------------------------------------===//
-
-LogicalResult AliasScopeMetadataOp::verify() {
-  Operation *domainOp = SymbolTable::lookupNearestSymbolFrom(
-      this->getOperation(), getDomainAttr());
-  if (!isa_and_nonnull<AliasScopeDomainMetadataOp>(domainOp)) {
-    return this->emitOpError()
-           << "expected '" << getDomain()
-           << "' to reference a domain operation in the same region";
-  }
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // OpAsmDialectInterface
 //===----------------------------------------------------------------------===//
@@ -3109,8 +3094,9 @@ struct LLVMOpAsmDialectInterface : public OpAsmDialectInterface {
 
   AliasResult getAlias(Attribute attr, raw_ostream &os) const override {
     return TypeSwitch<Attribute, AliasResult>(attr)
-        .Case<DIBasicTypeAttr, DICompileUnitAttr, DICompositeTypeAttr,
-              DIDerivedTypeAttr, DIFileAttr, DILabelAttr, DILexicalBlockAttr,
+        .Case<AliasScopeAttr, AliasScopeDomainAttr, DIBasicTypeAttr,
+              DICompileUnitAttr, DICompositeTypeAttr, DIDerivedTypeAttr,
+              DIFileAttr, DILabelAttr, DILexicalBlockAttr,
               DILexicalBlockFileAttr, DILocalVariableAttr, DINamespaceAttr,
               DINullTypeAttr, DISubprogramAttr, DISubroutineTypeAttr,
               LoopAnnotationAttr, LoopVectorizeAttr, LoopInterleaveAttr,

mlir/lib/Dialect/LLVMIR/IR/LLVMInterfaces.cpp

@@ -100,12 +100,6 @@ LogicalResult mlir::LLVM::detail::verifyAccessGroupOpInterface(Operation *op) {
 LogicalResult
 mlir::LLVM::detail::verifyAliasAnalysisOpInterface(Operation *op) {
   auto iface = cast<AliasAnalysisOpInterface>(op);
-  if (failed(verifySymbolRefsPointTo<LLVM::AliasScopeMetadataOp>(
-          iface, "alias scopes", iface.getAliasScopesOrNull())))
-    return failure();
-  if (failed(verifySymbolRefsPointTo<LLVM::AliasScopeMetadataOp>(
-          iface, "noalias scopes", iface.getNoAliasScopesOrNull())))
-    return failure();
   if (failed(verifySymbolRefsPointTo<LLVM::TBAATagOp>(
           iface, "tbaa tags", iface.getTBAATagsOrNull())))
     return failure();

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

@@ -191,13 +191,12 @@ static LogicalResult setLoopAttr(const llvm::MDNode *node, Operation *op,
       .Default([](auto) { return failure(); });
 }
 
-/// Looks up all the symbol references pointing to the alias scope operations
-/// that map to the alias scope nodes starting from the alias scope metadata
-/// `node`, and attaches all of them to the imported operation if the lookups
-/// succeed. Returns failure otherwise.
+/// Looks up all the alias scope attributes that map to the alias scope nodes
+/// starting from the alias scope metadata `node`, and attaches all of them to
+/// the imported operation if the lookups succeed. Returns failure otherwise.
 static LogicalResult setAliasScopesAttr(const llvm::MDNode *node, Operation *op,
                                         LLVM::ModuleImport &moduleImport) {
-  FailureOr<SmallVector<SymbolRefAttr>> aliasScopes =
+  FailureOr<SmallVector<AliasScopeAttr>> aliasScopes =
       moduleImport.lookupAliasScopeAttrs(node);
   if (failed(aliasScopes))
     return failure();
@@ -207,8 +206,7 @@ static LogicalResult setAliasScopesAttr(const llvm::MDNode *node, Operation *op,
     return failure();
 
   iface.setAliasScopes(ArrayAttr::get(
-      iface.getContext(),
-      SmallVector<Attribute>{aliasScopes->begin(), aliasScopes->end()}));
+      iface.getContext(), llvm::to_vector_of<Attribute>(*aliasScopes)));
   return success();
 }
 
@@ -219,7 +217,7 @@ static LogicalResult setAliasScopesAttr(const llvm::MDNode *node, Operation *op,
 static LogicalResult setNoaliasScopesAttr(const llvm::MDNode *node,
                                           Operation *op,
                                           LLVM::ModuleImport &moduleImport) {
-  FailureOr<SmallVector<SymbolRefAttr>> noAliasScopes =
+  FailureOr<SmallVector<AliasScopeAttr>> noAliasScopes =
       moduleImport.lookupAliasScopeAttrs(node);
   if (failed(noAliasScopes))
     return failure();
@@ -229,8 +227,7 @@ static LogicalResult setNoaliasScopesAttr(const llvm::MDNode *node,
     return failure();
 
   iface.setNoAliasScopes(ArrayAttr::get(
-      iface.getContext(),
-      SmallVector<Attribute>{noAliasScopes->begin(), noAliasScopes->end()}));
+      iface.getContext(), llvm::to_vector_of<Attribute>(*noAliasScopes)));
   return success();
 }