83 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-scf -lower-affine -convert-scf-to-cf -convert-vector-to-llvm="enable-amx" -finalize-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \2// RUN: mlir-translate -mlir-to-llvmir | \3// RUN: %lli --entry-function=entry --mattr="+amx-tile,+amx-int8,+amx-bf16" --dlopen=%mlir_c_runner_utils | \4// RUN: FileCheck %s5 6// Note: To run this test, your CPU must support AMX.7 8func.func @print(%arg0: memref<4x32xf32>) {9 %fu = arith.constant -1.0: f3210 %c0 = arith.constant 0: index11 %c1 = arith.constant 1: index12 %c4 = arith.constant 4: index13 scf.for %i = %c0 to %c4 step %c1 {14 %0 = vector.transfer_read %arg0[%i, %c0], %fu: memref<4x32xf32>, vector<32xf32>15 vector.print %0 : vector<32xf32>16 }17 return18}19 20func.func @kernel(%arg0: memref<4x32xf32>) {21 %c0 = arith.constant 0: index22 %c2 = arith.constant 2 : index23 %c4 = arith.constant 4 : index24 %c16 = arith.constant 16 : index25 %c32 = arith.constant 32 : index26 scf.for %i = %c0 to %c4 step %c2 {27 scf.for %j = %c0 to %c32 step %c16 {28 %0 = amx.tile_zero : vector<2x16xf32>29 amx.tile_store %arg0[%i, %j], %0 : memref<4x32xf32>, vector<2x16xf32>30 func.call @print(%arg0) : (memref<4x32xf32>) -> ()31 }32 }33 return34}35 36func.func @entry() -> i32 {37 %f1 = arith.constant 1.0: f3238 %c0 = arith.constant 0: index39 %c1 = arith.constant 1: index40 %c4 = arith.constant 4 : index41 %c32 = arith.constant 32 : index42 43 // Set up memory.44 %a = memref.alloc() : memref<4x32xf32>45 scf.for %i = %c0 to %c4 step %c1 {46 scf.for %j = %c0 to %c32 step %c1 {47 memref.store %f1, %a[%i, %j] : memref<4x32xf32>48 }49 }50 51 // Call kernel.52 func.call @kernel(%a) : (memref<4x32xf32>) -> ()53 54 // Verify progress of blocked tilezero.55 //56 // CHECK: ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )57 // CHECK-NEXT: ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )58 // CHECK-NEXT: ( 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )59 // CHECK-NEXT: ( 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )60 //61 // 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 )62 // 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 )63 // CHECK-NEXT: ( 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )64 // CHECK-NEXT: ( 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )65 //66 // 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 )67 // 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 )68 // CHECK-NEXT: ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )69 // CHECK-NEXT: ( 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1 )70 //71 // 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 )72 // 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 )73 // 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 )74 // 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 )75 //76 77 // Release resources.78 memref.dealloc %a : memref<4x32xf32>79 80 %i0 = arith.constant 0 : i3281 return %i0 : i3282}83