[mlir][llvm] Store memory op metadata using op attributes.

The revision introduces operation attributes to store tbaa metadata on
load and store operations rather than relying using dialect attributes.
At the same time, the change also ensures the provided getters and
setters instead are used instead of a string based lookup. The latter
is done for the tbaa, access groups, and alias scope attributes.

The goal of this change is to ensure the metadata attributes are only
placed on operations that have the corresponding operation attributes.
This is imported since only these operations later on translate these
attributes to LLVM IR. Dialect attributes placed on other operations
are lost during the translation.

Reviewed By: vzakhari, Dinistro

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

Author: gysit

flang/lib/Optimizer/CodeGen/TBAABuilder.cpp

@@ -12,6 +12,7 @@
 
 #include "TBAABuilder.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
@@ -159,9 +160,13 @@ void TBAABuilder::attachTBAATag(Operation *op, Type baseFIRType,
   else
     tbaaTagSym = getDataAccessTag(baseFIRType, accessFIRType, gep);
 
-  if (tbaaTagSym)
-    op->setAttr(LLVMDialect::getTBAAAttrName(),
-                ArrayAttr::get(op->getContext(), tbaaTagSym));
+  if (!tbaaTagSym)
+    return;
+
+  auto tbaaAttr = ArrayAttr::get(op->getContext(), tbaaTagSym);
+  llvm::TypeSwitch<Operation *>(op)
+      .Case<LoadOp, StoreOp>([&](auto memOp) { memOp.setTbaaAttr(tbaaAttr); })
+      .Default([](auto) { llvm_unreachable("expected LoadOp or StoreOp"); });
 }
 
 } // namespace fir

flang/test/Fir/tbaa.fir

