[msan] Don't use TLS slots of noundef args

[vitalybuka]

Summary

Transformations may strip the attribute from the
argument, e.g. for unused, which will result in
shadow offsets mismatch between caller and
callee.

Stripping noundef for used arguments can be
a problem, as TLS is not going to be set
by caller. However this is not the goal of the
patch and I am not aware if that's even
possible.

[kda]

LGTM

[kstoimenov]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3712–3713

Consider converting this to "else if (ByVal)" to reduce indentation level.

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3712–3713

Plain conversion will not work. There is IRB.CreateStore(getOrigin(A) which shared for both non EagerCheck branches

[kstoimenov]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3712–3713

That makes sense. Let's keep as is.

[eugenis]

LGTM

I agree it is a little bit scary that noundef can be removed from an argument - msan appears to put more requirements on the attribute than the rest of llvm (it's a call abi attribute for msan, not just optimization).

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

if (Argument *A = dyn_cast<Argument>(V)) {
             setOrigin(A, EntryIRB.CreateLoad(MS.OriginTy, OriginPtr));
           } else {
             setOrigin(A, getCleanOrigin());
           }
 
           break;
         }
 
-        if (!FArgEagerCheck)
-          ArgOffset += alignTo(Size, kShadowTLSAlignment);
       }
       assert(*ShadowPtr && "Could not find shadow for an argument");
       return *ShadowPtr;
     }
     // For everything else the shadow is zero.
     return getCleanShadow(V);
   }
 
for (auto ArgIt = CB.arg_begin(), End = CB.arg_end(); ArgIt != End;
       const DataLayout &DL = F.getParent()->getDataLayout();
 
       bool ByVal = CB.paramHasAttr(i, Attribute::ByVal);
       bool NoUndef = CB.paramHasAttr(i, Attribute::NoUndef);
       bool EagerCheck = MayCheckCall && !ByVal && NoUndef;
 
       if (EagerCheck) {
         insertShadowCheck(A, &CB);
-        continue;
-      }
+        Size = DL.getTypeAllocSize(A->getType());
+      } else {
         if (ByVal) {
           // ByVal requires some special handling as it's too big for a single
           // load
           assert(A->getType()->isPointerTy() &&
                  "ByVal argument is not a pointer!");
           Size = DL.getTypeAllocSize(CB.getParamByValType(i));
-        if (ArgOffset + Size > kParamTLSSize) break;
+          if (ArgOffset + Size > kParamTLSSize)
+            break;
           const MaybeAlign ParamAlignment(CB.getParamAlign(i));
           MaybeAlign Alignment = llvm::None;
           if (ParamAlignment)
             Alignment = std::min(*ParamAlignment, kShadowTLSAlignment);
           Value *AShadowPtr =
               getShadowOriginPtr(A, IRB, IRB.getInt8Ty(), Alignment,
                                  /*isStore*/ false)
                   .first;
 
           Store = IRB.CreateMemCpy(ArgShadowBase, Alignment, AShadowPtr,
                                    Alignment, Size);
           // TODO(glider): need to copy origins.
         } else {
-        // Any other parameters mean we need bit-grained tracking of uninit data
+          // Any other parameters mean we need bit-grained tracking of uninit
+          // data
           Size = DL.getTypeAllocSize(A->getType());
-        if (ArgOffset + Size > kParamTLSSize) break;
+          if (ArgOffset + Size > kParamTLSSize)
+            break;
           Store = IRB.CreateAlignedStore(ArgShadow, ArgShadowBase,
                                          kShadowTLSAlignment);
           Constant *Cst = dyn_cast<Constant>(ArgShadow);
-        if (Cst && Cst->isNullValue()) ArgIsInitialized = true;
+          if (Cst && Cst->isNullValue())
+            ArgIsInitialized = true;
         }
         if (MS.TrackOrigins && !ArgIsInitialized)
           IRB.CreateStore(getOrigin(A),
                           getOriginPtrForArgument(A, IRB, ArgOffset));
         (void)Store;
-      assert(Size != 0 && Store != nullptr);
+        assert(Store != nullptr);
         LLVM_DEBUG(dbgs() << "  Param:" << *Store << "\n");
+      }
+      assert(Size != 0);
       ArgOffset += alignTo(Size, kShadowTLSAlignment);
     }
     LLVM_DEBUG(dbgs() << "  done with call args\n");
 
     FunctionType *FT = CB.getFunctionType();
     if (FT->isVarArg()) {
       VAHelper->visitCallBase(CB, IRB);
     }

llvm/test/Instrumentation/MemorySanitizer/msan_eager.ll

 ;
   %p = inttoptr i64 0 to i32 *
   store i32 %a, i32 *%p
   ret void
 }
 
 define void @NormalArgAfterNoUndef(i32 noundef %a, i32 %b) nounwind uwtable sanitize_memory {
 ; CHECK-LABEL: @NormalArgAfterNoUndef(
-; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* bitcast ([100 x i64]* @__msan_param_tls to i32*), align 8
-; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__msan_param_origin_tls, i32 0, i32 0), align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__msan_param_tls to i64), i64 8) to i32*), align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* inttoptr (i64 add (i64 ptrtoint ([200 x i32]* @__msan_param_origin_tls to i64), i64 8) to i32*), align 4
 ; CHECK-NEXT:    call void @llvm.donothing()
 ; CHECK-NEXT:    [[P:%.*]] = inttoptr i64 0 to i32*
 ; CHECK-NEXT:    [[TMP3:%.*]] = ptrtoint i32* [[P]] to i64
 ; CHECK-NEXT:    [[TMP4:%.*]] = xor i64 [[TMP3]], 87960930222080
 ; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to i32*
 ; CHECK-NEXT:    [[TMP6:%.*]] = add i64 [[TMP4]], 17592186044416
 ; CHECK-NEXT:    [[TMP7:%.*]] = inttoptr i64 [[TMP6]] to i32*
 ; CHECK-NEXT:    store i32 [[TMP1]], i32* [[TMP5]], align 4
;
   call void @NormalArg(i32 %r) nounwind uwtable sanitize_memory
   ret void
 }
 
 define void @CallNormalArgAfterNoUndef() nounwind uwtable sanitize_memory {
 ; CHECK-LABEL: @CallNormalArgAfterNoUndef(
 ; CHECK-NEXT:    call void @llvm.donothing()
 ; CHECK-NEXT:    [[R:%.*]] = call i32 @NormalRet() #[[ATTR0]]
-; CHECK-NEXT:    store i32 0, i32* bitcast ([100 x i64]* @__msan_param_tls to i32*), align 8
+; CHECK-NEXT:    store i32 0, i32* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__msan_param_tls to i64), i64 8) to i32*), align 8
 ; CHECK-NEXT:    call void @NormalArgAfterNoUndef(i32 [[R]], i32 [[R]]) #[[ATTR0]]
 ; CHECK-NEXT:    ret void
 ;
   %r = call i32 @NormalRet() nounwind uwtable sanitize_memory
   call void @NormalArgAfterNoUndef(i32 %r, i32 %r) nounwind uwtable sanitize_memory
   ret void
 }