[Attributor] UB Attribute now handles all instructions that access memory through a pointer

[baziotis]

Summary

Follow-up on: https://reviews.llvm.org/D71435
We basically use checkForAllInstructions to loop through all the instructions in a function that access memory through a pointer: load, store, atomicrmw, atomiccmpxchg
Note that we can now use the getPointerOperand() that gets us the pointer operand for an instruction that belongs to the aforementioned set.

Question: This function returns nullptr if the instruction is volatile. Why?
Guess: Because if it is volatile, we don't want to do any transformation to it.

Another subtle point is that I had to add AtomicRMW, AtomicCmpXchg to initializeInformationCache(). Following checkAllInstructions() path, that
seemed the most reasonable place to add it and correct the fact that these instructions were ignored (they were not in OpcodeInstMap etc.). Is that ok?

[sstefan1]

Quick initial comment

llvm/lib/Transforms/IPO/Attributor.cpp

2035–2039

I think you should be able to use CheckForAllReadWriteInstructions here. Then you don't need to add 2 additional cases in initializeInformationCache.

[jdoerfert]

This looks generally fine to me. Let me know what you think about the two comments below.

llvm/lib/Transforms/IPO/Attributor.cpp

2035–2039

We can do that, but you'll see other instructions here too. I'm fine with doing it step by step as well.

llvm/test/Transforms/Attributor/undefined_behavior.ll

140

Maybe it makes sense to run the update_test_checks script on this file.

[baziotis]

llvm/lib/Transforms/IPO/Attributor.cpp

2035–2039

I think you should be able to use CheckForAllReadWriteInstructions here. Then you don't need to add 2 additional cases in initializeInformationCache.

Thanks Stefan! I hadn't noticed.

We can do that, but you'll see other instructions here too. I'm fine with doing it step by step as well.

Alright, I'll leave it as it is and re-introduce a TODO to not only check these instructions (i.e. make clear that other instructions will be added).

llvm/test/Transforms/Attributor/undefined_behavior.ll

140

Thanks, I didn't know about it.

[baziotis]

• llvm/utils/update_test_checks.py ran on undefined_behavior.ll test file.
• TODO to check more instructions in the future.

[jdoerfert]

LGTM. I'll commit them soon. You can continue on top if you find the time and motivation.

[sstefan1]

llvm/lib/Transforms/IPO/Attributor.cpp

2013

nit: I think you should update these comments (line 2021) as well.

[baziotis]

LGTM. I'll commit them soon. You can continue on top if you find the time and motivation.

Great, I'll continue with branches. Could you please answer this question:

