[DeviceSanitizer] Ignore target extension type (#15484)

We don't know exactly what size it is, so just ignore that type.

Author: zhaomaosu

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -1506,19 +1506,49 @@ static bool isUnsupportedAMDGPUAddrspace(Value *Addr) {
   return false;
 }
 
-static bool isUnsupportedSPIRAccess(Value *Addr, Function *Func) {
+static bool containsTargetExtType(const Type *Ty) {
+  if (isa<TargetExtType>(Ty))
+    return true;
+
+  if (Ty->isVectorTy())
+    return containsTargetExtType(Ty->getScalarType());
+
+  if (Ty->isArrayTy())
+    return containsTargetExtType(Ty->getArrayElementType());
+
+  if (auto *STy = dyn_cast<StructType>(Ty)) {
+    for (unsigned int i = 0; i < STy->getNumElements(); i++)
+      if (containsTargetExtType(STy->getElementType(i)))
+        return true;
+    return false;
+  }
+
+  return false;
+}
+
+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;
 
+  // Ignore load/store for target ext type since we can't know exactly what size
+  // it is.
+  if (isa<StoreInst>(Inst) &&
+      containsTargetExtType(
+          cast<StoreInst>(Inst)->getValueOperand()->getType()))
+    return true;
+
+  if (isa<LoadInst>(Inst) && containsTargetExtType(Inst->getType()))
+    return true;
+
   Type *PtrTy = cast<PointerType>(Addr->getType()->getScalarType());
   switch (PtrTy->getPointerAddressSpace()) {
   case kSpirOffloadPrivateAS: {
     if (!ClSpirOffloadPrivates)
       return true;
     // Skip kernel arguments
-    return Func->getCallingConv() == CallingConv::SPIR_KERNEL &&
+    return Inst->getFunction()->getCallingConv() == CallingConv::SPIR_KERNEL &&
            isa<Argument>(Addr);
   }
   case kSpirOffloadGlobalAS: {
@@ -1756,7 +1786,10 @@ bool AddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
        // swifterror allocas are register promoted by ISel
        !AI.isSwiftError() &&
        // safe allocas are not interesting
-       !(SSGI && SSGI->isSafe(AI)));
+       !(SSGI && SSGI->isSafe(AI)) &&
+       // ignore alloc contains target ext type since we can't know exactly what
+       // size it is.
+       !containsTargetExtType(AI.getAllocatedType()));
 
   ProcessedAllocas[&AI] = IsInteresting;
   return IsInteresting;
@@ -1765,7 +1798,7 @@ bool AddressSanitizer::isInterestingAlloca(const AllocaInst &AI) {
 bool AddressSanitizer::ignoreAccess(Instruction *Inst, Value *Ptr) {
   // SPIR has its own rules to filter the instrument accesses
   if (TargetTriple.isSPIROrSPIRV()) {
-    if (isUnsupportedSPIRAccess(Ptr, Inst->getFunction()))
+    if (isUnsupportedSPIRAccess(Ptr, Inst))
       return true;
   } else {
     // Instrument accesses from different address spaces only for AMDGPU.

llvm/test/Instrumentation/AddressSanitizer/SPIRV/ignore_target_ext_type.ll

@@ -0,0 +1,20 @@
+; RUN: opt < %s -passes=asan -asan-instrumentation-with-call-threshold=0 -asan-stack=0 -asan-globals=0 -asan-use-after-return=never -asan-stack-dynamic-alloca=0 -asan-mapping-scale=4 -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"
+
+%"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix" = type { target("spirv.JointMatrixINTEL", i16, 16, 32, 0, 3, 0, 1) }
+
+define spir_kernel void @_ZTS4multIN4sycl3_V13ext6oneapi8bfloat16ELm16ELm16ELm32EE() {
+entry:
+; CHECK-NOT: MyAlloc
+  %sub_a.i = alloca [2 x %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix"], i32 0, align 8
+  br label %for.cond10.i
+
+for.cond10.i:                                     ; preds = %for.cond10.i, %entry
+  %0 = load target("spirv.JointMatrixINTEL", i16, 16, 32, 0, 3, 0, 1), ptr null, align 8
+  store target("spirv.JointMatrixINTEL", float, 16, 16, 3, 3, 2) zeroinitializer, ptr null, align 8
+; CHECK-NOT: asan_load
+; CHECK-NOT: asan_store
+  br label %for.cond10.i
+}