[MLIR] Add SystemZ arg extensions for some tests (#116314)

The SystemZ ABI requires that i32 values should be extended when passed
between functions.

This patch fixes some tests that were lacking this, either by adding
some SystemZ specific inlinings of test functions or by disabling the
verification of this with the CL option controlling this.

Fixes #115564

Author: JonPsson1

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"

mlir/test/mlir-cpu-runner/simple.mlir

@@ -1,15 +1,22 @@
-// 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 %if target={{s390x-.*}} %{ -argext-abi-check=false %} \
+// RUN:   | FileCheck %s
+// RUN: mlir-cpu-runner %s -e foo %if target={{s390x-.*}} %{ -argext-abi-check=false %} \
+// RUN:   | FileCheck -check-prefix=NOMAIN %s
+// RUN: mlir-cpu-runner %s --entry-point-result=i32 -e int32_main %if target={{s390x-.*}} \
+// RUN:   %{ -argext-abi-check=false %} | FileCheck -check-prefix=INT32MAIN %s
+// RUN: mlir-cpu-runner %s --entry-point-result=i64 -e int64_main %if target={{s390x-.*}} \
+// RUN:   %{ -argext-abi-check=false %} | FileCheck -check-prefix=INT64MAIN %s
+// RUN: mlir-cpu-runner %s -O3 %if target={{s390x-.*}} %{ -argext-abi-check=false %} \
+// RUN:   | FileCheck %s
 
 // RUN: cp %s %t
-// RUN: mlir-cpu-runner %t -dump-object-file | FileCheck %t
+// RUN: mlir-cpu-runner %t -dump-object-file %if target={{s390x-.*}} \
+// RUN:   %{ -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 \
+// RUN:   %if target={{s390x-.*}} %{ -argext-abi-check=false %} | FileCheck %s
 // RUN: ls %T/test.o
 // RUN: rm %T/test.o
 

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);