[MemProf] Basic metadata support and verification

[tejohnson]

Summary

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.

[snehasish]

lgtm with a couple of minor comments.

llvm/lib/IR/Verifier.cpp

4536

Instruction& I is unused.

Perhaps add a check here that I is a CallInst? If you do then also add a test for it.

4537

Perhaps add a comment here about whats being verified. The transition from callsite verification to callstack may be hard to follow. How about something like --

// Here we verify the callstack profile annotated by memprof. This callsite is a part of a profiled allocation callstack.

[Enna1]

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

30

IIUC, does this !4 = !{i64 123, i64 789, i64 678} means we have a call stack: test3() -> test1() -> malloc()
If so, the callee in test3() should be test1() instead of test2() ? :

define i32* @test3() {
entry:
  %call = call noundef i32* @test1(), !callsite !7
  ret i32* %call
}

[tejohnson]

Address comments

llvm/lib/IR/Verifier.cpp

4536

Good idea, added this check both here and in visitMemProfMetadata, and added checks for both in the bad metadata test.

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

30

Indeed! This is a good catch, mistake in my hand crafted test case IR. Fixed. (Although it doesn't matter for verification and in general we can't detect this kind of issue since the call chains can cross module boundaries.)

llvm/include/llvm/IR/FixedMetadataKinds.def

 LLVM_FIXED_MD_KIND(MD_irr_loop, "irr_loop", 24)
 LLVM_FIXED_MD_KIND(MD_access_group, "llvm.access.group", 25)
 LLVM_FIXED_MD_KIND(MD_callback, "callback", 26)
 LLVM_FIXED_MD_KIND(MD_preserve_access_index, "llvm.preserve.access.index", 27)
 LLVM_FIXED_MD_KIND(MD_vcall_visibility, "vcall_visibility", 28)
 LLVM_FIXED_MD_KIND(MD_noundef, "noundef", 29)
 LLVM_FIXED_MD_KIND(MD_annotation, "annotation", 30)
 LLVM_FIXED_MD_KIND(MD_nosanitize, "nosanitize", 31)
+LLVM_FIXED_MD_KIND(MD_memprof, "memprof", 32)
+LLVM_FIXED_MD_KIND(MD_callsite, "callsite", 33)

llvm/lib/IR/Verifier.cpp

void visitModuleFlag(const MDNode *Op,
                        DenseMap<const MDString *, const MDNode *> &SeenIDs,
                        SmallVectorImpl<const MDNode *> &Requirements);
   void visitModuleFlagCGProfileEntry(const MDOperand &MDO);
   void visitFunction(const Function &F);
   void visitBasicBlock(BasicBlock &BB);
   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);
   void visitAccessGroupMetadata(const MDNode *MD);
 
   template <class Ty> bool isValidMetadataArray(const MDTuple &N);
 #define HANDLE_SPECIALIZED_MDNODE_LEAF(CLASS) void visit##CLASS(const CLASS &N);
 #include "llvm/IR/Metadata.def"
for (unsigned i = 1; i < MD->getNumOperands(); ++i) {
       auto &MDO = MD->getOperand(i);
       Check(MDO, "second operand should not be null", MD);
       Check(mdconst::dyn_extract<ConstantInt>(MDO),
             "!prof brunch_weights operand is not a const int");
     }
   }
 }
 
+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,
         "annotation must have at least one operand");
   for (const MDOperand &Op : Annotation->operands())
     Check(isa<MDString>(Op.get()), "operands must be strings");
 }
 
Check(isPowerOf2_64(Align), "align metadata value must be a power of 2!",
           &I);
     Check(Align <= Value::MaximumAlignment,
           "alignment is larger that implementation defined limit", &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);
 
   if (MDNode *N = I.getDebugLoc().getAsMDNode()) {
     CheckDI(isa<DILocation>(N), "invalid !dbg metadata attachment", &I, N);
     visitMDNode(*N, AreDebugLocsAllowed::Yes);
   }
 

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

+; 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

+; 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}