[mlir][llvm] Fuse MD_access_group & MD_loop import

[Dinistro]

Summary

This commit moves the importing logic of access group metadata into the
loop annotation importer. These two metadata imports can be grouped
because access groups are only used in combination with
llvm.loop.parallel_accesses.

As a nice side effect, this commit decouples the LoopAnnotationImporter
from the ModuleImport class.

[gysit]

nice! LGTM modulo doc comments.

mlir/lib/Target/LLVMIR/LoopAnnotationImporter.h

26

nit: I would drop the back ticks here since they are usually used to document function parameters, e.g.:

... and llvm.access.group metadata nodes into AccessGroupMetadataOps.

34–35

It sounds like we create the operations twice here, maybe:

Converts all LLVM access groups starting from node to MLIR access group operations nested in the region of metadataOp.

35

nit: node

[Dinistro]

address review comments

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

private:
   /// 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.
   DenseMap<const llvm::MDNode *, SymbolRefAttr> tbaaMapping;
-  /// Mapping between original LLVM access group metadata nodes and the symbol
-  /// references pointing to the imported MLIR access group operations.
-  DenseMap<const llvm::MDNode *, SymbolRefAttr> accessGroupMapping;
   /// 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/Target/LLVMIR/LoopAnnotationImporter.h

 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Target/LLVMIR/ModuleImport.h"
 
 namespace mlir {
 namespace LLVM {
 namespace detail {
 
-/// A helper class that converts a `llvm.loop` metadata node into a
-/// corresponding LoopAnnotationAttr.
+/// A helper class that converts llvm.loop metadata nodes into corresponding
+/// LoopAnnotationAttrs and llvm.access.group nodes into
+/// AccessGroupMetadataOps.
 class LoopAnnotationImporter {
 public:
-  explicit LoopAnnotationImporter(ModuleImport &moduleImport)
-      : moduleImport(moduleImport) {}
-  LoopAnnotationAttr translate(const llvm::MDNode *node, Location loc);
+  explicit LoopAnnotationImporter(OpBuilder &builder) : builder(builder) {}
+  LoopAnnotationAttr translateLoopAnnotation(const llvm::MDNode *node,
+                                             Location loc);
+
+  /// Converts all LLVM access groups starting from node to MLIR access group
+  /// operations mested in the region of metadataOp. It stores a mapping from
+  /// every nested access group nod to the symbol pointing to the translated
+  /// operation. Returns success if all conversions succeed and failure
+  /// otherwise.
+  LogicalResult translateAccessGroup(const llvm::MDNode *node, Location loc,
+                                     MetadataOp metadataOp);
+
+  /// Returns the symbol references pointing to the access group operations that
+  /// map to the access group nodes starting from the access group metadata
+  /// node. Returns failure, if any of the symbol references cannot be found.
+  FailureOr<SmallVector<SymbolRefAttr>>
+  lookupAccessGroupAttrs(const llvm::MDNode *node) const;
 
 private:
   /// Returns the LLVM metadata corresponding to a llvm loop metadata attribute.
   LoopAnnotationAttr lookupLoopMetadata(const llvm::MDNode *node) const {
     return loopMetadataMapping.lookup(node);
   }
 
   void mapLoopMetadata(const llvm::MDNode *metadata, LoopAnnotationAttr attr) {
     auto result = loopMetadataMapping.try_emplace(metadata, attr);
     (void)result;
     assert(result.second &&
            "attempting to map loop options that was already mapped");
   }
 
-  ModuleImport &moduleImport;
+  OpBuilder &builder;
   DenseMap<const llvm::MDNode *, LoopAnnotationAttr> loopMetadataMapping;
+  /// Mapping between original LLVM access group metadata nodes and the symbol
+  /// references pointing to the imported MLIR access group operations.
+  DenseMap<const llvm::MDNode *, SymbolRefAttr> accessGroupMapping;
 };
 
 } // namespace detail
 } // namespace LLVM
 } // namespace mlir
 
 #endif // MLIR_LIB_TARGET_LLVMIR_LOOPANNOTATIONIMPORTER_H_

