[DevTSAN] Ignore check for joint matrix access (#18773)

we can't know exactly what size of joint matrix type is, so we need to
ignore it.

Author: zhaomaosu

llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp

@@ -43,6 +43,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Instrumentation/SPIRVSanitizerCommonUtils.h"
 #include "llvm/Transforms/Utils/EscapeEnumerator.h"
 #include "llvm/Transforms/Utils/Instrumentation.h"
 #include "llvm/Transforms/Utils/Local.h"
@@ -52,13 +53,6 @@ using namespace llvm;
 
 #define DEBUG_TYPE "tsan"
 
-// Spir memory address space
-static constexpr unsigned kSpirOffloadPrivateAS = 0;
-static constexpr unsigned kSpirOffloadGlobalAS = 1;
-static constexpr unsigned kSpirOffloadConstantAS = 2;
-static constexpr unsigned kSpirOffloadLocalAS = 3;
-static constexpr unsigned kSpirOffloadGenericAS = 4;
-
 static cl::opt<bool> ClInstrumentMemoryAccesses(
     "tsan-instrument-memory-accesses", cl::init(true),
     cl::desc("Instrument memory accesses"), cl::Hidden);
@@ -127,6 +121,8 @@ struct ThreadSanitizerOnSpirv {
 
   void appendDebugInfoToArgs(Instruction *I, SmallVectorImpl<Value *> &Args);
 
+  bool isUnsupportedSPIRAccess(Value *Addr, Instruction *Inst);
+
 private:
   void instrumentGlobalVariables();
 
@@ -383,6 +379,38 @@ void ThreadSanitizerOnSpirv::appendDebugInfoToArgs(
   Args.push_back(ConstantExpr::getPointerCast(FuncNameGV, ConstASPtrTy));
 }
 
+bool ThreadSanitizerOnSpirv::isUnsupportedSPIRAccess(Value *Addr,
+                                                     Instruction *Inst) {
+  auto *OrigValue = getUnderlyingObject(Addr);
+  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 (auto *SI = dyn_cast<StoreInst>(Inst))
+    if (getTargetExtType(SI->getValueOperand()->getType()) ||
+        isJointMatrixAccess(SI->getPointerOperand()))
+      return true;
+
+  if (auto *LI = dyn_cast<LoadInst>(Inst))
+    if (getTargetExtType(Inst->getType()) ||
+        isJointMatrixAccess(LI->getPointerOperand()))
+      return true;
+
+  auto AddrAS = cast<PointerType>(Addr->getType()->getScalarType())
+                    ->getPointerAddressSpace();
+  switch (AddrAS) {
+  case kSpirOffloadPrivateAS:
+  case kSpirOffloadLocalAS:
+  case kSpirOffloadConstantAS:
+    return true;
+  case kSpirOffloadGlobalAS:
+  case kSpirOffloadGenericAS:
+    return false;
+  }
+  return false;
+}
+
 bool ThreadSanitizerOnSpirv::isSupportedSPIRKernel(Function &F) {
 
   if (!F.hasFnAttribute(Attribute::SanitizeThread) ||
@@ -709,30 +737,12 @@ static bool shouldInstrumentReadWriteFromAddress(const Module *M, Value *Addr) {
     }
   }
 
-  if (Triple(M->getTargetTriple()).isSPIROrSPIRV()) {
-    auto *OrigValue = getUnderlyingObject(Addr);
-    if (OrigValue->getName().starts_with("__spirv_BuiltIn"))
+  // Do not instrument accesses from different address spaces; we cannot deal
+  // with them.
+  if (Addr) {
+    Type *PtrTy = cast<PointerType>(Addr->getType()->getScalarType());
+    if (PtrTy->getPointerAddressSpace() != 0)
       return false;
-
-    auto AddrAS = cast<PointerType>(Addr->getType()->getScalarType())
-                      ->getPointerAddressSpace();
-    switch (AddrAS) {
-    case kSpirOffloadPrivateAS:
-    case kSpirOffloadLocalAS:
-    case kSpirOffloadConstantAS:
-      return false;
-    case kSpirOffloadGlobalAS:
-    case kSpirOffloadGenericAS:
-      return true;
-    }
-  } else {
-    // Do not instrument accesses from different address spaces; we cannot deal
-    // with them.
-    if (Addr) {
-      Type *PtrTy = cast<PointerType>(Addr->getType()->getScalarType());
-      if (PtrTy->getPointerAddressSpace() != 0)
-        return false;
-    }
   }
 
   return true;
@@ -781,7 +791,10 @@ void ThreadSanitizer::chooseInstructionsToInstrument(
     Value *Addr = IsWrite ? cast<StoreInst>(I)->getPointerOperand()
                           : cast<LoadInst>(I)->getPointerOperand();
 
-    if (!shouldInstrumentReadWriteFromAddress(I->getModule(), Addr))
+    if (Spirv) {
+      if (Spirv->isUnsupportedSPIRAccess(Addr, I))
+        continue;
+    } else if (!shouldInstrumentReadWriteFromAddress(I->getModule(), Addr))
       continue;
 
     if (!IsWrite) {
@@ -890,7 +903,8 @@ bool ThreadSanitizer::sanitizeFunction(Function &F,
       else if (isa<LoadInst>(Inst) || isa<StoreInst>(Inst))
         LocalLoadsAndStores.push_back(&Inst);
       else if (Spirv && isa<AllocaInst>(Inst) &&
-               cast<AllocaInst>(Inst).getAllocatedType()->isSized())
+               cast<AllocaInst>(Inst).getAllocatedType()->isSized() &&
+               !getTargetExtType(cast<AllocaInst>(Inst).getAllocatedType()))
         Allocas.push_back(&Inst);
       else if ((isa<CallInst>(Inst) && !isa<DbgInfoIntrinsic>(Inst)) ||
                isa<InvokeInst>(Inst)) {

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

@@ -0,0 +1,25 @@
+; RUN: opt < %s -passes='function(tsan),module(tsan-module)' -tsan-instrument-func-entry-exit=0 -tsan-instrument-memintrinsics=0 -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"
+
+%"class.sycl::_V1::ext::oneapi::bfloat16" = type { i16 }
+%"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix" = type { target("spirv.CooperativeMatrixKHR", i16, 3, 16, 32, 0) }
+
+declare dso_local spir_func noundef ptr addrspace(4) @_Z19__spirv_AccessChainIN4sycl3_V13ext6oneapi8bfloat16ES4_Lm16ELm32ELN5__spv9MatrixUseE0ELNS5_5Scope4FlagE3EEPT_PPNS5_28__spirv_CooperativeMatrixKHRIT0_XT4_EXT1_EXT2_EXT3_EEEm(ptr addrspace(4) noundef, i64 noundef)
+
+define weak_odr dso_local spir_kernel void @test() {
+; CHECK-LABEL: void @test
+; CHECK-NOT: call void @__tsan_write
+; CHECK-NOT: ptrtoint ptr %sub_a.i to i64
+entry:
+  %sub_a.i = alloca %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix", align 8
+  %element.i = alloca %"class.sycl::_V1::ext::oneapi::bfloat16", align 2
+  %0 = getelementptr inbounds { i16 }, ptr %element.i, i64 0, i32 0
+  %spvm.i = getelementptr inbounds %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix", ptr %sub_a.i, i64 0, i32 0
+  %addrcast = addrspacecast ptr %spvm.i to ptr addrspace(4)
+  %call.i67 = call spir_func noundef ptr addrspace(4) @_Z19__spirv_AccessChainIN4sycl3_V13ext6oneapi8bfloat16ES4_Lm16ELm32ELN5__spv9MatrixUseE0ELNS5_5Scope4FlagE3EEPT_PPNS5_28__spirv_CooperativeMatrixKHRIT0_XT4_EXT1_EXT2_EXT3_EEEm(ptr addrspace(4) noundef %addrcast, i64 1)
+  %gep = getelementptr inbounds nuw { i16 }, ptr addrspace(4) %call.i67, i64 0, i32 0
+  %val = load i16, ptr %0, align 2
+  store i16 %val, ptr addrspace(4) %gep, align 2
+  ret void
+}