[mlir] Support TBAA metadata in LLVMIR dialect.

This change introduces new LLVMIR dialect operations to represent
TBAA root, type descriptor and access tag metadata nodes.

For the purpose of importing TBAA metadata from LLVM IR it only
supports the current version of TBAA format described in
https://llvm.org/docs/LangRef.html#tbaa-metadata (i.e. size-aware
representation introduced in D41501 is not supported).

TBAA attribute support is only added for LLVM::LoadOp and LLVM::StoreOp.
Support for intrinsics operations (e.g. LLVM::MemcpyOp) may be added later.

The TBAA attribute is represented as an array of access tags, though,
LLVM IR supports only single access tag per memory accessing instruction.
I implemented it as an array anticipating similar support in LLVM IR
to combine TBAA graphs with different roots for Flang - one of the options
described in https://docs.google.com/document/d/16kKZVmI585wth01VSaJAqZMZpoX68rcdBmgfj0kNAt0/edit#heading=h.jzzheaz9vqac

It should be easy to restrict MLIR operation to a single access tag,
if we end up using a different approach for Flang.

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

Author: vzakhari

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

@@ -52,6 +52,7 @@ def LLVM_Dialect : Dialect {
     static StringRef getNoUndefAttrName() { return "llvm.noundef"; }
     static StringRef getSExtAttrName() { return "llvm.signext"; }
     static StringRef getZExtAttrName() { return "llvm.zeroext"; }
+    static StringRef getTBAAAttrName() { return "llvm.tbaa"; }
 
     /// Verifies if the attribute is a well-formed value for "llvm.struct_attrs"
     static LogicalResult verifyStructAttr(

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

@@ -205,6 +205,10 @@ class MemoryOpWithAlignmentAndAttributes : MemoryOpWithAlignmentBase {
   code setAliasScopeMetadataCode = [{
     moduleTranslation.setAliasScopeMetadata(op, inst);
   }];
+
+  code setTBAAMetadataCode = [{
+    moduleTranslation.setTBAAMetadata(op, inst);
+  }];
 }
 
 // Memory-related operations.
@@ -360,6 +364,7 @@ def LLVM_LoadOp : LLVM_Op<"load">, MemoryOpWithAlignmentAndAttributes {
      # setNonTemporalMetadataCode
      # setAccessGroupsMetadataCode
      # setAliasScopeMetadataCode
+     # setTBAAMetadataCode
      # [{
     $res = inst;
   }];
@@ -397,7 +402,8 @@ def LLVM_StoreOp : LLVM_Op<"store">, MemoryOpWithAlignmentAndAttributes {
   }] # setAlignmentCode
      # setNonTemporalMetadataCode
      # setAccessGroupsMetadataCode
