[MemProf] Basic metadata support and verification

Add basic support for the MemProf metadata (!memprof and !callsite)
which was initially described in "RFC: IR metadata format for MemProf"
(https://discourse.llvm.org/t/rfc-ir-metadata-format-for-memprof/59165).

The bulk of the patch is verification support, along with some tests.

There are a couple of changes to the format described in the original
RFC:

Initial measurements suggested that a tree format for the stack ids in
the contexts would be more efficient, but subsequent evaluation with
large applications showed that in fact the cost of the additional
metadata nodes required by this deduplication scheme overwhelmed the
benefit from sharing stack id nodes. Therefore, the implementation here
and in follow on patches utilizes a simpler scheme of lists of stack id
integers in the memprof profile contexts and callsite metadata. The
follow on matching patch employs context trimming optimizations to
reduce the cost.

Secondly, instead of verbosely listing all profiled fields in each
profiled context (memory info block or MIB), and deferring the
interpretation of the profile data, the profile data is evaluated and
converted into string tags specifying the behavior (e.g. "cold") during
profile matching. This reduces the verbosity of the profile metadata,
and allows additional context trimming optimizations. As a result, the
named metadata schema description is also no longer needed.

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

llvm/include/llvm/IR/FixedMetadataKinds.def

@@ -45,3 +45,5 @@ LLVM_FIXED_MD_KIND(MD_annotation, "annotation", 30)
 LLVM_FIXED_MD_KIND(MD_nosanitize, "nosanitize", 31)
 LLVM_FIXED_MD_KIND(MD_func_sanitize, "func_sanitize", 32)
 LLVM_FIXED_MD_KIND(MD_exclude, "exclude", 33)
+LLVM_FIXED_MD_KIND(MD_memprof, "memprof", 34)
+LLVM_FIXED_MD_KIND(MD_callsite, "callsite", 35)

llvm/lib/IR/Verifier.cpp

@@ -469,6 +469,9 @@ class Verifier : public InstVisitor<Verifier>, VerifierSupport {
   void visitRangeMetadata(Instruction &I, MDNode *Range, Type *Ty);
   void visitDereferenceableMetadata(Instruction &I, MDNode *MD);
   void visitProfMetadata(Instruction &I, MDNode *MD);
+  void visitCallStackMetadata(MDNode *MD);
+  void visitMemProfMetadata(Instruction &I, MDNode *MD);
+  void visitCallsiteMetadata(Instruction &I, MDNode *MD);
   void visitAnnotationMetadata(MDNode *Annotation);
   void visitAliasScopeMetadata(const MDNode *MD);
   void visitAliasScopeListMetadata(const MDNode *MD);
@@ -4487,6 +4490,55 @@ void Verifier::visitProfMetadata(Instruction &I, MDNode *MD) {
   }
 }
 
+void Verifier::visitCallStackMetadata(MDNode *MD) {
+  // Call stack metadata should consist of a list of at least 1 constant int
+  // (representing a hash of the location).
+  Check(MD->getNumOperands() >= 1,
+        "call stack metadata should have at least 1 operand", MD);
+
+  for (const auto &Op : MD->operands())
+    Check(mdconst::dyn_extract_or_null<ConstantInt>(Op),
+          "call stack metadata operand should be constant integer", Op);
+}
+
+void Verifier::visitMemProfMetadata(Instruction &I, MDNode *MD) {
+  Check(isa<CallBase>(I), "!memprof metadata should only exist on calls", &I);
+  Check(MD->getNumOperands() >= 1,
+        "!memprof annotations should have at least 1 metadata operand "
+        "(MemInfoBlock)",
+        MD);
+
+  // Check each MIB
+  for (auto &MIBOp : MD->operands()) {
+    MDNode *MIB = dyn_cast<MDNode>(MIBOp);
+    // The first operand of an MIB should be the call stack metadata.
+    // There rest of the operands should be MDString tags, and there should be
+    // at least one.
+    Check(MIB->getNumOperands() >= 2,
+          "Each !memprof MemInfoBlock should have at least 2 operands", MIB);
+
+    // Check call stack metadata (first operand).
+    Check(MIB->getOperand(0) != nullptr,
+          "!memprof MemInfoBlock first operand should not be null", MIB);
+    Check(isa<MDNode>(MIB->getOperand(0)),
+          "!memprof MemInfoBlock first operand should be an MDNode", MIB);
+    MDNode *StackMD = dyn_cast<MDNode>(MIB->getOperand(0));
+    visitCallStackMetadata(StackMD);
+
+    // Check that remaining operands are MDString.
+    Check(std::all_of(MIB->op_begin() + 1, MIB->op_end(),
+                      [](const MDOperand &Op) { return isa<MDString>(Op); }),
+          "Not all !memprof MemInfoBlock operands 1 to N are MDString", MIB);
+  }
+}
+
+void Verifier::visitCallsiteMetadata(Instruction &I, MDNode *MD) {
+  Check(isa<CallBase>(I), "!callsite metadata should only exist on calls", &I);
+  // Verify the partial callstack annotated from memprof profiles. This callsite
+  // is a part of a profiled allocation callstack.
+  visitCallStackMetadata(MD);
+}
+
 void Verifier::visitAnnotationMetadata(MDNode *Annotation) {
   Check(isa<MDTuple>(Annotation), "annotation must be a tuple");
   Check(Annotation->getNumOperands() >= 1,
@@ -4733,6 +4785,12 @@ void Verifier::visitInstruction(Instruction &I) {
   if (MDNode *MD = I.getMetadata(LLVMContext::MD_prof))
     visitProfMetadata(I, MD);
 
+  if (MDNode *MD = I.getMetadata(LLVMContext::MD_memprof))
+    visitMemProfMetadata(I, MD);
+
+  if (MDNode *MD = I.getMetadata(LLVMContext::MD_callsite))
+    visitCallsiteMetadata(I, MD);
+
   if (MDNode *Annotation = I.getMetadata(LLVMContext::MD_annotation))
     visitAnnotationMetadata(Annotation);
 

llvm/test/Verifier/memprof-metadata-bad.ll

@@ -0,0 +1,56 @@
+; Test that incorrect memprof and callsite metadata fail verification.
+; RUN: not llvm-as -disable-output < %s 2>&1 | FileCheck %s
+
+define i32* @test1() {
+entry:
+  %call1 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !0
+  %call2 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !1
+  %call3 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !3
+  %call4 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !5
+  %call5 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !7, !callsite !9
+  %0 = bitcast i8* %call5 to i32*
+  ret i32* %0
+}
+
+define i32* @test2() {
+entry:
+  %call = call noundef i32* @test1(), !callsite !10
+  ret i32* %call
+}
+
+define i32* @test3() {
+entry:
+  %call = call noundef i32* @test2(), !callsite !11
+  ret i32* %call
+}
+
+define void @wronginsttype() {
+  %1 = add i32 0, 1, !memprof !0, !callsite !9
+  ret void
+}
+
+declare dso_local noalias noundef i8* @malloc(i64 noundef)
+
+; CHECK: !memprof annotations should have at least 1 metadata operand (MemInfoBlock)
+!0 = !{}
+!1 = !{!2}
+; CHECK: !memprof MemInfoBlock first operand should not be null
+!2 = !{null, !"cold"}
+!3 = !{!4}
+; CHECK: !memprof MemInfoBlock first operand should be an MDNode
+!4 = !{i64 0, !"cold"}
+!5 = !{!6}
+; CHECK: Each !memprof MemInfoBlock should have at least 2 operands
+!6 = !{i64 0}
+!7 = !{!8}
+; CHECK: call stack metadata should have at least 1 operand
+; CHECK: Not all !memprof MemInfoBlock operands 1 to N are MDString
+!8 = !{!0, !"default", i64 0}
+!9 = !{i64 123}
+; CHECK: call stack metadata operand should be constant integer
+!10 = !{!"wrongtype"}
+!11 = !{i64 789, i64 678}
+
+; Errors from annotating incorrect instruction type in @wronginsttype.
+; CHECK: !memprof metadata should only exist on calls
+; CHECK: !callsite metadata should only exist on calls

llvm/test/Verifier/memprof-metadata-good.ll

@@ -0,0 +1,34 @@
+; Test that well-formed memprof and callsite metadata pass verification.
+; RUN: llvm-as -disable-output < %s 2>&1
+
+define i32* @test1() {
+entry:
+  %call = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !0, !callsite !5
+  %0 = bitcast i8* %call to i32*
+  ret i32* %0
+}
+
+define i32* @test2() {
+entry:
+  %call = call noundef i32* @test1(), !callsite !6
+  ret i32* %call
+}
+
+define i32* @test3() {
+entry:
+  %call = call noundef i32* @test1(), !callsite !7
+  ret i32* %call
+}
+
+declare dso_local noalias noundef i8* @malloc(i64 noundef)
+
+!0 = !{!1, !3}
+; !memprof metadata should be able to support an arbitrary list of string tags.
+!1 = !{!2, !"default", !"tag2"}
+!2 = !{i64 123, i64 456}
+!3 = !{!4, !"cold", !"tag3", !"tag4"}
+!4 = !{i64 123, i64 789, i64 678}
+!5 = !{i64 123}
+!6 = !{i64 456}
+; Inlined callsites will have more than one call stack id.
+!7 = !{i64 789, i64 678}