[MachineMemOperand] learn about ptr_provenance

[jeroen.dobbelaere]

Summary

For the machine instructions, MachineMemOperand is the better place to track the ptr_provenance.

[jeroen.dobbelaere]

• Rebased

[jeroen.dobbelaere]

Rebase to 49cb4fef02e635bf304906232214166c7531d753

[Zoxc]

Given that the provenance gets dropped in https://reviews.llvm.org/D107356, how does it end up in machine operands?

[jeroen.dobbelaere]

In D107354#3494239, @Zoxc wrote:

Given that the provenance gets dropped in https://reviews.llvm.org/D107356, how does it end up in machine operands?

Good question. The provenance embedded in the MachineOperand here originates from the ptr_provenance operand in load/store instructions.
Your question does point out that https://reviews.llvm.org/D107356 (aka CodeGen/IntrinsicLowering) is too early if we want to track the more complex forms
of provenance after lowering.

[jeroen.dobbelaere]

In D107354#3495878, @jeroen.dobbelaere wrote:

In D107354#3494239, @Zoxc wrote:

Given that the provenance gets dropped in https://reviews.llvm.org/D107356, how does it end up in machine operands?

Good question. The provenance embedded in the MachineOperand here originates from the ptr_provenance operand in load/store instructions.
Your question does point out that https://reviews.llvm.org/D107356 (aka CodeGen/IntrinsicLowering) is too early if we want to track the more complex forms
of provenance after lowering.

Having looked deeper into it: IntrinsicLowering.cpp seems only to be used for the interpreter. There it makes sense to eliminate the noalias/ptr_provenance intrinsics.
During code generation, SelectionDAG is taking care of seeing through the intrinsics while lowering. But, the original llvm-IR (with intrinsics) stays available for the alias analysis. So the machine operands must be able to see it.
(Although you might need to enable AA analysis on machine ir)

[jeroen.dobbelaere]

Rebased to ffe262a198a9f9030991df6d3ddd812e74fa3523

[jeroen.dobbelaere]

Rebased to ef545ef62a833152d8975ff16333b57cc41befcc (Jan 9, 2023)

[vzakhari]

LGTM. Can you please add a LIT test with a check that ptr_provenance is printed or is it useless?

[arsenm]

Needs some testing, especially MIR print and parse. Can you not simply look through the existing Value to find the provenance?

[jeroen.dobbelaere]

Rebased to: 8c7cfa357280dd93d33b10bbba0fe33797e27d63 (Feb 14, 2023)

Added MIR parser support and testcase.

