[Mlir-commits] [mlir] [mlir][llvm] Do not inline variadic functions (PR #77241)

[Tobias Gysi]

https://github.com/gysit created https://github.com/llvm/llvm-project/pull/77241

This revision updates the llvm dialect inliner to explicitly disallow the inlining of variadic functions. Already previously the inlining failed if the number of function arguments did not match the number of call arguments. After the change, inlining checks the function is not variadic and it does not contain a va_start intrinsic.

>From a2cf499574b06706320abc8022d14d9966d34935 Mon Sep 17 00:00:00 2001
From: Tobias Gysi <tobias.gysi at nextsilicon.com>
Date: Sun, 7 Jan 2024 14:11:14 +0000
Subject: [PATCH] [mlir][llvm] Do not inline variadic functions

This revision updates the llvm dialect inliner to explicitly disallow
the inlining of variadic functions. Already previously the inlining
failed if the number of function arguments did not match the number of
call arguments. After the change, inlining checks the function is not
variadic and it does not contain a va_start intrinsic.
---
 mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp |  7 +++++-
 mlir/test/Dialect/LLVMIR/inlining.mlir      | 25 +++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index 65c1daee6711ad..4a6154ea6d3004 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -663,6 +663,10 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
                  << "Cannot inline: callable is not an LLVM::LLVMFuncOp\n");
       return false;
     }
+    if (funcOp.isVarArg()) {
+      LLVM_DEBUG(llvm::dbgs() << "Cannot inline: callable is variadic\n");
+      return false;
+    }
     // TODO: Generate aliasing metadata from noalias argument/result attributes.
     if (auto attrs = funcOp.getArgAttrs()) {
       for (DictionaryAttr attrDict : attrs->getAsRange<DictionaryAttr>()) {
@@ -704,7 +708,8 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
   }
 
   bool isLegalToInline(Operation *op, Region *, bool, IRMapping &) const final {
-    return true;
+    // The inliner cannot handle variadic function arguments.
+    return !isa<LLVM::VaStartOp>(op);
   }
 
   /// Handle the given inlined return by replacing it with a branch. This
diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index 63e7a46f1bdb06..61fd88f133e5a8 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -644,3 +644,28 @@ llvm.func @caller(%ptr : !llvm.ptr) -> i32 {
   llvm.store %c5, %ptr { access_groups = [#caller] } : i32, !llvm.ptr
   llvm.return %0 : i32
 }
+
+// -----
+
+llvm.func @vararg_func(...) {
+  llvm.return
+}
+
+llvm.func @vararg_intrinrics() {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  %list = llvm.alloca %0 x !llvm.struct<"struct.va_list_opaque", (ptr)> : (i32) -> !llvm.ptr
+  // The vararg intinriscs should normally be part of a variadic function.
+  // However, this test uses a non-variadic function to ensure the presence of
+  // the intrinsic alone suffices to prevent inlining.
+  llvm.intr.vastart %list : !llvm.ptr
+  llvm.return
+}
+
+// CHECK-LABEL: @caller
+llvm.func @caller() {
+  // CHECK-NEXT: llvm.call @vararg_func()
+  llvm.call @vararg_func() vararg(!llvm.func<void (...)>) : () -> ()
+  // CHECK-NEXT: llvm.call @vararg_intrinrics()
+  llvm.call @vararg_intrinrics() : () -> ()
+  llvm.return
+}