[MSAN] Correct shadow type for atomicrmw instrumentation

[loladiro]

Summary

We were passing the type of Val to getShadowOriginPtr, rather
than the type of Val's shadow resulting in broken IR. The fix
is simple.

[loladiro]

fixup CHECKs

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

#endif
   }
 
   void handleCASOrRMW(Instruction &I) {
     assert(isa<AtomicRMWInst>(I) || isa<AtomicCmpXchgInst>(I));
 
     IRBuilder<> IRB(&I);
     Value *Addr = I.getOperand(0);
     Value *Val = I.getOperand(1);
-    Value *ShadowPtr = getShadowOriginPtr(Addr, IRB, Val->getType(), Align(1),
+    Value *ShadowPtr = getShadowOriginPtr(Addr, IRB, getShadowTy(Val), Align(1),
                                           /*isStore*/ true)
                            .first;
 
     if (ClCheckAccessAddress)
       insertShadowCheck(Addr, &I);
 
     // Only test the conditional argument of cmpxchg instruction.
     // The other argument can potentially be uninitialized, but we can not

llvm/test/Instrumentation/MemorySanitizer/atomics.ll

 }
 
 ; CHECK-LABEL: @AtomicRmwXchg
 ; CHECK: store i32 0,
 ; CHECK: atomicrmw xchg {{.*}} seq_cst
 ; CHECK: store i32 0, {{.*}} @__msan_retval_tls
 ; CHECK: ret i32
 
+; atomicrmw xchg: store clean shadow, return clean shadow
+
+define i32* @AtomicRmwXchgPtr(i32** %p, i32* %x) sanitize_memory {
+entry:
+  %0 = atomicrmw xchg i32** %p, i32* %x seq_cst
+  ret i32* %0
+}
+
+; CHECK-LABEL: @AtomicRmwXchg
+; CHECK: store i32 0,
+; CHECK: atomicrmw xchg {{.*}} seq_cst
+; CHECK: store i32 0, {{.*}} @__msan_retval_tls
+; CHECK: ret i32
+
 
 ; atomicrmw max: exactly the same as above
 
 define i32 @AtomicRmwMax(i32* %p, i32 %x) sanitize_memory {
 entry:
   %0 = atomicrmw max i32* %p, i32 %x seq_cst
   ret i32 %0
 }