[Asan] Teach FunctionStackPoisoner to filter out struct type with scalable vector type.

[yetingk]

FunctionStackPoisoner does not serve for AllocaInst with scalable vector type, but it does not filter out struct type with scalable vector introduced by c8eb535.
Currently, llvm does not allows an element of a struct type with scalable vector is an element of a struct type vector, so we only need to check the first layer of the struct type of the AllocaInst.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: Yeting Kuo (yetingk)

Changes

FunctionStackPoisoner does not serve for AllocaInst with scalable vector type, but it does not filter out struct type with scalable vector introduced by c8eb535.
Currently, llvm does not allows an element of a struct type with scalable vector is an element of a struct type vector, so we only need to check the first layer of the struct type of the AllocaInst.

---

Full diff: https://github.com/llvm/llvm-project/pull/93406.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp (+9-1)
• (added) llvm/test/Instrumentation/AddressSanitizer/asan-struct-scalable.ll (+11)

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index 9cc978dc6c16e..011262c5ee949 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -1139,8 +1139,16 @@ struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   /// Collect Alloca instructions we want (and can) handle.
   void visitAllocaInst(AllocaInst &AI) {
     // FIXME: Handle scalable vectors instead of ignoring them.
+    auto IsScalableVecTy = [&](const Type *Ty) {
+      if (const auto *STy = dyn_cast<StructType>(Ty))
+        return any_of(STy->elements(), [&](const Type *ElemTy) {
+          return isa<ScalableVectorType>(ElemTy);
+        });
+      return isa<ScalableVectorType>(Ty);
+    };
+
     if (!ASan.isInterestingAlloca(AI) ||
-        isa<ScalableVectorType>(AI.getAllocatedType())) {
+        IsScalableVecTy(AI.getAllocatedType())) {
       if (AI.isStaticAlloca()) {
         // Skip over allocas that are present *before* the first instrumented
         // alloca, we don't want to move those around.
diff --git a/llvm/test/Instrumentation/AddressSanitizer/asan-struct-scalable.ll b/llvm/test/Instrumentation/AddressSanitizer/asan-struct-scalable.ll
new file mode 100644
index 0000000000000..d03f70d808a53
--- /dev/null
+++ b/llvm/test/Instrumentation/AddressSanitizer/asan-struct-scalable.ll
@@ -0,0 +1,11 @@
+; RUN: opt -passes=asan -disable-output -S %s
+; Check not crash.
+
+define void @test() #0 {
+entry:
+  %t0 = alloca { <vscale x 2 x i32>, <vscale x 2 x i32> }, align 4
+  call void null(ptr null, ptr %t0, i64 0)
+  ret void
+}
+
+attributes #0 = { sanitize_address }

[preames]

See containsHomogeneousScalableVectorTypes on StructType.

[yetingk]

Done. Thank you for the suggestion.

[preames]

LGTM