@@ -35,11 +35,7 @@ def LLVM_Dialect : Dialect {
   let extraClassDeclaration = [{
     /// Name of the data layout attributes.
     static StringRef getDataLayoutAttrName() { return "llvm.data_layout"; }
-    static StringRef getNoAliasScopesAttrName() { return "noalias_scopes"; }
-    static StringRef getAliasScopesAttrName() { return "alias_scopes"; }
     static StringRef getLoopAttrName() { return "llvm.loop"; }
-    static StringRef getAccessGroupsAttrName() { return "access_groups"; }
-    static StringRef getTBAAAttrName() { return "llvm.tbaa"; }
 
     /// Names of llvm parameter attributes.
     static StringRef getAlignAttrName() { return "llvm.align"; }

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

@@ -350,6 +350,7 @@ def LLVM_LoadOp : LLVM_Op<"load">, MemoryOpBase {
                    OptionalAttr<SymbolRefArrayAttr>:$access_groups,
                    OptionalAttr<SymbolRefArrayAttr>:$alias_scopes,
                    OptionalAttr<SymbolRefArrayAttr>:$noalias_scopes,
+                   OptionalAttr<SymbolRefArrayAttr>:$tbaa,
                    OptionalAttr<I64Attr>:$alignment, UnitAttr:$volatile_,
                    UnitAttr:$nontemporal);
   let results = (outs LLVM_LoadableType:$res);
@@ -390,6 +391,7 @@ def LLVM_StoreOp : LLVM_Op<"store">, MemoryOpBase {
                    OptionalAttr<SymbolRefArrayAttr>:$access_groups,
                    OptionalAttr<SymbolRefArrayAttr>:$alias_scopes,
                    OptionalAttr<SymbolRefArrayAttr>:$noalias_scopes,
+                   OptionalAttr<SymbolRefArrayAttr>:$tbaa,
                    OptionalAttr<I64Attr>:$alignment, UnitAttr:$volatile_,
                    UnitAttr:$nontemporal);
   string llvmInstName = "Store";

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

@@ -120,15 +120,14 @@ class ModuleTranslation {
   /// in these blocks.
   void forgetMapping(Region &region);
 
-  /// Returns the LLVM metadata corresponding to a reference to an mlir LLVM
-  /// dialect access group operation.
-  llvm::MDNode *getAccessGroup(Operation &opInst,
+  /// 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 reference to an mlir LLVM
-  /// dialect alias scope operation
-  llvm::MDNode *getAliasScope(Operation &opInst,
-                              SymbolRefAttr aliasScopeRef) 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);
@@ -287,9 +286,9 @@ class ModuleTranslation {
   /// metadata nodes for them and their domains.
   LogicalResult createAliasScopeMetadata();
 
-  /// Returns the LLVM metadata corresponding to a reference to an mlir LLVM
-  /// dialect TBAATagOp operation.
-  llvm::MDNode *getTBAANode(Operation &memOp, SymbolRefAttr tagRef) const;
+  /// Returns the LLVM metadata corresponding to a symbol reference to an mlir
+  /// LLVM dialect TBAATagOp operation.
+  llvm::MDNode *getTBAANode(Operation *op, SymbolRefAttr tagRef) const;
 
   /// Process tbaa LLVM Metadata operations and create LLVM
   /// metadata nodes for them.

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

@@ -668,53 +668,60 @@ Type LLVM::GEPOp::getSourceElementType() {
 // Builder, printer and parser for for LLVM::LoadOp.
 //===----------------------------------------------------------------------===//
 
-LogicalResult verifySymbolAttribute(
-    Operation *op, StringRef attributeName,
+/// Verifies the given array attribute contains symbol references and checks the
+/// referenced symbol types using the provided verification function.
+LogicalResult verifyMemOpSymbolRefs(
+    Operation *op, StringRef name, ArrayAttr symbolRefs,
     llvm::function_ref<LogicalResult(Operation *, SymbolRefAttr)>
         verifySymbolType) {
-  if (Attribute attribute = op->getAttr(attributeName)) {
-    // Verify that the attribute is a symbol ref array attribute,
-    // because this constraint is not verified for all attribute
-    // names processed here (e.g. 'tbaa'). This verification
-    // is redundant in some cases.
-    if (!(attribute.isa<ArrayAttr>() &&
-          llvm::all_of(attribute.cast<ArrayAttr>(), [&](Attribute attr) {
-            return attr && attr.isa<SymbolRefAttr>();
-          })))
-      return op->emitOpError("attribute '")
-             << attributeName
-             << "' failed to satisfy constraint: symbol ref array attribute";
-
-    for (SymbolRefAttr symbolRef :
-         attribute.cast<ArrayAttr>().getAsRange<SymbolRefAttr>()) {
-      StringAttr metadataName = symbolRef.getRootReference();
-      StringAttr symbolName = symbolRef.getLeafReference();
-      // We want @metadata::@symbol, not just @symbol
-      if (metadataName == symbolName) {
-        return op->emitOpError() << "expected '" << symbolRef
-                                 << "' to specify a fully qualified reference";
-      }
-      auto metadataOp = SymbolTable::lookupNearestSymbolFrom<LLVM::MetadataOp>(
-          op->getParentOp(), metadataName);
-      if (!metadataOp)
-        return op->emitOpError()
-               << "expected '" << symbolRef << "' to reference a metadata op";
-      Operation *symbolOp =
-          SymbolTable::lookupNearestSymbolFrom(metadataOp, symbolName);
-      if (!symbolOp)
-        return op->emitOpError()
-               << "expected '" << symbolRef << "' to be a valid reference";
-      if (failed(verifySymbolType(symbolOp, symbolRef))) {
-        return failure();
-      }
+  assert(symbolRefs && "expected a non-null attribute");
+
+  // Verify that the attribute is a symbol ref array attribute,
+  // because this constraint is not verified for all attribute
+  // names processed here (e.g. 'tbaa'). This verification
+  // is redundant in some cases.
+  if (!llvm::all_of(symbolRefs, [](Attribute attr) {
+        return attr && attr.isa<SymbolRefAttr>();
+      }))
+    return op->emitOpError("attribute '")
+           << name
+           << "' failed to satisfy constraint: symbol ref array attribute";
+
+  for (SymbolRefAttr symbolRef : symbolRefs.getAsRange<SymbolRefAttr>()) {
+    StringAttr metadataName = symbolRef.getRootReference();
+    StringAttr symbolName = symbolRef.getLeafReference();
+    // We want @metadata::@symbol, not just @symbol
+    if (metadataName == symbolName) {
+      return op->emitOpError() << "expected '" << symbolRef
+                               << "' to specify a fully qualified reference";
+    }
+    auto metadataOp = SymbolTable::lookupNearestSymbolFrom<LLVM::MetadataOp>(
+        op->getParentOp(), metadataName);
+    if (!metadataOp)
+      return op->emitOpError()
+             << "expected '" << symbolRef << "' to reference a metadata op";
+    Operation *symbolOp =
+        SymbolTable::lookupNearestSymbolFrom(metadataOp, symbolName);
+    if (!symbolOp)
+      return op->emitOpError()
+             << "expected '" << symbolRef << "' to be a valid reference";
+    if (failed(verifySymbolType(symbolOp, symbolRef))) {
+      return failure();
     }
   }
+
   return success();
 }
 
-// Verifies that metadata ops are wired up properly.
+/// Verifies the given array attribute contains symbol references that point to
+/// metadata operations of the given type.
 template <typename OpTy>
-static LogicalResult verifyOpMetadata(Operation *op, StringRef attributeName) {
+static LogicalResult
+verifyMemOpSymbolRefsPointTo(Operation *op, StringRef name,
+                             std::optional<ArrayAttr> symbolRefs) {
+  if (!symbolRefs)
+    return success();
+
   auto verifySymbolType = [op](Operation *symbolOp,
                                SymbolRefAttr symbolRef) -> LogicalResult {
     if (!isa<OpTy>(symbolOp)) {
@@ -724,35 +731,33 @@ static LogicalResult verifyOpMetadata(Operation *op, StringRef attributeName) {
     }
     return success();
   };
-
-  return verifySymbolAttribute(op, attributeName, verifySymbolType);
+  return verifyMemOpSymbolRefs(op, name, *symbolRefs, verifySymbolType);
 }
 
-static LogicalResult verifyMemoryOpMetadata(Operation *op) {
-  // access_groups
-  if (failed(verifyOpMetadata<LLVM::AccessGroupMetadataOp>(
-          op, LLVMDialect::getAccessGroupsAttrName())))
+/// Verifies the types of the metadata operations referenced by aliasing and
+/// access group metadata.
+template <typename OpTy>
+LogicalResult verifyMemOpMetadata(OpTy memOp) {
+  if (failed(verifyMemOpSymbolRefsPointTo<LLVM::AccessGroupMetadataOp>(
+          memOp, memOp.getAccessGroupsAttrName(), memOp.getAccessGroups())))
     return failure();
 
-  // alias_scopes
-  if (failed(verifyOpMetadata<LLVM::AliasScopeMetadataOp>(
-          op, LLVMDialect::getAliasScopesAttrName())))
+  if (failed(verifyMemOpSymbolRefsPointTo<LLVM::AliasScopeMetadataOp>(
+          memOp, memOp.getAliasScopesAttrName(), memOp.getAliasScopes())))
     return failure();
 
-  // noalias_scopes
-  if (failed(verifyOpMetadata<LLVM::AliasScopeMetadataOp>(
-          op, LLVMDialect::getNoAliasScopesAttrName())))
+  if (failed(verifyMemOpSymbolRefsPointTo<LLVM::AliasScopeMetadataOp>(
+          memOp, memOp.getNoaliasScopesAttrName(), memOp.getNoaliasScopes())))
     return failure();
 
-  // tbaa
-  if (failed(verifyOpMetadata<LLVM::TBAATagOp>(op,
-                                               LLVMDialect::getTBAAAttrName())))
+  if (failed(verifyMemOpSymbolRefsPointTo<LLVM::TBAATagOp>(
+          memOp, memOp.getTbaaAttrName(), memOp.getTbaa())))
     return failure();
 
   return success();
 }
 
-LogicalResult LoadOp::verify() { return verifyMemoryOpMetadata(*this); }
+LogicalResult LoadOp::verify() { return verifyMemOpMetadata(*this); }
 
 void LoadOp::build(OpBuilder &builder, OperationState &result, Type t,
                    Value addr, unsigned alignment, bool isVolatile,
@@ -828,7 +833,7 @@ ParseResult LoadOp::parse(OpAsmParser &parser, OperationState &result) {
 // Builder, printer and parser for LLVM::StoreOp.
 //===----------------------------------------------------------------------===//
 
-LogicalResult StoreOp::verify() { return verifyMemoryOpMetadata(*this); }
+LogicalResult StoreOp::verify() { return verifyMemOpMetadata(*this); }
 
 void StoreOp::build(OpBuilder &builder, OperationState &result, Value value,
                     Value addr, unsigned alignment, bool isVolatile,

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

@@ -76,6 +76,30 @@ static ArrayRef<unsigned> getSupportedMetadataImpl() {
   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 {
+  AttributeSetter(Operation *op) : op(op) {}
+
+  /// Calls `attachFn` on the provided Operation if it has one of
+  /// the given operation types. Returns failure otherwise.
+  template <typename CallableT>
+  LogicalResult apply(CallableT &&attachFn) {
+    return llvm::TypeSwitch<Operation *, LogicalResult>(op)
+        .Case<OpTys...>([&attachFn](auto concreteOp) {
+          attachFn(concreteOp);
+          return success();
+        })
+        .Default([&](auto) { return failure(); });
+  }
+
+private:
+  Operation *op;
+};
+} // namespace
+
 /// Converts the given profiling metadata `node` to an MLIR profiling attribute
 /// and attaches it to the imported operation if the translation succeeds.
 /// Returns failure otherwise.
@@ -129,16 +153,10 @@ static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
     branchWeights.push_back(branchWeight->getZExtValue());
   }
 
-  // Attach the branch weights to the operations that support it.
-  return llvm::TypeSwitch<Operation *, LogicalResult>(op)
-      .Case<CondBrOp, SwitchOp, CallOp, InvokeOp>([&](auto branchWeightOp) {
+  return AttributeSetter<CondBrOp, SwitchOp, CallOp, InvokeOp>(op).apply(
+      [&](auto branchWeightOp) {
         branchWeightOp.setBranchWeightsAttr(
             builder.getI32VectorAttr(branchWeights));
-        return success();
-      })
-      .Default([op](auto) {
-        return op->emitWarning()
-               << op->getName() << " does not support branch weights";
       });
 }
 
@@ -151,9 +169,9 @@ static LogicalResult setTBAAAttr(const llvm::MDNode *node, Operation *op,
   if (!tbaaTagSym)
     return failure();
 
-  op->setAttr(LLVMDialect::getTBAAAttrName(),
-              ArrayAttr::get(op->getContext(), tbaaTagSym));
-  return success();
+  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
@@ -169,9 +187,10 @@ static LogicalResult setAccessGroupAttr(const llvm::MDNode *node, Operation *op,
 
   SmallVector<Attribute> accessGroupAttrs(accessGroups->begin(),
                                           accessGroups->end());
-  op->setAttr(LLVMDialect::getAccessGroupsAttrName(),
-              ArrayAttr::get(op->getContext(), accessGroupAttrs));
-  return success();
+  return AttributeSetter<LoadOp, StoreOp>(op).apply([&](auto memOp) {
+    memOp.setAccessGroupsAttr(
+        ArrayAttr::get(memOp.getContext(), accessGroupAttrs));
+  });
 }
 
 /// Converts the given loop metadata node to an MLIR loop annotation attribute

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

@@ -210,7 +210,7 @@ llvm::MDNode *LoopAnnotationConversion::convert() {
         llvm::MDString::get(ctx, "llvm.loop.parallel_accesses"));
     for (SymbolRefAttr accessGroupRef : parallelAccessGroups)
       parallelAccess.push_back(
-          moduleTranslation.getAccessGroup(*op, accessGroupRef));
+          moduleTranslation.getAccessGroup(op, accessGroupRef));
     metadataNodes.push_back(llvm::MDNode::get(ctx, parallelAccess));
   }