[mlir] Fix return values of AMX test

[aartbik]

Thanks for the change!

Do you need me to commit this or are you fully setup for that?

mlir/test/Integration/Dialect/Vector/CPU/AMX/test-mulf.mlir

func @kernel2(%arg0: memref<2x4xbf16>,
   %1 = amx.tile_load %arg0[%0, %0] : memref<2x4xbf16>  into vector<2x4xbf16>
   %2 = amx.tile_load %arg1[%0, %0] : memref<2x4xbf16>  into vector<2x4xbf16>
   %3 = amx.tile_load %arg2[%0, %0] : memref<2x2xf32> into vector<2x2xf32>
   %4 = amx.tile_mulf %1, %2, %3 : vector<2x4xbf16>, vector<2x4xbf16>, vector<2x2xf32>
   amx.tile_store %arg2[%0, %0], %4 : memref<2x2xf32>, vector<2x2xf32>
   return
 }
 
-func @entry() {
+func @entry() -> i32 {
   %f0 = constant 0.0: f32
   %c0 = constant 0: index
   %c1 = constant 1: index
   %c2 = constant 2: index
 
   // Set up memory.
   %a = memref.alloc() : memref<2x4xbf16>
   %b = memref.alloc() : memref<2x4xbf16>
scf.for %i = %c0 to %c2 step %c1 {
     vector.print %cv : vector<2xf32>
   }
 
   // Release resources.
   memref.dealloc %a : memref<2x4xbf16>
   memref.dealloc %b : memref<2x4xbf16>
   memref.dealloc %c : memref<2x2xf32>
 
-  return
+  %i0 = constant 0 : i32
+  return %i0 : i32
 }

mlir/test/Integration/Dialect/Vector/CPU/AMX/test-muli-ext.mlir

func @kernel4(%arg0: memref<16x16xi8>,
   %1 = amx.tile_load %arg0[%0, %0] : memref<16x16xi8>  into vector<16x16xi8>
   %2 = amx.tile_load %arg1[%0, %0] : memref<4x16xi8>  into vector<4x16xi8>
   %3 = amx.tile_zero : vector<16x4xi32>
   %4 = amx.tile_muli %1 zext, %2 zext, %3 : vector<16x16xi8>, vector<4x16xi8>, vector<16x4xi32>
   amx.tile_store %arg2[%0, %0], %4 : memref<16x4xi32>, vector<16x4xi32>
   return
 }
 
-func @entry() {
+func @entry() -> i32 {
   %c0 = constant 0: index
 
   // Set up memory.
   %a = memref.alloc() : memref<16x16xi8>
   %b = memref.alloc() : memref<4x16xi8>
   %c = memref.alloc() : memref<16x4xi32>
 
   %0 = std.constant dense<
func @entry() -> i32 {
   call @kernel4(%a, %b, %c) : (memref<16x16xi8>, memref<4x16xi8>, memref<16x4xi32>) -> ()
   call @print(%c) : (memref<16x4xi32>) -> ()
 
   // Release resources.
   memref.dealloc %a : memref<16x16xi8>
   memref.dealloc %b : memref<4x16xi8>
   memref.dealloc %c : memref<16x4xi32>
 
-  return
+  %i0 = constant 0 : i32
+  return %i0 : i32
 }

mlir/test/Integration/Dialect/Vector/CPU/AMX/test-muli.mlir

func @kernel2(%arg0: memref<2x8xi8>,
   %1 = amx.tile_load %arg0[%0, %0] : memref<2x8xi8>  into vector<2x8xi8>
   %2 = amx.tile_load %arg1[%0, %0] : memref<2x8xi8>  into vector<2x8xi8>
   %3 = amx.tile_load %arg2[%0, %0] : memref<2x2xi32> into vector<2x2xi32>
   %4 = amx.tile_muli %1 zext, %2 zext, %3 : vector<2x8xi8>, vector<2x8xi8>, vector<2x2xi32>
   amx.tile_store %arg2[%0, %0], %4 : memref<2x2xi32>, vector<2x2xi32>
   return
 }
 
-func @entry() {
+func @entry() -> i32 {
   %i0 = constant 0: i32
   %c0 = constant 0: index
   %c1 = constant 1: index
   %c2 = constant 2: index
 
   // Set up memory.
   %a = memref.alloc() : memref<2x8xi8>
   %b = memref.alloc() : memref<2x8xi8>
scf.for %i = %c0 to %c2 step %c1 {
     vector.print %cv : vector<2xi32>
   }
 
   // Release resources.
   memref.dealloc %a : memref<2x8xi8>
   memref.dealloc %b : memref<2x8xi8>
   memref.dealloc %c : memref<2x2xi32>
 
-  return
+  return %i0 : i32
 }

mlir/test/Integration/Dialect/Vector/CPU/AMX/test-tilezero-block.mlir

scf.for %j = %c0 to %c32 step %c16 {
       %0 = amx.tile_zero : vector<2x16xf32>
       amx.tile_store %arg0[%i, %j], %0 : memref<4x32xf32>, vector<2x16xf32>
       call @print(%arg0) : (memref<4x32xf32>) -> ()
     }
   }
   return
 }
 
-func @entry() {
+func @entry() -> i32 {
   %f1  = constant 1.0: f32
   %c0  = constant 0: index
   %c1  = constant 1: index
   %c4  = constant 4 : index
   %c32 = constant 32 : index
 
   // Set up memory.
   %a = memref.alloc() : memref<4x32xf32>
func @entry() -> i32 {
   // CHECK-NEXT: ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
   // CHECK-NEXT: ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
   // CHECK-NEXT: ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
   //
 
   // Release resources.
   memref.dealloc %a : memref<4x32xf32>
 
-  return
+  %i0 = constant 0 : i32
+  return %i0 : i32
 }

mlir/test/Integration/Dialect/Vector/CPU/AMX/test-tilezero.mlir

 // RUN: mlir-opt %s -convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm="enable-amx" -convert-std-to-llvm | \
 // RUN: mlir-translate -mlir-to-llvmir | \
 // RUN: %lli --jit-kind=mcjit --entry-function=entry --mattr="+amx-tile,+amx-int8,+amx-bf16" --dlopen=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 // TODO: drop lli's --jit-kind flag once PR#49906 (https://bugs.llvm.org/show_bug.cgi?id=49906) is fixed.
 
 // Note: To run this test, your CPU must support AMX.
 
 func @tilezero(%arg0: memref<?x?xi32>, %i: index, %j: index) {
   %1 = amx.tile_zero : vector<16x16xi32>
   amx.tile_store %arg0[%i, %j], %1 : memref<?x?xi32>, vector<16x16xi32>
   return
 }
 
-func @entry() {
+func @entry() -> i32 {
   %i0 = constant 0: i32
   %i1 = constant 1: i32
   %c0 = constant 0: index
   %c1 = constant 1: index
   %c3 = constant 3: index
   %c19 = constant 19: index
 
   // Set up memory.
func @entry() -> i32 {
   scf.for %i = %c0 to %c19 step %c1 {
     %av = vector.transfer_read %a[%i, %c0], %i0: memref<?x?xi32>, vector<19xi32>
     vector.print %av : vector<19xi32>
   }
 
   // Release resources.
   memref.dealloc %a : memref<?x?xi32>
 
-  return
+  return %i0 : i32
 }