[llvm] [FuzzMutate] Prevent UB caused by parameter ABI attributes (PR #139737)

[Manuel Carrasco via llvm-commits]

https://github.com/mgcarrasco updated https://github.com/llvm/llvm-project/pull/139737

>From f1534b0e3f94f5b92e6ef3204699008eb6c1dd14 Mon Sep 17 00:00:00 2001
From: Manuel Carrasco <Manuel.Carrasco at amd.com>
Date: Tue, 13 May 2025 03:23:14 -0700
Subject: [PATCH 1/2] [FuzzMutate] Refactor code determining whether the target
 function is unsupported.

---
 llvm/include/llvm/FuzzMutate/IRMutator.h |  3 +++
 llvm/lib/FuzzMutate/IRMutator.cpp        | 34 +++++++++++++++++-------
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/llvm/include/llvm/FuzzMutate/IRMutator.h b/llvm/include/llvm/FuzzMutate/IRMutator.h
index 4b02578792bc4..22681f4dcb612 100644
--- a/llvm/include/llvm/FuzzMutate/IRMutator.h
+++ b/llvm/include/llvm/FuzzMutate/IRMutator.h
@@ -144,6 +144,9 @@ class LLVM_ABI InsertFunctionStrategy : public IRMutationStrategy {
 
   using IRMutationStrategy::mutate;
   void mutate(BasicBlock &BB, RandomIRBuilder &IB) override;
+
+private:
+  bool isUnsupportedFunction(Function *F);
 };
 
 /// Strategy to split a random block and insert a random CFG in between.
diff --git a/llvm/lib/FuzzMutate/IRMutator.cpp b/llvm/lib/FuzzMutate/IRMutator.cpp
index 672c666d4e052..3725fc6ce229c 100644
--- a/llvm/lib/FuzzMutate/IRMutator.cpp
+++ b/llvm/lib/FuzzMutate/IRMutator.cpp
@@ -356,6 +356,28 @@ static uint64_t getUniqueCaseValue(SmallSet<uint64_t, 4> &CasesTaken,
   return tmp;
 }
 
+bool InsertFunctionStrategy::isUnsupportedFunction(Function *F) {
+  // Some functions accept metadata type or token type as arguments.
+  // We don't call those functions for now.
+  // For example, `@llvm.dbg.declare(metadata, metadata, metadata)`
+  // https://llvm.org/docs/SourceLevelDebugging.html#llvm-dbg-declare
+  auto IsUnsupportedTy = [](Type *T) {
+    return T->isMetadataTy() || T->isTokenTy();
+  };
+
+  if (IsUnsupportedTy(F->getReturnType()) ||
+      any_of(F->getFunctionType()->params(), IsUnsupportedTy)) {
+    return true;
+  }
+
+  // If it is not satisfied, the IR will be invalid.
+  if (!isCallableCC(F->getCallingConv())) {
+    return true;
+  }
+
+  return false;
+}
+
 void InsertFunctionStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
   Module *M = BB.getParent()->getParent();
   // If nullptr is selected, we will create a new function declaration.
@@ -366,16 +388,8 @@ void InsertFunctionStrategy::mutate(BasicBlock &BB, RandomIRBuilder &IB) {
 
   auto RS = makeSampler(IB.Rand, Functions);
   Function *F = RS.getSelection();
-  // Some functions accept metadata type or token type as arguments.
-  // We don't call those functions for now.
-  // For example, `@llvm.dbg.declare(metadata, metadata, metadata)`
-  // https://llvm.org/docs/SourceLevelDebugging.html#llvm-dbg-declare
-  auto IsUnsupportedTy = [](Type *T) {
-    return T->isMetadataTy() || T->isTokenTy();
-  };
-  if (!F || IsUnsupportedTy(F->getReturnType()) ||
-      any_of(F->getFunctionType()->params(), IsUnsupportedTy) ||
-      !isCallableCC(F->getCallingConv())) {
+
+  if (!F || isUnsupportedFunction(F)) {
     F = IB.createFunctionDeclaration(*M);
   }
 

>From 728c8bc64fcb387a95093ae74b54e327c8571980 Mon Sep 17 00:00:00 2001
From: Manuel Carrasco <Manuel.Carrasco at amd.com>
Date: Tue, 13 May 2025 07:00:55 -0700
Subject: [PATCH 2/2] [FuzzMutate] Prevent UB caused by parameter ABI
 attributes.

We make those cases unsupported as it happens now for functions accepting metadata or token types.
---
 llvm/include/llvm/FuzzMutate/IRMutator.h |  3 --
 llvm/lib/FuzzMutate/IRMutator.cpp        | 36 +++++++++++++++++++++++-
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/llvm/include/llvm/FuzzMutate/IRMutator.h b/llvm/include/llvm/FuzzMutate/IRMutator.h
index 22681f4dcb612..4b02578792bc4 100644
--- a/llvm/include/llvm/FuzzMutate/IRMutator.h
+++ b/llvm/include/llvm/FuzzMutate/IRMutator.h
@@ -144,9 +144,6 @@ class LLVM_ABI InsertFunctionStrategy : public IRMutationStrategy {
 
   using IRMutationStrategy::mutate;
   void mutate(BasicBlock &BB, RandomIRBuilder &IB) override;
-
-private:
-  bool isUnsupportedFunction(Function *F);
 };
 
 /// Strategy to split a random block and insert a random CFG in between.
diff --git a/llvm/lib/FuzzMutate/IRMutator.cpp b/llvm/lib/FuzzMutate/IRMutator.cpp
index 3725fc6ce229c..2ab79d25ce6f3 100644
--- a/llvm/lib/FuzzMutate/IRMutator.cpp
+++ b/llvm/lib/FuzzMutate/IRMutator.cpp
@@ -356,7 +356,13 @@ static uint64_t getUniqueCaseValue(SmallSet<uint64_t, 4> &CasesTaken,
   return tmp;
 }
 
-bool InsertFunctionStrategy::isUnsupportedFunction(Function *F) {
+/// Determines whether a function is unsupported by the current mutator's
+/// implementation. The function returns true if any of the following criteria
+/// are met:
+///   * The function accepts metadata or token types as arguments.
+///   * The function has ABI attributes that could cause UB.
+///   * The function uses a non-callable CC that may result in UB.
+static bool isUnsupportedFunction(Function *F) {
   // Some functions accept metadata type or token type as arguments.
   // We don't call those functions for now.
   // For example, `@llvm.dbg.declare(metadata, metadata, metadata)`
@@ -370,6 +376,34 @@ bool InsertFunctionStrategy::isUnsupportedFunction(Function *F) {
     return true;
   }
 
+  // ABI attributes must be specified both at the function
+  // declaration/definition and call-site, otherwise the
+  // behavior may be undefined.
+  // We don't call those functions for now to prevent UB from happening.
+  auto IsABIAttribute = [](AttributeSet A) {
+    static const Attribute::AttrKind ABIAttrs[] = {
+        Attribute::StructRet,      Attribute::ByVal,
+        Attribute::InAlloca,       Attribute::InReg,
+        Attribute::StackAlignment, Attribute::SwiftSelf,
+        Attribute::SwiftAsync,     Attribute::SwiftError,
+        Attribute::Preallocated,   Attribute::ByRef,
+        Attribute::ZExt,           Attribute::SExt};
+
+    return std::any_of(
+        std::begin(ABIAttrs), std::end(ABIAttrs),
+        [&](Attribute::AttrKind kind) { return A.hasAttribute(kind); });
+  };
+
+  auto FuncAttrs = F->getAttributes();
+  if (IsABIAttribute(FuncAttrs.getRetAttrs())) {
+    return true;
+  }
+  for (size_t i = 0; i < F->arg_size(); i++) {
+    if (IsABIAttribute(FuncAttrs.getParamAttrs(i))) {
+      return true;
+    }
+  }
+
   // If it is not satisfied, the IR will be invalid.
   if (!isCallableCC(F->getCallingConv())) {
     return true;