[DeviceSanitizer] Do sanitize for device globals in AddressSanitizer pass (#13678)

zhaomaosu · callumfare · web-flow · commit 33fe64c3973d · 2024-11-19T15:07:31.000+01:00
Now UR already implemented API "urProgramGetGlobalVariablePointer", so we can use it to query the size of device globals and remove "__AsanDeviceGlobalCount". UR Part: oneapi-src/unified-runtime#1584 --------- Co-authored-by: Callum Fare <callum@codeplay.com>
diff --git a/clang/lib/Driver/OffloadBundler.cpp b/clang/lib/Driver/OffloadBundler.cpp
@@ -688,8 +688,11 @@ class ObjectFileHandler final : public FileHandler {
           return std::move(Err);
 
         // If we are dealing with a bitcode file do not add special globals
-        // llvm.used and llvm.compiler.used to the list of defined symbols.
-        if (SF->isIR() && (Name == "llvm.used" || Name == "llvm.compiler.used"))
+        // llvm.used and llvm.compiler.used and __AsanDeviceGlobalMetadata to
+        // the list of defined symbols.
+        if (SF->isIR() &&
+            (Name == "llvm.used" || Name == "llvm.compiler.used" ||
+             Name == "__AsanDeviceGlobalMetadata"))
           continue;
 
         // Add symbol name with the target prefix to the buffer.
diff --git a/llvm/include/llvm/SYCLLowerIR/SanitizeDeviceGlobal.h b/llvm/include/llvm/SYCLLowerIR/SanitizeDeviceGlobal.h
diff --git a/llvm/lib/SYCLLowerIR/CMakeLists.txt b/llvm/lib/SYCLLowerIR/CMakeLists.txt
@@ -69,7 +69,6 @@ add_llvm_component_library(LLVMSYCLLowerIR
   SYCLPropagateJointMatrixUsage.cpp
   SYCLVirtualFunctionsAnalysis.cpp
   SYCLUtils.cpp
-  SanitizeDeviceGlobal.cpp
 
   LocalAccessorToSharedMemory.cpp
   GlobalOffset.cpp
diff --git a/llvm/lib/SYCLLowerIR/SanitizeDeviceGlobal.cpp b/llvm/lib/SYCLLowerIR/SanitizeDeviceGlobal.cpp
diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/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"
@@ -465,6 +466,10 @@ static cl::opt<bool>
                           cl::desc("instrument generic pointer"), cl::Hidden,
                           cl::init(true));
 
+static cl::opt<bool> ClDeviceGlobals("asan-device-globals",
+                                     cl::desc("instrument device globals"),
+                                     cl::Hidden, cl::init(true));
+
 // Debug flags.
 
 static cl::opt<int> ClDebug("asan-debug", cl::desc("debug"), cl::Hidden,
@@ -970,6 +975,7 @@ class ModuleAddressSanitizer {
 private:
   void initializeCallbacks();
 
+  void instrumentDeviceGlobal(IRBuilder<> &IRB);
   void instrumentGlobals(IRBuilder<> &IRB, bool *CtorComdat);
   void InstrumentGlobalsCOFF(IRBuilder<> &IRB,
                              ArrayRef<GlobalVariable *> ExtendedGlobals,
@@ -1556,12 +1562,23 @@ static bool isJointMatrixAccess(Value *V) {
   return false;
 }
 
+static bool isUnsupportedDeviceGlobal(GlobalVariable *G) {
+  // 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.
+  return (!isDeviceGlobalVariable(*G) || !hasDeviceImageScopeProperty(*G));
+}
+
 static bool isUnsupportedSPIRAccess(Value *Addr, Instruction *Inst) {
   // Skip SPIR-V built-in varibles
   auto *OrigValue = Addr->stripInBoundsOffsets();
   if (OrigValue->getName().starts_with("__spirv_BuiltIn"))
     return true;
 
+  GlobalVariable *GV = dyn_cast<GlobalVariable>(OrigValue);
+  if (GV && isUnsupportedDeviceGlobal(GV))
+    return true;
+
   // Ignore load/store for target ext type since we can't know exactly what size
   // it is.
   if (auto *SI = dyn_cast<StoreInst>(Inst))
@@ -2766,6 +2783,71 @@ Instruction *ModuleAddressSanitizer::CreateAsanModuleDtor() {
   return ReturnInst::Create(*C, AsanDtorBB);
 }
 
+void ModuleAddressSanitizer::instrumentDeviceGlobal(IRBuilder<> &IRB) {
+  auto &DL = M.getDataLayout();
+  SmallVector<GlobalVariable *, 8> GlobalsToRemove;
+  SmallVector<Constant *, 8> DeviceGlobalMetadata;
+
+  Type *IntptrTy = M.getDataLayout().getIntPtrType(*C, kSpirOffloadGlobalAS);
+
+  // 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(IntptrTy, IntptrTy, IntptrTy);
+
+  for (auto &G : M.globals()) {
+    if (isUnsupportedDeviceGlobal(&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);
+    DeviceGlobalMetadata.push_back(ConstantStruct::get(
+        StructTy, ConstantInt::get(IntptrTy, SizeInBytes),
+        ConstantInt::get(IntptrTy, SizeInBytes + RightRedzoneSize),
+        ConstantExpr::getPointerCast(NewGlobal, IntptrTy)));
+  }
+
+  if (GlobalsToRemove.empty())
+    return;
+
+  // Create meta data global to record device globals' information
+  ArrayType *ArrayTy = ArrayType::get(StructTy, DeviceGlobalMetadata.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, ArrayRef<GlobalVariable *> ExtendedGlobals,
     ArrayRef<Constant *> MetadataInitializers) {
@@ -3234,6 +3316,11 @@ bool ModuleAddressSanitizer::instrumentModule() {
     auto *MD = M.getOrInsertNamedMetadata("device.sanitizer");
     Metadata *MDVals[] = {MDString::get(Ctx, "asan")};
     MD->addOperand(MDNode::get(Ctx, MDVals));
+
+    if (ClDeviceGlobals) {
+      IRBuilder<> IRB(*C);
+      instrumentDeviceGlobal(IRB);
+    }
   }
 
   const uint64_t Priority = GetCtorAndDtorPriority(TargetTriple);
diff --git a/llvm/lib/Transforms/Instrumentation/CMakeLists.txt b/llvm/lib/Transforms/Instrumentation/CMakeLists.txt
@@ -39,6 +39,7 @@ add_llvm_component_library(LLVMInstrumentation
   Core
   Demangle
   MC
+  SYCLLowerIR
   Support
   TargetParser
   TransformUtils
diff --git a/llvm/test/Instrumentation/AddressSanitizer/SPIRV/device_global.ll b/llvm/test/Instrumentation/AddressSanitizer/SPIRV/device_global.ll
@@ -0,0 +1,13 @@
+; RUN: opt < %s -passes=asan -asan-instrumentation-with-call-threshold=0 -asan-stack=0 -asan-globals=0 -asan-constructor-kind=none -S | FileCheck %s
+
+; check that image scope device globals can be correctly instrumented.
+
+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"
+target triple = "spir64-unknown-unknown"
+
+@dev_global = addrspace(1) global { [4 x i32] } zeroinitializer #0
+
+; CHECK: @dev_global = addrspace(1) global { { [4 x i32] }, [16 x i8] }
+; CHECK: @__AsanDeviceGlobalMetadata = appending local_unnamed_addr addrspace(1) global [1 x { i64, i64, i64 }] [{ i64, i64, i64 } { i64 16, i64 32, i64 ptrtoint (ptr addrspace(1) @dev_global to i64) }]
+
+attributes #0 = { "sycl-device-global-size"="16" "sycl-device-image-scope" }
diff --git a/llvm/test/Instrumentation/AddressSanitizer/SPIRV/device_global_non_image_scope.ll b/llvm/test/Instrumentation/AddressSanitizer/SPIRV/device_global_non_image_scope.ll
@@ -0,0 +1,11 @@
+; RUN: opt < %s -passes=asan -asan-instrumentation-with-call-threshold=0 -asan-stack=0 -asan-globals=0 -asan-constructor-kind=none -S | FileCheck %s
+
+; check non image scope device globals will not be instrumented.
+
+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"
+target triple = "spir64-unknown-unknown"
+
+@dev_global = addrspace(1) global { ptr addrspace(1), [4 x i32] } zeroinitializer
+
+; CHECK: @dev_global = addrspace(1) global { ptr addrspace(1), [4 x i32] }
+; CHECK-NOT: @__AsanDeviceGlobalMetadata
diff --git a/llvm/tools/sycl-post-link/sycl-post-link.cpp b/llvm/tools/sycl-post-link/sycl-post-link.cpp
@@ -40,7 +40,6 @@
 #include "llvm/SYCLLowerIR/ModuleSplitter.h"
 #include "llvm/SYCLLowerIR/SYCLJointMatrixTransform.h"
 #include "llvm/SYCLLowerIR/SYCLUtils.h"
-#include "llvm/SYCLLowerIR/SanitizeDeviceGlobal.h"
 #include "llvm/SYCLLowerIR/SpecConstants.h"
 #include "llvm/SYCLLowerIR/Support.h"
 #include "llvm/Support/CommandLine.h"
@@ -791,11 +790,6 @@ processInputModule(std::unique_ptr<Module> M) {
   if (M->getTargetTriple().find("spir") != std::string::npos)
     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);
-
   // Transform Joint Matrix builtin calls to align them with SPIR-V friendly
   // LLVM IR specification.
   Modified |= runModulePass<SYCLJointMatrixTransformPass>(*M);
diff --git a/sycl/cmake/modules/UnifiedRuntimeTag.cmake b/sycl/cmake/modules/UnifiedRuntimeTag.cmake
@@ -1,7 +1,7 @@
-# commit 30391c65d2d2ccc7ee3688a14815804bfb7fdf05
-# Merge: 5e6d79b3 58dabfe8
+# commit 0ea47d7c70b9a21a3d90612a0a0e7525034e62f7
+# Merge: e3247c23 e36941cb
 # Author: Callum Fare <callum@codeplay.com>
-# Date:   Fri Nov 15 15:13:20 2024 +0000
-#     Merge pull request #2222 from RossBrunton/ross/cfi
-#     Enable -flto and -fsanitize=cfi in clang
-set(UNIFIED_RUNTIME_TAG 30391c65d2d2ccc7ee3688a14815804bfb7fdf05)
+# Date:   Tue Nov 19 10:24:08 2024 +0000
+#     Merge pull request #1584 from zhaomaosu/simplify-device-global
+#     [DeviceSanitizer] Remove device global "__AsanDeviceGlobalCount"
+set(UNIFIED_RUNTIME_TAG 0ea47d7c70b9a21a3d90612a0a0e7525034e62f7)
