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

[Yeting Kuo via llvm-commits]

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

>From 129e4a7ef410f5c63f9bc11864acb9c3bfed9793 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..88c49b2ef72e8 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 }