[DeviceMSAN] Fix gpu crashed on device global variable

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -767,50 +767,110 @@ Constant *getOrCreateGlobalString(Module &M, StringRef Name, StringRef Value,
   });
 }
 
-static void extendSpirKernelArgs(Module &M) {
-  SmallVector<Constant *, 8> SpirKernelsMetadata;
+static bool isUnsupportedDeviceGlobal(const GlobalVariable *G) {
+  // Skip instrumenting on "__MsanKernelMetadata" etc.
+  if (G->getName().starts_with("__Msan"))
+    return true;
+  if (G->getName().starts_with("__spirv_BuiltIn"))
+    return true;
+  if (G->getName().starts_with("__usid_str"))
+    return true;
+  if (G->getAddressSpace() == kSpirOffloadLocalAS ||
+      G->getAddressSpace() == kSpirOffloadConstantAS)
+    return true;
+  return false;
+}
+
+static void instrumentSPIRModule(Module &M) {
 
   const auto &DL = M.getDataLayout();
   Type *IntptrTy = DL.getIntPtrType(M.getContext());
 
-  // SpirKernelsMetadata only saves fixed kernels, and is described by
-  // following structure:
-  //  uptr unmangled_kernel_name
-  //  uptr unmangled_kernel_name_size
-  StructType *StructTy = StructType::get(IntptrTy, IntptrTy);
-  for (Function &F : M) {
-    if (F.getCallingConv() != CallingConv::SPIR_KERNEL)
-      continue;
+  // Instrument __MsanKernelMetadata, which records information of sanitized
+  // kernel
+  {
+    SmallVector<Constant *, 8> SpirKernelsMetadata;
+
+    // SpirKernelsMetadata only saves fixed kernels, and is described by
+    // following structure:
+    //  uptr unmangled_kernel_name
+    //  uptr unmangled_kernel_name_size
+    StructType *StructTy = StructType::get(IntptrTy, IntptrTy);
+    for (Function &F : M) {
+      if (F.getCallingConv() != CallingConv::SPIR_KERNEL)
+        continue;
 
-    if (!F.hasFnAttribute(Attribute::SanitizeMemory) ||
-        F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation))
-      continue;
+      if (!F.hasFnAttribute(Attribute::SanitizeMemory) ||
+          F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation))
+        continue;
 
-    auto KernelName = F.getName();
-    auto *KernelNameGV = getOrCreateGlobalString(M, "__msan_kernel", KernelName,
-                                                 kSpirOffloadConstantAS);
-    SpirKernelsMetadata.emplace_back(ConstantStruct::get(
-        StructTy, ConstantExpr::getPointerCast(KernelNameGV, IntptrTy),
-        ConstantInt::get(IntptrTy, KernelName.size())));
-  }
-
-  // Create global variable to record spirv kernels' information
-  ArrayType *ArrayTy = ArrayType::get(StructTy, SpirKernelsMetadata.size());
-  Constant *MetadataInitializer =
-      ConstantArray::get(ArrayTy, SpirKernelsMetadata);
-  GlobalVariable *MsanSpirKernelMetadata = new GlobalVariable(
-      M, MetadataInitializer->getType(), false, GlobalValue::AppendingLinkage,
-      MetadataInitializer, "__MsanKernelMetadata", nullptr,
-      GlobalValue::NotThreadLocal, 1);
-  MsanSpirKernelMetadata->setUnnamedAddr(GlobalValue::UnnamedAddr::Local);
-  // Add device global attributes
-  MsanSpirKernelMetadata->addAttribute(
-      "sycl-device-global-size", std::to_string(DL.getTypeAllocSize(ArrayTy)));
-  MsanSpirKernelMetadata->addAttribute("sycl-device-image-scope");
-  MsanSpirKernelMetadata->addAttribute("sycl-host-access", "0"); // read only
-  MsanSpirKernelMetadata->addAttribute("sycl-unique-id",
-                                       "_Z20__MsanKernelMetadata");
-  MsanSpirKernelMetadata->setDSOLocal(true);
+      auto KernelName = F.getName();
+      auto *KernelNameGV = getOrCreateGlobalString(
+          M, "__msan_kernel", KernelName, kSpirOffloadConstantAS);
+      SpirKernelsMetadata.emplace_back(ConstantStruct::get(
+          StructTy, ConstantExpr::getPointerCast(KernelNameGV, IntptrTy),
+          ConstantInt::get(IntptrTy, KernelName.size())));
+    }
+
+    // Create global variable to record spirv kernels' information
+    ArrayType *ArrayTy = ArrayType::get(StructTy, SpirKernelsMetadata.size());
+    Constant *MetadataInitializer =
+        ConstantArray::get(ArrayTy, SpirKernelsMetadata);
+    GlobalVariable *MsanSpirKernelMetadata = new GlobalVariable(
+        M, MetadataInitializer->getType(), false, GlobalValue::AppendingLinkage,
+        MetadataInitializer, "__MsanKernelMetadata", nullptr,
+        GlobalValue::NotThreadLocal, 1);
+    MsanSpirKernelMetadata->setUnnamedAddr(GlobalValue::UnnamedAddr::Local);
+    // Add device global attributes
+    MsanSpirKernelMetadata->addAttribute(
+        "sycl-device-global-size",
+        std::to_string(DL.getTypeAllocSize(ArrayTy)));
+    MsanSpirKernelMetadata->addAttribute("sycl-device-image-scope");
+    MsanSpirKernelMetadata->addAttribute("sycl-host-access",
+                                         "0"); // read only
+    MsanSpirKernelMetadata->addAttribute("sycl-unique-id",
+                                         "_Z20__MsanKernelMetadata");
+    MsanSpirKernelMetadata->setDSOLocal(true);
+  }
+
+  // Handle global variables:
+  //   - Skip sanitizing unsupported variables
+  //   - Instrument __MsanDeviceGlobalMetadata for device globals
+  do {
+    SmallVector<Constant *, 8> DeviceGlobalMetadata;
+
+    // Device global meta data is described by a structure
+    //  size_t device_global_size
+    //  size_t beginning address of the device global
+    StructType *StructTy = StructType::get(IntptrTy, IntptrTy);
+
+    for (auto &G : M.globals()) {
+      if (isUnsupportedDeviceGlobal(&G)) {
+        for (auto *User : G.users())
+          if (auto *Inst = dyn_cast<Instruction>(User))
+            Inst->setNoSanitizeMetadata();
+        continue;
+      }
+
+      DeviceGlobalMetadata.push_back(ConstantStruct::get(
+          StructTy,
+          ConstantInt::get(IntptrTy, DL.getTypeAllocSize(G.getValueType())),
+          ConstantExpr::getPointerCast(&G, IntptrTy)));
+    }
+
+    if (DeviceGlobalMetadata.empty())
+      break;
+
+    // Create meta data global to record device globals' information
+    ArrayType *ArrayTy = ArrayType::get(StructTy, DeviceGlobalMetadata.size());
+    Constant *MetadataInitializer =
+        ConstantArray::get(ArrayTy, DeviceGlobalMetadata);
+    GlobalVariable *MsanDeviceGlobalMetadata = new GlobalVariable(
+        M, MetadataInitializer->getType(), false, GlobalValue::AppendingLinkage,
+        MetadataInitializer, "__MsanDeviceGlobalMetadata", nullptr,
+        GlobalValue::NotThreadLocal, 1);
+    MsanDeviceGlobalMetadata->setUnnamedAddr(GlobalValue::UnnamedAddr::Local);
+  } while (false);
 }
 
 PreservedAnalyses MemorySanitizerPass::run(Module &M,
@@ -827,7 +887,7 @@ PreservedAnalyses MemorySanitizerPass::run(Module &M,
   }
 
   if (TargetTriple.isSPIROrSPIRV()) {
-    extendSpirKernelArgs(M);
+    instrumentSPIRModule(M);
     Modified = true;
   }
 

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

@@ -0,0 +1,10 @@
+; 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"
+
+@.str = external addrspace(1) constant [59 x i8]
+@__spirv_BuiltInGlobalInvocationId = external addrspace(1) constant <3 x i64>
+
+; CHECK: @__MsanDeviceGlobalMetadata
+; CHECK-NOT: @__spirv_BuiltInGlobalInvocationId
+; CHECK-SAME: @.str

sycl/cmake/modules/UnifiedRuntimeTag.cmake

@@ -1,7 +1,7 @@
-# commit afbb289aa8d4f3b27b1536ba33ca618b0aba65c7
-# Merge: ef70004f d7c33f88
-# Author: Kenneth Benzie (Benie) <k.benzie@codeplay.com>
-# Date:   Wed Jan 15 11:54:25 2025 +0000
-#     Merge pull request #2520 from zhaomaosu/fix-buffer-shadow
-#     [DevMSAN] Propagate shadow memory in buffer related APIs
-set(UNIFIED_RUNTIME_TAG afbb289aa8d4f3b27b1536ba33ca618b0aba65c7)
+# commit 9e48f543b8dd39d45563169433bb529583625dfe
+# Merge: 6a3fece6 1a1108b3
+# Author: Martin Grant <martin.morrisongrant@codeplay.com>
+# Date:   Wed Jan 15 14:33:29 2025 +0000
+#     Merge pull request #2540 from martygrant/martin/program-info-unswitch
+#     Move urProgramGetInfo success test from a switch to individual tests.
+set(UNIFIED_RUNTIME_TAG 9e48f543b8dd39d45563169433bb529583625dfe)

sycl/test-e2e/MemorySanitizer/check_device_global.cpp

@@ -0,0 +1,58 @@
+// REQUIRES: linux, cpu || (gpu && level_zero)
+// RUN: %{build} %device_msan_flags -O0 -g -o %t1.out
+// RUN: %{run} not %t1.out 2>&1 | FileCheck %s
+// RUN: %{build} %device_msan_flags -O1 -g -o %t2.out
+// RUN: %{run} not %t2.out 2>&1 | FileCheck %s
+// RUN: %{build} %device_msan_flags -O2 -g -o %t3.out
+// RUN: %{run} not %t3.out 2>&1 | FileCheck %s
+
+#include <sycl/detail/core.hpp>
+#include <sycl/ext/oneapi/device_global/device_global.hpp>
+#include <sycl/usm.hpp>
+
+using namespace sycl;
+using namespace sycl::ext::oneapi;
+using namespace sycl::ext::oneapi::experimental;
+
+sycl::ext::oneapi::experimental::device_global<
+    int[4], decltype(properties(device_image_scope, host_access_read_write))>
+    dev_global;
+
+__attribute__((noinline)) int check(int data) { return data + 1; }
+
+int main() {
+  sycl::queue Q;
+  int *array = sycl::malloc_device<int>(4, Q);
+
+  Q.submit([&](sycl::handler &h) {
+     h.single_task<class Test1>([=]() {
+       dev_global[0] = 42;
+       array[0] = check(dev_global[1]);
+       array[1] = dev_global[1];
+     });
+   }).wait();
+
+  int val[4];
+  Q.copy(dev_global, val).wait();
+  assert(val[0] == 42);
+
+  Q.submit([&](sycl::handler &h) {
+     h.single_task<class Test2>([=]() {
+       array[0] = check(array[1]);
+       dev_global[1] = array[2]; // uninitialzed value
+     });
+   }).wait();
+
+  Q.submit([&](sycl::handler &h) {
+     h.single_task<class Test3>([=]() {
+       array[0] = dev_global[1];
+       check(array[0]);
+     });
+   }).wait();
+  // CHECK: use-of-uninitialized-value
+  // CHECK-NEXT: kernel <{{.*Test3}}>
+
+  sycl::free(array, Q);
+
+  return 0;
+}