333 lines · plain
1// RUN: mlir-opt --xevm-attach-target='module=xevm_* chip=pvc' -xegpu-subgroup-distribute \2// RUN: -allow-unregistered-dialect -canonicalize -cse -split-input-file %s | FileCheck %s3 4// CHECK-LABEL: gpu.func @load_dpas_postop_store5// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<16x16xf16>,6// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {7// CHECK: %[[T2:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>8// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T2]][%{{.*}}] : !xegpu.tensor_desc<8x16xf16> -> vector<8xf16>9// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16>10// CHECK: %[[T1:.*]] = xegpu.load_nd %[[T0]][%{{.*}}] <{packed}> : !xegpu.tensor_desc<16x16xf16> -> vector<16xf16>11// CHECK-DAG: %[[T4:.*]] = xegpu.dpas %[[T3]], %[[T1]] : vector<8xf16>, vector<16xf16> -> vector<8xf32>12// CHECK: %[[T5:.*]] = vector.shape_cast %[[T4]] : vector<8xf32> to vector<8x1xf32>13// CHECK: %[[T6:.*]] = math.exp %[[T5]] {{{.*}}} : vector<8x1xf32>14// CHECK-DAG: %[[T8:.*]] = vector.shape_cast %[[T6]] : vector<8x1xf32> to vector<8xf32>15// CHECK-DAG: %[[T7:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>16// CHECK: xegpu.store_nd %[[T8]], %[[T7]][{{.*}}] : vector<8xf32>, !xegpu.tensor_desc<8x16xf32>17gpu.module @xevm_module{18 gpu.func @load_dpas_postop_store(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<8x16xf32>) {19 %c0 = arith.constant 0 : index20 %0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16>21 -> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>22 %1 = xegpu.load_nd %0[%c0, %c0]23 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :24 !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>25 26 %2 = xegpu.create_nd_tdesc %arg1: memref<16x16xf16>27 -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>28 %3 = xegpu.load_nd %2[%c0, %c0]29 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}30 : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>31 -> vector<16x16xf16>32 33 %4 = xegpu.dpas %1, %334 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}35 : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>36 37 %5 = math.exp %438 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}39 : vector<8x16xf32>40 41 %6 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> ->42 !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>43 xegpu.store_nd %5, %6[%c0, %c0] : vector<8x16xf32>,44 !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>45 gpu.return46 }47}48 49// -----50// CHECK-LABEL: gpu.func @gemm51// CHECK: (%[[ARG0:[0-9a-zA-Z]+]]: memref<1024x1024xbf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<1024x1024xbf16>,52// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]+]]: memref<1024x1024xf32>) {53// CHECK-DAG: %[[BLOCK_ID_X:.*]] = gpu.block_id x54// CHECK-DAG: %[[BLOCK_ID_Y:.*]] = gpu.block_id y55// CHECK-DAG: %[[Y_COORD:.*]] = arith.muli %[[BLOCK_ID_Y]], %c16 : index56// CHECK-DAG: %[[X_COORD:.*]] = arith.muli %[[BLOCK_ID_X]], %c8 : index57// CHECK: %[[T2:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<1024x1024xf32> -> !xegpu.tensor_desc<8x16xf32>58// CHECK-NEXT: %[[T3:.*]] = xegpu.load_nd %[[T2]][%[[X_COORD]], %[[Y_COORD]]] : !xegpu.tensor_desc<8x16xf32> -> vector<8xf32>59// CHECK-NEXT: %[[T4:.*]] = vector.shape_cast %[[T3]] : vector<8xf32> to vector<8x1xf32>60// CHECK: %[[T5:.*]] = scf.for %[[K:.*]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[ARG4:.*]] = %[[T4]])61// CHECK-SAME: -> (vector<8x1xf32>) {62// CHECK-DAG: %[[T10:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<1024x1024xbf16> -> !xegpu.tensor_desc<16x16xbf16>63// CHECK-DAG: %[[T11:.*]] = xegpu.load_nd %[[T10]][%[[K]], %[[Y_COORD]]] <{packed}> : !xegpu.tensor_desc<16x16xbf16> -> vector<16xbf16>64// CHECK-DAG: %[[T12:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<1024x1024xbf16> -> !xegpu.tensor_desc<8x16xbf16>65// CHECK-DAG: %[[T13:.*]] = xegpu.load_nd %[[T12]][%[[X_COORD]], %[[K]]] : !xegpu.tensor_desc<8x16xbf16> -> vector<8xbf16>66// CHECK-DAG: %[[T14:.*]] = vector.shape_cast %[[ARG4]] : vector<8x1xf32> to vector<8xf32>67// CHECK-NEXT: %[[T15:.*]] = xegpu.dpas %[[T13]], %[[T11]], %[[T14]]68// CHECK-SAME: : vector<8xbf16>, vector<16xbf16>, vector<8xf32> -> vector<8xf32>69// CHECK-NEXT: %[[T16:.*]] = vector.shape_cast %[[T15]] : vector<8xf32> to vector<8x1xf32>70// CHECK-NEXT: scf.yield %[[T16]] : vector<8x1xf32>71// CHECK-NEXT: }72// CHECK-NEXT: %[[T9:.*]] = vector.shape_cast %[[T5]] : vector<8x1xf32> to vector<8xf32>73// CHECK-NEXT: xegpu.store_nd %[[T9]], %[[T2]][%[[X_COORD]], %[[Y_COORD]]] : vector<8xf32>, !xegpu.tensor_desc<8x16xf32>74gpu.module @xevm_module{75gpu.func @gemm(%arg0: memref<1024x1024xbf16>, %arg1: memref<1024x1024xbf16>, %arg2: memref<1024x1024xf32>){76 %c0 = arith.constant 0 : index77 %c16 = arith.constant 16 : index78 %c8 = arith.constant 8 : index79 %c1024 = arith.constant 1024 : index80 %block_id_x = gpu.block_id x81 %block_id_y = gpu.block_id y82 %0 = arith.muli %block_id_x, %c8 : index83 %1 = arith.muli %block_id_y, %c16 : index84 %2 = xegpu.create_nd_tdesc %arg2 : memref<1024x1024xf32> ->85 !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>86 %3 = xegpu.load_nd %2[%0, %1]87 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}88 : !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf32>89 90 %4 = scf.for %arg3 = %c0 to %c1024 step %c16 iter_args(%arg4 = %3) -> (vector<8x16xf32>) {91 92 %5 = xegpu.create_nd_tdesc %arg0: memref<1024x1024xbf16>93 -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>94 %6 = xegpu.create_nd_tdesc %arg1 : memref<1024x1024xbf16>95 -> !xegpu.tensor_desc<16x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>96 97 %7 = xegpu.load_nd %5[%0, %arg3]98 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}99 : !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xbf16>100 %8 = xegpu.load_nd %6[%arg3, %1]101 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}102 : !xegpu.tensor_desc<16x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xbf16>103 104 %9 = xegpu.dpas %7, %8, %arg4105 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}106 : vector<8x16xbf16>, vector<16x16xbf16>, vector<8x16xf32> -> vector<8x16xf32>107 108 scf.yield %9 : vector<8x16xf32>109 } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}110 111 xegpu.store_nd %4, %2[%0, %1] : vector<8x16xf32>,112 !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>113 gpu.return114}115}116 117// -----118// CHECK-LABEL: gpu.func @scatter_ops_scf_yield119// CHECK: (%{{.*}}: memref<256xf16>, %[[PREDICATE:[a-zA-Z0-9]+]]: i1) {120// CHECK-DAG: %[[CST:.*]] = arith.constant dense<1.200000e+01> : vector<1x8xf16>121// CHECK-DAG: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>122// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>123// CHECK: %[[IF:.*]] = scf.if %[[PREDICATE]] -> (vector<1x8xf16>) {124// CHECK-NEXT: %[[LD:.*]] = xegpu.load %{{.*}}[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>125// CHECK-SAME: : memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>126// CHECK-NEXT: %[[LD_CAST:.*]] = vector.shape_cast %[[LD]] : vector<8xf16> to vector<1x8xf16>127// CHECK-NEXT: scf.yield %[[LD_CAST]] : vector<1x8xf16>128// CHECK-NEXT: } else {129// CHECK-NEXT: scf.yield %[[CST]] : vector<1x8xf16>130// CHECK-NEXT: }131// CHECK-NEXT: %[[IF_CAST:.*]] = vector.shape_cast %[[IF]] : vector<1x8xf16> to vector<8xf16>132// CHECK-NEXT: xegpu.store %[[IF_CAST]], %{{.*}}[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>133// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>134gpu.module @xevm_module{135 gpu.func @scatter_ops_scf_yield(%src: memref<256xf16>, %pred : i1) {136 %1 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1>: vector<16xi1>137 %offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16xindex>138 %loaded = scf.if %pred -> (vector<16x8xf16>) {139 %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {140 layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>141 } : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>142 scf.yield %3 : vector<16x8xf16>143 } else {144 %3 = arith.constant {145 layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>146 } dense<12.> : vector<16x8xf16>147 scf.yield %3 : vector<16x8xf16>148 } { layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]> }149 xegpu.store %loaded, %src[%offset], %1 <{chunk_size=8}> : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>150 gpu.return151 }152}153 154// -----155// CHECK-LABEL: gpu.func @scatter_ops_scf_non_yield({{.*}}) {156// CHECK: %[[OFFSET:.*]] = arith.constant dense<12> : vector<1xindex>157// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<1xi1>158// CHECK: %[[PREDICATE:.*]] = llvm.mlir.poison : i1159// CHECK: scf.if %[[PREDICATE]] {160// CHECK-NEXT: %[[LOADED:.*]] = xegpu.load %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>161// CHECK-SAME: memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>162// CHECK-NEXT: xegpu.store %[[LOADED]], %arg0[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>163// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>164// CHECK-NEXT: }165gpu.module @xevm_module{166 gpu.func @scatter_ops_scf_non_yield(%src: memref<256xf16>) {167 %pred = llvm.mlir.poison : i1168 %1 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1>: vector<16xi1>169 %offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16xindex>170 scf.if %pred {171 %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {172 layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>173 } : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>174 xegpu.store %3, %src[%offset], %1 <{chunk_size=8}> : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>175 }176 gpu.return177 }178}179 180// -----181// CHECK-LABEL: gpu.func @mma_transpose_b(182// CHECK: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<16x8xi32>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {183// CHECK-DAG: %[[ADESC:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>184// CHECK-DAG: %[[BDESC:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<16x8xi32> -> !xegpu.tensor_desc<16x8xi32>185// CHECK-DAG: %[[A:.*]] = xegpu.load_nd %[[ADESC]][%{{.*}}] : !xegpu.tensor_desc<8x16xf16> -> vector<8xf16>186// CHECK-DAG: %[[B:.*]] = xegpu.load_nd %[[BDESC]][%{{.*}}] <{transpose = array<i64: 1, 0>}>187// CHECK-SAME: !xegpu.tensor_desc<16x8xi32> -> vector<8xi32>188// CHECK-NEXT: %[[BCAST0:.*]] = vector.shape_cast %[[B]] : vector<8xi32> to vector<1x8xi32>189// CHECK-NEXT: %[[BCAST1:.*]] = vector.bitcast %[[BCAST0]] : vector<1x8xi32> to vector<1x16xf16>190// CHECK-NEXT: %[[BCAST2:.*]] = vector.shape_cast %[[BCAST1]] : vector<1x16xf16> to vector<16xf16>191// CHECK-NEXT: %[[C:.*]] = xegpu.dpas %[[A]], %[[BCAST2]] : vector<8xf16>, vector<16xf16> -> vector<8xf32>192gpu.module @xevm_module{193 gpu.func @mma_transpose_b(%arg0: memref<8x16xf16>, %arg1: memref<16x8xi32>, %arg2: memref<8x16xf32>) {194 %c0 = arith.constant 0 : index195 %0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16>196 -> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>197 %1 = xegpu.load_nd %0[%c0, %c0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}198 : !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>199 %2 = xegpu.create_nd_tdesc %arg1 : memref<16x8xi32>200 -> !xegpu.tensor_desc<16x8xi32, #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>>201 %3 = xegpu.load_nd %2[%c0, %c0] {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}202 : !xegpu.tensor_desc<16x8xi32, #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>> -> vector<16x8xi32>203 %4 = vector.bitcast %3 {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}204 : vector<16x8xi32> to vector<16x16xf16>205 %5 = vector.transpose %4, [1, 0] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}206 : vector<16x16xf16> to vector<16x16xf16>207 %6 = xegpu.dpas %1, %5 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}208 : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>209 %7 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32>210 -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>211 xegpu.store_nd %6, %7[%c0, %c0] : vector<8x16xf32>,212 !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>213 gpu.return214 215 }216}217 218// -----219// CHECK-LABEL: gpu.func @warp_scf_for_unused_uniform_for_result(220// CHECK: %[[W:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16] args(%{{.*}} : index,221// CHECK-SAME: !xegpu.tensor_desc<16x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,222// CHECK-SAME: memref<16x16xf32>) -> (vector<16x1xf32>, vector<16x1xf32>) {223// CHECK: gpu.yield %{{.*}}, {{.*}} : vector<16x16xf32>, vector<16x1xf32>224// CHECK: }225// CHECK: %{{.*}}:2 = scf.for {{.*}} to %{{.*}} step %{{.*}} iter_args226// CHECK-SAME: (%{{.*}} = %[[W]]#0, %{{.*}} = %[[W]]#1) -> (vector<16x1xf32>, vector<16x1xf32>) {227// CHECK: %[[W1:.*]]:2 = gpu.warp_execute_on_lane_0(%{{.*}})[16]228// CHECK-SAME: args(%{{.*}} : vector<16x1xf32>, vector<16x1xf32>) -> (vector<16x1xf32>, vector<16x1xf32>) {229// CHECK: gpu.yield %{{.*}}, %{{.*}} : vector<16x16xf32>, vector<16x1xf32>230// CHECK: }231// CHECK: scf.yield %[[W1]]#0, %[[W1]]#1 : vector<16x1xf32>, vector<16x1xf32>232// CHECK: }233gpu.module @xevm_module{234 gpu.func @warp_scf_for_unused_uniform_for_result(%arg0: index,235 %arg1: !xegpu.tensor_desc<16x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,236 %arg2: memref<16x16xf32>) {237 %c128 = arith.constant 128 : index238 %c1 = arith.constant 1 : index239 %c0 = arith.constant 0 : index240 %ini = "some_def"() {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}241 : () -> (vector<16x1xf32>)242 %ini2 = "some_def"() {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}243 : () -> (vector<16x16xf32>)244 %3:2 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %ini2, %arg5 = %ini) -> (vector<16x16xf32>, vector<16x1xf32>) {245 %1 = "some_def"(%arg5)246 {247 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,248 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>249 }250 : (vector<16x1xf32>) -> (vector<16x1xf32>)251 %acc = "some_def"(%arg4, %1)252 {253 layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,254 layout_operand_1 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,255 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>256 }257 : (vector<16x16xf32>, vector<16x1xf32>) -> (vector<16x16xf32>)258 scf.yield %acc, %1 : vector<16x16xf32>, vector<16x1xf32>259 }260 {261 layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>262 }263 xegpu.store_nd %3#0, %arg1[%c0, %c0]264 : vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>265 gpu.return266 }267}268 269// -----270// CHECK-LABEL: gpu.func @load_store_matrix_1({{.*}}) {271// CHECK: %[[C2:.*]] = arith.constant 2 : index272// CHECK: %[[C8:.*]] = arith.constant 8 : index273// CHECK: %[[LANE_ID:.*]] = gpu.lane_id274// CHECK: %[[REMU1:.*]] = index.remu %[[LANE_ID]], %[[C8]]275// CHECK: %[[DIVU:.*]] = index.divu %[[LANE_ID]], %[[C8]]276// CHECK: %[[REMU2:.*]] = index.remu %[[DIVU]], %[[C2]]277// CHECK: %[[REMU3:.*]] = index.remu %[[REMU2]], %[[C2]]278// CHECK: %[[REMU4:.*]] = index.remu %[[REMU1]], %[[C8]]279// CHECK: %[[MAT:.*]] = xegpu.load_matrix %arg0[%[[REMU3]], %[[REMU4]]] : !xegpu.mem_desc<32x32xf32>, index, index -> vector<1x1xf32>280// CHECK: xegpu.store_matrix %[[MAT]], %arg0[%[[REMU3]], %[[REMU4]]] : vector<1x1xf32>, !xegpu.mem_desc<32x32xf32>, index, index281gpu.module @xevm_module{282 gpu.func @load_store_matrix_1(%arg0: !xegpu.mem_desc<32x32xf32>) {283 %c0 = arith.constant 0 : index284 %1 = xegpu.load_matrix %arg0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}> : !xegpu.mem_desc<32x32xf32>, index, index -> vector<2x8xf32>285 xegpu.store_matrix %1, %arg0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}> : vector<2x8xf32>, !xegpu.mem_desc<32x32xf32>, index, index286 gpu.return287 }288}289 290// -----291// CHECK-LABEL: gpu.func @load_store_matrix_2({{.*}}) {292// CHECK: %[[C8:.*]] = arith.constant 8 : index293// CHECK: %[[C2:.*]] = arith.constant 2 : index294// CHECK: %[[C4:.*]] = arith.constant 4 : index295// CHECK: %[[C1:.*]] = arith.constant 1 : index296// CHECK: %[[LANE_ID:.*]] = gpu.lane_id297// CHECK: %[[REMU1:.*]] = index.remu %[[LANE_ID]], %[[C4]]298// CHECK: %[[DIVU:.*]] = index.divu %[[LANE_ID]], %[[C4]]299// CHECK: %[[REMU2:.*]] = index.remu %[[DIVU]], %[[C4]]300// CHECK: %[[MUL:.*]] = index.mul %[[REMU2]], %[[C2]]301// CHECK: %[[REMU3:.*]] = index.remu %[[MUL]], %[[C8]]302// CHECK: %[[REMU4:.*]] = index.remu %[[REMU1]], %[[C4]]303// CHECK: %[[ADD:.*]] = index.add %[[REMU4]], %[[C1]]304// CHECK: %[[MAT:.*]] = xegpu.load_matrix %arg0[%[[REMU3]], %[[ADD]]] : !xegpu.mem_desc<32x32xf32>, index, index -> vector<2x1xf32>305// CHECK: xegpu.store_matrix %[[MAT]], %arg0[%[[REMU3]], %[[ADD]]] : vector<2x1xf32>, !xegpu.mem_desc<32x32xf32>, index, index306gpu.module @xevm_module{307 gpu.func @load_store_matrix_2(%arg0: !xegpu.mem_desc<32x32xf32>) {308 %c0 = arith.constant 0 : index309 %c1 = arith.constant 1 : index310 %1 = xegpu.load_matrix %arg0[%c0, %c1] <{layout = #xegpu.layout<lane_layout = [4, 4], lane_data = [2, 1]>}> : !xegpu.mem_desc<32x32xf32>, index, index -> vector<8x4xf32>311 xegpu.store_matrix %1, %arg0[%c0, %c1] <{layout = #xegpu.layout<lane_layout = [4, 4], lane_data = [2, 1]>}> : vector<8x4xf32>, !xegpu.mem_desc<32x32xf32>, index, index312 gpu.return313 }314}315 316// -----317// CHECK-LABEL: gpu.func @load_store_matrix_3({{.*}}) {318// CHECK: %[[MAT:.*]] = xegpu.load_matrix %arg0[%{{.*}}, %{{.*}}] <{subgroup_block_io}>:319// CHECK-SAME: !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<block = [16, 1], stride = [1, 32]>>, index, index -> vector<1x2xf32>320// CHECK: xegpu.store_matrix %[[MAT]], %arg0[%{{.*}}, %{{.*}}] <{subgroup_block_io}>:321// CHECK-SAME: vector<1x2xf32>, !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<block = [16, 1], stride = [1, 32]>>, index, index322gpu.module @xevm_module{323 gpu.func @load_store_matrix_3(%arg0: !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [16, 1]>>) {324 %c0 = arith.constant 0 : index325 %c1 = arith.constant 1 : index326 %1 = xegpu.load_matrix %arg0[%c0, %c1] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>} :327 !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [16, 1]>>, index, index -> vector<16x2xf32>328 xegpu.store_matrix %1, %arg0[%c0, %c1] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>} :329 vector<16x2xf32>, !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [16, 1]>>, index, index330 gpu.return331 }332}333