llvm/include/llvm/CodeGen/MachineMemOperand.h

 
 /// This class contains a discriminated union of information about pointers in
 /// memory operands, relating them back to LLVM IR or to virtual locations (such
 /// as frame indices) that are exposed during codegen.
 struct MachinePointerInfo {
   /// This is the IR pointer value for the access, or it is null if unknown.
   PointerUnion<const Value *, const PseudoSourceValue *> V;
 
+  /// The provenance of the pointer. When UnknownProvenance, the provenance can
+  /// be any object.
+  /// FIXME: nullptr here means : use V as provenance. This will change in
+  /// future.
+  const Value *PtrProvenance = nullptr;
+
   /// Offset - This is an offset from the base Value*.
   int64_t Offset;
 
   unsigned AddrSpace = 0;
 
   uint8_t StackID;
 
   explicit MachinePointerInfo(const Value *v, int64_t offset = 0,
-                              uint8_t ID = 0)
-      : V(v), Offset(offset), StackID(ID) {
+                              uint8_t ID = 0,
+                              const Value *ptrProvenance = nullptr)
+      : V(v), PtrProvenance(ptrProvenance), Offset(offset), StackID(ID) {
     AddrSpace = v ? v->getType()->getPointerAddressSpace() : 0;
   }
 
   explicit MachinePointerInfo(const PseudoSourceValue *v, int64_t offset = 0,
                               uint8_t ID = 0)
       : V(v), Offset(offset), StackID(ID) {
     AddrSpace = v ? v->getAddressSpace() : 0;
   }
 
-  explicit MachinePointerInfo(unsigned AddressSpace = 0, int64_t offset = 0)
-      : V((const Value *)nullptr), Offset(offset), AddrSpace(AddressSpace),
-        StackID(0) {}
+  explicit MachinePointerInfo(unsigned AddressSpace = 0, int64_t offset = 0,
+                              const Value *ptrProvenance = nullptr)
+      : V((const Value *)nullptr), PtrProvenance(ptrProvenance), Offset(offset),
+        AddrSpace(AddressSpace), StackID(0) {}
 
   explicit MachinePointerInfo(
       PointerUnion<const Value *, const PseudoSourceValue *> v,
-    int64_t offset = 0,
-    uint8_t ID = 0)
-    : V(v), Offset(offset), StackID(ID) {
+      int64_t offset = 0, uint8_t ID = 0, const Value *ptrProvenance = nullptr)
+      : V(v), PtrProvenance(ptrProvenance), Offset(offset), StackID(ID) {
     if (V) {
       if (const auto *ValPtr = V.dyn_cast<const Value*>())
         AddrSpace = ValPtr->getType()->getPointerAddressSpace();
       else
         AddrSpace = V.get<const PseudoSourceValue*>()->getAddressSpace();
     }
   }
 
public:
   const Value *getValue() const { return PtrInfo.V.dyn_cast<const Value*>(); }
 
   const PseudoSourceValue *getPseudoValue() const {
     return PtrInfo.V.dyn_cast<const PseudoSourceValue*>();
   }
 
   const void *getOpaqueValue() const { return PtrInfo.V.getOpaqueValue(); }
 
+  const Value *getPtrProvenance() const { return PtrInfo.PtrProvenance; }
+
   /// Return the raw flags of the source value, \see Flags.
   Flags getFlags() const { return FlagVals; }
 
   /// Bitwise OR the current flags with the given flags.
   void setFlags(Flags f) { FlagVals |= f; }
 
   /// For normal values, this is a byte offset added to the base address.
   /// For PseudoSourceValue::FPRel values, this is the FrameIndex number.

llvm/lib/CodeGen/MIRParser/MILexer.h

enum TokenKind {
     kw_ppc_fp128,
     kw_target_flags,
     kw_volatile,
     kw_non_temporal,
     kw_invariant,
     kw_align,
     kw_basealign,
     kw_addrspace,
+    kw_ptr_provenance,
     kw_stack,
     kw_got,
     kw_jump_table,
     kw_constant_pool,
     kw_call_entry,
     kw_custom,
     kw_liveout,
     kw_landing_pad,

llvm/lib/CodeGen/MIRParser/MILexer.cpp

return StringSwitch<MIToken::TokenKind>(Identifier)
       .Case("target-flags", MIToken::kw_target_flags)
       .Case("volatile", MIToken::kw_volatile)
       .Case("non-temporal", MIToken::kw_non_temporal)
       .Case("dereferenceable", MIToken::kw_dereferenceable)
       .Case("invariant", MIToken::kw_invariant)
       .Case("align", MIToken::kw_align)
       .Case("basealign", MIToken::kw_basealign)
       .Case("addrspace", MIToken::kw_addrspace)
+      .Case("ptr_provenance", MIToken::kw_ptr_provenance)
       .Case("stack", MIToken::kw_stack)
       .Case("got", MIToken::kw_got)
       .Case("jump-table", MIToken::kw_jump_table)
       .Case("constant-pool", MIToken::kw_constant_pool)
       .Case("call-entry", MIToken::kw_call_entry)
       .Case("custom", MIToken::kw_custom)
       .Case("liveout", MIToken::kw_liveout)
       .Case("landing-pad", MIToken::kw_landing_pad)

llvm/lib/CodeGen/MIRParser/MIParser.cpp

public:
   bool parseOffset(int64_t &Offset);
   bool parseIRBlockAddressTaken(BasicBlock *&BB);
   bool parseAlignment(uint64_t &Alignment);
   bool parseAddrspace(unsigned &Addrspace);
   bool parseSectionID(std::optional<MBBSectionID> &SID);
   bool parseBBID(std::optional<unsigned> &BBID);
   bool parseOperandsOffset(MachineOperand &Op);
   bool parseIRValue(const Value *&V);
+  bool parseIRValueAsPointer(const Value *&V);
   bool parseMemoryOperandFlag(MachineMemOperand::Flags &Flags);
   bool parseMemoryPseudoSourceValue(const PseudoSourceValue *&PSV);
   bool parseMachinePointerInfo(MachinePointerInfo &Dest);
   bool parseOptionalScope(LLVMContext &Context, SyncScope::ID &SSID);
   bool parseOptionalAtomicOrdering(AtomicOrdering &Order);
   bool parseMachineMemoryOperand(MachineMemOperand *&Dest);
   bool parsePreOrPostInstrSymbol(MCSymbol *&Symbol);
   bool parseHeapAllocMarker(MDNode *&Node);

 
 bool MIParser::parseIRValue(const Value *&V) {
   return ::parseIRValue(
       Token, PFS, V, [this](StringRef::iterator Loc, const Twine &Msg) -> bool {
         return error(Loc, Msg);
       });
 }
 
+bool MIParser::parseIRValueAsPointer(const Value *&V) {
+  if (Token.isNot(MIToken::NamedIRValue) && Token.isNot(MIToken::IRValue) &&
+      Token.isNot(MIToken::GlobalValue) &&
+      Token.isNot(MIToken::NamedGlobalValue) &&
+      Token.isNot(MIToken::QuotedIRValue) &&
+      Token.isNot(MIToken::kw_unknown_address)) {
+    return error("expected an IR value reference");
+  }
+  V = nullptr;
+  if (parseIRValue(V))
+    return true;
+  if (V && !V->getType()->isPointerTy())
+    return error("expected a pointer IR value");
+  lex();
+  return false;
+}
+
 bool MIParser::getUint64(uint64_t &Result) {
   if (Token.hasIntegerValue()) {
     if (Token.integerValue().getActiveBits() > 64)
       return error("expected 64-bit integer (too large)");
     Result = Token.integerValue().getZExtValue();
     return false;
   }
   if (Token.is(MIToken::HexLiteral)) {
if (Token.is(MIToken::kw_constant_pool) || Token.is(MIToken::kw_stack) ||
     if (parseMemoryPseudoSourceValue(PSV))
       return true;
     int64_t Offset = 0;
     if (parseOffset(Offset))
       return true;
     Dest = MachinePointerInfo(PSV, Offset);
     return false;
   }
-  if (Token.isNot(MIToken::NamedIRValue) && Token.isNot(MIToken::IRValue) &&
-      Token.isNot(MIToken::GlobalValue) &&
-      Token.isNot(MIToken::NamedGlobalValue) &&
-      Token.isNot(MIToken::QuotedIRValue) &&
-      Token.isNot(MIToken::kw_unknown_address))
-    return error("expected an IR value reference");
   const Value *V = nullptr;
-  if (parseIRValue(V))
+  if (parseIRValueAsPointer(V))
     return true;
-  if (V && !V->getType()->isPointerTy())
-    return error("expected a pointer IR value");
-  lex();
   int64_t Offset = 0;
   if (parseOffset(Offset))
     return true;
   Dest = MachinePointerInfo(V, Offset);
   return false;
 }
 
 bool MIParser::parseOptionalScope(LLVMContext &Context,
if (parseLowLevelType(Token.location(), MemoryType))
       return true;
     if (expectAndConsume(MIToken::rparen))
       return true;
 
     Size = MemoryType.getSizeInBytes();
   }
 
   MachinePointerInfo Ptr = MachinePointerInfo();
+  const Value *PtrProvenancePtr = nullptr;
   if (Token.is(MIToken::Identifier)) {
     const char *Word =
         ((Flags & MachineMemOperand::MOLoad) &&
          (Flags & MachineMemOperand::MOStore))
             ? "on"
             : Flags & MachineMemOperand::MOLoad ? "from" : "into";
     if (Token.stringValue() != Word)
       return error(Twine("expected '") + Word + "'");
case MIToken::kw_basealign:
       // basealign is printed if it is different than align.
       if (parseAlignment(BaseAlignment))
         return true;
       break;
     case MIToken::kw_addrspace:
       if (parseAddrspace(Ptr.AddrSpace))
         return true;
       break;
+    case MIToken::kw_ptr_provenance: {
+      if (PtrProvenancePtr)
+        return error("ptr_provenance must occur at most once");
+      lex();
+      if (parseIRValueAsPointer(PtrProvenancePtr))
+        return true;
+      Ptr.PtrProvenance = PtrProvenancePtr;
+      break;
+    }
     case MIToken::md_tbaa:
       lex();
       if (parseMDNode(AAInfo.TBAA))
         return true;
       break;
     case MIToken::md_alias_scope:
       lex();
       if (parseMDNode(AAInfo.Scope))

llvm/lib/CodeGen/MachineOperand.cpp

if (const Value *Val = getValue()) {
     }
   } else if (getOpaqueValue() == nullptr && getOffset() != 0) {
     OS << ((isLoad() && isStore()) ? " on "
            : isLoad()              ? " from "
                                    : " into ")
        << "unknown-address";
   }
   MachineOperand::printOperandOffset(OS, getOffset());
+  if (getPtrProvenance()) {
+    OS << ", ptr_provenance ";
+    MIRFormatter::printIRValue(OS, *getPtrProvenance(), MST);
+  }
   if (getSize() > 0 && getAlign() != getSize())
     OS << ", align " << getAlign().value();
   if (getAlign() != getBaseAlign())
     OS << ", basealign " << getBaseAlign().value();
   auto AAInfo = getAAInfo();
   if (AAInfo.TBAA) {
     OS << ", !tbaa ";
     AAInfo.TBAA->printAsOperand(OS, MST);

llvm/test/CodeGen/X86/ptr_provenance_memoperand.ll

+; NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -stop-after=finalize-isel 2>&1 | FileCheck %s
+
+; NOTE: change 'RUN' to 'R U N' in following lines before regenerating the results; afterwards, put it back.
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown -stop-before=finalize-isel -o %t.mir
+; RUN: llc %t.mir -mtriple=x86_64-unknown-unknown -run-pass=finalize-isel -o - 2>&1 | FileCheck %s
+
+; This test validates that the ptr_provenance is added to the MachineMemOperand.
+; It also validates that the MIR parser can read it back in.
+
+; Function Attrs: nounwind
+define i32 @test01(ptr %pa, ptr %pb) #0 {
+  ; CHECK-LABEL: name: test01
+  ; CHECK: bb.0 (%ir-block.0):
+  ; CHECK-NEXT:   liveins: $rdi, $rsi
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gr64 = COPY $rsi
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:gr64 = COPY $rdi
+  ; CHECK-NEXT:   MOV32mi [[COPY]], 1, $noreg, 0, $noreg, 43 :: (store (s32) into %ir.pb, ptr_provenance `ptr unknown_provenance`)
+  ; CHECK-NEXT:   MOV32mi [[COPY1]], 1, $noreg, 0, $noreg, 42 :: (store (s32) into %ir.pa, ptr_provenance %ir.pb)
+  ; CHECK-NEXT:   [[MOV32rm:%[0-9]+]]:gr32 = MOV32rm [[COPY]], 1, $noreg, 0, $noreg :: (load (s32) from %ir.pb, ptr_provenance %ir.pa)
+  ; CHECK-NEXT:   $eax = COPY [[MOV32rm]]
+  ; CHECK-NEXT:   RET 0, $eax
+  store i32 43, ptr %pb, ptr_provenance ptr unknown_provenance
+  store i32 42, ptr %pa, ptr_provenance ptr %pb
+  %retval = load i32, ptr %pb, ptr_provenance ptr %pa
+  ret i32 %retval
+}
+
+; Function Attrs: nounwind
+define i32 @test02(ptr %pa, ptr %pb) #0 {
+  ; CHECK-LABEL: name: test02
+  ; CHECK: bb.0 (%ir-block.0):
+  ; CHECK-NEXT:   liveins: $rdi, $rsi
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gr64 = COPY $rsi
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:gr64 = COPY $rdi
+  ; CHECK-NEXT:   MOV32mi [[COPY1]], 1, $noreg, 0, $noreg, 43 :: (store (s32) into %ir.pa)
+  ; CHECK-NEXT:   [[MOV32rm:%[0-9]+]]:gr32 = MOV32rm [[COPY]], 1, $noreg, 0, $noreg :: (load (s32) from %ir.pb)
+  ; CHECK-NEXT:   $eax = COPY [[MOV32rm]]
+  ; CHECK-NEXT:   RET 0, $eax
+  store i32 43, ptr %pa
+  %retval = load i32, ptr %pb
+  ret i32 %retval
+}
+
+attributes #0 = { nounwind }