Question: getPointerOperand() (meaning, Attributor's getPointerOperand()) returns nullptr if the instruction is volatile. Why?
Guess: Because if it is volatile, we don't want to do any transformation to it.

llvm/lib/Transforms/IPO/Attributor.cpp

2013

Yep, I missed that. I'm uploading a new diff, thanks again.

[baziotis]

Comments update.

llvm/lib/Transforms/IPO/Attributor.cpp

 };
 
 /// -------------------- Undefined-Behavior Attributes ------------------------
 
 struct AAUndefinedBehaviorImpl : public AAUndefinedBehavior {
   AAUndefinedBehaviorImpl(const IRPosition &IRP) : AAUndefinedBehavior(IRP) {}
 
   /// See AbstractAttribute::updateImpl(...).
+  // TODO: We should not only check instructions that access memory
+  // through a pointer (i.e. also branches etc.)
   ChangeStatus updateImpl(Attributor &A) override {
-    size_t PrevSize = NoUBLoads.size();
+    const size_t PrevSize = NoUBMemAccessInsts.size();
 
-    // TODO: We should not only check for load instructions.
-    auto InspectLoadForUB = [&](Instruction &I) {
+    auto InspectMemAccessInstForUB = [&](Instruction &I) {
       // Skip instructions that are already saved.
-      if (NoUBLoads.count(&I) || UBLoads.count(&I))
+      if (NoUBMemAccessInsts.count(&I) || UBMemAccessInsts.count(&I))
         return true;
 
-      Value *PtrOp = cast<LoadInst>(&I)->getPointerOperand();
+      // `InspectMemAccessInstForUB` is only called on instructions
+      // for which getPointerOperand() should give us their
+      // pointer operand unless they're volatile.
+      const Value *PtrOp = getPointerOperand(&I);
+      if (!PtrOp)
+        return true;
 
       // A load is considered UB only if it dereferences a constant
       // null pointer.
       if (!isa<ConstantPointerNull>(PtrOp)) {
-        NoUBLoads.insert(&I);
+        NoUBMemAccessInsts.insert(&I);
         return true;
       }
-      Type *PtrTy = PtrOp->getType();
+      const Type *PtrTy = PtrOp->getType();
 
       // Because we only consider loads inside functions,
       // assume that a parent function exists.
       const Function *F = I.getFunction();
 
       // A dereference on constant null is only considered UB
       // if null dereference is _not_ defined for the target platform.
       // TODO: Expand it to not only check constant values.
       if (!llvm::NullPointerIsDefined(F, PtrTy->getPointerAddressSpace()))
-        UBLoads.insert(&I);
+        UBMemAccessInsts.insert(&I);
       else
-        NoUBLoads.insert(&I);
+        NoUBMemAccessInsts.insert(&I);
       return true;
     };
 
-    A.checkForAllInstructions(InspectLoadForUB, *this, {Instruction::Load});
-    if (PrevSize != NoUBLoads.size())
+    A.checkForAllInstructions(InspectMemAccessInstForUB, *this,
+                              {Instruction::Load, Instruction::Store,
+                               Instruction::AtomicCmpXchg,
+                               Instruction::AtomicRMW});
+    if (PrevSize != NoUBMemAccessInsts.size())
       return ChangeStatus::CHANGED;
     return ChangeStatus::UNCHANGED;
   }
 
   bool isAssumedToCauseUB(Instruction *I) const override {
-    return UBLoads.count(I);
+    return UBMemAccessInsts.count(I);
   }
 
   ChangeStatus manifest(Attributor &A) override {
-    if (!UBLoads.size())
+    if (!UBMemAccessInsts.size())
       return ChangeStatus::UNCHANGED;
-    for (Instruction *I : UBLoads)
+    for (Instruction *I : UBMemAccessInsts)
       changeToUnreachable(I, /* UseLLVMTrap */ false);
     return ChangeStatus::CHANGED;
   }
 
   /// See AbstractAttribute::getAsStr()
   const std::string getAsStr() const override {
     return getAssumed() ? "undefined-behavior" : "no-ub";
   }
 
 protected:
-  // A set of all the (live) load instructions that _are_ assumed to cause UB.
-  SmallPtrSet<Instruction *, 8> UBLoads;
+  // A set of all the (live) memory accessing instructions that _are_ assumed to
+  // cause UB.
+  SmallPtrSet<Instruction *, 8> UBMemAccessInsts;
 
 private:
-  // A set of all the (live) load instructions that are _not_ assumed to cause
-  // UB.
+  // A set of all the (live) memory accessing instructions
+  // that are _not_ assumed to cause UB.
   //   Note: The correctness of the procedure depends on the fact that this
   //   set stops changing after some point. "Change" here means that the size
   //   of the set changes. The size of this set is monotonically increasing
-  //   (we only add items to it) and is upper bounded by the number of load
-  //   instructions in the processed function (we can never save more elements
-  //   in this set than this number). Hence, the size of this set, at some
-  //   point, will stop increasing, effectively reaching a fixpoint.
-  SmallPtrSet<Instruction *, 8> NoUBLoads;
+  //   (we only add items to it) and is upper bounded by the number of memory
+  //   accessing instructions in the processed function (we can never save more
+  //   elements in this set than this number). Hence, the size of this set, at
+  //   some point, will stop increasing, effectively reaching a fixpoint.
+  SmallPtrSet<Instruction *, 8> NoUBMemAccessInsts;
 };
 
 struct AAUndefinedBehaviorFunction final : AAUndefinedBehaviorImpl {
   AAUndefinedBehaviorFunction(const IRPosition &IRP)
       : AAUndefinedBehaviorImpl(IRP) {}
 
   /// See AbstractAttribute::trackStatistics()
   void trackStatistics() const override {
     STATS_DECL(UndefinedBehaviorInstruction, Instruction,
                "Number of instructions known to have UB");
     BUILD_STAT_NAME(UndefinedBehaviorInstruction, Instruction) +=
-        UBLoads.size();
+        UBMemAccessInsts.size();
   }
 };
 
 /// ------------------------ Will-Return Attributes ----------------------------
 
 // Helper function that checks whether a function has any cycle.
 // TODO: Replace with more efficent code
 static bool containsCycle(Function &F) {
case Instruction::Load:
       // The alignment of a pointer is interesting for loads.
     case Instruction::Store:
       // The alignment of a pointer is interesting for stores.
     case Instruction::Call:
     case Instruction::CallBr:
     case Instruction::Invoke:
     case Instruction::CleanupRet:
     case Instruction::CatchSwitch:
+    case Instruction::AtomicRMW:
+    case Instruction::AtomicCmpXchg:
     case Instruction::Resume:
     case Instruction::Ret:
       IsInterestingOpcode = true;
     }
     if (IsInterestingOpcode)
       InstOpcodeMap[I.getOpcode()].push_back(&I);
     if (I.mayReadOrWriteMemory())
       ReadOrWriteInsts.push_back(&I);

 
 raw_ostream &llvm::operator<<(raw_ostream &OS, const IRPosition &Pos) {
   const Value &AV = Pos.getAssociatedValue();
   return OS << "{" << Pos.getPositionKind() << ":" << AV.getName() << " ["
             << Pos.getAnchorValue().getName() << "@" << Pos.getArgNo() << "]}";
 }
 
 template <typename base_ty, base_ty BestState, base_ty WorstState>
-raw_ostream &
-llvm::operator<<(raw_ostream &OS,
-                 const IntegerStateBase<base_ty, BestState, WorstState> &S) {
+raw_ostream &llvm::
+operator<<(raw_ostream &OS,
   return OS << "(" << S.getKnown() << "-" << S.getAssumed() << ")"
             << static_cast<const AbstractState &>(S);
 }
 
 raw_ostream &llvm::operator<<(raw_ostream &OS, const AbstractState &S) {
   return OS << (!S.isValidState() ? "top" : (S.isAtFixpoint() ? "fix" : ""));
 }
 

llvm/test/Transforms/Attributor/undefined_behavior.ll

 ; RUN: opt --attributor --attributor-disable=false -S < %s | FileCheck %s --check-prefix=ATTRIBUTOR
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 
 ; Test cases specifically designed for the "undefined behavior" abstract function attribute.
 ; We want to verify that whenever undefined behavior is assumed, the code becomes unreachable.
 ; We use FIXME's to indicate problems and missing attributes.
 
-; ATTRIBUTOR: define void @wholly_unreachable()
+; -- Load tests --
+
+; ATTRIBUTOR-LABEL: define void @load_wholly_unreachable()
+define void @load_wholly_unreachable() {
 ; ATTRIBUTOR-NEXT: unreachable
-define void @wholly_unreachable() {
   %a = load i32, i32* null
   ret void
 }
 
-; ATTRIBUTOR: define void @single_bb_unreachable(i1 %cond)
-; ATTRIBUTOR-NEXT: br i1 %cond, label %t, label %e
-; ATTRIBUTOR-EMPTY: 
-; ATTRIBUTOR-NEXT: t:
-; ATTRIBUTOR-NEXT: unreachable
-; ATTRIBUTOR-EMPTY:
-; ATTRIBUTOR-NEXT: e:
-; ATTRIBUTOR-NEXT: ret void
-define void @single_bb_unreachable(i1 %cond) {
+define void @load_single_bb_unreachable(i1 %cond) {
+; ATTRIBUTOR-LABEL: @load_single_bb_unreachable(
+; ATTRIBUTOR-NEXT:    br i1 [[COND:%.*]], label [[T:%.*]], label [[E:%.*]]
+; ATTRIBUTOR:       t:
+; ATTRIBUTOR-NEXT:    unreachable
+; ATTRIBUTOR:       e:
+; ATTRIBUTOR-NEXT:    ret void
+;
   br i1 %cond, label %t, label %e
 t:
   %b = load i32, i32* null
   br label %e
 e:
   ret void
 }
 
-; ATTRIBUTOR: define void @null_pointer_is_defined()
-; ATTRIBUTOR-NEXT: %a = load i32, i32* null
-define void @null_pointer_is_defined() "null-pointer-is-valid"="true" {
+define void @load_null_pointer_is_defined() "null-pointer-is-valid"="true" {
+; ATTRIBUTOR-LABEL: @load_null_pointer_is_defined(
+; ATTRIBUTOR-NEXT:    [[A:%.*]] = load i32, i32* null
+; ATTRIBUTOR-NEXT:    ret void
+;
   %a = load i32, i32* null
   ret void
 }
+
+; -- Store tests --
+
+define void @store_wholly_unreachable() {
+; ATTRIBUTOR-LABEL: @store_wholly_unreachable(
+; ATTRIBUTOR-NEXT:    unreachable
+;
+  store i32 5, i32* null
+  ret void
+}
+
+define void @store_single_bb_unreachable(i1 %cond) {
+; ATTRIBUTOR-LABEL: @store_single_bb_unreachable(
+; ATTRIBUTOR-NEXT:    br i1 [[COND:%.*]], label [[T:%.*]], label [[E:%.*]]
+; ATTRIBUTOR:       t:
+; ATTRIBUTOR-NEXT:    unreachable
+; ATTRIBUTOR:       e:
+; ATTRIBUTOR-NEXT:    ret void
+;
+  br i1 %cond, label %t, label %e
+t:
+  store i32 5, i32* null
+  br label %e
+e:
+  ret void
+}
+
+define void @store_null_pointer_is_defined() "null-pointer-is-valid"="true" {
+; ATTRIBUTOR-LABEL: @store_null_pointer_is_defined(
+; ATTRIBUTOR-NEXT:    store i32 5, i32* null
+; ATTRIBUTOR-NEXT:    ret void
+;
+  store i32 5, i32* null
+  ret void
+}
+
+; -- AtomicRMW tests --
+
+define void @atomicrmw_wholly_unreachable() {
+; ATTRIBUTOR-LABEL: @atomicrmw_wholly_unreachable(
+; ATTRIBUTOR-NEXT:    unreachable
+;
+  %a = atomicrmw add i32* null, i32 1 acquire
+  ret void
+}
+
+define void @atomicrmw_single_bb_unreachable(i1 %cond) {
+; ATTRIBUTOR-LABEL: @atomicrmw_single_bb_unreachable(
+; ATTRIBUTOR-NEXT:    br i1 [[COND:%.*]], label [[T:%.*]], label [[E:%.*]]
+; ATTRIBUTOR:       t:
+; ATTRIBUTOR-NEXT:    unreachable
+; ATTRIBUTOR:       e:
+; ATTRIBUTOR-NEXT:    ret void
+;
+  br i1 %cond, label %t, label %e
+t:
+  %a = atomicrmw add i32* null, i32 1 acquire
+  br label %e
+e:
+  ret void
+}
+
+define void @atomicrmw_null_pointer_is_defined() "null-pointer-is-valid"="true" {
+; ATTRIBUTOR-LABEL: @atomicrmw_null_pointer_is_defined(
+; ATTRIBUTOR-NEXT:    [[A:%.*]] = atomicrmw add i32* null, i32 1 acquire
+; ATTRIBUTOR-NEXT:    ret void
+;
+  %a = atomicrmw add i32* null, i32 1 acquire
+  ret void
+}
+
+; -- AtomicCmpXchg tests --
+
+define void @atomiccmpxchg_wholly_unreachable() {
+; ATTRIBUTOR-LABEL: @atomiccmpxchg_wholly_unreachable(
+; ATTRIBUTOR-NEXT:    unreachable
+;
+  %a = cmpxchg i32* null, i32 2, i32 3 acq_rel monotonic
+  ret void
+}
+
+define void @atomiccmpxchg_single_bb_unreachable(i1 %cond) {
+; ATTRIBUTOR-LABEL: @atomiccmpxchg_single_bb_unreachable(
+; ATTRIBUTOR-NEXT:    br i1 [[COND:%.*]], label [[T:%.*]], label [[E:%.*]]
+; ATTRIBUTOR:       t:
+; ATTRIBUTOR-NEXT:    unreachable
+; ATTRIBUTOR:       e:
+; ATTRIBUTOR-NEXT:    ret void
+;
+  br i1 %cond, label %t, label %e
+t:
+  %a = cmpxchg i32* null, i32 2, i32 3 acq_rel monotonic
+  br label %e
+e:
+  ret void
+}
+
+define void @atomiccmpxchg_null_pointer_is_defined() "null-pointer-is-valid"="true" {
+; ATTRIBUTOR-LABEL: @atomiccmpxchg_null_pointer_is_defined(
+; ATTRIBUTOR-NEXT:    [[A:%.*]] = cmpxchg i32* null, i32 2, i32 3 acq_rel monotonic
+; ATTRIBUTOR-NEXT:    ret void
+;
+  %a = cmpxchg i32* null, i32 2, i32 3 acq_rel monotonic
+  ret void
+}