[mlir][llvm] Translate alias scopes lazily

[gysit]

Summary

Change the LLVM dialect to LLVM IR translation to convert the alias
scope attributes lazily to LLVM IR metadata. Previously, the alias
scopes have been translated upfront walking the alias scopes of
operations that implement the AliasAnalysisOpInterface. As a result,
the translation of a module that contains only a noalias scope
intrinsic failed, since its alias scope attribute has not been
translated due to the intrinsic not implementing
AliasAnalysisOpInterface.

[gysit]

Avoid using the walker.

[zero9178]

LGTM

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

 }
 
 def LLVM_NoAliasScopeDeclOp
     : LLVM_ZeroResultIntrOp<"experimental.noalias.scope.decl"> {
   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({$scope});
+    llvm::MDNode* node = moduleTranslation.getOrCreateAliasScopes({$scope});
     builder.CreateNoAliasScopeDeclaration(node);
   }];
   string mlirBuilder = [{
     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
     // used to convert the alias scope metadata.

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

 
 namespace LLVM {
 
 namespace detail {
 class DebugTranslation;
 class LoopAnnotationTranslation;
 } // namespace detail
 
+class AliasScopeAttr;
+class AliasScopeDomainAttr;
 class DINodeAttr;
 class LLVMFuncOp;
 class ComdatSelectorOp;
 
 /// Implementation class for module translation. Holds a reference to the module
 /// being translated, and the mappings between the original and the translated
 /// functions, basic blocks and values. It is practically easier to hold these
 /// mappings in one class since the conversion of control flow operations
llvm::CallInst *lookupCall(Operation *op) const {
     return callMapping.lookup(op);
   }
 
   /// Removes the mapping for blocks contained in the region and values defined
   /// in these blocks.
   void forgetMapping(Region &region);
 
   /// Returns the LLVM metadata corresponding to a mlir LLVM dialect alias scope
-  /// attribute.
-  llvm::MDNode *getAliasScope(AliasScopeAttr aliasScopeAttr) const;
+  /// attribute. Creates the metadata node if it has not been converted before.
+  llvm::MDNode *getOrCreateAliasScope(AliasScopeAttr aliasScopeAttr);
 
   /// Returns the LLVM metadata corresponding to an array of mlir LLVM dialect
-  /// alias scope attributes.
-  llvm::MDNode *getAliasScopes(ArrayRef<AliasScopeAttr> aliasScopeAttrs) const;
+  /// alias scope attributes. Creates the metadata nodes if they have not been
+  /// converted before.
+  llvm::MDNode *
+  getOrCreateAliasScopes(ArrayRef<AliasScopeAttr> aliasScopeAttrs);
 
   // Sets LLVM metadata for memory operations that are in a parallel loop.
   void setAccessGroupsMetadata(AccessGroupOpInterface op,
                                llvm::Instruction *inst);
 
   // Sets LLVM metadata for memory operations that have alias scope information.
   void setAliasScopeMetadata(AliasAnalysisOpInterface op,
                              llvm::Instruction *inst);
private:
   /// Converts individual components.
   LogicalResult convertOperation(Operation &op, llvm::IRBuilderBase &builder);
   LogicalResult convertFunctionSignatures();
   LogicalResult convertFunctions();
   LogicalResult convertComdats();
   LogicalResult convertGlobals();
   LogicalResult convertOneFunction(LLVMFuncOp func);
 
-  /// Process alias.scope LLVM Metadata operations and create LLVM
-  /// metadata nodes for them and their domains.
-  LogicalResult createAliasScopeMetadata();
-
   /// Returns the LLVM metadata corresponding to the given mlir LLVM dialect
   /// TBAATagAttr.
   llvm::MDNode *getTBAANode(TBAATagAttr tbaaAttr) const;
 
   /// Process tbaa LLVM Metadata operations and create LLVM
   /// metadata nodes for them.
   LogicalResult createTBAAMetadata();
 
private:
   /// values after all operations are converted.
   DenseMap<Operation *, llvm::Instruction *> branchMapping;
 
   /// A mapping between MLIR LLVM dialect call operations and LLVM IR call
   /// instructions. This allows for adding branch weights after the operations
   /// have been converted.
   DenseMap<Operation *, llvm::CallInst *> callMapping;
 
-  /// Mapping from an alias scope metadata operation to its LLVM metadata.
-  /// This map is populated on module entry.
-  DenseMap<Attribute, llvm::MDNode *> aliasScopeMetadataMapping;
+  /// Mapping from an alias scope attribute to its LLVM metadata.
+  /// This map is populated lazily.
+  DenseMap<AliasScopeAttr, llvm::MDNode *> aliasScopeMetadataMapping;
+
+  /// Mapping from an alias scope domain attribute to its LLVM metadata.
+  /// This map is populated lazily.
+  DenseMap<AliasScopeDomainAttr, llvm::MDNode *> aliasDomainMetadataMapping;
 
   /// Mapping from a tbaa attribute to its LLVM metadata.
   /// This map is populated on module entry.
   DenseMap<Attribute, llvm::MDNode *> tbaaMetadataMapping;
 
   /// Mapping from a comdat selector operation to its LLVM comdat struct.
   /// This map is populated on module entry.
   DenseMap<ComdatSelectorOp, llvm::Comdat *> comdatMapping;

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 }
 
 void ModuleTranslation::setAccessGroupsMetadata(AccessGroupOpInterface op,
                                                 llvm::Instruction *inst) {
   if (llvm::MDNode *node = loopAnnotationTranslation->getAccessGroups(op))
     inst->setMetadata(llvm::LLVMContext::MD_access_group, node);
 }
 
-LogicalResult ModuleTranslation::createAliasScopeMetadata() {
-  DenseMap<Attribute, llvm::MDNode *> aliasScopeDomainMetadataMapping;
-
-  AttrTypeWalker walker;
-  walker.addWalk([&](LLVM::AliasScopeDomainAttr op) {
-    llvm::LLVMContext &ctx = llvmModule->getContext();
+llvm::MDNode *
+ModuleTranslation::getOrCreateAliasScope(AliasScopeAttr aliasScopeAttr) {
+  auto [scopeIt, scopeInserted] =
+      aliasScopeMetadataMapping.try_emplace(aliasScopeAttr, nullptr);
+  if (!scopeInserted)
+    return scopeIt->second;
+  // Convert the domain metadata node if necessary.
+  auto [domainIt, insertedDomain] = aliasDomainMetadataMapping.try_emplace(
+      aliasScopeAttr.getDomain(), nullptr);
+  if (insertedDomain) {
     llvm::SmallVector<llvm::Metadata *, 2> operands;
-    operands.push_back({}); // Placeholder for self-reference
-    if (StringAttr description = op.getDescription())
+    // Placeholder for self-reference.
+    operands.push_back({});
+    if (StringAttr description = aliasScopeAttr.getDomain().getDescription())
       operands.push_back(llvm::MDString::get(ctx, description));
-    llvm::MDNode *domain = llvm::MDNode::get(ctx, operands);
-    domain->replaceOperandWith(0, domain); // Self-reference for uniqueness
-    aliasScopeDomainMetadataMapping.insert({op, domain});
-  });
-
-  walker.addWalk([&](LLVM::AliasScopeAttr op) {
-    llvm::LLVMContext &ctx = llvmModule->getContext();
-    llvm::MDNode *domain =
-        aliasScopeDomainMetadataMapping.lookup(op.getDomain());
-    assert(domain && "Scope's domain should already be valid");
-    llvm::SmallVector<llvm::Metadata *, 3> operands;
-    operands.push_back({}); // Placeholder for self-reference
-    operands.push_back(domain);
-    if (StringAttr description = op.getDescription())
-      operands.push_back(llvm::MDString::get(ctx, description));
-    llvm::MDNode *scope = llvm::MDNode::get(ctx, operands);
-    scope->replaceOperandWith(0, scope); // Self-reference for uniqueness
-    aliasScopeMetadataMapping.insert({op, scope});
-  });
-
-  mlirModule->walk([&](AliasAnalysisOpInterface op) {
-    if (auto aliasScopes = op.getAliasScopesOrNull())
-      walker.walk(aliasScopes);
-    if (auto noAliasScopes = op.getNoAliasScopesOrNull())
-      walker.walk(noAliasScopes);
-  });
-
-  return success();
+    domainIt->second = llvm::MDNode::get(ctx, operands);
+    // Self-reference for uniqueness.
+    domainIt->second->replaceOperandWith(0, domainIt->second);
+  }
+  // Convert the scope metadata node.
+  assert(domainIt->second && "Scope's domain should already be valid");
+  // Placeholder for self-reference.
+  operands.push_back({});
+  operands.push_back(domainIt->second);
+  if (StringAttr description = aliasScopeAttr.getDescription())
+  scopeIt->second = llvm::MDNode::get(ctx, operands);
+  // Self-reference for uniqueness.
+  scopeIt->second->replaceOperandWith(0, scopeIt->second);
+  return scopeIt->second;
 }
 
-llvm::MDNode *
-ModuleTranslation::getAliasScope(AliasScopeAttr aliasScopeAttr) const {
-  return aliasScopeMetadataMapping.lookup(aliasScopeAttr);
-}
-
-llvm::MDNode *ModuleTranslation::getAliasScopes(
-    ArrayRef<AliasScopeAttr> aliasScopeAttrs) const {
+llvm::MDNode *ModuleTranslation::getOrCreateAliasScopes(
+    ArrayRef<AliasScopeAttr> aliasScopeAttrs) {
   SmallVector<llvm::Metadata *> nodes;
   nodes.reserve(aliasScopeAttrs.size());
   for (AliasScopeAttr aliasScopeAttr : aliasScopeAttrs)
-    nodes.push_back(getAliasScope(aliasScopeAttr));
+    nodes.push_back(getOrCreateAliasScope(aliasScopeAttr));
   return llvm::MDNode::get(getLLVMContext(), nodes);
 }
 
 void ModuleTranslation::setAliasScopeMetadata(AliasAnalysisOpInterface op,
                                               llvm::Instruction *inst) {
   auto populateScopeMetadata = [&](ArrayAttr aliasScopeAttrs, unsigned kind) {
     if (!aliasScopeAttrs || aliasScopeAttrs.empty())
       return;
-    llvm::MDNode *node = getAliasScopes(
+    llvm::MDNode *node = getOrCreateAliasScopes(
         llvm::to_vector(aliasScopeAttrs.getAsRange<AliasScopeAttr>()));
     inst->setMetadata(kind, node);
   };
 
   populateScopeMetadata(op.getAliasScopesOrNull(),
                         llvm::LLVMContext::MD_alias_scope);
   populateScopeMetadata(op.getNoAliasScopesOrNull(),
                         llvm::LLVMContext::MD_noalias);
mlir::translateModuleToLLVMIR(Operation *module, llvm::LLVMContext &llvmContext,
 
   ModuleTranslation translator(module, std::move(llvmModule));
   if (failed(translator.convertComdats()))
     return nullptr;
   if (failed(translator.convertFunctionSignatures()))
     return nullptr;
   if (failed(translator.convertGlobals()))
     return nullptr;
-  if (failed(translator.createAliasScopeMetadata()))
-    return nullptr;
   if (failed(translator.createTBAAMetadata()))
     return nullptr;
 
   // Convert other top-level operations if possible.
   llvm::IRBuilder<> llvmBuilder(llvmContext);
   for (Operation &o : getModuleBody(module).getOperations()) {
     if (!isa<LLVM::LLVMFuncOp, LLVM::GlobalOp, LLVM::GlobalCtorsOp,
              LLVM::GlobalDtorsOp, LLVM::ComdatOp>(&o) &&

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

+// RUN: mlir-translate -mlir-to-llvmir -split-input-file %s | FileCheck %s
+
+llvm.func @foo(%arg0: !llvm.ptr)
+
+#alias_scope_domain = #llvm.alias_scope_domain<id = distinct[0]<>, description = "The domain">
+#alias_scope1 = #llvm.alias_scope<id = distinct[1]<>, domain = #alias_scope_domain, description = "The first scope">
+#alias_scope2 = #llvm.alias_scope<id = distinct[2]<>, domain = #alias_scope_domain>
+#alias_scope3 = #llvm.alias_scope<id = distinct[3]<>, domain = #alias_scope_domain>
+
+// CHECK-LABEL: @alias_scopes
+llvm.func @alias_scopes(%arg1 : !llvm.ptr) {
+  %0 = llvm.mlir.constant(0 : i32) : i32
+  // CHECK:  call void @llvm.experimental.noalias.scope.decl(metadata ![[SCOPES1:[0-9]+]])
+  llvm.intr.experimental.noalias.scope.decl #alias_scope1
+  // CHECK:  store {{.*}}, !alias.scope ![[SCOPES1]], !noalias ![[SCOPES23:[0-9]+]]
+  llvm.store %0, %arg1 {alias_scopes = [#alias_scope1], noalias_scopes = [#alias_scope2, #alias_scope3]} : i32, !llvm.ptr
+  // CHECK:  load {{.*}}, !alias.scope ![[SCOPES2:[0-9]+]], !noalias ![[SCOPES13:[0-9]+]]
+  %1 = llvm.load %arg1 {alias_scopes = [#alias_scope2], noalias_scopes = [#alias_scope1, #alias_scope3]} : !llvm.ptr -> i32
+  // CHECK:  atomicrmw {{.*}}, !alias.scope ![[SCOPES3:[0-9]+]], !noalias ![[SCOPES12:[0-9]+]]
+  %2 = llvm.atomicrmw add %arg1, %0 monotonic {alias_scopes = [#alias_scope3], noalias_scopes = [#alias_scope1, #alias_scope2]} : !llvm.ptr, i32
+  // CHECK:  cmpxchg {{.*}}, !alias.scope ![[SCOPES3]]
+  %3 = llvm.cmpxchg %arg1, %1, %2 acq_rel monotonic {alias_scopes = [#alias_scope3]} : !llvm.ptr, i32
+  %5 = llvm.mlir.constant(42 : i8) : i8
+  // CHECK:  llvm.memcpy{{.*}}, !alias.scope ![[SCOPES3]]
+  "llvm.intr.memcpy"(%arg1, %arg1, %0) <{isVolatile = false}> {alias_scopes = [#alias_scope3]} : (!llvm.ptr, !llvm.ptr, i32) -> ()
+  // CHECK:  llvm.memset{{.*}}, !noalias ![[SCOPES3]]
+  "llvm.intr.memset"(%arg1, %5, %0) <{isVolatile = false}> {noalias_scopes = [#alias_scope3]} : (!llvm.ptr, i8, i32) -> ()
+  // CHECK: call void @foo({{.*}} !alias.scope ![[SCOPES3]]
+  llvm.call @foo(%arg1) {alias_scopes = [#alias_scope3]} : (!llvm.ptr) -> ()
+  // CHECK: call void @foo({{.*}} !noalias ![[SCOPES3]]
+  llvm.call @foo(%arg1) {noalias_scopes = [#alias_scope3]} : (!llvm.ptr) -> ()
+  llvm.return
+}
+
+// Check the intrinsic declarations.
+// CHECK-DAG: declare void @llvm.experimental.noalias.scope.decl(metadata)
+// CHECK-DAG: declare void @llvm.memcpy.p0.p0.i32(ptr noalias nocapture writeonly, ptr noalias nocapture readonly, i32, i1 immarg)
+// CHECK-DAG: declare void @llvm.memset.p0.i32(ptr nocapture writeonly, i8, i32, i1 immarg)
+
+// Check the translated metadata.
+// CHECK-DAG: ![[DOMAIN:[0-9]+]] = distinct !{![[DOMAIN]], !"The domain"}
+// CHECK-DAG: ![[SCOPE1:[0-9]+]] = distinct !{![[SCOPE1]], ![[DOMAIN]], !"The first scope"}
+// CHECK-DAG: ![[SCOPE2:[0-9]+]] = distinct !{![[SCOPE2]], ![[DOMAIN]]}
+// CHECK-DAG: ![[SCOPE3:[0-9]+]] = distinct !{![[SCOPE3]], ![[DOMAIN]]}
+// CHECK-DAG: ![[SCOPES1]] = !{![[SCOPE1]]}
+// CHECK-DAG: ![[SCOPES2]] = !{![[SCOPE2]]}
+// CHECK-DAG: ![[SCOPES3]] = !{![[SCOPE3]]}
+// CHECK-DAG: ![[SCOPES12]] = !{![[SCOPE1]], ![[SCOPE2]]}
+// CHECK-DAG: ![[SCOPES13]] = !{![[SCOPE1]], ![[SCOPE3]]}
+// CHECK-DAG: ![[SCOPES23]] = !{![[SCOPE2]], ![[SCOPE3]]}
+
+// -----
+
+// This test verifies the noalias scope intrinsice can be translated in
+// isolation. It is the only operation using alias scopes attributes without
+// implementing AliasAnalysisOpInterface.
+
+#alias_scope_domain = #llvm.alias_scope_domain<id = distinct[0]<>, description = "The domain">
+#alias_scope1 = #llvm.alias_scope<id = distinct[1]<>, domain = #alias_scope_domain>
+
+// CHECK-LABEL: @noalias_intr_only
+llvm.func @noalias_intr_only(%arg1 : !llvm.ptr) {
+  %0 = llvm.mlir.constant(0 : i32) : i32
+  // CHECK:  call void @llvm.experimental.noalias.scope.decl(metadata ![[SCOPES1:[0-9]+]])
+  llvm.intr.experimental.noalias.scope.decl #alias_scope1
+  llvm.return
+}
+
+// Check the translated metadata.
+// CHECK-DAG: ![[DOMAIN:[0-9]+]] = distinct !{![[DOMAIN]], !"The domain"}
+// CHECK-DAG: ![[SCOPE1:[0-9]+]] = distinct !{![[SCOPE1]], ![[DOMAIN]]}
+// CHECK-DAG: ![[SCOPES1]] = !{![[SCOPE1]]}

mlir/test/Target/LLVMIR/llvmir.mlir

 ^bb3: // pred: ^bb0
   llvm.return %arg0 : i32
 }
 
 // CHECK: ![[SWITCH_WEIGHT_NODE]] = !{!"branch_weights", i32 13, i32 17, i32 19}
 
 // -----
 
-llvm.func @foo(%arg0: !llvm.ptr)
-
-#alias_scope_domain = #llvm.alias_scope_domain<id = distinct[0]<>, description = "The domain">
-#alias_scope1 = #llvm.alias_scope<id = distinct[1]<>, domain = #alias_scope_domain, description = "The first scope">
-#alias_scope2 = #llvm.alias_scope<id = distinct[2]<>, domain = #alias_scope_domain>
-#alias_scope3 = #llvm.alias_scope<id = distinct[3]<>, domain = #alias_scope_domain>
-
-// CHECK-LABEL: aliasScope
-llvm.func @aliasScope(%arg1 : !llvm.ptr) {
-  %0 = llvm.mlir.constant(0 : i32) : i32
-  // CHECK:  call void @llvm.experimental.noalias.scope.decl(metadata ![[SCOPES1:[0-9]+]])
-  llvm.intr.experimental.noalias.scope.decl #alias_scope1
-  // CHECK:  store {{.*}}, !alias.scope ![[SCOPES1]], !noalias ![[SCOPES23:[0-9]+]]
-  llvm.store %0, %arg1 {alias_scopes = [#alias_scope1], noalias_scopes = [#alias_scope2, #alias_scope3]} : i32, !llvm.ptr
-  // CHECK:  load {{.*}}, !alias.scope ![[SCOPES2:[0-9]+]], !noalias ![[SCOPES13:[0-9]+]]
-  %1 = llvm.load %arg1 {alias_scopes = [#alias_scope2], noalias_scopes = [#alias_scope1, #alias_scope3]} : !llvm.ptr -> i32
-  // CHECK:  atomicrmw {{.*}}, !alias.scope ![[SCOPES3:[0-9]+]], !noalias ![[SCOPES12:[0-9]+]]
-  %2 = llvm.atomicrmw add %arg1, %0 monotonic {alias_scopes = [#alias_scope3], noalias_scopes = [#alias_scope1, #alias_scope2]} : !llvm.ptr, i32
-  // CHECK:  cmpxchg {{.*}}, !alias.scope ![[SCOPES3]]
-  %3 = llvm.cmpxchg %arg1, %1, %2 acq_rel monotonic {alias_scopes = [#alias_scope3]} : !llvm.ptr, i32
-  %5 = llvm.mlir.constant(42 : i8) : i8
-  // CHECK:  llvm.memcpy{{.*}}, !alias.scope ![[SCOPES3]]
-  "llvm.intr.memcpy"(%arg1, %arg1, %0) <{isVolatile = false}> {alias_scopes = [#alias_scope3]} : (!llvm.ptr, !llvm.ptr, i32) -> ()
-  // CHECK:  llvm.memset{{.*}}, !noalias ![[SCOPES3]]
-  "llvm.intr.memset"(%arg1, %5, %0) <{isVolatile = false}> {noalias_scopes = [#alias_scope3]} : (!llvm.ptr, i8, i32) -> ()
-  // CHECK: call void @foo({{.*}} !alias.scope ![[SCOPES3]]
-  llvm.call @foo(%arg1) {alias_scopes = [#alias_scope3]} : (!llvm.ptr) -> ()
-  // CHECK: call void @foo({{.*}} !noalias ![[SCOPES3]]
-  llvm.call @foo(%arg1) {noalias_scopes = [#alias_scope3]} : (!llvm.ptr) -> ()
-  llvm.return
-}
-
-// Check the intrinsic declarations.
-// CHECK-DAG: declare void @llvm.experimental.noalias.scope.decl(metadata)
-// CHECK-DAG: declare void @llvm.memcpy.p0.p0.i32(ptr noalias nocapture writeonly, ptr noalias nocapture readonly, i32, i1 immarg)
-// CHECK-DAG: declare void @llvm.memset.p0.i32(ptr nocapture writeonly, i8, i32, i1 immarg)
-
-// Check the translated metadata.
-// CHECK-DAG: ![[DOMAIN:[0-9]+]] = distinct !{![[DOMAIN]], !"The domain"}
-// CHECK-DAG: ![[SCOPE1:[0-9]+]] = distinct !{![[SCOPE1]], ![[DOMAIN]], !"The first scope"}
-// CHECK-DAG: ![[SCOPE2:[0-9]+]] = distinct !{![[SCOPE2]], ![[DOMAIN]]}
-// CHECK-DAG: ![[SCOPE3:[0-9]+]] = distinct !{![[SCOPE3]], ![[DOMAIN]]}
-// CHECK-DAG: ![[SCOPES1]] = !{![[SCOPE1]]}
-// CHECK-DAG: ![[SCOPES2]] = !{![[SCOPE2]]}
-// CHECK-DAG: ![[SCOPES3]] = !{![[SCOPE3]]}
-// CHECK-DAG: ![[SCOPES12]] = !{![[SCOPE1]], ![[SCOPE2]]}
-// CHECK-DAG: ![[SCOPES13]] = !{![[SCOPE1]], ![[SCOPE3]]}
-// CHECK-DAG: ![[SCOPES23]] = !{![[SCOPE2]], ![[SCOPE3]]}
-
-// -----
-
 // It is okay to have repeated successors if they have no arguments.
 
 // CHECK-LABEL: @duplicate_block_in_switch
 // CHECK-SAME: float %[[FIRST:.*]],
 // CHECK-SAME: float %[[SECOND:.*]])
 // CHECK:   switch i32 %{{.*}}, label %[[DEFAULT:.*]] [
 // CHECK:     i32 105, label %[[DUPLICATE:.*]]
 // CHECK:     i32 108, label %[[BLOCK:.*]]