[CodeGen][i386] Move -mregparm storage earlier and fix Runtime calls (#89707)

When building the Linux kernel for i386, the -mregparm=3 option is
enabled. Crashes were observed in the sanitizer handler functions, and
the problem was found to be mismatched calling convention.

As was fixed in commit c167c0a4dcdb ("[BuildLibCalls] infer inreg param
attrs from NumRegisterParameters"), call arguments need to be marked as
"in register" when -mregparm is set. Use the same helper developed there
to update the function arguments.

Since CreateRuntimeFunction() is actually part of CodeGenModule, storage
of the -mregparm value is also moved to the constructor, as doing this
in Release() is too late.

Fixes: https://github.com/llvm/llvm-project/issues/89670

4 files changed, 36 insertions, 6 deletions

diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index d085e73..c8898ce 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -73,6 +73,7 @@
 #include "llvm/TargetParser/RISCVISAInfo.h"
 #include "llvm/TargetParser/Triple.h"
 #include "llvm/TargetParser/X86TargetParser.h"
+#include "llvm/Transforms/Utils/BuildLibCalls.h"
 #include <optional>
 
 using namespace clang;
@@ -442,6 +443,11 @@ CodeGenModule::CodeGenModule(ASTContext &C,
       }
     ModuleNameHash = llvm::getUniqueInternalLinkagePostfix(Path);
   }
+
+  // Record mregparm value now so it is visible through all of codegen.
+  if (Context.getTargetInfo().getTriple().getArch() == llvm::Triple::x86)
+    getModule().addModuleFlag(llvm::Module::Error, "NumRegisterParameters",
+                              CodeGenOpts.NumRegisterParameters);
 }
 
 CodeGenModule::~CodeGenModule() {}
@@ -980,11 +986,6 @@ void CodeGenModule::Release() {
       NMD->addOperand(MD);
   }
 
-  // Record mregparm value now so it is visible through rest of codegen.
-  if (Context.getTargetInfo().getTriple().getArch() == llvm::Triple::x86)
-    getModule().addModuleFlag(llvm::Module::Error, "NumRegisterParameters",
-                              CodeGenOpts.NumRegisterParameters);
-
   if (CodeGenOpts.DwarfVersion) {
     getModule().addModuleFlag(llvm::Module::Max, "Dwarf Version",
                               CodeGenOpts.DwarfVersion);
@@ -4781,6 +4782,10 @@ CodeGenModule::CreateRuntimeFunction(llvm::FunctionType *FTy, StringRef Name,
         }
       }
       setDSOLocal(F);
+      // FIXME: We should use CodeGenModule::SetLLVMFunctionAttributes() instead
+      // of trying to approximate the attributes using the LLVM function
+      // signature. This requires revising the API of CreateRuntimeFunction().
+      markRegisterParameterAttributes(F);
     }
   }
 
diff --git a/clang/test/CodeGen/regparm-flag.c b/clang/test/CodeGen/regparm-flag.c
index c35b53c..d888c1e 100644
--- a/clang/test/CodeGen/regparm-flag.c
+++ b/clang/test/CodeGen/regparm-flag.c
@@ -1,4 +1,8 @@
 // RUN: %clang_cc1 -triple i386-unknown-unknown -mregparm 4 %s -emit-llvm -o - | FileCheck %s
+// RUN: %clang_cc1 -triple i386-unknown-unknown -fsanitize=array-bounds %s -emit-llvm -o - | FileCheck %s --check-prefix=RUNTIME0
+// RUN: %clang_cc1 -triple i386-unknown-unknown -mregparm 1 -fsanitize=array-bounds %s -emit-llvm -o - | FileCheck %s --check-prefix=RUNTIME1
+// RUN: %clang_cc1 -triple i386-unknown-unknown -mregparm 2 -fsanitize=array-bounds %s -emit-llvm -o - | FileCheck %s --check-prefix=RUNTIME2
+// RUN: %clang_cc1 -triple i386-unknown-unknown -mregparm 3 -fsanitize=array-bounds %s -emit-llvm -o - | FileCheck %s --check-prefix=RUNTIME2
 
 void f1(int a, int b, int c, int d,
         int e, int f, int g, int h);
@@ -13,7 +17,21 @@ void f0(void) {
   f2(1, 2);
 }
 
+struct has_array {
+  int a;
+  int b[4];
+  int c;
+};
+
+int access(struct has_array *p, int index)
+{
+  return p->b[index];
+}
+
 // CHECK: declare void @f1(i32 inreg noundef, i32 inreg noundef, i32 inreg noundef, i32 inreg noundef,
 // CHECK: i32 noundef, i32 noundef, i32 noundef, i32 noundef)
 // CHECK: declare void @f2(i32 noundef, i32 noundef)
 
+// RUNTIME0: declare void @__ubsan_handle_out_of_bounds_abort(ptr, i32)
+// RUNTIME1: declare void @__ubsan_handle_out_of_bounds_abort(ptr inreg, i32)
+// RUNTIME2: declare void @__ubsan_handle_out_of_bounds_abort(ptr inreg, i32 inreg)
diff --git a/llvm/include/llvm/Transforms/Utils/BuildLibCalls.h b/llvm/include/llvm/Transforms/Utils/BuildLibCalls.h
index 9ebb950..429d6a2 100644
--- a/llvm/include/llvm/Transforms/Utils/BuildLibCalls.h
+++ b/llvm/include/llvm/Transforms/Utils/BuildLibCalls.h
@@ -62,6 +62,13 @@ namespace llvm {
                      LibFunc TheLibFunc, AttributeList AttributeList,
                      FunctionType *Invalid, ArgsTy... Args) = delete;
 
+  // Handle -mregparm for the given function.
+  // Note that this function is a rough approximation that only works for simple
+  // function signatures; it does not apply other relevant attributes for
+  // function signatures, including sign/zero-extension for arguments and return
+  // values.
+  void markRegisterParameterAttributes(Function *F);
+
   /// Check whether the library function is available on target and also that
   /// it in the current Module is a Function with the right type.
   bool isLibFuncEmittable(const Module *M, const TargetLibraryInfo *TLI,
diff --git a/llvm/lib/Transforms/Utils/BuildLibCalls.cpp b/llvm/lib/Transforms/Utils/BuildLibCalls.cpp
index ed0ed34..e97506b 100644
--- a/llvm/lib/Transforms/Utils/BuildLibCalls.cpp
+++ b/llvm/lib/Transforms/Utils/BuildLibCalls.cpp
@@ -1255,7 +1255,7 @@ static void setRetExtAttr(Function &F,
 }
 
 // Modeled after X86TargetLowering::markLibCallAttributes.
-static void markRegisterParameterAttributes(Function *F) {
+void llvm::markRegisterParameterAttributes(Function *F) {
   if (!F->arg_size() || F->isVarArg())
     return;