[DeviceSanitizer] Do sanitize for device globals in AddressSanitizer pass

llvm/lib/SYCLLowerIR/CMakeLists.txt

@@ -70,7 +70,6 @@ add_llvm_component_library(LLVMSYCLLowerIR
   SYCLPropagateAspectsUsage.cpp
   SYCLPropagateJointMatrixUsage.cpp
   SYCLUtils.cpp
-  SanitizeDeviceGlobal.cpp
 
   LocalAccessorToSharedMemory.cpp
   GlobalOffset.cpp

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -64,6 +64,7 @@
 #include "llvm/IR/Use.h"
 #include "llvm/IR/Value.h"
 #include "llvm/MC/MCSectionMachO.h"
+#include "llvm/SYCLLowerIR/DeviceGlobals.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -923,6 +924,7 @@ class ModuleAddressSanitizer {
   void initializeCallbacks(Module &M);
 
   void instrumentGlobals(IRBuilder<> &IRB, Module &M, bool *CtorComdat);
+  void instrumentDeviceGlobal(IRBuilder<> &IRB, Module &M);
   void InstrumentGlobalsCOFF(IRBuilder<> &IRB, Module &M,
                              ArrayRef<GlobalVariable *> ExtendedGlobals,
                              ArrayRef<Constant *> MetadataInitializers);
@@ -2449,6 +2451,77 @@ Instruction *ModuleAddressSanitizer::CreateAsanModuleDtor(Module &M) {
   return ReturnInst::Create(*C, AsanDtorBB);
 }
 
+void ModuleAddressSanitizer::instrumentDeviceGlobal(IRBuilder<> &IRB,
+                                                    Module &M) {
+  auto &DL = M.getDataLayout();
+  SmallVector<GlobalVariable *, 8> GlobalsToRemove;
+  SmallVector<GlobalVariable *, 8> NewDeviceGlobals;
+  SmallVector<Constant *, 8> DeviceGlobalMetadata;
+
+  Type *IntTy = Type::getIntNTy(M.getContext(), DL.getPointerSizeInBits());
+
+  // Device global meta data is described by a structure
+  //  size_t device_global_size
+  //  size_t device_global_size_with_red_zone
+  //  size_t beginning address of the device global
+  StructType *StructTy = StructType::get(IntTy, IntTy, IntTy);
+
+  for (auto &G : M.globals()) {
+    // Non image scope device globals are implemented by device USM, and the
+    // out-of-bounds check for them will be done by sanitizer USM part. So we
+    // exclude them here.
+    if (!isDeviceGlobalVariable(G) || !hasDeviceImageScopeProperty(G))
+      continue;
+
+    Type *Ty = G.getValueType();
+    const uint64_t SizeInBytes = DL.getTypeAllocSize(Ty);
+    const uint64_t RightRedzoneSize = getRedzoneSizeForGlobal(SizeInBytes);
+    Type *RightRedZoneTy = ArrayType::get(IRB.getInt8Ty(), RightRedzoneSize);
+    StructType *NewTy = StructType::get(Ty, RightRedZoneTy);
+    Constant *NewInitializer = ConstantStruct::get(
+        NewTy, G.getInitializer(), Constant::getNullValue(RightRedZoneTy));
+
+    // Create a new global variable with enough space for a redzone.
+    GlobalVariable *NewGlobal = new GlobalVariable(
+        M, NewTy, G.isConstant(), G.getLinkage(), NewInitializer, "", &G,
+        G.getThreadLocalMode(), G.getAddressSpace());
+    NewGlobal->copyAttributesFrom(&G);
+    NewGlobal->setComdat(G.getComdat());
+    NewGlobal->setAlignment(Align(getMinRedzoneSizeForGlobal()));
+    NewGlobal->copyMetadata(&G, 0);
+
+    Value *Indices2[2];
+    Indices2[0] = IRB.getInt32(0);
+    Indices2[1] = IRB.getInt32(0);
+
+    G.replaceAllUsesWith(
+        ConstantExpr::getGetElementPtr(NewTy, NewGlobal, Indices2, true));
+    NewGlobal->takeName(&G);
+    GlobalsToRemove.push_back(&G);
+    NewDeviceGlobals.push_back(NewGlobal);
+    DeviceGlobalMetadata.push_back(ConstantStruct::get(
+        StructTy, ConstantInt::get(IntTy, SizeInBytes),
+        ConstantInt::get(IntTy, SizeInBytes + RightRedzoneSize),
+        ConstantExpr::getPointerCast(NewGlobal, IntTy)));
+  }
+
+  if (GlobalsToRemove.empty())
+    return;
+
+  // Create meta data global to record device globals' information
+  ArrayType *ArrayTy = ArrayType::get(StructTy, NewDeviceGlobals.size());
+  Constant *MetadataInitializer =
+      ConstantArray::get(ArrayTy, DeviceGlobalMetadata);
+  GlobalVariable *AsanDeviceGlobalMetadata = new GlobalVariable(
+      M, MetadataInitializer->getType(), false, GlobalValue::AppendingLinkage,
+      MetadataInitializer, "__AsanDeviceGlobalMetadata", nullptr,
+      GlobalValue::NotThreadLocal, 1);
+  AsanDeviceGlobalMetadata->setUnnamedAddr(GlobalValue::UnnamedAddr::Local);
+
+  for (auto *G : GlobalsToRemove)
+    G->eraseFromParent();
+}
+
 void ModuleAddressSanitizer::InstrumentGlobalsCOFF(
     IRBuilder<> &IRB, Module &M, ArrayRef<GlobalVariable *> ExtendedGlobals,
     ArrayRef<Constant *> MetadataInitializers) {
@@ -2922,6 +2995,9 @@ bool ModuleAddressSanitizer::instrumentModule(Module &M) {
     auto *MD = M.getOrInsertNamedMetadata("device.sanitizer");
     Metadata *MDVals[] = {MDString::get(Ctx, "asan")};
     MD->addOperand(MDNode::get(Ctx, MDVals));
+
+    IRBuilder<> IRB(*C);
+    instrumentDeviceGlobal(IRB, M);
   }
 
   const uint64_t Priority = GetCtorAndDtorPriority(TargetTriple);

llvm/lib/Transforms/Instrumentation/CMakeLists.txt

@@ -36,6 +36,7 @@ add_llvm_component_library(LLVMInstrumentation
   Core
   Demangle
   MC
+  SYCLLowerIR
   Support
   TargetParser
   TransformUtils

llvm/tools/sycl-post-link/sycl-post-link.cpp

@@ -43,7 +43,6 @@
 #include "llvm/SYCLLowerIR/ModuleSplitter.h"
 #include "llvm/SYCLLowerIR/SYCLDeviceRequirements.h"
 #include "llvm/SYCLLowerIR/SYCLUtils.h"
-#include "llvm/SYCLLowerIR/SanitizeDeviceGlobal.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/InitLLVM.h"
@@ -1088,11 +1087,6 @@ processInputModule(std::unique_ptr<Module> M) {
   // "llvm.compiler.used" they can be erased safely.
   Modified |= removeDeviceGlobalFromCompilerUsed(*M.get());
 
-  // Instrument each image scope device globals if the module has been
-  // instrumented by sanitizer pass.
-  if (isModuleUsingAsan(*M))
-    Modified |= runModulePass<SanitizeDeviceGlobalPass>(*M);
-
   // Do invoke_simd processing before splitting because this:
   // - saves processing time (the pass is run once, even though on larger IR)
   // - doing it before SYCL/ESIMD splitting is required for correctness

sycl/plugins/unified_runtime/CMakeLists.txt

@@ -94,14 +94,8 @@ if(SYCL_PI_UR_USE_FETCH_CONTENT)
       CACHE PATH "Path to external '${name}' adapter source dir" FORCE)
   endfunction()
 
-  set(UNIFIED_RUNTIME_REPO "https://github.com/oneapi-src/unified-runtime.git")
-  # commit ebf873fb5996c9ddca32bbb7c9330d3ffe15473c
-  # Merge: 633ec408 8f375039
-  # Author: Kenneth Benzie (Benie) <[email protected]>
-  # Date:   Thu May 2 10:56:55 2024 +0100
-  #    Merge pull request #1535 from przemektmalon/przemek/sampled-image-fetch
-  #    [Bindless][Exp] Add device queries for sampled image fetch
-  set(UNIFIED_RUNTIME_TAG ebf873fb5996c9ddca32bbb7c9330d3ffe15473c)
+  set(UNIFIED_RUNTIME_REPO "https://github.com/zhaomaosu/unified-runtime.git")
+  set(UNIFIED_RUNTIME_TAG simplify-device-global)
 
   fetch_adapter_source(level_zero
     "https://github.com/oneapi-src/unified-runtime.git"