mlir/lib/Target/LLVMIR/LoopAnnotationImporter.cpp

 
 using namespace mlir;
 using namespace mlir::LLVM;
 using namespace mlir::LLVM::detail;
 
 namespace {
 /// Helper class that keeps the state of one metadata to attribute conversion.
 struct LoopMetadataConversion {
-  LoopMetadataConversion(const llvm::MDNode *node, ModuleImport &moduleImport,
-                         Location loc,
+  LoopMetadataConversion(const llvm::MDNode *node, Location loc,
                          LoopAnnotationImporter &loopAnnotationImporter)
-      : node(node), moduleImport(moduleImport), loc(loc),
-        loopAnnotationImporter(loopAnnotationImporter),
+      : node(node), loc(loc), loopAnnotationImporter(loopAnnotationImporter),
         ctx(loc->getContext()){};
   /// Converts this structs loop metadata node into a LoopAnnotationAttr.
   LoopAnnotationAttr convert();
 
   LogicalResult initPropertyMap();
 
   /// Helper function to get and erase a property.
   const llvm::MDNode *lookupAndEraseProperty(StringRef name);
struct LoopMetadataConversion {
   FailureOr<LoopUnrollAndJamAttr> convertUnrollAndJamAttr();
   FailureOr<LoopLICMAttr> convertLICMAttr();
   FailureOr<LoopDistributeAttr> convertDistributeAttr();
   FailureOr<LoopPipelineAttr> convertPipelineAttr();
   FailureOr<SmallVector<SymbolRefAttr>> convertParallelAccesses();
 
   llvm::StringMap<const llvm::MDNode *> propertyMap;
   const llvm::MDNode *node;
-  ModuleImport &moduleImport;
   Location loc;
   LoopAnnotationImporter &loopAnnotationImporter;
   MLIRContext *ctx;
 };
 } // namespace
 
 LogicalResult LoopMetadataConversion::initPropertyMap() {
   // Check if it's a valid node.

 FailureOr<LoopAnnotationAttr>
 LoopMetadataConversion::lookupFollowupNode(StringRef name) {
   auto node = lookupMDNode(name);
   if (failed(node))
     return failure();
   if (*node == nullptr)
     return LoopAnnotationAttr(nullptr);
 
-  return loopAnnotationImporter.translate(*node, loc);
+  return loopAnnotationImporter.translateLoopAnnotation(*node, loc);
 }
 
 static bool isEmptyOrNull(const Attribute attr) { return !attr; }
 
 template <typename T>
 static bool isEmptyOrNull(const SmallVectorImpl<T> &vec) {
   return vec.empty();
 }

 LoopMetadataConversion::convertParallelAccesses() {
   FailureOr<SmallVector<llvm::MDNode *>> nodes =
       lookupMDNodes("llvm.loop.parallel_accesses");
   if (failed(nodes))
     return failure();
   SmallVector<SymbolRefAttr> refs;
   for (llvm::MDNode *node : *nodes) {
     FailureOr<SmallVector<SymbolRefAttr>> accessGroups =
-        moduleImport.lookupAccessGroupAttrs(node);
+        loopAnnotationImporter.lookupAccessGroupAttrs(node);
     if (failed(accessGroups))
       return emitWarning(loc) << "could not lookup access group";
     llvm::append_range(refs, *accessGroups);
   }
   return refs;
 }
 
 LoopAnnotationAttr LoopMetadataConversion::convert() {
LoopAnnotationAttr LoopMetadataConversion::convert() {
   }
 
   return createIfNonNull<LoopAnnotationAttr>(
       ctx, disableNonForced, vecAttr, interleaveAttr, unrollAttr,
       unrollAndJamAttr, licmAttr, distributeAttr, pipelineAttr, mustProgress,
       parallelAccesses);
 }
 
-LoopAnnotationAttr LoopAnnotationImporter::translate(const llvm::MDNode *node,
-                                                     Location loc) {
+LoopAnnotationAttr
+LoopAnnotationImporter::translateLoopAnnotation(const llvm::MDNode *node,
   if (!node)
     return {};
 
   // Note: This check is necessary to distinguish between failed translations
   // and not yet attempted translations.
   auto it = loopMetadataMapping.find(node);
   if (it != loopMetadataMapping.end())
     return it->getSecond();
 
-  LoopAnnotationAttr attr =
-      LoopMetadataConversion(node, moduleImport, loc, *this).convert();
+  LoopAnnotationAttr attr = LoopMetadataConversion(node, loc, *this).convert();
 
   mapLoopMetadata(node, attr);
   return attr;
 }
+
+LogicalResult LoopAnnotationImporter::translateAccessGroup(
+    const llvm::MDNode *node, Location loc, MetadataOp metadataOp) {
+  SmallVector<const llvm::MDNode *> accessGroups;
+  if (!node->getNumOperands())
+    accessGroups.push_back(node);
+  for (const llvm::MDOperand &operand : node->operands()) {
+    auto *childNode = dyn_cast<llvm::MDNode>(operand);
+    if (!childNode)
+      return emitWarning(loc)
+             << "expected access group operands to be metadata nodes";
+    accessGroups.push_back(cast<llvm::MDNode>(operand.get()));
+  }
+
+  // Convert all entries of the access group list to access group operations.
+  for (const llvm::MDNode *accessGroup : accessGroups) {
+    if (accessGroupMapping.count(accessGroup))
+      continue;
+    // Verify the access group node is distinct and empty.
+    if (accessGroup->getNumOperands() != 0 || !accessGroup->isDistinct())
+      return emitWarning(loc)
+             << "expected an access group node to be empty and distinct";
+
+    OpBuilder::InsertionGuard guard(builder);
+    builder.setInsertionPointToEnd(&metadataOp.getBody().back());
+    auto groupOp = builder.create<AccessGroupMetadataOp>(
+        loc, llvm::formatv("group_{0}", accessGroupMapping.size()).str());
+    // Add a mapping from the access group node to the symbol reference pointing
+    // to the newly created operation.
+    accessGroupMapping[accessGroup] = SymbolRefAttr::get(
+        builder.getContext(), metadataOp.getSymName(),
+        FlatSymbolRefAttr::get(builder.getContext(), groupOp.getSymName()));
+  }
+  return success();
+}
+
+FailureOr<SmallVector<SymbolRefAttr>>
+LoopAnnotationImporter::lookupAccessGroupAttrs(const llvm::MDNode *node) const {
+  // An access group node is either a single access group or an access group
+  // list.
+  SmallVector<SymbolRefAttr> accessGroups;
+  if (!node->getNumOperands())
+    accessGroups.push_back(accessGroupMapping.lookup(node));
+  for (const llvm::MDOperand &operand : node->operands()) {
+    auto *node = cast<llvm::MDNode>(operand.get());
+    accessGroups.push_back(accessGroupMapping.lookup(node));
+  }
+  // Exit if one of the access group node lookups failed.
+  if (llvm::is_contained(accessGroups, nullptr))
+    return failure();
+  return accessGroups;
+}

mlir/lib/Target/LLVMIR/ModuleImport.cpp

 
 ModuleImport::ModuleImport(ModuleOp mlirModule,
                            std::unique_ptr<llvm::Module> llvmModule)
     : builder(mlirModule->getContext()), context(mlirModule->getContext()),
       mlirModule(mlirModule), llvmModule(std::move(llvmModule)),
       iface(mlirModule->getContext()),
       typeTranslator(*mlirModule->getContext()),
       debugImporter(std::make_unique<DebugImporter>(mlirModule)),
-      loopAnnotationImporter(std::make_unique<LoopAnnotationImporter>(*this)) {
+      loopAnnotationImporter(
+          std::make_unique<LoopAnnotationImporter>(builder)) {
   builder.setInsertionPointToStart(mlirModule.getBody());
 }
 
 MetadataOp ModuleImport::getGlobalMetadataOp() {
   if (globalMetadataOp)
     return globalMetadataOp;
 
   OpBuilder::InsertionGuard guard(builder);
for (const auto *current : nodesToConvert) {
     }
   }
 
   return success();
 }
 
 LogicalResult
 ModuleImport::processAccessGroupMetadata(const llvm::MDNode *node) {
-  // An access group node is either access group or an access group list. Start
-  // by collecting all access groups to translate.
-  SmallVector<const llvm::MDNode *> accessGroups;
-  if (!node->getNumOperands())
-    accessGroups.push_back(node);
-  for (const llvm::MDOperand &operand : node->operands())
-    accessGroups.push_back(cast<llvm::MDNode>(operand.get()));
-
-  // Convert all entries of the access group list to access group operations.
-  for (const llvm::MDNode *accessGroup : accessGroups) {
-    if (accessGroupMapping.count(accessGroup))
-      continue;
-    // Verify the access group node is distinct and empty.
-    Location loc = mlirModule.getLoc();
-    if (accessGroup->getNumOperands() != 0 || !accessGroup->isDistinct())
-      return emitError(loc) << "unsupported access group node: "
-                            << diagMD(accessGroup, llvmModule.get());
-
-    MetadataOp metadataOp = getGlobalMetadataOp();
-    OpBuilder::InsertionGuard guard(builder);
-    builder.setInsertionPointToEnd(&metadataOp.getBody().back());
-    auto groupOp = builder.create<AccessGroupMetadataOp>(
-        loc, (Twine("group_") + Twine(accessGroupMapping.size())).str());
-    // Add a mapping from the access group node to the symbol reference pointing
-    // to the newly created operation.
-    accessGroupMapping[accessGroup] = SymbolRefAttr::get(
-        builder.getContext(), metadataOp.getSymName(),
-        FlatSymbolRefAttr::get(builder.getContext(), groupOp.getSymName()));
-  }
+  if (failed(loopAnnotationImporter->translateAccessGroup(
+          node, loc, getGlobalMetadataOp())))
+                          << diagMD(node, llvmModule.get());
   return success();
 }
 
 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)) {
if (Operation *op = lookupOperation(&inst)) {
       emitWarning(loc) << "dropped instruction: " << diag(inst);
     }
   }
   return success();
 }
 
 FailureOr<SmallVector<SymbolRefAttr>>
 ModuleImport::lookupAccessGroupAttrs(const llvm::MDNode *node) const {
-  // An access group node is either a single access group or an access group
-  // list.
-  SmallVector<SymbolRefAttr> accessGroups;
-  if (!node->getNumOperands())
-    accessGroups.push_back(accessGroupMapping.lookup(node));
-  for (const llvm::MDOperand &operand : node->operands()) {
-    auto *node = cast<llvm::MDNode>(operand.get());
-    accessGroups.push_back(accessGroupMapping.lookup(node));
-  }
-  // Exit if one of the access group node lookups failed.
-  if (llvm::is_contained(accessGroups, nullptr))
-    return failure();
-  return accessGroups;
+  return loopAnnotationImporter->lookupAccessGroupAttrs(node);
 }
 
 LoopAnnotationAttr
 ModuleImport::translateLoopAnnotationAttr(const llvm::MDNode *node,
                                           Location loc) const {
-  return loopAnnotationImporter->translate(node, loc);
+  return loopAnnotationImporter->translateLoopAnnotation(node, loc);
 }
 
 OwningOpRef<ModuleOp>
 mlir::translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
                               MLIRContext *context) {
   // Preload all registered dialects to allow the import to iterate the
   // registered LLVMImportDialectInterface implementations and query the
   // supported LLVM IR constructs before starting the translation. Assumes the

mlir/test/Target/LLVMIR/Import/import-failure.ll

 !2 = !{!3, !3, i64 0, i64 4}
 !3 = !{!4, i64 4, !"int"}
 !4 = !{!5, i64 1, !"omnipotent char"}
 !5 = !{!"Simple C++ TBAA"}
 
 ; // -----
 
 ; CHECK:      import-failure.ll
-; CHECK-SAME: error: unsupported access group node: !0 = !{}
+; CHECK-SAME: warning: expected an access group node to be empty and distinct
+; CHECK:      error: unsupported access group node: !0 = !{}
 define void @access_group(ptr %arg1) {
   %1 = load i32, ptr %arg1, !llvm.access.group !0
   ret void
 }
 
 !0 = !{}
 
 ; // -----
 
 ; CHECK:      import-failure.ll
-; CHECK-SAME: error: unsupported access group node: !1 = distinct !{!"unsupported access group"}
+; CHECK-SAME: warning: expected an access group node to be empty and distinct
+; CHECK:      error: unsupported access group node: !0 = !{!1}
 define void @access_group(ptr %arg1) {
   %1 = load i32, ptr %arg1, !llvm.access.group !0
   ret void
 }
 
 !0 = !{!1}
 !1 = distinct !{!"unsupported access group"}
 
 ; // -----
 
 ; CHECK:      import-failure.ll
+; CHECK-SAME: warning: expected access group operands to be metadata nodes
+; CHECK:      error: unsupported access group node: !0 = !{i1 false}
+define void @access_group(ptr %arg1) {
+  %1 = load i32, ptr %arg1, !llvm.access.group !0
+  ret void
+}
+
+!0 = !{i1 false}
+
+; // -----
+
+; CHECK:      import-failure.ll
 ; CHECK-SAME: warning: expected all loop properties to be either debug locations or metadata nodes
 ; CHECK:      import-failure.ll
 ; CHECK-SAME: warning: unhandled metadata: !0 = distinct !{!0, i32 42}
 define void @invalid_loop_node(i64 %n, ptr %A) {
 entry:
   br label %end, !llvm.loop !0
 end:
   ret void