[mlir][llvm] Import alias scope metadata from LLVM IR.

[gysit]

Summary

The revision adds support for importing alias.scope and noalias metadata
from LLVM IR into LLVM dialect. It also adds a verifier to the
AliasScopeMetadataOp to check that the associated domain exists
and is of type AliasScopeDomainMetadataOp.

[gysit]

Improve error handling and add tets.

[Dinistro]

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

2660

NIT: isa_and_nonnull<>() could be used here.

[gysit]

Address review comments.

[Dinistro]

LGTM!

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

Example:
           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
 ]> {
   let arguments = (ins
     SymbolNameAttr:$sym_name
   );

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

public:
   FailureOr<SmallVector<SymbolRefAttr>>
   lookupAccessGroupAttrs(const llvm::MDNode *node) const;
 
   /// Returns the loop annotation attribute that corresponds to the given LLVM
   /// loop metadata `node`.
   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>>
+  lookupAliasScopeAttrs(const llvm::MDNode *node) const;
+
 private:
   /// Clears the block and value mapping before processing a new region.
   void clearBlockAndValueMapping() {
     valueMapping.clear();
     noResultOpMapping.clear();
     blockMapping.clear();
   }
   /// Sets the constant insertion point to the start of the given block.
private:
   /// (unless all dependencies have been converted by a previous
   /// invocation of this function).
   LogicalResult processTBAAMetadata(const llvm::MDNode *node);
   /// Converts all LLVM access groups starting from `node` to MLIR access group
   /// operations and stores a mapping from every nested access group node to the
   /// symbol pointing to the translated operation. Returns success if all
   /// conversions succeed and failure otherwise.
   LogicalResult processAccessGroupMetadata(const llvm::MDNode *node);
+  /// Converts all LLVM alias scopes and domains starting from `node` to MLIR
+  /// alias scope and domain operations and stores a mapping from every nested
+  /// alias scope or alias domain node to the symbol pointing to the translated
+  /// operation. Returns success if all conversions succeed and failure
+  /// otherwise.
+  LogicalResult processAliasScopeMetadata(const llvm::MDNode *node);
 
   /// Builder pointing at where the next instruction should be generated.
   OpBuilder builder;
   /// Block to insert the next constant into.
   Block *constantInsertionBlock = nullptr;
   /// Operation to insert the next constant after.
   Operation *constantInsertionOp = nullptr;
   /// Operation to insert the next global after.
private:
   /// Function-local mapping between original and imported block.
   DenseMap<llvm::BasicBlock *, Block *> blockMapping;
   /// Function-local mapping between original and imported values.
   DenseMap<llvm::Value *, Value> valueMapping;
   /// Function-local mapping between original instructions and imported
   /// 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 TBAA metadata nodes and symbol references
-  /// to the LLVMIR dialect TBAA operations corresponding to these
-  /// nodes.
+  /// 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 TBAA metadata nodes and symbol references to the LLVM
+  /// dialect TBAA operations corresponding to these nodes.
   DenseMap<const llvm::MDNode *, SymbolRefAttr> tbaaMapping;
   /// The stateful type translator (contains named structs).
   LLVM::TypeFromLLVMIRTranslator typeTranslator;
   /// Stateful debug information importer.
   std::unique_ptr<detail::DebugImporter> debugImporter;
   /// Loop annotation importer.
   std::unique_ptr<detail::LoopAnnotationImporter> loopAnnotationImporter;
 };
 
 } // namespace LLVM
 } // namespace mlir
 
 #endif // MLIR_TARGET_LLVMIR_MODULEIMPORT_H

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

