[mlir] [LLVM IR] Introduce VaArgOp

[ChuanqiXu9]

I find there is no LLVMOp corresponding to LLVM's va_arg instruction so I tried to add one. This is helpful for clangir (llvm/clangir#865).

New to MLIR and not sure who are the appropriate reviewers. Appreciated in ahead for reviewing and triaging.

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Chuanqi Xu (ChuanqiXu9)

Changes

I find there is no LLVMOp corresponding to LLVM's va_arg instruction so I tried to add one. This is helpful for clangir (llvm/clangir#865).

New to MLIR and not sure who are the appropriate reviewers. Appreciated in ahead for reviewing and triaging.

---

Full diff: https://github.com/llvm/llvm-project/pull/109260.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td (+20)
• (modified) mlir/test/Target/LLVMIR/llvmir.mlir (+6-4)

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
 }
 
 // -----

[zero9178]

Looks great, thank you so much (and welcome to MLIR 😄)!

The only nit I have is whether you could also add a test for the import functionality from LLVM IR to the LLVM Dialect (e.g. here:

llvm-project/mlir/test/Target/LLVMIR/Import/basic.ll

Line 80 in 912e821

define void @variadic_function(i32 %X, ...) {

).

[zero9178]

Suggested change

	let assemblyFormat = "$arg attr-dict `:` `(` type($arg) `)` `->` type($res)";
	let assemblyFormat = "$arg attr-dict `:` functional-type($arg, $res)";

[ChuanqiXu9]

Done

[gysit]

Thanks!

In addition to the import test, it could also make sense to add it to the roundtrip tests:

llvm-project/mlir/test/Dialect/LLVMIR/roundtrip.mlir

Line 634 in 04ccbe6

llvm.func @vararg_func(%arg0: i32, ...) {

[ChuanqiXu9]

Done by adding test in [mlir/test/Dialect/LLVMIR/roundtrip.mlir](https://github.com/llvm/llvm-project/blob/04ccbe6e70cf11e846da3fbc800832c6e56b573f/mlir/test/Dialect/LLVMIR/roundtrip.mlir#L634) and [mlir/test/Target/LLVMIR/Import/basic.ll](https://github.com/llvm/llvm-project/blob/912e821ab3d415dc8421f53738ebabcc7f1dac2d/mlir/test/Target/LLVMIR/Import/basic.ll#L80)