[mlir][llvm] Fuse access_group & loop export (NFC)

[Dinistro]

Summary

This commit moves the access group translation into the
LoopAnnotationTranslation class as these two metadata kinds only appear
together.

Drops the access group cleanup from ModuleTranslation::forgetMapping
as this is only used on function regions. Access groups only appear in the
region of a global metadata operation and will thus not be cleaned here.

Analogous to https://reviews.llvm.org/D143577

[gysit]

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

335–338

Is this still used? I would have expected this can be deleted now?

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

452

can this be deleted since the access group metadata ops only exist in metadata operation regions?

[gysit]

Also should this be an NFC commit?

[Dinistro]

In D144253#4134704, @gysit wrote:

Also should this be an NFC commit?

That depends on the definition of NFC. It removes public member functions, so in theory it has a functional change for potential downstream users. On the other hand, this does not change any upstream LLVMIR export behavior.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

452

Exactly. Should I mention this in the commit message?

[gysit]

In D144253#4134708, @Dinistro wrote:

In D144253#4134704, @gysit wrote:

Also should this be an NFC commit?

That depends on the definition of NFC. It removes public member functions, so in theory it has a functional change for potential downstream users. On the other hand, this does not change any upstream LLVMIR export behavior.

The standard definition is usually if there is a test that changed, which is not the case here. However, it is a corner case since the public API changed.

Exactly. Should I mention this in the commit message?

Yes that would be good.

[Dinistro]

address review comments

[Dinistro]

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

335–338

Good catch

[gysit]

LGTM

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

