[mlir][ArmSME][test] Cleanup printing scalable vectors in vector-load-store.mlir (NFC)

[benmxwl-arm]

Summary

This replaces the manual print loops with vector.print which now supports
scalable vectors.

[awarzynski]

Beautiful refactor, thank you Ben! LGTM :)

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector-load-store.mlir

 
 // RUN: %{compile} | %{run} | FileCheck %s --check-prefix=CHECK-ZA0_D
 
 // REDEFINE: %{entry_point} = load_store_two_za_s_tiles
 // RUN: %{compile} | %{run} | FileCheck %s
 
 // Integration tests demonstrating load/store to/from SME ZA tile.
 
-llvm.func @printF64(f64)
-llvm.func @printI64(i64)
-llvm.func @printOpen()
-llvm.func @printClose()
-llvm.func @printComma()
-llvm.func @printNewline()
 llvm.func @printCString(!llvm.ptr<i8>)
 
 // This test verifies a 64-bit element ZA with FP64 data is correctly
 // loaded/stored to/from memory.
 func.func @za0_d_f64() -> i32 {
   %c0 = arith.constant 0 : index
   %c0_f64 = arith.constant 0.0 : f64
   %c1_f64 = arith.constant 1.0 : f64
func.func @za0_d_f64() -> i32 {
 
   // Dump "mem1". The smallest SVL is 128-bits so the tile will be at least
   // 2x2xi64.
   //
   // CHECK-ZA0_D:      ( 0.1, 0.1
   // CHECK-ZA0_D-NEXT: ( 1.1, 1.1
   scf.for %i = %c0 to %tilesize step %svl_d {
     %tileslice = vector.load %mem1[%i] : memref<?xf64>, vector<[2]xf64>
-
-    llvm.call @printOpen() : () -> ()
-    scf.for %i2 = %c0 to %svl_d step %c1_index {
-      %elem = vector.extractelement %tileslice[%i2 : index] : vector<[2]xf64>
-      llvm.call @printF64(%elem) : (f64) -> ()
-      %last_i = arith.subi %svl_d, %c1_index : index
-      %isNotLastIter = arith.cmpi ult, %i2, %last_i : index
-      scf.if %isNotLastIter {
-        llvm.call @printComma() : () -> ()
-      }
-    }
-    llvm.call @printClose() : () -> ()
-    llvm.call @printNewline() : () -> ()
+    vector.print %tileslice : vector<[2]xf64>
   }
 
   // Load ZA0.D from "mem1"
   %za0_d = vector.load %mem1[%c0] : memref<?xf64>, vector<[2]x[2]xf64>
 
   // Allocate "mem2" to store ZA0.D to
   %mem2 = memref.alloca(%tilesize) : memref<?xf64>
 
func.func @za0_d_f64() -> i32 {
 
   // Dump zeroed "mem2". The smallest SVL is 128-bits so the tile will be at
   // least 2x2xi64.
   //
   // CHECK-ZA0_D-NEXT: ( 0, 0
   // CHECK-ZA0_D-NEXT: ( 0, 0
   scf.for %i = %c0 to %tilesize step %svl_d {
     %tileslice = vector.load %mem2[%i] : memref<?xf64>, vector<[2]xf64>
-
-    llvm.call @printOpen() : () -> ()
-    scf.for %i2 = %c0 to %svl_d step %c1_index {
-      %elem = vector.extractelement %tileslice[%i2 : index] : vector<[2]xf64>
-      llvm.call @printF64(%elem) : (f64) -> ()
-      %last_i = arith.subi %svl_d, %c1_index : index
-      %isNotLastIter = arith.cmpi ult, %i2, %last_i : index
-      scf.if %isNotLastIter {
-        llvm.call @printComma() : () -> ()
-      }
-    }
-    llvm.call @printClose() : () -> ()
-    llvm.call @printNewline() : () -> ()
+    vector.print %tileslice : vector<[2]xf64>
   }
 
   // Verify "mem1" != "mem2"
   %init_1 = arith.constant 1 : i64
   %mul_reduce_0 = scf.for %vnum = %c0 to %tilesize step %svl_d iter_args(%iter = %init_1) -> (i64) {
     %row_1 = vector.load %mem1[%vnum] : memref<?xf64>, vector<[2]xf64>
     %row_2 = vector.load %mem2[%vnum] : memref<?xf64>, vector<[2]xf64>
     %cmp = arith.cmpf one, %row_1, %row_2 : vector<[2]xf64>
func.func @za0_d_f64() -> i32 {
 
   // Dump "mem2". The smallest SVL is 128-bits so the tile will be at least
   // 2x2xi64.
   //
   // CHECK-ZA0_D-NEXT: ( 0.1, 0.1
   // CHECK-ZA0_D-NEXT: ( 1.1, 1.1
   scf.for %i = %c0 to %tilesize step %svl_d {
     %tileslice = vector.load %mem2[%i] : memref<?xf64>, vector<[2]xf64>
-
-    llvm.call @printOpen() : () -> ()
-    scf.for %i2 = %c0 to %svl_d step %c1_index {
-      %elem = vector.extractelement %tileslice[%i2 : index] : vector<[2]xf64>
-      llvm.call @printF64(%elem) : (f64) -> ()
-      %last_i = arith.subi %svl_d, %c1_index : index
-      %isNotLastIter = arith.cmpi ult, %i2, %last_i : index
-      scf.if %isNotLastIter {
-        llvm.call @printComma() : () -> ()
-      }
-    }
-    llvm.call @printClose() : () -> ()
-    llvm.call @printNewline() : () -> ()
+    vector.print %tileslice : vector<[2]xf64>
   }
 
   %c0_i32 = arith.constant 0 : i32
   return %c0_i32 : i32
 }
 
 func.func @printTileBegin() {
   %0 = llvm.mlir.addressof @str_tile_begin : !llvm.ptr<array<11 x i8>>
func.func @load_store_two_za_s_tiles() -> i32 {
   // CHECK-NEXT: ( 1, 1, 1, 1
   // CHECK-NEXT: ( 1, 1, 1, 1
   // CHECK:      ( 2, 2, 2, 2
   // CHECK-NEXT: ( 2, 2, 2, 2
   // CHECK-NEXT: ( 2, 2, 2, 2
   // CHECK-NEXT: ( 2, 2, 2, 2
   scf.for %i = %c0 to %size_of_two_tiles step %svl_s {
     %tileslice = vector.load %mem1[%i] : memref<?xi32>, vector<[4]xi32>
-
-    llvm.call @printOpen() : () -> ()
-    scf.for %i2 = %c0 to %svl_s step %c1_index {
-      %elem = vector.extractelement %tileslice[%i2 : index] : vector<[4]xi32>
-      %elem_i64 = llvm.zext %elem : i32 to i64
-      llvm.call @printI64(%elem_i64) : (i64) -> ()
-      %last_i = arith.subi %svl_s, %c1_index : index
-      %isNotLastIter = arith.cmpi ult, %i2, %last_i : index
-      scf.if %isNotLastIter {
-        llvm.call @printComma() : () -> ()
-      }
-    }
-    llvm.call @printClose() : () -> ()
-    llvm.call @printNewline() : () -> ()
+    vector.print %tileslice : vector<[4]xi32>
   }
 
   // Load tile 1 from memory
   %za0_s = vector.load %mem1[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
 
   // Load tile 2 from memory
   %za1_s = vector.load %mem1[%size_of_tile] : memref<?xi32>, vector<[4]x[4]xi32>
 
func.func @load_store_two_za_s_tiles() -> i32 {
   // CHECK-NEXT: ( 1, 1, 1, 1
   // CHECK-NEXT: ( 1, 1, 1, 1
   // CHECK-NEXT: ( 1, 1, 1, 1
   // CHECK-NEXT: ( 1, 1, 1, 1
   // CHECK:      TILE END
   func.call @printTileBegin() : () -> ()
   scf.for %i = %c0 to %size_of_two_tiles step %svl_s {
     %av = vector.load %mem2[%i] : memref<?xi32>, vector<[4]xi32>
-
-    llvm.call @printOpen() : () -> ()
-    scf.for %i2 = %c0 to %svl_s step %c1_index {
-      %elem = vector.extractelement %av[%i2 : index] : vector<[4]xi32>
-      %elem_i64 = llvm.zext %elem : i32 to i64
-      llvm.call @printI64(%elem_i64) : (i64) -> ()
-      %last_i = arith.subi %svl_s, %c1_index : index
-      %isNotLastIter = arith.cmpi ult, %i2, %last_i : index
-      scf.if %isNotLastIter {
-        llvm.call @printComma() : () -> ()
-      }
-    }
-    llvm.call @printClose() : () -> ()
-    llvm.call @printNewline() : () -> ()
+    vector.print %av : vector<[4]xi32>
 
     %tileSizeMinusStep = arith.subi %size_of_tile, %svl_s : index
     %isNextTile = arith.cmpi eq, %i, %tileSizeMinusStep : index
     scf.if %isNextTile {
       func.call @printTileEnd() : () -> ()
       func.call @printTileBegin() : () -> ()
     }
   }
   func.call @printTileEnd() : () -> ()
 
   return %c0_i32 : i32
 }
 
 llvm.mlir.global internal constant @str_tile_begin("TILE BEGIN\0A")
 llvm.mlir.global internal constant @str_tile_end("TILE END\0A")