147 lines · plain
1// RUN: mlir-opt %s \2// RUN: | mlir-opt -pass-pipeline='builtin.module(cse,func.func(gpu-async-region),xevm-attach-target,gpu.module(convert-gpu-to-llvm-spv{use-64bit-index=true},convert-xevm-to-llvm,cse))' \3// RUN: | mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-vector-to-llvm -convert-arith-to-llvm \4// RUN: | mlir-opt -gpu-to-llvm -reconcile-unrealized-casts -cse -gpu-module-to-binary \5// RUN: | mlir-runner \6// RUN: --shared-libs=%mlir_levelzero_runtime \7// RUN: --shared-libs=%mlir_runner_utils \8// RUN: --shared-libs=%mlir_c_runner_utils \9// RUN: --entry-point-result=void \10// RUN: | FileCheck %s11 12module @gemm attributes {gpu.container_module} {13 gpu.module @kernel {14 // - Sets of `matrix_mad` intrinsics can differ based on device's *minimal* supported sub-group size.15 // The *minimum supported* sub-group size should be used to call `matrix_mad` intrinsics.16 // https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_matrix_multiply_accumulate.html17 18 gpu.func @block_dpas(%a: !llvm.ptr<1>, %b: !llvm.ptr<1>, %c: !llvm.ptr<1>) kernel {19 %base_width_a = arith.constant 32 : i3220 %base_height_a = arith.constant 8 : i3221 %base_pitch_a = arith.constant 32 : i3222 %x = arith.constant 0 : i3223 %y = arith.constant 0 : i3224 %loaded_a = xevm.blockload2d %a, %base_width_a, %base_height_a, %base_pitch_a, %x, %y25 <{elem_size_in_bits=16 : i32, tile_width=16 : i32, tile_height=8 : i32, v_blocks=1 : i32,26 transpose=false, pack_register=false}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi16>27 28 %base_width_b = arith.constant 32 : i3229 %base_height_b = arith.constant 16 : i3230 %base_pitch_b = arith.constant 32 : i3231 %loaded_b1 = xevm.blockload2d %b, %base_width_b, %base_height_b, %base_pitch_b, %x, %y32 <{elem_size_in_bits=16 : i32, tile_width=16 : i32, tile_height=16 : i32, v_blocks=1 : i32,33 transpose=false, pack_register=false}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<16xi16>34 %loaded_b_casted = vector.bitcast %loaded_b1 : vector<16xi16> to vector<8xi32>35 36 %base_width_c = arith.constant 64 : i3237 %base_height_c = arith.constant 8 : i3238 %base_pitch_c = arith.constant 64 : i3239 %loaded_c = xevm.blockload2d %c, %base_width_c, %base_height_c, %base_pitch_c, %x, %y40 <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32, v_blocks=1 : i32,41 transpose=false, pack_register=false}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>42 43 %loaded_c_casted = vector.bitcast %loaded_c : vector<8xi32> to vector<8xf32>44 %c_result = xevm.mma %loaded_a, %loaded_b_casted, %loaded_c_casted45 {shape=<m=8, n=16, k=16>, types=<d=f32, a=f16, b=f16, c=f32>}46 : (vector<8xi16>, vector<8xi32>, vector<8xf32>) -> vector<8xf32>47 %c_result_casted = vector.bitcast %c_result : vector<8xf32> to vector<8xi32>48 49 xevm.blockstore2d %c, %base_width_c, %base_height_c, %base_pitch_c, %x, %y, %c_result_casted50 <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32}>51 : (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)52 gpu.return53 }54 }55 56 func.func @test(%a : memref<8x16xf16>, %b : memref<16x16xf16>, %c : memref<8x16xf32>) -> memref<8x16xf32> attributes {llvm.emit_c_interface} {57 %c1 = arith.constant 1 : index58 %c16 = arith.constant 16 : index59 60 %memref_a = gpu.alloc() : memref<8x16xf16>61 gpu.memcpy %memref_a, %a : memref<8x16xf16>, memref<8x16xf16>62 %a_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_a : memref<8x16xf16> -> index63 %a_ptr_as_i64 = arith.index_cast %a_ptr_as_idx : index to i6464 %a_ptr = llvm.inttoptr %a_ptr_as_i64 : i64 to !llvm.ptr65 %a_ptr_casted = llvm.addrspacecast %a_ptr : !llvm.ptr to !llvm.ptr<1>66 67 %memref_b = gpu.alloc() : memref<16x16xf16>68 gpu.memcpy %memref_b, %b : memref<16x16xf16>, memref<16x16xf16>69 %b_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_b : memref<16x16xf16> -> index70 %b_ptr_as_i64 = arith.index_cast %b_ptr_as_idx : index to i6471 %b_ptr = llvm.inttoptr %b_ptr_as_i64 : i64 to !llvm.ptr72 %b_ptr_casted = llvm.addrspacecast %b_ptr : !llvm.ptr to !llvm.ptr<1>73 74 %memref_c = gpu.alloc() : memref<8x16xf32>75 gpu.memcpy %memref_c, %c : memref<8x16xf32>, memref<8x16xf32>76 %c_ptr_as_idx = memref.extract_aligned_pointer_as_index %memref_c : memref<8x16xf32> -> index77 %c_ptr_as_i64 = arith.index_cast %c_ptr_as_idx : index to i6478 %c_ptr = llvm.inttoptr %c_ptr_as_i64 : i64 to !llvm.ptr79 %c_ptr_casted = llvm.addrspacecast %c_ptr : !llvm.ptr to !llvm.ptr<1>80 81 gpu.launch_func @kernel::@block_dpas blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)82 args(%a_ptr_casted : !llvm.ptr<1>, %b_ptr_casted : !llvm.ptr<1>, %c_ptr_casted : !llvm.ptr<1>)83 gpu.dealloc %memref_a : memref<8x16xf16>84 gpu.dealloc %memref_b : memref<16x16xf16>85 %res = memref.alloc() : memref<8x16xf32>86 gpu.memcpy %res, %memref_c : memref<8x16xf32>, memref<8x16xf32>87 gpu.dealloc %memref_c : memref<8x16xf32>88 return %res : memref<8x16xf32>89 }90 91 func.func @main() attributes {llvm.emit_c_interface} {92 %A = memref.alloc() : memref<8x16xf16>93 %c0 = arith.constant 0 : index94 %c1 = arith.constant 1 : index95 %c8 = arith.constant 8 : index96 %c16 = arith.constant 16 : index97 98 scf.for %i = %c0 to %c8 step %c1 {99 scf.for %j = %c0 to %c16 step %c1 {100 %row_idx = arith.index_cast %i : index to i32101 %row = arith.sitofp %row_idx : i32 to f16102 memref.store %row, %A[%i, %j] : memref<8x16xf16>103 }104 }105 %B = memref.alloc() : memref<16x16xf16>106 scf.for %i = %c0 to %c16 step %c1 {107 scf.for %j = %c0 to %c16 step %c1 {108 %col_idx = arith.index_cast %j : index to i32109 %col = arith.sitofp %col_idx : i32 to f16110 memref.store %col, %B[%i, %j] : memref<16x16xf16>111 }112 }113 114 %C = memref.alloc() : memref<8x16xf32>115 %c0_f16 = arith.constant 0.0 : f32116 scf.for %i = %c0 to %c8 step %c1 {117 scf.for %j = %c0 to %c16 step %c1 {118 memref.store %c0_f16, %C[%i, %j] : memref<8x16xf32>119 }120 }121 122 %C_res = call @test(%A, %B, %C) : (memref<8x16xf16>, memref<16x16xf16>, memref<8x16xf32>) -> memref<8x16xf32>123 %C_cast = memref.cast %C_res : memref<8x16xf32> to memref<*xf32>124 %A_cast = memref.cast %A : memref<8x16xf16> to memref<*xf16>125 call @printMemrefF32(%C_cast) : (memref<*xf32>) -> ()126 127 // CHECK: Unranked Memref base@ = 0x{{[0-9a-f]+}}128 // CHECK-NEXT: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]129 // CHECK-NEXT: [0, 16, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240]130 // CHECK-NEXT: [0, 32, 64, 96, 128, 160, 192, 224, 256, 288, 320, 352, 384, 416, 448, 480]131 // CHECK-NEXT: [0, 48, 96, 144, 192, 240, 288, 336, 384, 432, 480, 528, 576, 624, 672, 720]132 // CHECK-NEXT: [0, 64, 128, 192, 256, 320, 384, 448, 512, 576, 640, 704, 768, 832, 896, 960]133 // CHECK-NEXT: [0, 80, 160, 240, 320, 400, 480, 560, 640, 720, 800, 880, 960, 1040, 1120, 1200]134 // CHECK-NEXT: [0, 96, 192, 288, 384, 480, 576, 672, 768, 864, 960, 1056, 1152, 1248, 1344, 1440]135 // CHECK-NEXT: [0, 112, 224, 336, 448, 560, 672, 784, 896, 1008, 1120, 1232, 1344, 1456, 1568, 1680]136 137 memref.dealloc %A : memref<8x16xf16>138 memref.dealloc %B : memref<16x16xf16>139 memref.dealloc %C : memref<8x16xf32>140 memref.dealloc %C_res : memref<8x16xf32>141 return142 }143 func.func private @printMemrefF16(%ptr : memref<*xf16>) attributes { llvm.emit_c_interface }144 func.func private @printMemrefF32(%ptr : memref<*xf32>) attributes { llvm.emit_c_interface }145 146}147