404 lines · plain
1//2// NOTE: this test requires gpu-sm803//4// RUN: mlir-opt \5// RUN: --pass-pipeline="builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,convert-nvgpu-to-nvvm,affine-expand-index-ops,lower-affine,arith-expand,convert-arith-to-llvm),convert-vector-to-llvm,canonicalize,cse)" \6// RUN: %s \7// RUN: | mlir-opt --gpu-lower-to-nvvm-pipeline="cubin-chip=sm_80 cubin-features=+ptx71 cubin-format=%gpu_compilation_format" \8// RUN: | mlir-runner \9// RUN: --shared-libs=%mlir_cuda_runtime \10// RUN: --shared-libs=%mlir_c_runner_utils \11// RUN: --e main --entry-point-result=void \12// RUN: | FileCheck %s13 14module attributes {gpu.container_module} {15 16 // Kernels that run on the device.17 18 gpu.module @kernels {19 20 //21 // An NVidia GPU kernel to compute22 // C = A x B23 // (or, technically, D = A x B + C)24 // using 2:4 structured sparsity for A.25 //26 // This kernel provides building block for sparse compilation of a larger27 // enveloping matrix multiplication computation on a GPU.28 //29 // Operand A values (2:4 sparse): row major format, logically "16x32xf16"30 // but "16x16xf16" after compression31 //32 // Operand A metadata:33 // - The metadata is logically "16x16xi2". Each 2-bit value indicates34 // the position of a non-zero value within the respective group of 4 elements.35 // - However, we represent it as "16x2xi16".36 // - Each sparse instruction type specifies how the metadata should be distributed37 // among threads. In this case, within each quad (group of 4 consecutive threads38 // starting with a thread ID which is a multiple of 4), thread 4i and 4i+139 // will require to hold the metadata different metadata. For uniformity below,40 // we just have all threads load metadata, and the way they determine which metadata41 // to load is given below.42 // - Thread map for the 16x32x16 instruction is:43 // 2i -> col 044 // 2i + 1 -> col 145 //46 // Operand B (dense): column major format.47 //48 // Operand C (accum): assumed zero on entry, used as output.49 //50 gpu.func @mma_sp_sync_f16_16832(51 %argA: memref<16x16xf16>,52 %argA_meta: memref<16x2xi16>,53 %argB: memref<8x32xf16>,54 %argC: memref<16x8xf16>) kernel {55 %f0 = arith.constant 0.0 : f1656 %c4 = arith.constant 4 : index57 %c8 = arith.constant 8 : index58 59 // Assume we have a linear thread id and the kernel launches 32 threads (1 warp).60 // So CUDA launch would be threadblock = (32, 1, 1), grid = (1, 1, 1)61 %lane_id = gpu.thread_id x62 // Which group of 4 threads do we belong to?63 %quad_id = affine.apply affine_map<()[s0]->(s0 floordiv 4)>()[%lane_id]64 // Are we even group or odd group?65 %pair_id = affine.apply affine_map<()[s0]->(s0 mod 2)>()[%lane_id]66 67 // Now we have68 // MMA lane=0 quad=0 pair=069 // MMA lane=1 quad=0 pair=170 // MMA lane=2 quad=0 pair=071 // MMA lane=3 quad=0 pair=172 // MMA lane=4 quad=1 pair=073 // MMA lane=5 quad=1 pair=174 // ...75 // MMA lane=30 quad=7 pair=276 // MMA lane=31 quad=7 pair=177 //78 // gpu.printf "MMA lane=%lld quad=%lld pair=%lld\n" %lane_id, %quad_id, %pair_id : index, index, index79 80 //===----------------------------------------------------------------------===//81 // Load the operandA metadata82 // (https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#sparse-mma-metadata-16832-f16bf16)83 //===----------------------------------------------------------------------===//84 85 // For the 16x2xi16 metadata, all threads that load metadata will load one86 // i16 value from the first 8 rows, and one i16 value from the second 8 rows.87 // The i16 values are then put into a i32 with the value from the first 8 rows88 // going in the lower bits.89 //90 // The below IR loads and combines the two pieces of i16 metadata required.91 // Obviously, it's possible to re-pack the metadata before launching the kernel in92 // order to eliminate this cost and load a single i32 operand. This just shows93 // how to put them together if you do the naive load per the diagram in94 // the PTX docs. Technically only the first two threads in each quad need95 // to do this, but for simplicity we just have all threads participate since96 // it can't hurt.97 //98 // The mapping is99 // Lower i16 bits <- (thread_id) -> load A_meta[ quad_id , pair_id]100 // Lower i16 bits <- (thread_id) -> load A_meta[ quad_id + 8, pair_id]101 102 %quad_id_plus_8 = affine.apply affine_map<()[s0]->(s0 + 8)>()[%quad_id]103 %meta_A_per_thread0 = memref.load %argA_meta[%quad_id , %pair_id] : memref<16x2xi16>104 %meta_A_per_thread1 = memref.load %argA_meta[%quad_id_plus_8, %pair_id] : memref<16x2xi16>105 106 %low_i32 = arith.extui %meta_A_per_thread0 : i16 to i32107 %high_i32 = arith.extui %meta_A_per_thread1 : i16 to i32108 109 %meta_init = arith.constant dense<0> : vector<2xi16>110 %meta_low = vector.insert %meta_A_per_thread0, %meta_init[0] : i16 into vector<2xi16>111 %meta = vector.insert %meta_A_per_thread1, %meta_low[1] : i16 into vector<2xi16>112 113 //===----------------------------------------------------------------------===//114 // Load operandA115 //===----------------------------------------------------------------------===//116 117 // Load the actual fragments for the sparse values. This can be done using ldmatrix,118 // but here we just do naive individual loads, which would also be required for119 // a layout/element type that is not compatible with ldmatrix (e.g. i8 transpose load).120 //121 // The thread map here is different that operandA metadata. Each thread will122 // load one 2xf16 vector from each of the four (8x8xf16) quadrants fo the 16x16xf16123 // operand.124 //125 // The (thread_id)->(row, col) map within each 8x4x(2xf16) quadrant is (t)->(t/4, t%4). We126 // can use "affine.delinearize_index" which means the same thing.127 128 %quad_row, %col_8x4 = affine.delinearize_index %lane_id into (%c8, %c4) : index, index129 %quad_col = affine.apply affine_map<()[s0]->(s0 * 2)>()[%col_8x4] // account for 2xf16/col130 131 // Load quad (0, 0)132 %A_quad00 = vector.transfer_read %argA[%quad_row, %quad_col], %f0 {in_bounds = [true]} : memref<16x16xf16>, vector<2xf16>133 134 // Load quad (1, 0). Just shift row down 8.135 %quad_row_plus_8 = affine.apply affine_map<(d0)[]->(d0+8)>(%quad_row)[]136 %A_quad10 = vector.transfer_read %argA[%quad_row_plus_8, %quad_col], %f0 {in_bounds = [true]} : memref<16x16xf16>, vector<2xf16>137 138 // Load quad (0, 1). Just shift col right 8 (4 2xf16 values)139 %quad_col_plus_8 = affine.apply affine_map<(d0)[]->(d0+8)>(%quad_col)[]140 %A_quad01 = vector.transfer_read %argA[%quad_row, %quad_col_plus_8], %f0 {in_bounds = [true]} : memref<16x16xf16>, vector<2xf16>141 142 // Load quad (1, 1)143 %A_quad11 = vector.transfer_read %argA[%quad_row_plus_8, %quad_col_plus_8], %f0 {in_bounds = [true]} : memref<16x16xf16>, vector<2xf16>144 145 // Assemble the elements into a vector146 %A_init0 = arith.constant dense<0.0> : vector<4x2xf16>147 %A_data0 = vector.insert %A_quad00, %A_init0[0] : vector<2xf16> into vector<4x2xf16>148 %A_data1 = vector.insert %A_quad10, %A_data0[1] : vector<2xf16> into vector<4x2xf16>149 %A_data2 = vector.insert %A_quad01, %A_data1[2] : vector<2xf16> into vector<4x2xf16>150 %A_data = vector.insert %A_quad11, %A_data2[3] : vector<2xf16> into vector<4x2xf16>151 152 //===----------------------------------------------------------------------===//153 // Load operand B154 //===----------------------------------------------------------------------===//155 156 // Load the actual fragments for the dense values. This can be done using ldmatrix,157 // but here we just do naive individual loads, as would be required if we could158 // not use ldmatrix.159 //160 // The thread map here is different from operandA. This operand is in the form161 // memref<8x32xf16> (col major). Each thread load a 2xf16 vector from a162 // 8x8xf16 quadrant.163 //164 // The (thread_id)->(col, row) map within each 8x4x(2xf16) quadrant is165 // (t) -> (t/4, t % 4). So we can re-use some of the calculation from A.166 167 // Load quad (0, 0)168 %B_quad0 = vector.transfer_read %argB[%quad_row, %quad_col], %f0 {in_bounds = [true]} : memref<8x32xf16>, vector<2xf16>169 170 // Load quad (0, 1)171 %B_quad1 = vector.transfer_read %argB[%quad_row, %quad_col_plus_8], %f0 {in_bounds = [true]} : memref<8x32xf16>, vector<2xf16>172 173 // Load quad (0, 2)174 %quad_col_plus_16 = affine.apply affine_map<()[s0]->(s0 + 16)>()[%quad_col]175 %B_quad2 = vector.transfer_read %argB[%quad_row, %quad_col_plus_16], %f0 {in_bounds = [true]} : memref<8x32xf16>, vector<2xf16>176 177 // Load quad (0, 3)178 %quad_col_plus_24 = affine.apply affine_map<()[s0]->(s0 + 24)>()[%quad_col]179 %B_quad3 = vector.transfer_read %argB[%quad_row, %quad_col_plus_24], %f0 {in_bounds = [true]} : memref<8x32xf16>, vector<2xf16>180 181 // Assemble into vector182 %B_init0 = arith.constant dense<0.0> : vector<4x2xf16>183 %B_data0 = vector.insert %B_quad0, %B_init0[0] : vector<2xf16> into vector<4x2xf16>184 %B_data1 = vector.insert %B_quad1, %B_data0[1] : vector<2xf16> into vector<4x2xf16>185 %B_data2 = vector.insert %B_quad2, %B_data1[2] : vector<2xf16> into vector<4x2xf16>186 %B_data = vector.insert %B_quad3, %B_data2[3] : vector<2xf16> into vector<4x2xf16>187 188 // For now just say accum is a zero-d register189 %accum = arith.constant dense<0.0> : vector<2x2xf16>190 191 gpu.barrier192 193 // Sparsity selector. For 16x8x32, the default "0" means threads T0/T1194 // within each group of four threads contribute metadata.195 %d = nvgpu.mma.sp.sync(%A_data, %B_data, %accum)196 metadata(%meta)197 {mmaShape = [16, 8, 32]} : (vector<4x2xf16>, vector<4x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>198 199 //===----------------------------------------------------------------------===//200 // Write back results to gpu global memory201 //===----------------------------------------------------------------------===//202 203 // The mma instruction gave us two 2xf16 vectors per thread. These values204 // correspond to different positions in the 16x8xf16 result matrix. Each value belongs205 // to one of the 8x4x(2xf16) halves. The halves are indexed as follows (as you might guess):206 // vector0: (tid) -> (tid / 4 , tid %4)207 // vector1: (tid) -> (tid / 4 + 8, tid %4)208 %C_0 = vector.extract %d[0] : vector<2xf16> from vector<2x2xf16>209 %C_1 = vector.extract %d[1] : vector<2xf16> from vector<2x2xf16>210 vector.transfer_write %C_0, %argC[%quad_row, %quad_col] {in_bounds = [true]} : vector<2xf16>, memref<16x8xf16>211 vector.transfer_write %C_1, %argC[%quad_row_plus_8, %quad_col] {in_bounds = [true]} : vector<2xf16>, memref<16x8xf16>212 213 gpu.return214 }215 }216 217 // Code than runs on the host.218 219 //220 // This test performs a matrix multiplication221 // C = A x B222 // using NVidia 2:4 structured sparsity for A.223 //224 func.func @main() {225 %f0 = arith.constant 0.0 : f16226 %c0 = arith.constant 0 : index227 %c1 = arith.constant 1 : index228 %c2 = arith.constant 2 : index229 %c8 = arith.constant 8 : index230 %c16 = arith.constant 16 : index231 %c32 = arith.constant 32 : index232 %c64 = arith.constant 64 : index233 234 // Matrices A, B, C (16x32, 32x8, 16x8).235 %a = memref.alloc() : memref<16x16xf16> // 16x32 but 2:4, row-major236 %b = memref.alloc() : memref<8x32xf16> // regular dense column-major237 %c = memref.alloc() : memref<16x8xf16> // accumulator row-major238 239 // Metadata for A.240 %m = memref.alloc() : memref<16x2xi16>241 242 //243 // Setup matrix A.244 //245 scf.for %ai = %c0 to %c16 step %c1 {246 scf.for %aj = %c0 to %c16 step %c1 {247 %a0 = arith.addi %ai, %aj : index248 %a1 = arith.addi %a0, %c1 : index249 %a2 = arith.index_cast %a1 : index to i32250 %a3 = arith.sitofp %a2 : i32 to f16251 memref.store %a3, %a[%ai, %aj] : memref<16x16xf16>252 }253 }254 255 //256 // Setup metadata for matrix A.257 //258 // Here we assume that all 2:4 elements are in pos 0 and 2,259 // viz. in matrix260 // | A 0 B 0 |261 // { 0 2 }262 //263 // Note that within each i16, we need little-endian264 // storage of the indices, as follows:265 //266 // 10 00 10 00 10 00 10 00 10 00 10 00 = 0x8888267 //268 %bits = arith.constant 0x8888 : i16269 scf.for %mi = %c0 to %c16 step %c1 {270 memref.store %bits, %m[%mi, %c0] : memref<16x2xi16>271 memref.store %bits, %m[%mi, %c1] : memref<16x2xi16>272 }273 274 //275 // Setup matrix B.276 //277 scf.for %bi = %c0 to %c8 step %c1 {278 scf.for %bj = %c0 to %c32 step %c1 {279 %b0 = arith.subi %bi, %bj : index280 %b1 = arith.index_cast %b0 : index to i32281 %b2 = arith.sitofp %b1 : i32 to f16282 memref.store %b2, %b[%bi, %bj] : memref<8x32xf16>283 }284 }285 286 //287 // Reset matrix C.288 //289 scf.for %ci = %c0 to %c16 step %c1 {290 scf.for %cj = %c0 to %c8 step %c1 {291 memref.store %f0, %c[%ci, %cj] : memref<16x8xf16>292 }293 }294 295 //296 // Sanity check on **compressed** input matrix A.297 //298 // Note that it really is a 16x32 matrix:299 // | 1 0 2 0 3 0 ...300 // | 2 0 3 0 4 0 ...301 // etc.302 //303 // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16 )304 // CHECK-NEXT: ( 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17 )305 // CHECK-NEXT: ( 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18 )306 // CHECK-NEXT: ( 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19 )307 // CHECK-NEXT: ( 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 )308 // CHECK-NEXT: ( 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21 )309 // CHECK-NEXT: ( 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22 )310 // CHECK-NEXT: ( 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23 )311 // CHECK-NEXT: ( 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24 )312 // CHECK-NEXT: ( 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25 )313 // CHECK-NEXT: ( 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26 )314 // CHECK-NEXT: ( 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27 )315 // CHECK-NEXT: ( 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28 )316 // CHECK-NEXT: ( 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29 )317 // CHECK-NEXT: ( 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30 )318 // CHECK-NEXT: ( 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31 )319 //320 scf.for %pai = %c0 to %c16 step %c1 {321 %pa0 = vector.transfer_read %a[%pai, %c0], %f0 : memref<16x16xf16>, vector<16xf16>322 vector.print %pa0 : vector<16xf16>323 }324 325 //326 // Sanity check on input matrix 32x8 B.327 // Note that this is really shown as B^T328 //329 // CHECK-NEXT: ( 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12, -13, -14, -15, -16, -17, -18, -19, -20, -21, -22, -23, -24, -25, -26, -27, -28, -29, -30, -31 )330 // CHECK-NEXT: ( 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12, -13, -14, -15, -16, -17, -18, -19, -20, -21, -22, -23, -24, -25, -26, -27, -28, -29, -30 )331 // CHECK-NEXT: ( 2, 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12, -13, -14, -15, -16, -17, -18, -19, -20, -21, -22, -23, -24, -25, -26, -27, -28, -29 )332 // CHECK-NEXT: ( 3, 2, 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12, -13, -14, -15, -16, -17, -18, -19, -20, -21, -22, -23, -24, -25, -26, -27, -28 )333 // CHECK-NEXT: ( 4, 3, 2, 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12, -13, -14, -15, -16, -17, -18, -19, -20, -21, -22, -23, -24, -25, -26, -27 )334 // CHECK-NEXT: ( 5, 4, 3, 2, 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12, -13, -14, -15, -16, -17, -18, -19, -20, -21, -22, -23, -24, -25, -26 )335 // CHECK-NEXT: ( 6, 5, 4, 3, 2, 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12, -13, -14, -15, -16, -17, -18, -19, -20, -21, -22, -23, -24, -25 )336 // CHECK-NEXT: ( 7, 6, 5, 4, 3, 2, 1, 0, -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12, -13, -14, -15, -16, -17, -18, -19, -20, -21, -22, -23, -24 )337 //338 //339 scf.for %pbi = %c0 to %c8 step %c1 {340 %pb0 = vector.transfer_read %b[%pbi, %c0], %f0 : memref<8x32xf16>, vector<32xf16>341 vector.print %pb0 : vector<32xf16>342 }343 344 // Maps the provided host buffers into the device address space.345 // Writes from the host are guaranteed to be visible to device346 // kernels that are launched afterwards. Writes from the device347 // are guaranteed to be visible on the host after synchronizing348 // with the device kernel completion.349 %cast_a = memref.cast %a : memref<16x16xf16> to memref<*xf16>350 gpu.host_register %cast_a : memref<*xf16>351 %cast_m = memref.cast %m : memref<16x2xi16> to memref<*xi16>352 gpu.host_register %cast_m : memref<*xi16>353 %cast_b = memref.cast %b : memref<8x32xf16> to memref<*xf16>354 gpu.host_register %cast_b : memref<*xf16>355 %cast_c = memref.cast %c : memref<16x8xf16> to memref<*xf16>356 gpu.host_register %cast_c : memref<*xf16>357 358 // Call the kernel, using a single warp of 32 threads.359 %t1 = arith.constant 1 : index360 %t32 = arith.constant 32 : index361 gpu.launch_func362 @kernels::@mma_sp_sync_f16_16832363 blocks in (%t1, %t1, %t1) // gridSizeX,Y,Z364 threads in (%t32, %t1, %t1) // blockSizeX,Y,Z365 args(%a : memref<16x16xf16>,366 %m : memref<16x2xi16>,367 %b : memref<8x32xf16>,368 %c : memref<16x8xf16>)369 370 // Unmaps the host buffers.371 gpu.host_unregister %cast_a : memref<*xf16>372 gpu.host_unregister %cast_m : memref<*xi16>373 gpu.host_unregister %cast_b : memref<*xf16>374 gpu.host_unregister %cast_c : memref<*xf16>375 376 //377 // Verify computed matrix C.378 //379 // CHECK-NEXT: ( -2720, -2584, -2448, -2312, -2176, -2040, -1904, -1768 )380 // CHECK-NEXT: ( -2960, -2808, -2656, -2504, -2352, -2200, -2048, -1896 )381 // CHECK-NEXT: ( -3200, -3032, -2864, -2696, -2528, -2360, -2192, -2024 )382 // CHECK-NEXT: ( -3440, -3256, -3072, -2888, -2704, -2520, -2336, -2152 )383 // CHECK-NEXT: ( -3680, -3480, -3280, -3080, -2880, -2680, -2480, -2280 )384 // CHECK-NEXT: ( -3920, -3704, -3488, -3272, -3056, -2840, -2624, -2408 )385 // CHECK-NEXT: ( -4160, -3928, -3696, -3464, -3232, -3000, -2768, -2536 )386 // CHECK-NEXT: ( -4400, -4152, -3904, -3656, -3408, -3160, -2912, -2664 )387 // CHECK-NEXT: ( -4640, -4376, -4112, -3848, -3584, -3320, -3056, -2792 )388 // CHECK-NEXT: ( -4880, -4600, -4320, -4040, -3760, -3480, -3200, -2920 )389 // CHECK-NEXT: ( -5120, -4824, -4528, -4232, -3936, -3640, -3344, -3048 )390 // CHECK-NEXT: ( -5360, -5048, -4736, -4424, -4112, -3800, -3488, -3176 )391 // CHECK-NEXT: ( -5600, -5272, -4944, -4616, -4288, -3960, -3632, -3304 )392 // CHECK-NEXT: ( -5840, -5496, -5152, -4808, -4464, -4120, -3776, -3432 )393 // CHECK-NEXT: ( -6080, -5720, -5360, -5000, -4640, -4280, -3920, -3560 )394 // CHECK-NEXT: ( -6320, -5944, -5568, -5192, -4816, -4440, -4064, -3688 )395 //396 scf.for %pci = %c0 to %c16 step %c1 {397 %pc0 = vector.transfer_read %c[%pci, %c0], %f0 : memref<16x8xf16>, vector<8xf16>398 vector.print %pc0 : vector<8xf16>399 }400 401 return402 }403}404