llvm::Instruction *lookupBranch(Operation *op) const {
     return branchMapping.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 symbol reference to an mlir
-  /// LLVM dialect access group operation.
-  llvm::MDNode *getAccessGroup(Operation *op,
-                               SymbolRefAttr accessGroupRef) const;
-
-  /// 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;
 
   // Sets LLVM metadata for memory operations that are in a parallel loop.
   void setAccessGroupsMetadata(Operation *op, llvm::Instruction *inst);
 
   // Sets LLVM metadata for memory operations that have alias scope information.
   void setAliasScopeMetadata(Operation *op, llvm::Instruction *inst);
private:
   DenseMap<Value, llvm::Value *> valueMapping;
   DenseMap<Block *, llvm::BasicBlock *> blockMapping;
 
   /// A mapping between MLIR LLVM dialect terminators and LLVM IR terminators
   /// they are converted to. This allows for connecting PHI nodes to the source
   /// values after all operations are converted.
   DenseMap<Operation *, llvm::Instruction *> branchMapping;
 
-  /// Mapping from an access group metadata operation to its LLVM metadata.
-  /// This map is populated on module entry and is used to annotate loops (as
-  /// identified via their branches) and contained memory accesses.
-  DenseMap<Operation *, llvm::MDNode *> accessGroupMetadataMapping;
-
   /// Mapping from an alias scope metadata operation to its LLVM metadata.
   /// This map is populated on module entry.
   DenseMap<Operation *, llvm::MDNode *> aliasScopeMetadataMapping;
 
   /// Mapping from a tbaa metadata operation to its LLVM metadata.
   /// This map is populated on module entry.
   DenseMap<const Operation *, llvm::MDNode *> tbaaMetadataMapping;
 

mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.h

 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"
 
 namespace mlir {
 namespace LLVM {
 namespace detail {
 
-/// A helper class that converts a LoopAnnotationAttr into a corresponding
-/// llvm::MDNode.
+/// A helper class that converts LoopAnnotationAttrs and AccessGroupMetadataOps
+/// into a corresponding llvm::MDNodes.
 class LoopAnnotationTranslation {
 public:
-  LoopAnnotationTranslation(LLVM::ModuleTranslation &moduleTranslation)
-      : moduleTranslation(moduleTranslation) {}
+  LoopAnnotationTranslation(Operation *mlirModule, llvm::Module &llvmModule)
+      : mlirModule(mlirModule), llvmModule(llvmModule) {}
 
-  llvm::MDNode *translate(LoopAnnotationAttr attr, Operation *op);
+  llvm::MDNode *translateLoopAnnotation(LoopAnnotationAttr attr, Operation *op);
+
+  /// Traverses the global access group metadata operation in the `mlirModule`
+  /// and creates corresponding LLVM metadata nodes.
+  LogicalResult createAccessGroupMetadata();
+
+  /// Returns the LLVM metadata corresponding to a symbol reference to an mlir
+  /// LLVM dialect access group operation.
+  llvm::MDNode *getAccessGroup(Operation *op,
+                               SymbolRefAttr accessGroupRef) const;
+
+  /// Returns the LLVM metadata corresponding to a list of symbol reference to
+  /// an mlir LLVM dialect access group operation. Returns nullptr if
+  /// `accessGroupRefs` is null or empty.
+  llvm::MDNode *getAccessGroups(Operation *op, ArrayAttr accessGroupRefs) const;
 
 private:
   /// Returns the LLVM metadata corresponding to a llvm loop metadata attribute.
   llvm::MDNode *lookupLoopMetadata(Attribute options) const {
     return loopMetadataMapping.lookup(options);
   }
 
   void mapLoopMetadata(Attribute options, llvm::MDNode *metadata) {
     auto result = loopMetadataMapping.try_emplace(options, metadata);
     (void)result;
     assert(result.second &&
            "attempting to map loop options that was already mapped");
   }
 
   /// Mapping from an attribute describing loop metadata to its LLVM metadata.
   /// The metadata is attached to Latch block branches with this attribute.
   DenseMap<Attribute, llvm::MDNode *> loopMetadataMapping;
 
-  LLVM::ModuleTranslation &moduleTranslation;
+  /// Mapping from an access group metadata operation to its LLVM metadata.
+  /// This map is populated on module entry and is used to annotate loops (as
+  /// identified via their branches) and contained memory accesses.
+  DenseMap<Operation *, llvm::MDNode *> accessGroupMetadataMapping;
+
+  Operation *mlirModule;
+  llvm::Module &llvmModule;
 };
 
 } // namespace detail
 } // namespace LLVM
 } // namespace mlir
 
 #endif // MLIR_LIB_TARGET_LLVMIR_LOOPANNOTATIONTRANSLATION_H_

mlir/lib/Target/LLVMIR/LoopAnnotationTranslation.cpp

 
 using namespace mlir;
 using namespace mlir::LLVM;
 using namespace mlir::LLVM::detail;
 
 namespace {
 /// Helper class that keeps the state of one attribute to metadata conversion.
 struct LoopAnnotationConversion {
-  LoopAnnotationConversion(LoopAnnotationAttr attr,
-                           ModuleTranslation &moduleTranslation, Operation *op,
-                           LoopAnnotationTranslation &loopAnnotationTranslation)
-      : attr(attr), moduleTranslation(moduleTranslation), op(op),
-        loopAnnotationTranslation(loopAnnotationTranslation),
-        ctx(moduleTranslation.getLLVMContext()) {}
+  LoopAnnotationConversion(LoopAnnotationAttr attr, Operation *op,
+                           LoopAnnotationTranslation &loopAnnotationTranslation,
+                           llvm::LLVMContext &ctx)
+      : attr(attr), op(op),
+        loopAnnotationTranslation(loopAnnotationTranslation), ctx(ctx) {}
 
   /// Converts this struct's loop annotation into a corresponding LLVMIR
   /// metadata representation.
   llvm::MDNode *convert();
 
   /// Conversion functions for different payload attribute kinds.
   void addUnitNode(StringRef name);
   void addUnitNode(StringRef name, BoolAttr attr);
struct LoopAnnotationConversion {
   void convertLoopOptions(LoopUnrollAndJamAttr options);
   void convertLoopOptions(LoopLICMAttr options);
   void convertLoopOptions(LoopDistributeAttr options);
   void convertLoopOptions(LoopPipelineAttr options);
   void convertLoopOptions(LoopPeeledAttr options);
   void convertLoopOptions(LoopUnswitchAttr options);
 
   LoopAnnotationAttr attr;
-  ModuleTranslation &moduleTranslation;
   Operation *op;
   LoopAnnotationTranslation &loopAnnotationTranslation;
   llvm::LLVMContext &ctx;
   llvm::SmallVector<llvm::Metadata *> metadataNodes;
 };
 } // namespace
 
 void LoopAnnotationConversion::addUnitNode(StringRef name) {
void LoopAnnotationConversion::convertI32Node(StringRef name,
   addI32NodeWithVal(name, attr.getInt());
 }
 
 void LoopAnnotationConversion::convertFollowupNode(StringRef name,
                                                    LoopAnnotationAttr attr) {
   if (!attr)
     return;
 
-  llvm::MDNode *node = loopAnnotationTranslation.translate(attr, op);
+  llvm::MDNode *node =
+      loopAnnotationTranslation.translateLoopAnnotation(attr, op);
 
   metadataNodes.push_back(
       llvm::MDNode::get(ctx, {llvm::MDString::get(ctx, name), node}));
 }
 
 void LoopAnnotationConversion::convertLoopOptions(LoopVectorizeAttr options) {
   convertBoolNode("llvm.loop.vectorize.enable", options.getDisable(), true);
   convertBoolNode("llvm.loop.vectorize.predicate.enable",
llvm::MDNode *LoopAnnotationConversion::convert() {
 
   ArrayRef<SymbolRefAttr> parallelAccessGroups = attr.getParallelAccesses();
   if (!parallelAccessGroups.empty()) {
     SmallVector<llvm::Metadata *> parallelAccess;
     parallelAccess.push_back(
         llvm::MDString::get(ctx, "llvm.loop.parallel_accesses"));
     for (SymbolRefAttr accessGroupRef : parallelAccessGroups)
       parallelAccess.push_back(
-          moduleTranslation.getAccessGroup(op, accessGroupRef));
+          loopAnnotationTranslation.getAccessGroup(op, accessGroupRef));
     metadataNodes.push_back(llvm::MDNode::get(ctx, parallelAccess));
   }
 
   // Create loop options and set the first operand to itself.
   llvm::MDNode *loopMD = llvm::MDNode::get(ctx, metadataNodes);
   loopMD->replaceOperandWith(0, loopMD);
 
   return loopMD;
 }
 
-llvm::MDNode *LoopAnnotationTranslation::translate(LoopAnnotationAttr attr,
+llvm::MDNode *
+LoopAnnotationTranslation::translateLoopAnnotation(LoopAnnotationAttr attr,
                                                    Operation *op) {
   if (!attr)
     return nullptr;
 
   llvm::MDNode *loopMD = lookupLoopMetadata(attr);
   if (loopMD)
     return loopMD;
 
   loopMD =
-      LoopAnnotationConversion(attr, moduleTranslation, op, *this).convert();
+      LoopAnnotationConversion(attr, op, *this, this->llvmModule.getContext())
+          .convert();
   // Store a map from this Attribute to the LLVM metadata in case we
   // encounter it again.
   mapLoopMetadata(attr, loopMD);
   return loopMD;
 }
+
+LogicalResult LoopAnnotationTranslation::createAccessGroupMetadata() {
+  mlirModule->walk([&](LLVM::MetadataOp metadatas) {
+    metadatas.walk([&](LLVM::AccessGroupMetadataOp op) {
+      llvm::MDNode *accessGroup =
+          llvm::MDNode::getDistinct(llvmModule.getContext(), {});
+      accessGroupMetadataMapping.insert({op, accessGroup});
+    });
+  });
+  return success();
+}
+
+llvm::MDNode *
+LoopAnnotationTranslation::getAccessGroup(Operation *op,
+                                          SymbolRefAttr accessGroupRef) const {
+  auto metadataName = accessGroupRef.getRootReference();
+  auto accessGroupName = accessGroupRef.getLeafReference();
+  auto metadataOp = SymbolTable::lookupNearestSymbolFrom<LLVM::MetadataOp>(
+      op->getParentOp(), metadataName);
+  auto *accessGroupOp =
+      SymbolTable::lookupNearestSymbolFrom(metadataOp, accessGroupName);
+  return accessGroupMetadataMapping.lookup(accessGroupOp);
+}
+
+llvm::MDNode *
+LoopAnnotationTranslation::getAccessGroups(Operation *op,
+                                           ArrayAttr accessGroupRefs) const {
+  if (!accessGroupRefs || accessGroupRefs.empty())
+    return nullptr;
+
+  SmallVector<llvm::Metadata *> groupMDs;
+  for (SymbolRefAttr groupRef : accessGroupRefs.getAsRange<SymbolRefAttr>())
+    groupMDs.push_back(getAccessGroup(op, groupRef));
+  if (groupMDs.size() == 1)
+    return llvm::cast<llvm::MDNode>(groupMDs.front());
+  return llvm::MDNode::get(llvmModule.getContext(), groupMDs);
+}

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

llvm::Constant *mlir::LLVM::detail::getLLVMConstant(
   return nullptr;
 }
 
 ModuleTranslation::ModuleTranslation(Operation *module,
                                      std::unique_ptr<llvm::Module> llvmModule)
     : mlirModule(module), llvmModule(std::move(llvmModule)),
       debugTranslation(
           std::make_unique<DebugTranslation>(module, *this->llvmModule)),
-      loopAnnotationTranslation(
-          std::make_unique<LoopAnnotationTranslation>(*this)),
+      loopAnnotationTranslation(std::make_unique<LoopAnnotationTranslation>(
+          module, *this->llvmModule)),
       typeTranslator(this->llvmModule->getContext()),
       iface(module->getContext()) {
   assert(satisfiesLLVMModule(mlirModule) &&
          "mlirModule should honor LLVM's module semantics.");
 }
 ModuleTranslation::~ModuleTranslation() {
   if (ompBuilder)
     ompBuilder->finalize();
for (Block &block : *current) {
         valueMapping.erase(arg);
       for (Operation &op : block) {
         for (Value value : op.getResults())
           valueMapping.erase(value);
         if (op.hasSuccessors())
           branchMapping.erase(&op);
         if (isa<LLVM::GlobalOp>(op))
           globalsMapping.erase(&op);
-        accessGroupMetadataMapping.erase(&op);
         llvm::append_range(
             toProcess,
             llvm::map_range(op.getRegions(), [](Region &r) { return &r; }));
       }
     }
   }
 }
 
for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
 
     if (failed(convertOneFunction(function)))
       return failure();
   }
 
   return success();
 }
 
-llvm::MDNode *
-ModuleTranslation::getAccessGroup(Operation *op,
-                                  SymbolRefAttr accessGroupRef) const {
-  auto metadataName = accessGroupRef.getRootReference();
-  auto accessGroupName = accessGroupRef.getLeafReference();
-  auto metadataOp = SymbolTable::lookupNearestSymbolFrom<LLVM::MetadataOp>(
-      op->getParentOp(), metadataName);
-  auto *accessGroupOp =
-      SymbolTable::lookupNearestSymbolFrom(metadataOp, accessGroupName);
-  return accessGroupMetadataMapping.lookup(accessGroupOp);
-}
-
 LogicalResult ModuleTranslation::createAccessGroupMetadata() {
-  mlirModule->walk([&](LLVM::MetadataOp metadatas) {
-    metadatas.walk([&](LLVM::AccessGroupMetadataOp op) {
-      llvm::LLVMContext &ctx = llvmModule->getContext();
-      llvm::MDNode *accessGroup = llvm::MDNode::getDistinct(ctx, {});
-      accessGroupMetadataMapping.insert({op, accessGroup});
-    });
-  });
-  return success();
+  return loopAnnotationTranslation->createAccessGroupMetadata();
 }
 
 void ModuleTranslation::setAccessGroupsMetadata(Operation *op,
                                                 llvm::Instruction *inst) {
   auto populateGroupsMetadata = [&](ArrayAttr groupRefs) {
-    if (!groupRefs || groupRefs.empty())
-      return;
-
-    llvm::Module *module = inst->getModule();
-    SmallVector<llvm::Metadata *> groupMDs;
-    for (SymbolRefAttr groupRef : groupRefs.getAsRange<SymbolRefAttr>())
-      groupMDs.push_back(getAccessGroup(op, groupRef));
-
-    llvm::MDNode *node = nullptr;
-    if (groupMDs.size() == 1)
-      node = llvm::cast<llvm::MDNode>(groupMDs.front());
-    else if (groupMDs.size() >= 2)
-      node = llvm::MDNode::get(module->getContext(), groupMDs);
-
+    if (llvm::MDNode *node =
+            loopAnnotationTranslation->getAccessGroups(op, groupRefs))
       inst->setMetadata(llvm::LLVMContext::MD_access_group, node);
   };
 
   auto groupRefs =
       op->getAttrOfType<ArrayAttr>(LLVMDialect::getAccessGroupsAttrName());
   populateGroupsMetadata(groupRefs);
 }
 
 LogicalResult ModuleTranslation::createAliasScopeMetadata() {

 void ModuleTranslation::setLoopMetadata(Operation *op,
                                         llvm::Instruction *inst) {
   LoopAnnotationAttr attr =
       TypeSwitch<Operation *, LoopAnnotationAttr>(op)
           .Case<LLVM::BrOp, LLVM::CondBrOp>(
               [](auto branchOp) { return branchOp.getLoopAnnotationAttr(); });
   if (!attr)
     return;
-  llvm::MDNode *loopMD = loopAnnotationTranslation->translate(attr, op);
+  llvm::MDNode *loopMD =
+      loopAnnotationTranslation->translateLoopAnnotation(attr, op);
   inst->setMetadata(llvm::LLVMContext::MD_loop, loopMD);
 }
 
 llvm::Type *ModuleTranslation::convertType(Type type) {
   return typeTranslator.translateType(type);
 }
 
 /// A helper to look up remapped operands in the value remapping table.