[llvm] [Asan] Teach FunctionStackPoisoner to filter out struct type with sclable vector type. (PR #93406)

[Yeting Kuo via llvm-commits]

https://github.com/yetingk created https://github.com/llvm/llvm-project/pull/93406

FunctionStackPoisoner does not serve for `AllocaInst` with scalable vector type, but it does not filter out struct type with scalable vector introduced by c8eb535aed0368c20b25fe05bca563ab38dd91e9.
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`.

>From 0e00497bb73156f9f65887c1147882edd9035cb6 Mon Sep 17 00:00:00 2001
From: Yeting Kuo <yeting.kuo at sifive.com>
Date: Sun, 26 May 2024 01:31:46 -0700
Subject: [PATCH] [Asan] Teach FunctionStackPoisoner to filter out struct type
 with sclable vector type.

FunctionStackPoisoner does not serve for AllocaInst with scalable vector type,
but it does not filter out struct type with scalable vector introduced
by c8eb535aed0368c20b25fe05bca563ab38dd91e9.
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 AllocaInst.
---
 .../Transforms/Instrumentation/AddressSanitizer.cpp   | 10 +++++++++-
 .../AddressSanitizer/asan-struct-scalable.ll          | 11 +++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Instrumentation/AddressSanitizer/asan-struct-scalable.ll

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 }