From 3addf3c48af8513798c39c6a10f348a0cfaa81c2 Mon Sep 17 00:00:00 2001
From: Jonas Paulsson <paulson1@linux.ibm.com>
Date: Fri, 15 Nov 2024 01:44:36 +0100
Subject: [PATCH] Fix arg exts for tests

---
 mlir/test/CAPI/execution_engine.c         |  4 ++++
 mlir/test/mlir-cpu-runner/simple.mlir     | 14 +++++++-------
 mlir/unittests/ExecutionEngine/Invoke.cpp | 21 +++++++++++++++++++++
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/mlir/test/CAPI/execution_engine.c b/mlir/test/CAPI/execution_engine.c
index 18120c6ec80280..4751288c3ee4bd 100644
--- a/mlir/test/CAPI/execution_engine.c
+++ b/mlir/test/CAPI/execution_engine.c
@@ -55,7 +55,11 @@ void testSimpleExecution(void) {
       ctx, mlirStringRefCreateFromCString(
                // clang-format off
 "module {                                                                    \n"
+#ifdef __s390__
+"  func.func @add(%arg0 : i32) -> (i32 {llvm.signext}) attributes { llvm.emit_c_interface } {     \n"
+#else
 "  func.func @add(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {     \n"
+#endif
 "    %res = arith.addi %arg0, %arg0 : i32                                        \n"
 "    return %res : i32                                                           \n"
 "  }                                                                             \n"
diff --git a/mlir/test/mlir-cpu-runner/simple.mlir b/mlir/test/mlir-cpu-runner/simple.mlir
index 38d9dcaf553714..dc283e76219f09 100644
--- a/mlir/test/mlir-cpu-runner/simple.mlir
+++ b/mlir/test/mlir-cpu-runner/simple.mlir
@@ -1,15 +1,15 @@
-// RUN: mlir-cpu-runner %s | FileCheck %s
-// RUN: mlir-cpu-runner %s -e foo | FileCheck -check-prefix=NOMAIN %s
-// RUN: mlir-cpu-runner %s --entry-point-result=i32 -e int32_main | FileCheck -check-prefix=INT32MAIN %s
-// RUN: mlir-cpu-runner %s --entry-point-result=i64 -e int64_main | FileCheck -check-prefix=INT64MAIN %s
-// RUN: mlir-cpu-runner %s -O3 | FileCheck %s
+// RUN: mlir-cpu-runner %s -argext-abi-check=false | FileCheck %s
+// RUN: mlir-cpu-runner %s -e foo  -argext-abi-check=false | FileCheck -check-prefix=NOMAIN %s
+// RUN: mlir-cpu-runner %s --entry-point-result=i32 -e int32_main -argext-abi-check=false | FileCheck -check-prefix=INT32MAIN %s
+// RUN: mlir-cpu-runner %s --entry-point-result=i64 -e int64_main -argext-abi-check=false | FileCheck -check-prefix=INT64MAIN %s
+// RUN: mlir-cpu-runner %s -O3  -argext-abi-check=false | FileCheck %s
 
 // RUN: cp %s %t
-// RUN: mlir-cpu-runner %t -dump-object-file | FileCheck %t
+// RUN: mlir-cpu-runner %t -dump-object-file -argext-abi-check=false | FileCheck %t
 // RUN: ls %t.o
 // RUN: rm %t.o
 
-// RUN: mlir-cpu-runner %s -dump-object-file -object-filename=%T/test.o | FileCheck %s
+// RUN: mlir-cpu-runner %s -dump-object-file -object-filename=%T/test.o -argext-abi-check=false | FileCheck %s
 // RUN: ls %T/test.o
 // RUN: rm %T/test.o
 
diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
index ff87fc9fad805a..887db227cfc4b2 100644
--- a/mlir/unittests/ExecutionEngine/Invoke.cpp
+++ b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -61,12 +61,21 @@ static LogicalResult lowerToLLVMDialect(ModuleOp module) {
 }
 
 TEST(MLIRExecutionEngine, SKIP_WITHOUT_JIT(AddInteger)) {
+#ifdef __s390__
+  std::string moduleStr = R"mlir(
+  func.func @foo(%arg0 : i32 {llvm.signext}) -> (i32 {llvm.signext}) attributes { llvm.emit_c_interface } {
+    %res = arith.addi %arg0, %arg0 : i32
+    return %res : i32
+  }
+  )mlir";
+#else
   std::string moduleStr = R"mlir(
   func.func @foo(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {
     %res = arith.addi %arg0, %arg0 : i32
     return %res : i32
   }
   )mlir";
+#endif
   DialectRegistry registry;
   registerAllDialects(registry);
   registerBuiltinDialectTranslation(registry);
@@ -259,6 +268,16 @@ TEST(NativeMemRefJit, MAYBE_JITCallback) {
   for (float &elt : *a)
     elt = count++;
 
+#ifdef __s390__
+  std::string moduleStr = R"mlir(
+  func.func private @callback(%arg0: memref<?x?xf32>, %coefficient: i32 {llvm.signext})  attributes { llvm.emit_c_interface }
+  func.func @caller_for_callback(%arg0: memref<?x?xf32>, %coefficient: i32 {llvm.signext}) attributes { llvm.emit_c_interface } {
+    %unranked = memref.cast %arg0: memref<?x?xf32> to memref<*xf32>
+    call @callback(%arg0, %coefficient) : (memref<?x?xf32>, i32) -> ()
+    return
+  }
+  )mlir";
+#else
   std::string moduleStr = R"mlir(
   func.func private @callback(%arg0: memref<?x?xf32>, %coefficient: i32)  attributes { llvm.emit_c_interface }
   func.func @caller_for_callback(%arg0: memref<?x?xf32>, %coefficient: i32) attributes { llvm.emit_c_interface } {
@@ -267,6 +286,8 @@ TEST(NativeMemRefJit, MAYBE_JITCallback) {
     return
   }
   )mlir";
+#endif
+
   DialectRegistry registry;
   registerAllDialects(registry);
   registerBuiltinDialectTranslation(registry);