[DevMSAN] Always make sure poison shuffle mask point to clean shadow (#18191)

For Spirv target, we need to make sure poison mask always point to a
clean shadow to avoid pollution of launch info clean shadow.

Signed-off-by: Zhao, Maosu <[email protected]>

Author: zhaomaosu

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -3152,8 +3152,24 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     IRBuilder<> IRB(&I);
     auto *Shadow0 = getShadow(&I, 0);
     auto *Shadow1 = getShadow(&I, 1);
-    setShadow(&I, IRB.CreateShuffleVector(Shadow0, Shadow1, I.getShuffleMask(),
-                                          "_msprop"));
+    // For Spirv target, we need to make sure poison mask always point to a
+    // clean shadow to avoid pollution of launch info clean shadow.
+    if (SpirOrSpirv &&
+        any_of(I.getShuffleMask(), [](int Mask) { return Mask == -1; })) {
+      auto *V = IRB.CreateShuffleVector(Shadow0, Shadow1, I.getShuffleMask(),
+                                        "_msprop");
+      auto *Shadow2 = getCleanShadow(V);
+      SmallVector<int, 4> NewMask;
+      I.getShuffleMask(NewMask);
+      for (auto &Mask : NewMask) {
+        if (Mask == -1)
+          Mask = NewMask.size() + 1;
+      }
+      setShadow(&I, IRB.CreateShuffleVector(V, Shadow2, NewMask, "_msprop"));
+    } else {
+      setShadow(&I, IRB.CreateShuffleVector(Shadow0, Shadow1,
+                                            I.getShuffleMask(), "_msprop"));
+    }
     setOriginForNaryOp(I);
   }
 

llvm/test/Instrumentation/MemorySanitizer/SPIRV/instrument_shuffle_vector.ll

@@ -0,0 +1,17 @@
+; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -S | FileCheck %s
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
+target triple = "spir64-unknown-unknown"
+
+define spir_kernel void @MyKernel(ptr addrspace(1) noundef align 4 %_arg_array) sanitize_memory {
+entry:
+  %arrayidx3 = getelementptr inbounds <3 x i16>, ptr addrspace(1) %_arg_array, i64 0
+  %extractVec4 = shufflevector <3 x i16> <i16 0, i16 0, i16 0>, <3 x i16> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 poison>
+; CHECK: [[REG1:%[0-9]+]] = call i64 @__msan_get_shadow
+; CHECK: [[REG2:%[0-9]+]] = inttoptr i64 [[REG1]] to ptr addrspace(1)
+; CHECK: store <4 x i16> zeroinitializer, ptr addrspace(1) [[REG2]]
+  store <4 x i16> %extractVec4, ptr addrspace(1) %arrayidx3, align 8
+  br label %exit
+
+exit:
+  ret void
+}