-     # setAliasScopeMetadataCode;
+     # setAliasScopeMetadataCode
+     # setTBAAMetadataCode;
   // FIXME: Import attributes.
   string mlirBuilder = [{
     $_op = $_builder.create<LLVM::StoreOp>($_location, $value, $addr);
@@ -992,6 +998,7 @@ def LLVM_MetadataOp : LLVM_Op<"metadata", [
   }];
   let regions = (region SizedRegion<1>:$body);
   let assemblyFormat = "$sym_name attr-dict-with-keyword $body";
+  let hasRegionVerifier = 1;
 }
 
 def LLVM_AliasScopeDomainMetadataOp : LLVM_Op<"alias_scope_domain", [
@@ -1074,6 +1081,152 @@ def LLVM_AccessGroupMetadataOp : LLVM_Op<"access_group", [
   let assemblyFormat = "$sym_name attr-dict";
 }
 
+def LLVM_TBAARootMetadataOp : LLVM_Op<"tbaa_root", [
+  HasParent<"MetadataOp">, Symbol
+]> {
+  let arguments = (ins
+    SymbolNameAttr:$sym_name,
+    StrAttr:$identity
+  );
+  let summary = "LLVM dialect TBAA root node metadata.";
+  let description = [{
+    Defines a TBAA root node.
+
+    Example:
+      llvm.metadata @tbaa {
+        llvm.tbaa_root @tbaa_root_0 {identity = "Simple C/C++ TBAA"}
+        llvm.return
+      }
+
+    See the following link for more details:
+    https://llvm.org/docs/LangRef.html#tbaa-metadata
+  }];
+  let assemblyFormat = [{
+    $sym_name ` ` `{` `id` `=` $identity `}` attr-dict
+  }];
+  let hasVerifier = 1;
+}
+
+def LLVM_TBAATypeDescriptorOp : LLVM_Op<"tbaa_type_desc", [
+  HasParent<"MetadataOp">, Symbol
+]> {
+  let arguments = (ins
+    SymbolNameAttr:$sym_name,
+    OptionalAttr<StrAttr>:$identity,
+    FlatSymbolRefArrayAttr:$members,
+    DenseI64ArrayAttr:$offsets
+  );
+  let summary = "LLVM dialect TBAA node describing a type.";
+  let description = [{
+    Defines a TBAA node describing a type.
+
+    Example:
+      llvm.metadata @tbaa {
+        llvm.tbaa_root @tbaa_root_0 {identity = "Simple C/C++ TBAA"}
+        llvm.tbaa_type_desc @tbaa_type_desc_1 {
+            identity = "omnipotent char",
+            members = [@tbaa_root_0],
+            offsets = array<i64: 0>
+        }
+        llvm.tbaa_type_desc @tbaa_type_desc_2 {
+            identity = "long long",
+            members = [@tbaa_type_desc_1],
+            offsets = array<i64: 0>
+        }
+        llvm.tbaa_type_desc @tbaa_type_desc_3 {
+            identity = "agg2_t",
+            members = [@tbaa_type_desc_2, @tbaa_type_desc_2],
+            offsets = array<i64: 0, 8>
+        }
+        llvm.tbaa_type_desc @tbaa_type_desc_5 {
+            identity = "int",
+            members = [@tbaa_type_desc_1],
+            offsets = array<i64: 0>
+        }
+        llvm.tbaa_type_desc @tbaa_type_desc_6 {
+            identity = "agg1_t",
+            members = [@tbaa_type_desc_5, @tbaa_type_desc_5],
+            offsets = array<i64: 0, 4>
+        }
+        llvm.return
+      }
+
+    See the following link for more details:
+    https://llvm.org/docs/LangRef.html#tbaa-metadata
+  }];
+
+  // Interleave member types and offsets for better matching
+  // LLVM IR text representation.
+  let assemblyFormat = [{
+    $sym_name ` ` `{`
+      ( `id` `=` $identity^ )? `,`
+      `members` `=` `{` custom<TBAAMembers>($members, $offsets) `}`
+    `}` attr-dict
+  }];
+  let hasVerifier = 1;
+}
+
+def LLVM_TBAATagOp : LLVM_Op<"tbaa_tag", [
+  HasParent<"MetadataOp">, Symbol
+]> {
+  let arguments = (ins
+    SymbolNameAttr:$sym_name,
+    FlatSymbolRefAttr:$base_type,
+    FlatSymbolRefAttr:$access_type,
+    I64Attr:$offset,
+    UnitAttr:$constant
+  );
+  let summary = "LLVM dialect TBAA node describing a memory access.";
+  let description = [{
+    Defines a TBAA node describing a memory access.
+
+    Example:
+      llvm.metadata @tbaa {
+        llvm.tbaa_root @tbaa_root_0 {identity = "Simple C/C++ TBAA"}
+        llvm.tbaa_type_desc @tbaa_type_desc_1 {
+            identity = "omnipotent char",
+            members = [@tbaa_root_0],
+            offsets = array<i64: 0>
+        }
+        llvm.tbaa_type_desc @tbaa_type_desc_2 {
+            identity = "long long",
+            members = [@tbaa_type_desc_1],
+            offsets = array<i64: 0>
+        }
+        llvm.tbaa_type_desc @tbaa_type_desc_3 {
+            identity = "agg2_t",
+            members = [@tbaa_type_desc_2, @tbaa_type_desc_2],
+            offsets = array<i64: 0, 8>
+        }
+        llvm.tbaa_tag @tbaa_tag_4 {
+            access_type = @tbaa_type_desc_2,
+            base_type = @tbaa_type_desc_3,
+            offset = 8 : i64
+        }
+        llvm.tbaa_type_desc @tbaa_type_desc_5 {
+            identity = "int",
+            members = [@tbaa_type_desc_1],
+            offsets = array<i64: 0>
+        }
+        llvm.tbaa_type_desc @tbaa_type_desc_6 {
+            identity = "agg1_t",
+            members = [@tbaa_type_desc_5, @tbaa_type_desc_5],
+            offsets = array<i64: 0, 4>
+        }
+        llvm.tbaa_tag @tbaa_tag_7 {
+            access_type = @tbaa_type_desc_5,
+            base_type = @tbaa_type_desc_6,
+            offset = 0 : i64
+        }
+        llvm.return
+      }
+
+    See the following link for more details:
+    https://llvm.org/docs/LangRef.html#tbaa-metadata
+  }];
+  let assemblyFormat = "$sym_name attr-dict";
+}
+
 def LLVM_GlobalOp : LLVM_Op<"mlir.global",
     [IsolatedFromAbove, SingleBlockImplicitTerminator<"ReturnOp">, Symbol]> {
   let arguments = (ins

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

@@ -149,6 +149,17 @@ class ModuleImport {
   /// implement the fastmath interface.
   void setFastmathFlagsAttr(llvm::Instruction *inst, Operation *op) const;
 
+  /// Converts LLVM metadata to corresponding MLIR representation,
+  /// e.g. metadata nodes referenced via !tbaa are converted to
+  /// TBAA operations hosted inside a MetadataOp.
+  LogicalResult convertMetadata();
+
+  /// Returns SymbolRefAttr representing TBAA metadata `node`
+  /// in `tbaaMapping`.
+  SymbolRefAttr lookupTBAAAttr(const llvm::MDNode *node) {
+    return tbaaMapping.lookup(node);
+  }
+
 private:
   /// Clears the block and value mapping before processing a new region.
   void clearBlockAndValueMapping() {
@@ -219,6 +230,25 @@ class ModuleImport {
   /// them fails. All operations are inserted at the start of the current
   /// function entry block.
   FailureOr<Value> convertConstantExpr(llvm::Constant *constant);
+  /// Returns symbol name to be used for MetadataOp containing
+  /// TBAA metadata operations. It must not conflict with the user
+  /// name space.
+  StringRef getTBAAMetadataOpName() const { return "__tbaa"; }
+  /// Returns a terminated MetadataOp into which TBAA metadata
+  /// operations can be placed. The MetadataOp is created
+  /// on the first invocation of this function.
+  MetadataOp getTBAAMetadataOp();
+  /// Performs conversion of LLVM TBAA metadata starting from
+  /// `node`. On exit from this function all nodes reachable
+  /// from `node` are converted, and tbaaMapping map is updated
+  /// (unless all dependencies have been converted by a previous
+  /// invocation of this function).
+  LogicalResult processTBAAMetadata(const llvm::MDNode *node);
+  /// Returns unique string name of a symbol that may be used
+  /// for a TBAA metadata operation. The name will contain
+  /// the provided `basename` and will be uniqued via
+  /// tbaaNodeCounter (see below).
+  std::string getNewTBAANodeName(StringRef basename);
 
   /// Builder pointing at where the next instruction should be generated.
   OpBuilder builder;
@@ -251,6 +281,16 @@ class ModuleImport {
   LLVM::TypeFromLLVMIRTranslator typeTranslator;
   /// Stateful debug information importer.
   std::unique_ptr<detail::DebugImporter> debugImporter;
+  /// A terminated MetadataOp where TBAA metadata operations
+  /// can be inserted.
+  MetadataOp tbaaMetadataOp{};
+  /// 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;
+  /// A counter to be used as a unique suffix for symbols
+  /// defined by TBAA operations.
+  unsigned tbaaNodeCounter = 0;
 };
 
 } // namespace LLVM

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

@@ -148,6 +148,10 @@ class ModuleTranslation {
   // Sets LLVM metadata for memory operations that have alias scope information.
   void setAliasScopeMetadata(Operation *op, llvm::Instruction *inst);
 
+  /// Sets LLVM TBAA metadata for memory operations that have
+  /// TBAA attributes.
+  void setTBAAMetadata(Operation *op, llvm::Instruction *inst);
+
   /// Converts the type from MLIR LLVM dialect to LLVM.
   llvm::Type *convertType(Type type);
 
@@ -291,6 +295,14 @@ 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;
+
+  /// Process tbaa LLVM Metadata operations and create LLVM
+  /// metadata nodes for them.
+  LogicalResult createTBAAMetadata();
+
   /// Translates dialect attributes attached to the given operation.
   LogicalResult convertDialectAttributes(Operation *op);
 
@@ -333,10 +345,14 @@ class ModuleTranslation {
   /// attribute.
   DenseMap<Attribute, llvm::MDNode *> loopOptionsMetadataMapping;
 
-  /// Mapping from an access scope metadata operation to its LLVM metadata.
+  /// 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;
+
   /// Stack of user-specified state elements, useful when translating operations
   /// with regions.
   SmallVector<std::unique_ptr<StackFrame>> stack;