[DeviceSanitizer] Move isUnsupportedSPIRAccess to SPIRVSanitizerCommonUtils.cpp. (#17597)

Both AddressSanitizer and MemorySanitizer ignore target extension type.

Author: haonanya1

.github/CODEOWNERS

@@ -202,9 +202,11 @@ llvm/include/llvm/Transforms/Instrumentation/AddressSanitizer.h @intel/dpcpp-san
 llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h @intel/dpcpp-sanitizers-review
 llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerOptions.h @intel/dpcpp-sanitizers-review
 llvm/include/llvm/Transforms/Instrumentation/MemorySanitizer.h @intel/dpcpp-sanitizers-review
+llvm/include/llvm/Transforms/Instrumentation/SPIRVSanitizerCommonUtils.h @intel/dpcpp-sanitizers-review
 llvm/include/llvm/Transforms/Instrumentation/ThreadSanitizer.h @intel/dpcpp-sanitizers-review
 llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp @intel/dpcpp-sanitizers-review
 llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp @intel/dpcpp-sanitizers-review
+llvm/lib/Transforms/Instrumentation/SPIRVSanitizerCommonUtils.cpp @intel/dpcpp-sanitizers-review
 llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp @intel/dpcpp-sanitizers-review
 llvm/test/Instrumentation/AddressSanitizer/ @intel/dpcpp-sanitizers-review
 llvm/test/Instrumentation/MemorySanitizer/ @intel/dpcpp-sanitizers-review

llvm/include/llvm/Transforms/Instrumentation/SPIRVSanitizerCommonUtils.h

@@ -0,0 +1,31 @@
+//===- SPIRVSanitizerCommonUtils.h - Commnon utils --------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares common infrastructure for SPIRV Sanitizer.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_TRANSFORMS_INSTRUMENTATION_SPIRVSANITIZERCOMMONUTILS_H
+#define LLVM_TRANSFORMS_INSTRUMENTATION_SPIRVSANITIZERCOMMONUTILS_H
+
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Type.h"
+#include "llvm/IR/Value.h"
+
+namespace llvm {
+// Spir memory address space
+constexpr unsigned kSpirOffloadPrivateAS = 0;
+constexpr unsigned kSpirOffloadGlobalAS = 1;
+constexpr unsigned kSpirOffloadConstantAS = 2;
+constexpr unsigned kSpirOffloadLocalAS = 3;
+constexpr unsigned kSpirOffloadGenericAS = 4;
+
+TargetExtType *getTargetExtType(Type *Ty);
+bool isJointMatrixAccess(Value *V);
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_INSTRUMENTATION_SPIRVSANITIZERCOMMONUTILS_H

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -74,6 +74,7 @@
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerCommon.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerOptions.h"
+#include "llvm/Transforms/Instrumentation/SPIRVSanitizerCommonUtils.h"
 #include "llvm/Transforms/Utils/ASanStackFrameLayout.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Instrumentation.h"
@@ -195,13 +196,6 @@ constexpr size_t kAccessSizeIndexMask = 0xf;
 constexpr size_t kIsWriteShift = 5;
 constexpr size_t kIsWriteMask = 0x1;
 
-// 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;
-
 // Command-line flags.
 
 static cl::opt<bool> ClEnableKasan(
@@ -1653,49 +1647,6 @@ static bool isUnsupportedAMDGPUAddrspace(Value *Addr) {
   return false;
 }
 
-static TargetExtType *getTargetExtType(Type *Ty) {
-  if (auto *TargetTy = dyn_cast<TargetExtType>(Ty))
-    return TargetTy;
-
-  if (Ty->isVectorTy())
-    return getTargetExtType(Ty->getScalarType());
-
-  if (Ty->isArrayTy())
-    return getTargetExtType(Ty->getArrayElementType());
-
-  if (auto *STy = dyn_cast<StructType>(Ty)) {
-    for (unsigned int i = 0; i < STy->getNumElements(); i++)
-      if (auto *TargetTy = getTargetExtType(STy->getElementType(i)))
-        return TargetTy;
-    return nullptr;
-  }
-
-  return nullptr;
-}
-
-// Skip pointer operand that is sycl joint matrix access since it isn't from
-// user code, e.g. %call:
-// clang-format off
-// %a = alloca %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix", align 8
-// %0 = getelementptr inbounds %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix", ptr %a, i64 0, i32 0
-// %call = call spir_func ptr
-// @_Z19__spirv_AccessChainIfN4sycl3_V13ext6oneapi12experimental6matrix9precision4tf32ELm8ELm8ELN5__spv9MatrixUseE0ELNS8_5Scope4FlagE3EEPT_PPNS8_28__spirv_CooperativeMatrixKHRIT0_XT4_EXT1_EXT2_EXT3_EEEm(ptr %0, i64 0)
-// %1 = load float, ptr %call, align 4
-// store float %1, ptr %call, align 4
-// clang-format on
-static bool isJointMatrixAccess(Value *V) {
-  auto *ActualV = V->stripInBoundsOffsets();
-  if (auto *CI = dyn_cast<CallInst>(ActualV)) {
-    for (Value *Op : CI->args()) {
-      if (auto *AI = dyn_cast<AllocaInst>(Op->stripInBoundsOffsets()))
-        if (auto *TargetTy = getTargetExtType(AI->getAllocatedType()))
-          return TargetTy->getName().starts_with("spirv.") &&
-                 TargetTy->getName().contains("Matrix");
-    }
-  }
-  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

llvm/lib/Transforms/Instrumentation/CMakeLists.txt

@@ -26,6 +26,7 @@ add_llvm_component_library(LLVMInstrumentation
   TypeSanitizer.cpp
   HWAddressSanitizer.cpp
   RealtimeSanitizer.cpp
+  SPIRVSanitizerCommonUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -200,6 +200,7 @@
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TargetParser/Triple.h"
+#include "llvm/Transforms/Instrumentation/SPIRVSanitizerCommonUtils.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Instrumentation.h"
 #include "llvm/Transforms/Utils/Local.h"
@@ -588,13 +589,6 @@ static const PlatformMemoryMapParams Intel_SPIR_MemoryMapParams = {
     &Intel_SPIR64_MemoryMapParams,
 };
 
-// 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;
-
 namespace {
 
 class MemorySanitizerOnSpirv;
@@ -1746,6 +1740,18 @@ static bool isUnsupportedSPIRAccess(const Value *Addr, Instruction *I) {
   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>(I))
+    if (getTargetExtType(SI->getValueOperand()->getType()) ||
+        isJointMatrixAccess(SI->getPointerOperand()))
+      return true;
+
+  if (auto *LI = dyn_cast<LoadInst>(I))
+    if (getTargetExtType(I->getType()) ||
+        isJointMatrixAccess(LI->getPointerOperand()))
+      return true;
+
   Type *PtrTy = cast<PointerType>(Addr->getType()->getScalarType());
   switch (PtrTy->getPointerAddressSpace()) {
   case kSpirOffloadPrivateAS:

llvm/lib/Transforms/Instrumentation/SPIRVSanitizerCommonUtils.cpp

@@ -0,0 +1,61 @@
+//===- SPIRVSanitizerCommonUtils.cpp- SPIRV Sanitizer commnon utils ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines common infrastructure for SPIRV Sanitizer.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Instrumentation/SPIRVSanitizerCommonUtils.h"
+#include "llvm/IR/Instructions.h"
+
+using namespace llvm;
+
+namespace llvm {
+TargetExtType *getTargetExtType(Type *Ty) {
+  if (auto *TargetTy = dyn_cast<TargetExtType>(Ty))
+    return TargetTy;
+
+  if (Ty->isVectorTy())
+    return getTargetExtType(Ty->getScalarType());
+
+  if (Ty->isArrayTy())
+    return getTargetExtType(Ty->getArrayElementType());
+
+  if (auto *STy = dyn_cast<StructType>(Ty)) {
+    for (unsigned int i = 0; i < STy->getNumElements(); i++)
+      if (auto *TargetTy = getTargetExtType(STy->getElementType(i)))
+        return TargetTy;
+    return nullptr;
+  }
+
+  return nullptr;
+}
+
+// Skip pointer operand that is sycl joint matrix access since it isn't from
+// user code, e.g. %call:
+// clang-format off
+// %a = alloca %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix", align 8
+// %0 = getelementptr inbounds %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix", ptr %a, i64 0, i32 0
+// %call = call spir_func ptr
+// @_Z19__spirv_AccessChainIfN4sycl3_V13ext6oneapi12experimental6matrix9precision4tf32ELm8ELm8ELN5__spv9MatrixUseE0ELNS8_5Scope4FlagE3EEPT_PPNS8_28__spirv_CooperativeMatrixKHRIT0_XT4_EXT1_EXT2_EXT3_EEEm(ptr %0, i64 0)
+// %1 = load float, ptr %call, align 4
+// store float %1, ptr %call, align 4
+// clang-format on
+bool isJointMatrixAccess(Value *V) {
+  auto *ActualV = V->stripInBoundsOffsets();
+  if (auto *CI = dyn_cast<CallInst>(ActualV)) {
+    for (Value *Op : CI->args()) {
+      if (auto *AI = dyn_cast<AllocaInst>(Op->stripInBoundsOffsets()))
+        if (auto *TargetTy = getTargetExtType(AI->getAllocatedType()))
+          return TargetTy->getName().starts_with("spirv.") &&
+                 TargetTy->getName().contains("Matrix");
+    }
+  }
+  return false;
+}
+} // namespace llvm

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

@@ -0,0 +1,25 @@
+; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -msan-spir-privates=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) }
+
+; CHECK-LABEL: @test
+; CHECK-NOT: call i64 @__msan_get_shadow
+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() {
+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
+}