[Mlir-commits] [mlir] [mlir] [LLVM IR] Introduce VaArgOp (PR #109260)

[Chuanqi Xu]

https://github.com/ChuanqiXu9 updated https://github.com/llvm/llvm-project/pull/109260

>From a004a265c81b98a994c886f52f31e3a7eede86ed Mon Sep 17 00:00:00 2001
From: Chuanqi Xu <yedeng.yd at linux.alibaba.com>
Date: Thu, 19 Sep 2024 17:43:01 +0800
Subject: [PATCH 1/2] [mlir] [LLVM IR] Introduce VaArgOp

---
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td | 20 ++++++++++++++++++++
 mlir/test/Target/LLVMIR/llvmir.mlir         | 10 ++++++----
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index d956d7f27f784d..862180ac940aac 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -582,6 +582,26 @@ def LLVM_LandingpadOp : LLVM_Op<"landingpad"> {
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// VAArgOp
+//===----------------------------------------------------------------------===//
+
+def LLVM_VaArgOp : LLVM_Op<"va_arg"> {
+  let arguments = (ins LLVM_AnyPointer:$arg);
+  let results = (outs LLVM_Type:$res);
+
+  let builders = [LLVM_OneResultOpBuilder];
+
+  let llvmBuilder = " $res = builder.CreateVAArg($arg, $_resultType); ";
+  let assemblyFormat = "$arg attr-dict `:` `(` type($arg) `)` `->` type($res)";
+
+  string llvmInstName = "VAArg";
+  string mlirBuilder = [{
+    $res = $_builder.create<mlir::LLVM::VaArgOp>(
+      $_location, $_resultType, $arg);
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // CallOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 7eca1a40373054..de0ab1b3442671 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -2275,8 +2275,8 @@ llvm.func @qux(f32)
 
 // CHECK: %struct.va_list = type { ptr }
 
-// CHECK: define void @vararg_function(i32 %{{.*}}, ...)
-llvm.func @vararg_function(%arg0: i32, ...) {
+// CHECK: define i32 @vararg_function(i32 %{{.*}}, ...)
+llvm.func @vararg_function(%arg0: i32, ...) -> i32 {
   %0 = llvm.mlir.constant(1 : i32) : i32
   %1 = llvm.mlir.constant(1 : i32) : i32
   // CHECK: %[[ALLOCA0:.+]] = alloca %struct.va_list, align 8
@@ -2287,12 +2287,14 @@ llvm.func @vararg_function(%arg0: i32, ...) {
   %4 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
   // CHECK: call void @llvm.va_copy.p0(ptr %[[ALLOCA1]], ptr %[[ALLOCA0]])
   llvm.intr.vacopy %2 to %4 : !llvm.ptr, !llvm.ptr
+  // CHECK: %[[RET:.+]] = va_arg ptr %[[ALLOCA1]], i32
+  %ret = llvm.va_arg %4 : (!llvm.ptr) -> i32
   // CHECK: call void @llvm.va_end.p0(ptr %[[ALLOCA1]])
   // CHECK: call void @llvm.va_end.p0(ptr %[[ALLOCA0]])
   llvm.intr.vaend %4 : !llvm.ptr
   llvm.intr.vaend %2 : !llvm.ptr
-  // CHECK: ret void
-  llvm.return
+  // CHECK: ret i32 %[[RET]]
+  llvm.return %ret : i32
 }
 
 // -----

>From 16d574087b60b0f56e3df05d39cead390416d69c Mon Sep 17 00:00:00 2001
From: Chuanqi Xu <yedeng.yd at linux.alibaba.com>
Date: Fri, 20 Sep 2024 10:46:00 +0800
Subject: [PATCH 2/2] Address comments

---
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td | 2 +-
 mlir/test/Dialect/LLVMIR/roundtrip.mlir     | 3 +++
 mlir/test/Target/LLVMIR/Import/basic.ll     | 8 +++++---
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 862180ac940aac..8584a25f8b3d6c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -593,7 +593,7 @@ def LLVM_VaArgOp : LLVM_Op<"va_arg"> {
   let builders = [LLVM_OneResultOpBuilder];
 
   let llvmBuilder = " $res = builder.CreateVAArg($arg, $_resultType); ";
-  let assemblyFormat = "$arg attr-dict `:` `(` type($arg) `)` `->` type($res)";
+  let assemblyFormat = "$arg attr-dict `:` functional-type($arg, $res)";
 
   string llvmInstName = "VAArg";
   string mlirBuilder = [{
diff --git a/mlir/test/Dialect/LLVMIR/roundtrip.mlir b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
index 0b251b81e97870..89d303fcac8ff2 100644
--- a/mlir/test/Dialect/LLVMIR/roundtrip.mlir
+++ b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
@@ -644,6 +644,9 @@ llvm.func @vararg_func(%arg0: i32, ...) {
   %list2 = llvm.alloca %1 x !llvm.struct<"struct.va_list_opaque", (ptr)> : (i32) -> !llvm.ptr
   llvm.intr.vacopy %list to %list2 : !llvm.ptr, !llvm.ptr
 
+  // CHECK: %[[RET:.+]] = llvm.va_arg %[[LIST2]] : (!llvm.ptr) -> i32
+  %ret = llvm.va_arg %list2 : (!llvm.ptr) -> i32
+
   // CHECK: llvm.intr.vaend %[[LIST]] : !llvm.ptr{{$}}
   // CHECK: llvm.intr.vaend %[[LIST2]] : !llvm.ptr{{$}}
   llvm.intr.vaend %list : !llvm.ptr
diff --git a/mlir/test/Target/LLVMIR/Import/basic.ll b/mlir/test/Target/LLVMIR/Import/basic.ll
index 448b0ebe25746c..332b583b5ce3b1 100644
--- a/mlir/test/Target/LLVMIR/Import/basic.ll
+++ b/mlir/test/Target/LLVMIR/Import/basic.ll
@@ -76,8 +76,8 @@ declare void @llvm.va_start.p0(ptr)
 declare void @llvm.va_copy.p0(ptr, ptr)
 declare void @llvm.va_end.p0(ptr)
 
-; CHECK-LABEL: llvm.func @variadic_function
-define void @variadic_function(i32 %X, ...) {
+; CHECK-LABEL: llvm.func @variadic_function(%arg0: i32, ...) -> i32
+define i32 @variadic_function(i32 %X, ...) {
   ; CHECK: %[[ALLOCA0:.+]] = llvm.alloca %{{.*}} x !llvm.struct<"struct.va_list", (ptr)> {{.*}} : (i32) -> !llvm.ptr
   %ap = alloca %struct.va_list
   ; CHECK: llvm.intr.vastart %[[ALLOCA0]]
@@ -87,11 +87,13 @@ define void @variadic_function(i32 %X, ...) {
   %aq = alloca ptr
   ; CHECK: llvm.intr.vacopy %[[ALLOCA0]] to %[[ALLOCA1]]
   call void @llvm.va_copy.p0(ptr %aq, ptr %ap)
+
+  %ret = va_arg ptr %aq, i32
   ; CHECK: llvm.intr.vaend %[[ALLOCA1]]
   call void @llvm.va_end.p0(ptr %aq)
 
   ; CHECK: llvm.intr.vaend %[[ALLOCA0]]
   call void @llvm.va_end.p0(ptr %ap)
   ; CHECK: llvm.return
-  ret void
+  ret i32 %ret
 }