return emitOpError() << "expected the same number of elements in "
                          << membersName << " and " << getOffsetsAttrName()
                          << ": " << members.size()
                          << " != " << getOffsets().size();
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// AliasScopeMetadataOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult AliasScopeMetadataOp::verify() {
+  Operation *domainOp = SymbolTable::lookupNearestSymbolFrom(
+      this->getOperation(), getDomainAttr());
+  if (!domainOp || !isa<AliasScopeDomainMetadataOp>(domainOp)) {
+    return this->emitOpError()
+           << "expected '" << getDomain()
+           << "' to reference a domain operation in the same region";
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // OpAsmDialectInterface
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct LLVMOpAsmDialectInterface : public OpAsmDialectInterface {
   using OpAsmDialectInterface::OpAsmDialectInterface;
 
   AliasResult getAlias(Attribute attr, raw_ostream &os) const override {

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

#include "mlir/Dialect/LLVMIR/LLVMIntrinsicFromLLVMIRConversions.inc"
 
   return failure();
 }
 
 /// Returns the list of LLVM IR metadata kinds that are convertible to MLIR LLVM
 /// dialect attributes.
 static ArrayRef<unsigned> getSupportedMetadataImpl() {
   static const SmallVector<unsigned> convertibleMetadata = {
-      llvm::LLVMContext::MD_prof, llvm::LLVMContext::MD_tbaa,
-      llvm::LLVMContext::MD_access_group, llvm::LLVMContext::MD_loop};
+      llvm::LLVMContext::MD_prof,         llvm::LLVMContext::MD_tbaa,
+      llvm::LLVMContext::MD_access_group, llvm::LLVMContext::MD_loop,
+      llvm::LLVMContext::MD_noalias,      llvm::LLVMContext::MD_alias_scope};
   return convertibleMetadata;
 }
 
 namespace {
 /// Helper class to attach metadata attributes to specific operation types. It
 /// specializes TypeSwitch to take an Operation and return a LogicalResult.
 template <typename... OpTys>
 struct AttributeSetter {
return AttributeSetter<LoadOp, StoreOp>(op).apply([&](auto memOp) {
     memOp.setTbaaAttr(ArrayAttr::get(memOp.getContext(), tbaaTagSym));
   });
 }
 
 /// Looks up all the symbol references pointing to the access group operations
 /// that map to the access group nodes starting from the access group metadata
 /// `node`, and attaches all of them to the imported operation if the lookups
 /// succeed. Returns failure otherwise.
-static LogicalResult setAccessGroupAttr(const llvm::MDNode *node, Operation *op,
+static LogicalResult setAccessGroupsAttr(const llvm::MDNode *node,
+                                         Operation *op,
                                          LLVM::ModuleImport &moduleImport) {
   FailureOr<SmallVector<SymbolRefAttr>> accessGroups =
       moduleImport.lookupAccessGroupAttrs(node);
   if (failed(accessGroups))
     return failure();
 
   SmallVector<Attribute> accessGroupAttrs(accessGroups->begin(),
                                           accessGroups->end());
   return AttributeSetter<LoadOp, StoreOp>(op).apply([&](auto memOp) {
LoopAnnotationAttr attr =
       moduleImport.translateLoopAnnotationAttr(node, op->getLoc());
   if (!attr)
     return failure();
 
   op->setAttr(LLVMDialect::getLoopAttrName(), attr);
   return success();
 }
 
+/// 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.
+static LogicalResult setAliasScopesAttr(const llvm::MDNode *node, Operation *op,
+                                        LLVM::ModuleImport &moduleImport) {
+  FailureOr<SmallVector<SymbolRefAttr>> aliasScopes =
+      moduleImport.lookupAliasScopeAttrs(node);
+  if (failed(aliasScopes))
+    return failure();
+
+  SmallVector<Attribute> aliasScopeAttrs(aliasScopes->begin(),
+                                         aliasScopes->end());
+  return AttributeSetter<LoadOp, StoreOp>(op).apply([&](auto memOp) {
+    memOp.setAliasScopesAttr(
+        ArrayAttr::get(memOp.getContext(), aliasScopeAttrs));
+  });
+}
+
+/// Looks up all the symbol references pointing to the alias scope operations
+/// that map to the alias scope nodes starting from the noalias metadata `node`,
+/// and attaches all of them to the imported operation if the lookups succeed.
+/// Returns failure otherwise.
+static LogicalResult setNoaliasScopesAttr(const llvm::MDNode *node,
+                                          Operation *op,
+                                          LLVM::ModuleImport &moduleImport) {
+  FailureOr<SmallVector<SymbolRefAttr>> noaliasScopes =
+      moduleImport.lookupAliasScopeAttrs(node);
+  if (failed(noaliasScopes))
+    return failure();
+
+  SmallVector<Attribute> noaliasScopeAttrs(noaliasScopes->begin(),
+                                           noaliasScopes->end());
+  return AttributeSetter<LoadOp, StoreOp>(op).apply([&](auto memOp) {
+    memOp.setNoaliasScopesAttr(
+        ArrayAttr::get(memOp.getContext(), noaliasScopeAttrs));
+  });
+}
+
 namespace {
 
 /// Implementation of the dialect interface that converts operations belonging
 /// to the LLVM dialect to LLVM IR.
 class LLVMDialectLLVMIRImportInterface : public LLVMImportDialectInterface {
 public:
   using LLVMImportDialectInterface::LLVMImportDialectInterface;
 
LogicalResult setMetadataAttrs(OpBuilder &builder, unsigned kind,
                                  llvm::MDNode *node, Operation *op,
                                  LLVM::ModuleImport &moduleImport) const final {
     // Call metadata specific handlers.
     if (kind == llvm::LLVMContext::MD_prof)
       return setProfilingAttr(builder, node, op, moduleImport);
     if (kind == llvm::LLVMContext::MD_tbaa)
       return setTBAAAttr(node, op, moduleImport);
     if (kind == llvm::LLVMContext::MD_access_group)
-      return setAccessGroupAttr(node, op, moduleImport);
+      return setAccessGroupsAttr(node, op, moduleImport);
     if (kind == llvm::LLVMContext::MD_loop)
       return setLoopAttr(node, op, moduleImport);
+    if (kind == llvm::LLVMContext::MD_alias_scope)
+      return setAliasScopesAttr(node, op, moduleImport);
+    if (kind == llvm::LLVMContext::MD_noalias)
+      return setNoaliasScopesAttr(node, op, moduleImport);
 
     // A handler for a supported metadata kind is missing.
     llvm_unreachable("unknown metadata type");
   }
 
   /// Returns the list of LLVM IR intrinsic identifiers that are convertible to
   /// MLIR LLVM dialect intrinsics.
   ArrayRef<unsigned> getSupportedIntrinsics() const final {

mlir/lib/Target/LLVMIR/ModuleImport.cpp

ModuleImport::processAccessGroupMetadata(const llvm::MDNode *node) {
   Location loc = mlirModule.getLoc();
   if (failed(loopAnnotationImporter->translateAccessGroup(
           node, loc, getGlobalMetadataOp())))
     return emitError(loc) << "unsupported access group node: "
                           << diagMD(node, llvmModule.get());
   return success();
 }
 
+LogicalResult
+ModuleImport::processAliasScopeMetadata(const llvm::MDNode *node) {
+  Location loc = mlirModule.getLoc();
+  // Helpers that verify the structure of the metadata.
+  auto hasSelfReference = [](const llvm::MDNode *node) {
+    return node->getNumOperands() != 0 &&
+           node == dyn_cast<llvm::MDNode>(node->getOperand(0));
+  };
+  auto hasOptionalDescription = [](const llvm::MDNode *node, unsigned idx) {
+    return idx >= node->getNumOperands() ||
+           isa<llvm::MDString>(node->getOperand(idx));
+  };
+  auto isDomainNode = [&](const llvm::MDNode *node) {
+    return node && hasSelfReference(node) && hasOptionalDescription(node, 1);
+  };
+  auto isScopeNode = [&](const llvm::MDNode *node) {
+    llvm::AliasScopeNode aliasScope(node);
+    return hasSelfReference(node) && isDomainNode(aliasScope.getDomain()) &&
+           hasOptionalDescription(node, 2);
+  };
+
+  // Helper that creates an alias scope domain operation.
+  auto createAliasScopeDomainOp = [&](const llvm::MDNode *aliasDomain) {
+    StringAttr description = nullptr;
+    if (aliasDomain->getNumOperands() >= 2)
+      if (auto *desc = dyn_cast<llvm::MDString>(aliasDomain->getOperand(1)))
+        description = builder.getStringAttr(desc->getString());
+    std::string name = llvm::formatv("domain_{0}", aliasScopeMapping.size());
+    return builder.create<AliasScopeDomainMetadataOp>(loc, name, description);
+  };
+
+  // Collect the alias scopes and and domains and translate them.
+  for (const llvm::MDOperand &operand : node->operands()) {
+    if (const auto *node = dyn_cast<llvm::MDNode>(operand)) {
+      if (!isScopeNode(node))
+        return emitError(loc) << "unsupported alias scope node: "
+                              << diagMD(node, llvmModule.get());
+
+      llvm::AliasScopeNode aliasScope(node);
+      if (aliasScopeMapping.count(node))
+        continue;
+
+      // Setup the insertion point to the end of the global metadata operation.
+      MetadataOp metadataOp = getGlobalMetadataOp();
+      StringAttr metadataOpName =
+          SymbolTable::getSymbolName(getGlobalMetadataOp());
+      OpBuilder::InsertionGuard guard(builder);
+      builder.setInsertionPointToEnd(&metadataOp.getBody().back());
+
+      // Convert the domain metadata node if it has not been translated before.
+      auto it = aliasScopeMapping.find(aliasScope.getDomain());
+      if (it == aliasScopeMapping.end()) {
+        const llvm::MDNode *domain = aliasScope.getDomain();
+        auto aliasScopeDomainOp = createAliasScopeDomainOp(domain);
+        auto symbolRef = SymbolRefAttr::get(
+            builder.getContext(), metadataOpName,
+            FlatSymbolRefAttr::get(builder.getContext(),
+                                   aliasScopeDomainOp.getSymName()));
+        it = aliasScopeMapping.try_emplace(domain, symbolRef).first;
+      }
+
+      // Convert the scope metadata node if it has not been converted before.
+      StringAttr description = nullptr;
+      if (!aliasScope.getName().empty())
+        description = builder.getStringAttr(aliasScope.getName());
+      std::string name = llvm::formatv("scope_{0}", aliasScopeMapping.size());
+      auto aliasScopeOp = builder.create<AliasScopeMetadataOp>(
+          loc, name, it->getSecond().getLeafReference().getValue(),
+          description);
+      auto symbolRef =
+          SymbolRefAttr::get(builder.getContext(), metadataOpName,
+                             FlatSymbolRefAttr::get(builder.getContext(),
+                                                    aliasScopeOp.getSymName()));
+      aliasScopeMapping.try_emplace(aliasScope.getNode(), symbolRef);
+    }
+  }
+  return success();
+}
+
+FailureOr<SmallVector<SymbolRefAttr>>
+ModuleImport::lookupAliasScopeAttrs(const llvm::MDNode *node) const {
+  SmallVector<SymbolRefAttr> aliasScopes;
+  aliasScopes.reserve(node->getNumOperands());
+  for (const llvm::MDOperand &operand : node->operands()) {
+    auto *node = cast<llvm::MDNode>(operand.get());
+    aliasScopes.push_back(aliasScopeMapping.lookup(node));
+  }
+  // Return failure if one of the alias scope lookups failed.
+  if (llvm::is_contained(aliasScopes, nullptr))
+    return failure();
+  return aliasScopes;
+}
+
 LogicalResult ModuleImport::convertMetadata() {
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointToEnd(mlirModule.getBody());
   for (const llvm::Function &func : llvmModule->functions()) {
     for (const llvm::Instruction &inst : llvm::instructions(func)) {
       // Convert access group metadata nodes.
       if (llvm::MDNode *node =
               inst.getMetadata(llvm::LLVMContext::MD_access_group))
         if (failed(processAccessGroupMetadata(node)))
           return failure();
 
       // Convert alias analysis metadata nodes.
       llvm::AAMDNodes aliasAnalysisNodes = inst.getAAMetadata();
       if (!aliasAnalysisNodes)
         continue;
       if (aliasAnalysisNodes.TBAA)
         if (failed(processTBAAMetadata(aliasAnalysisNodes.TBAA)))
           return failure();
-
-      // TODO: Support noalias and scope metadata nodes.
+      if (aliasAnalysisNodes.Scope)
+        if (failed(processAliasScopeMetadata(aliasAnalysisNodes.Scope)))
+          return failure();
+      if (aliasAnalysisNodes.NoAlias)
+        if (failed(processAliasScopeMetadata(aliasAnalysisNodes.NoAlias)))
+          return failure();
     }
   }
   return success();
 }
 
 LogicalResult ModuleImport::convertGlobals() {
   for (llvm::GlobalVariable &globalVar : llvmModule->globals()) {
     if (globalVar.getName() == getGlobalCtorsVarName() ||

mlir/test/Dialect/LLVMIR/invalid.mlir

module {
   }
   llvm.metadata @metadata {
     llvm.access_group @group
   }
 }
 
 // -----
 
+module {
+  llvm.metadata @metadata {
+    llvm.access_group @group
+    // expected-error@below {{expected 'group' to reference a domain operation in the same region}}
+    llvm.alias_scope @scope { domain = @group }
+  }
+}
+
+// -----
+
+module {
+  llvm.metadata @metadata {
+    // expected-error@below {{expected 'domain' to reference a domain operation in the same region}}
+    llvm.alias_scope @scope { domain = @domain }
+  }
+  llvm.metadata @other_metadata {
+    llvm.alias_scope_domain @domain
+  }
+}
+
+// -----
+
 llvm.func @wmmaLoadOp_invalid_mem_space(%arg0: !llvm.ptr<i32, 5>, %arg1: i32) {
   // expected-error@+1 {{'nvvm.wmma.load' op expected source pointer in memory space 0, 1, 3}}
   %0 = nvvm.wmma.load %arg0, %arg1
     {eltype = #nvvm.mma_type<f16>, frag = #nvvm.mma_frag<a>, k = 16 : i32, layout = #nvvm.mma_layout<row>, m = 16 : i32, n = 16 : i32}
     : (!llvm.ptr<i32, 5>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
   llvm.return
 }
 

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

+; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
+
+; CHECK: llvm.metadata @__llvm_global_metadata {
+; CHECK:   llvm.alias_scope_domain @[[DOMAIN:.*]] {description = "The domain"}
+; CHECK:   llvm.alias_scope @[[$SCOPE0:.*]] {description = "The first scope", domain = @[[DOMAIN]]}
+; CHECK:   llvm.alias_scope @[[$SCOPE1:.*]] {domain = @[[DOMAIN]]}
+; CHECK:   llvm.alias_scope @[[$SCOPE2:.*]] {domain = @[[DOMAIN]]}
+; CHECK: }
+
+; CHECK-LABEL: llvm.func @alias_scope
+define void @alias_scope(ptr %arg1) {
+  ; CHECK: llvm.load
+  ; CHECK-SAME:  alias_scopes = [@__llvm_global_metadata::@[[$SCOPE0]]]
+  ; CHECK-SAME:  noalias_scopes = [@__llvm_global_metadata::@[[$SCOPE1]], @__llvm_global_metadata::@[[$SCOPE2]]]
+  %1 = load i32, ptr %arg1, !alias.scope !4, !noalias !7
+  ; CHECK: llvm.load
+  ; CHECK-SAME:  alias_scopes = [@__llvm_global_metadata::@[[$SCOPE1]]]
+  ; CHECK-SAME:  noalias_scopes = [@__llvm_global_metadata::@[[$SCOPE0]], @__llvm_global_metadata::@[[$SCOPE2]]]
+  %2 = load i32, ptr %arg1, !alias.scope !5, !noalias !8
+  ; CHECK: llvm.load
+  ; CHECK-SAME:  alias_scopes = [@__llvm_global_metadata::@[[$SCOPE2]]]
+  ; CHECK-SAME:  noalias_scopes = [@__llvm_global_metadata::@[[$SCOPE0]], @__llvm_global_metadata::@[[$SCOPE1]]]
+  %3 = load i32, ptr %arg1, !alias.scope !6, !noalias !9
+  ret void
+}
+
+!0 = distinct !{!0, !"The domain"}
+!1 = distinct !{!1, !0, !"The first scope"}
+!2 = distinct !{!2, !0}
+!3 = distinct !{!3, !0}
+!4 = !{!1}
+!5 = !{!2}
+!6 = !{!3}
+!7 = !{!2, !3}
+!8 = !{!1, !3}
+!9 = !{!1, !2}
+
+; // -----
+
+; CHECK: llvm.metadata @__llvm_global_metadata {
+; CHECK:   llvm.alias_scope_domain @[[DOMAIN0:.*]] {description = "The domain"}
+; CHECK:   llvm.alias_scope @[[$SCOPE0:.*]] {domain = @[[DOMAIN0]]}
+; CHECK:   llvm.alias_scope_domain @[[DOMAIN1:.*]]
+; CHECK:   llvm.alias_scope @[[$SCOPE1:.*]] {domain = @[[DOMAIN1]]}
+; CHECK: }
+
+; CHECK-LABEL: llvm.func @two_domains
+define void @two_domains(ptr %arg1) {
+  ; CHECK: llvm.load
+  ; CHECK-SAME:  alias_scopes = [@__llvm_global_metadata::@[[$SCOPE0]]]
+  ; CHECK-SAME:  noalias_scopes = [@__llvm_global_metadata::@[[$SCOPE1]]]
+  %1 = load i32, ptr %arg1, !alias.scope !4, !noalias !5
+  ret void
+}
+
+!0 = distinct !{!0, !"The domain"}
+!1 = distinct !{!1}
+!2 = !{!2, !0}
+!3 = !{!3, !1}
+!4 = !{!2}
+!5 = !{!3}