225 lines · plain
1// RUN: mlir-opt -split-input-file -convert-xegpu-to-xevm -cse %s | FileCheck %s2 3gpu.module @test_kernel [#xevm.target<chip = "pvc">] {4 5 // e.g. for mem_desc<32x32xf16, @strides=[1, 16]>6 // its memory layout tuple is (blocked shape = [1,1,32,32],strides=[1024,1024,32,1])7 //CHECK-LABEL: load_store_matrix_plain8 gpu.func @load_store_matrix_plain(%arg0: memref<4096xi8, 3>) -> f32 {9 %0 = xegpu.create_mem_desc %arg0 : memref<4096xi8, 3> -> !xegpu.mem_desc<32x32xf32>10 11 //CHECK: %[[TID:.*]] = gpu.thread_id x12 //CHECK: %[[C1:.*]] = arith.constant 1 : index13 //CHECK: %[[MUL1:.*]] = arith.muli %[[TID]], %[[C1]] : index14 //CHECK: %[[C4:.*]] = arith.constant 4 : i3215 //CHECK: %[[MUL2:.*]] = arith.muli {{.*}}, %[[C4]] : i3216 //CHECK: llvm.load {{.*}} : !llvm.ptr<3> -> f3217 18 %tid_x = gpu.thread_id x19 %c0 = arith.constant 0 : index20 %1 = xegpu.load_matrix %0[%c0, %tid_x]: !xegpu.mem_desc<32x32xf32>, index, index -> f3221 22 //CHECK: llvm.store {{.*}}, {{.*}} : f32, !llvm.ptr<3>23 24 xegpu.store_matrix %1, %0[%c0, %tid_x]: f32, !xegpu.mem_desc<32x32xf32>, index, index25 26 gpu.return %1: f3227 }28 29 //CHECK-LABEL: load_store_matrix_plain_2d_input30 gpu.func @load_store_matrix_plain_2d_input(%arg0: memref<8192xi8, 3>) -> f32 {31 %c0 = arith.constant 0 : index32 %view = memref.view %arg0[%c0][]: memref<8192xi8, 3> to memref<64x32xf32, 3>33 34 %subview = memref.subview %view[32, 0] [32, 32] [1, 1] : memref<64x32xf32, 3> to memref<32x32xf32, strided<[32, 1], offset: 1024>, 3>35 36 %0 = xegpu.create_mem_desc %subview : memref<32x32xf32, strided<[32, 1], offset: 1024>, 3> -> !xegpu.mem_desc<32x32xf32>37 38 //CHECK: %[[TID:.*]] = gpu.thread_id x39 //CHECK: %[[C1:.*]] = arith.constant 1 : index40 //CHECK: %[[MUL1:.*]] = arith.muli %[[TID]], %[[C1]] : index41 //CHECK: %[[C4:.*]] = arith.constant 4 : i3242 //CHECK: %[[MUL2:.*]] = arith.muli {{.*}}, %[[C4]] : i3243 //CHECK: llvm.load {{.*}} : !llvm.ptr<3> -> f3244 45 %tid_x = gpu.thread_id x46 47 %1 = xegpu.load_matrix %0[%c0, %tid_x]: !xegpu.mem_desc<32x32xf32>, index, index -> f3248 49 //CHECK: llvm.store {{.*}}, {{.*}} : f32, !llvm.ptr<3>50 51 xegpu.store_matrix %1, %0[%c0, %tid_x]: f32, !xegpu.mem_desc<32x32xf32>, index, index52 53 gpu.return %1: f3254 }55 56 57// e.g. for mem_desc<32x64xf16, @block=[16, 16], @strides=[1, 32]>58 // its memory layout tuple is ([2,4,16,16],[256,512,1,16])59 //CHECK-LABEL: load_store_matrix_blocked_strided60 gpu.func @load_store_matrix_blocked_strided(%arg0: memref<4096xi8, 3>) -> f16 {61 %0 = xegpu.create_mem_desc %arg0 : memref<4096xi8, 3> -> !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>62 63 //CHECK: %[[tid_x:.*]] = gpu.thread_id x64 //CHECK: %[[c13:.*]] = arith.constant 13 : index65 //CHECK: %[[c16:.*]] = arith.constant 16 : index66 //CHECK: %[[offsetx_0:.*]] = arith.divsi %[[c13]], %[[c16]] : index67 //CHECK: %[[offsetx_1:.*]] = arith.remsi %[[c13]], %[[c16]] : index68 //CHECK: %[[offsety_0:.*]] = arith.divsi %[[tid_x]], %[[c16]] : index69 //CHECK: %[[offsety_1:.*]] = arith.remsi %[[tid_x]], %[[c16]] : index70 //CHECK: %[[c0:.*]] = arith.constant 0 : index71 //CHECK: %[[c256:.*]] = arith.constant 256 : index72 //CHECK: %[[mul0:.*]] = arith.muli %[[offsetx_0]], %[[c256]] : index73 //CHECK: %[[add0:.*]] = arith.addi %[[mul0]], %[[c0]] : index74 //CHECK: %[[c512:.*]] = arith.constant 512 : index75 //CHECK: %[[mul1:.*]] = arith.muli %[[offsety_0]], %[[c512]] : index76 //CHECK: %[[add1:.*]] = arith.addi %[[mul1]], %[[add0]] : index77 //CHECK: %[[c1:.*]] = arith.constant 1 : index78 //CHECK: %[[mul2:.*]] = arith.muli %[[offsetx_1]], %[[c1]] : index79 //CHECK: %[[add2:.*]] = arith.addi %[[mul2]], %[[add1]] : index80 //CHECK: %[[mul3:.*]] = arith.muli %[[offsety_1]], %[[c16]] : index81 //CHECK: %[[add3:.*]] = arith.addi %[[mul3]], %[[add2]] : index82 83 //CHECK: %[[loaded:.*]] = llvm.load {{.*}}: !llvm.ptr<3> -> f1684 85 86 %tid_x = gpu.thread_id x87 %c13 = arith.constant 13 : index88 %1 = xegpu.load_matrix %0[%c13, %tid_x]: !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index -> f1689 90 //CHECK: llvm.store %[[loaded]], {{.*}} : f16, !llvm.ptr<3>91 92 xegpu.store_matrix %1, %0[%c13, %tid_x]: f16, !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index 93 gpu.return %1: f1694 }95 96 97 // e.g. for mem_desc<32x64xf16, @block=[16, 16]>98 // its memory layout tuple is ([2,4,16,16],[1024,256,16,1])99 //CHECK-LABEL: load_store_matrix_blocked_nostride100 gpu.func @load_store_matrix_blocked_nostride(%arg0: memref<4096xi8, 3>) -> f16 {101 102 //CHECK: %[[intptr:.*]] = memref.extract_aligned_pointer_as_index %arg0 : memref<4096xi8, 3> -> index103 //CHECK: %[[basePtrI64:.*]] = arith.index_castui %[[intptr]] : index to i32104 %0 = xegpu.create_mem_desc %arg0 : memref<4096xi8, 3> -> !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>105 106 //CHECK: %[[tid_x:.*]] = gpu.thread_id x107 //CHECK: %[[c19:.*]] = arith.constant 19 : index108 %tid_x = gpu.thread_id x109 %c19 = arith.constant 19: index110 111 //CHECK: %[[c16:.*]] = arith.constant 16 : index112 //CHECK: %[[offsetx_0:.*]] = arith.divsi %[[c19]], %[[c16]] : index113 //CHECK: %[[offsetx_1:.*]] = arith.remsi %[[c19]], %[[c16]] : index114 //CHECK: %[[offsety_0:.*]] = arith.divsi %[[tid_x]], %[[c16]] : index115 //CHECK: %[[offsety_1:.*]] = arith.remsi %[[tid_x]], %[[c16]] : index116 //CHECK: %[[c0:.*]] = arith.constant 0 : index117 //CHECK: %[[c1024:.*]] = arith.constant 1024 : index118 //CHECK: %[[mul0:.*]] = arith.muli %[[offsetx_0]], %[[c1024]] : index119 //CHECK: %[[add0:.*]] = arith.addi %[[mul0]], %[[c0]] : index120 //CHECK: %[[c256:.*]] = arith.constant 256 : index121 //CHECK: %[[mul1:.*]] = arith.muli %[[offsety_0]], %[[c256]] : index122 //CHECK: %[[add1:.*]] = arith.addi %[[mul1]], %[[add0]] : index123 //CHECK: %[[mul2:.*]] = arith.muli %[[offsetx_1]], %[[c16]] : index124 //CHECK: %[[add2:.*]] = arith.addi %[[mul2]], %[[add1]] : index125 //CHECK: %[[c1:.*]] = arith.constant 1 : index126 //CHECK: %[[mul3:.*]] = arith.muli %[[offsety_1]], %[[c1]] : index127 //CHECK: %[[add3:.*]] = arith.addi %[[mul3]], %[[add2]] : index128 //CHECK: %[[loaded:.*]] = llvm.load {{.*}} : !llvm.ptr<3> -> f16129 %1 = xegpu.load_matrix %0[%c19, %tid_x]: !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>, index, index -> f16130 131 //CHECK: llvm.store %[[loaded]], {{.*}} : f16, !llvm.ptr<3>132 xegpu.store_matrix %1, %0[%c19, %tid_x]: f16, !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>, index, index133 134 //CHECK: gpu.return %[[loaded]] : f16135 gpu.return %1: f16136 }137 138 // e.g. for mem_desc<32x64xf16, @block=[16, 16], @strides=[1, 16]>139 // its memory layout tuple is ([2,4,16,16],[256,512,1,16])140 //CHECK-LABEL: load_store_matrix_blocked_strided_return_vector141 gpu.func @load_store_matrix_blocked_strided_return_vector(%arg0: memref<4096xi8, 3>) -> vector<8xf16> {142 %0 = xegpu.create_mem_desc %arg0 : memref<4096xi8, 3> -> !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>143 144 //CHECK: %[[tid_x:.*]] = gpu.thread_id x145 //CHECK: %[[c16:.*]] = arith.constant 16 : index146 //CHECK: %[[offsetx_0:.*]] = arith.divsi %[[c16]], %[[c16]] : index147 //CHECK: %[[offsetx_1:.*]] = arith.remsi %[[c16]], %[[c16]] : index148 //CHECK: %[[offsety_0:.*]] = arith.divsi %[[tid_x]], %[[c16]] : index149 //CHECK: %[[offsety_1:.*]] = arith.remsi %[[tid_x]], %[[c16]] : index150 //CHECK: %[[c0:.*]] = arith.constant 0 : index151 //CHECK: %[[c256:.*]] = arith.constant 256 : index152 //CHECK: %[[mul0:.*]] = arith.muli %[[offsetx_0]], %[[c256]] : index153 //CHECK: %[[add0:.*]] = arith.addi %[[mul0]], %[[c0]] : index154 //CHECK: %[[c512:.*]] = arith.constant 512 : index155 //CHECK: %[[mul1:.*]] = arith.muli %[[offsety_0]], %[[c512]] : index156 //CHECK: %[[add1:.*]] = arith.addi %[[mul1]], %[[add0]] : index157 //CHECK: %[[c1:.*]] = arith.constant 1 : index158 //CHECK: %[[mul2:.*]] = arith.muli %[[offsetx_1]], %[[c1]] : index159 //CHECK: %[[add2:.*]] = arith.addi %[[mul2]], %[[add1]] : index160 //CHECK: %[[mul3:.*]] = arith.muli %[[offsety_1]], %[[c16]] : index161 //CHECK: %[[add3:.*]] = arith.addi %[[mul3]], %[[add2]] : index162 163 //CHECK: %[[loaded:.*]] = llvm.load {{.*}}: !llvm.ptr<3> -> vector<8xf16>164 165 %tid_x = gpu.thread_id x166 %c16 = arith.constant 16 : index167 %1 = xegpu.load_matrix %0[%c16, %tid_x] : !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index -> vector<8xf16>168 169 //CHECK: llvm.store %[[loaded]], {{.*}} : vector<8xf16>, !llvm.ptr<3>170 xegpu.store_matrix %1, %0[%c16, %tid_x] : vector<8xf16>, !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<stride = [1, 32], block = [16, 16]>>, index, index171 172 gpu.return %1: vector<8xf16>173 }174 175 176 // e.g. for mem_desc<32x64xf16, @block=[16, 16]>177 // its memory layout tuple is ([2,4,16,16],[1024,256,16,1])178 //CHECK-LABEL: load_store_matrix_blocked_subgroupblockio179 gpu.func @load_store_matrix_blocked_subgroupblockio(%arg0: memref<4096xi8, 3>) -> vector<8xf16> {180 181 //CHECK: %[[intptr:.*]] = memref.extract_aligned_pointer_as_index %arg0 : memref<4096xi8, 3> -> index182 //CHECK: %[[basePtrI32:.*]] = arith.index_castui %[[intptr]] : index to i32183 %0 = xegpu.create_mem_desc %arg0 : memref<4096xi8, 3> -> !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>184 185 //CHECK: %[[c16:.*]] = arith.constant 16 : index186 //CHECK: %[[c48:.*]] = arith.constant 48 : index187 %c16 = arith.constant 16 : index188 %c48 = arith.constant 48 : index189 190 //CHECK: %[[offset0:.*]] = arith.divsi %[[c16]], %[[c16]] : index191 //CHECK: %[[offset1:.*]] = arith.remsi %[[c16]], %[[c16]] : index192 //CHECK: %[[offset2:.*]] = arith.divsi %[[c48]], %[[c16]] : index193 //CHECK: %[[offset3:.*]] = arith.remsi %[[c48]], %[[c16]] : index194 //CHECK: %[[c0:.*]] = arith.constant 0 : index195 //CHECK: %[[c1024:.*]] = arith.constant 1024 : index196 //CHECK: %[[mul0:.*]] = arith.muli %[[offset0]], %[[c1024]] : index197 //CHECK: %[[add0:.*]] = arith.addi %[[mul0]], %[[c0]] : index198 //CHECK: %[[c256:.*]] = arith.constant 256 : index199 //CHECK: %[[mul1:.*]] = arith.muli %[[offset2]], %[[c256]] : index200 //CHECK: %[[add1:.*]] = arith.addi %[[mul1]], %[[add0]] : index201 //CHECK: %[[mul2:.*]] = arith.muli %[[offset1]], %[[c16]] : index202 //CHECK: %[[add2:.*]] = arith.addi %[[mul2]], %[[add1]] : index203 //CHECK: %[[c1:.*]] = arith.constant 1 : index204 //CHECK: %[[mul3:.*]] = arith.muli %[[offset3]], %[[c1]] : index205 //CHECK: %[[linearOffset:.*]] = arith.addi %[[mul3]], %[[add2]] : index206 //CHECK: %[[linearOffsetI64:.*]] = arith.index_castui %[[linearOffset]] : index to i32207 //CHECK: %[[c2:.*]] = arith.constant 2 : i32208 //CHECK: %[[byteOffset:.*]] = arith.muli %[[linearOffsetI64]], %[[c2]] : i32209 //CHECK: %[[finalPtr:.*]] = arith.addi %[[basePtrI32]], %[[byteOffset]] : i32210 //CHECK: %[[ptr:.*]] = llvm.inttoptr %[[finalPtr]] : i32 to !llvm.ptr<3>211 //CHECK: %[[loadedI16:.*]] = xevm.blockload %[[ptr]] : (!llvm.ptr<3>) -> vector<8xi16>212 //CHECK: %[[loaded:.*]] = vector.bitcast %[[loadedI16]] : vector<8xi16> to vector<8xf16>213 214 %1 = xegpu.load_matrix %0[%c16, %c48] {subgroup_block_io}: !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>, index, index -> vector<8xf16>215 216 //CHECK: %[[storeDataI16:.*]] = vector.bitcast %[[loaded]] : vector<8xf16> to vector<8xi16>217 //CHECK: xevm.blockstore %[[ptr]], %[[storeDataI16]] : (!llvm.ptr<3>, vector<8xi16>) 218 219 xegpu.store_matrix %1, %0[%c16, %c48] {subgroup_block_io}: vector<8xf16>, !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>, index, index220 221 gpu.return %1: vector<8xf16>222 }223 224}225