[llvm] [msan] Generalize handleIntrinsicByApplyingToShadow to allow alternative intrinsic for shadows (PR #124831)

[Thurston Dang via llvm-commits]

https://github.com/thurstond updated https://github.com/llvm/llvm-project/pull/124831

>From fa429558ffd95492eebe5a6578b962c948a18e50 Mon Sep 17 00:00:00 2001
From: Thurston Dang <thurston at google.com>
Date: Tue, 28 Jan 2025 20:10:34 +0000
Subject: [PATCH 1/2] [msan] Generalize handleIntrinsicByApplyingToShadow to
 allow alternative intrinsic for shadows

https://github.com/llvm/llvm-project/pull/124159 uses
handleIntrinsicByApplyingToShadow for horizontal add/sub, but Vitaly
recommends always using the add version to avoid false negatives for
fully uninitialized data (https://github.com/llvm/llvm-project/issues/124662).

This patch lays the groundwork by generalizing handleIntrinsicByApplyingToShadow to allow using
a different intrinsic (of the same type as the original intrinsic) for the shadow. Planned work will apply it to horizontal sub.
---
 .../Instrumentation/MemorySanitizer.cpp       | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 0d4be09846b604..d2d6a2391a686a 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -4049,7 +4049,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     //   consider this an acceptable tradeoff for performance.
     // To make shadow propagation precise, we want the equivalent of
     // "horizontal OR", but this is not available.
-    return handleIntrinsicByApplyingToShadow(I, /* trailingVerbatimArgs */ 0);
+    return handleIntrinsicByApplyingToShadow(
+        I, /*trailingVerbatimArgs*/ 0, /*shadowIntrinsicID=*/std::nullopt);
   }
 
   /// Handle Arm NEON vector store intrinsics (vst{2,3,4}, vst1x_{2,3,4},
@@ -4156,6 +4157,9 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   ///     shadow[out] =
   ///         intrinsic(shadow[var1], shadow[var2], opType) | shadow[opType]
   ///
+  /// Optionally, the intrinsic for the shadow can be replaced with another
+  /// intrinsic of the same type.
+  ///
   /// CAUTION: this assumes that the intrinsic will handle arbitrary
   ///          bit-patterns (for example, if the intrinsic accepts floats for
   ///          var1, we require that it doesn't care if inputs are NaNs).
@@ -4164,8 +4168,9 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   /// (tbl{1,2,3,4}).
   ///
   /// The origin is approximated using setOriginForNaryOp.
-  void handleIntrinsicByApplyingToShadow(IntrinsicInst &I,
-                                         unsigned int trailingVerbatimArgs) {
+  void handleIntrinsicByApplyingToShadow(
+      IntrinsicInst &I, unsigned int trailingVerbatimArgs,
+      std::optional<Intrinsic::ID> shadowIntrinsicID) {
     IRBuilder<> IRB(&I);
 
     assert(trailingVerbatimArgs < I.arg_size());
@@ -4187,8 +4192,9 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       ShadowArgs.push_back(Arg);
     }
 
-    CallInst *CI =
-        IRB.CreateIntrinsic(I.getType(), I.getIntrinsicID(), ShadowArgs);
+    CallInst *CI = IRB.CreateIntrinsic(
+        I.getType(), shadowIntrinsicID.value_or(I.getIntrinsicID()),
+        ShadowArgs);
     Value *CombinedShadow = CI;
 
     // Combine the computed shadow with the shadow of trailing args
@@ -4664,7 +4670,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     case Intrinsic::aarch64_neon_tbx3:
     case Intrinsic::aarch64_neon_tbx4: {
       // The last trailing argument (index register) should be handled verbatim
-      handleIntrinsicByApplyingToShadow(I, 1);
+      handleIntrinsicByApplyingToShadow(I, /*trailingVerbatimArgs*/ 1,
+                                        /*shadowIntrinsicID=*/std::nullopt);
       break;
     }
 

>From e15c1f19b4375e3e0a33a1179ae044fb0e58f14a Mon Sep 17 00:00:00 2001
From: Thurston Dang <thurston at google.com>
Date: Tue, 28 Jan 2025 20:32:43 +0000
Subject: [PATCH 2/2] Make shadowIntrinsicID mandatory

---
 .../Instrumentation/MemorySanitizer.cpp       | 22 ++++++++++---------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index d2d6a2391a686a..fadebc3ae42660 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -4050,7 +4050,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     // To make shadow propagation precise, we want the equivalent of
     // "horizontal OR", but this is not available.
     return handleIntrinsicByApplyingToShadow(
-        I, /*trailingVerbatimArgs*/ 0, /*shadowIntrinsicID=*/std::nullopt);
+        I, /*shadowIntrinsicID=*/I.getIntrinsicID(),
+        /*trailingVerbatimArgs*/ 0);
   }
 
   /// Handle Arm NEON vector store intrinsics (vst{2,3,4}, vst1x_{2,3,4},
@@ -4157,7 +4158,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   ///     shadow[out] =
   ///         intrinsic(shadow[var1], shadow[var2], opType) | shadow[opType]
   ///
-  /// Optionally, the intrinsic for the shadow can be replaced with another
+  /// Typically, shadowIntrinsicID will be specified by the caller to be
+  /// I.getIntrinsicID(), but the caller can choose to replace it with another
   /// intrinsic of the same type.
   ///
   /// CAUTION: this assumes that the intrinsic will handle arbitrary
@@ -4168,9 +4170,9 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   /// (tbl{1,2,3,4}).
   ///
   /// The origin is approximated using setOriginForNaryOp.
-  void handleIntrinsicByApplyingToShadow(
-      IntrinsicInst &I, unsigned int trailingVerbatimArgs,
-      std::optional<Intrinsic::ID> shadowIntrinsicID) {
+  void handleIntrinsicByApplyingToShadow(IntrinsicInst &I,
+                                         Intrinsic::ID shadowIntrinsicID,
+                                         unsigned int trailingVerbatimArgs) {
     IRBuilder<> IRB(&I);
 
     assert(trailingVerbatimArgs < I.arg_size());
@@ -4192,9 +4194,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       ShadowArgs.push_back(Arg);
     }
 
-    CallInst *CI = IRB.CreateIntrinsic(
-        I.getType(), shadowIntrinsicID.value_or(I.getIntrinsicID()),
-        ShadowArgs);
+    CallInst *CI =
+        IRB.CreateIntrinsic(I.getType(), shadowIntrinsicID, ShadowArgs);
     Value *CombinedShadow = CI;
 
     // Combine the computed shadow with the shadow of trailing args
@@ -4670,8 +4671,9 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     case Intrinsic::aarch64_neon_tbx3:
     case Intrinsic::aarch64_neon_tbx4: {
       // The last trailing argument (index register) should be handled verbatim
-      handleIntrinsicByApplyingToShadow(I, /*trailingVerbatimArgs*/ 1,
-                                        /*shadowIntrinsicID=*/std::nullopt);
+      handleIntrinsicByApplyingToShadow(
+          I, /*shadowIntrinsicID=*/I.getIntrinsicID(),
+          /*trailingVerbatimArgs*/ 1);
       break;
     }