247 lines · plain
1// RUN: mlir-opt %s -split-input-file -canonicalize="test-convergence" | FileCheck %s2 3// CHECK-LABEL: func @known_oob_load4func.func @known_oob_load(%arg0: memref<4xf32>) -> f32 {5 // CHECK: %[[zero:.*]] = arith.constant 0.000000e+00 : f326 // CHECK: return %[[zero]]7 %c4_i32 = arith.constant 4 : i328 %0 = amdgpu.raw_buffer_load {boundsCheck = true} %arg0[%c4_i32] : memref<4xf32>, i32 -> f329 func.return %0 : f3210}11 12// -----13 14// CHECK-LABEL: func @known_oob_load_2d15func.func @known_oob_load_2d(%arg0: memref<4x4xf32>) -> f32 {16 // CHECK: %[[zero:.*]] = arith.constant 0.000000e+00 : f3217 // CHECK: return %[[zero]]18 %c0_i32 = arith.constant 0 : i3219 %c4_i32 = arith.constant 4 : i3220 %0 = amdgpu.raw_buffer_load {boundsCheck = true} %arg0[%c4_i32, %c0_i32] : memref<4x4xf32>, i32, i32 -> f3221 func.return %0 : f3222}23 24// -----25 26// CHECK-LABEL: func @known_oob_load_2d_on_last27func.func @known_oob_load_2d_on_last(%arg0: memref<4x4xf32>) -> f32 {28 // CHECK: %[[zero:.*]] = arith.constant 0.000000e+00 : f3229 // CHECK: return %[[zero]]30 %c0_i32 = arith.constant 0 : i3231 %c16_i32 = arith.constant 16 : i3232 %0 = amdgpu.raw_buffer_load {boundsCheck = true} %arg0[%c0_i32, %c16_i32] : memref<4x4xf32>, i32, i32 -> f3233 func.return %0 : f3234}35 36// -----37 38// CHECK-LABEL: func @known_oob_load_index39func.func @known_oob_load_index(%arg0: memref<4xf32>) -> f32 {40 // CHECK: %[[zero:.*]] = arith.constant 0.000000e+00 : f3241 // CHECK: return %[[zero]]42 %c0_i32 = arith.constant 0 : i3243 %0 = amdgpu.raw_buffer_load {boundsCheck = true, indexOffset = 4 : i32} %arg0[%c0_i32] : memref<4xf32>, i32 -> f3244 func.return %0 : f3245}46 47// -----48 49// CHECK-LABEL: func @known_oob_load_sgproffset50func.func @known_oob_load_sgproffset(%arg0: memref<4xf32>) -> f32 {51 // CHECK: %[[zero:.*]] = arith.constant 0.000000e+00 : f3252 // CHECK: return %[[zero]]53 %c2_i32 = arith.constant 2 : i3254 %0 = amdgpu.raw_buffer_load {boundsCheck = true} %arg0[%c2_i32] sgprOffset %c2_i32 : memref<4xf32>, i32 -> f3255 func.return %0 : f3256}57 58// -----59 60// CHECK-LABEL: func @unknown_load61func.func @unknown_load(%arg0: memref<4xf32>, %arg1: i32) -> f32 {62 // CHECK: %[[loaded:.*]] = amdgpu.raw_buffer_load63 // CHECK: return %[[loaded]]64 %c4_i32 = arith.constant 4 : i3265 %0 = amdgpu.raw_buffer_load {boundsCheck = true} %arg0[%arg1] sgprOffset %c4_i32 : memref<4xf32>, i32 -> f3266 func.return %0 : f3267}68 69// -----70 71// CHECK-LABEL: func @unknown_load_sgproffset72func.func @unknown_load_sgproffset(%arg0: memref<4xf32>, %arg1: i32) -> f32 {73 // CHECK: %[[loaded:.*]] = amdgpu.raw_buffer_load74 // CHECK: return %[[loaded]]75 %c4_i32 = arith.constant 4 : i3276 %0 = amdgpu.raw_buffer_load {boundsCheck = true} %arg0[%c4_i32] sgprOffset %arg1 : memref<4xf32>, i32 -> f3277 func.return %0 : f3278}79 80// -----81 82// CHECK-LABEL: func @unranked83func.func @unranked(%arg0: memref<?xf32>) -> f32 {84 // CHECK: %[[loaded:.*]] = amdgpu.raw_buffer_load85 // CHECK: return %[[loaded]]86 %c4_i32 = arith.constant 4 : i3287 %0 = amdgpu.raw_buffer_load {boundsCheck = true} %arg0[%c4_i32] : memref<?xf32>, i32 -> f3288 func.return %0 : f3289}90 91// -----92 93// CHECK-LABEL: func @no_oob_check94func.func @no_oob_check(%arg0: memref<4xf32>) -> f32 {95 // CHECK: %[[loaded:.*]] = amdgpu.raw_buffer_load96 // CHECK: return %[[loaded]]97 %c4_i32 = arith.constant 4 : i3298 %0 = amdgpu.raw_buffer_load {boundsCheck = false} %arg0[%c4_i32] : memref<4xf32>, i32 -> f3299 func.return %0 : f32100}101 102// -----103 104// CHECK-LABEL: func @in_bounds_overall105func.func @in_bounds_overall(%arg0: memref<4x4xf32>) -> f32 {106 // CHECK: %[[loaded:.*]] = amdgpu.raw_buffer_load107 // CHECK: return %[[loaded]]108 %c0_i32 = arith.constant 0 : i32109 %c15_i32 = arith.constant 15 : i32110 %0 = amdgpu.raw_buffer_load {boundsCheck = true} %arg0[%c0_i32, %c15_i32] : memref<4x4xf32>, i32, i32 -> f32111 func.return %0 : f32112}113 114// -----115 116// CHECK-LABEL: func @dead_store117func.func @dead_store(%arg0: memref<4xf32>, %arg1: f32) {118 // CHECK-NOT: amdgpu.raw_buffer_store119 %c4_i32 = arith.constant 4 : i32120 amdgpu.raw_buffer_store {boundsCheck = true} %arg1 -> %arg0[%c4_i32] : f32 -> memref<4xf32>, i32121 func.return122}123 124// -----125 126// CHECK-LABEL: func @dead_atomic_add127func.func @dead_atomic_add(%arg0: memref<4xf32>, %arg1: f32) {128 // CHECK-NOT: amdgpu.raw_buffer_atomic_fadd129 %c4_i32 = arith.constant 4 : i32130 amdgpu.raw_buffer_atomic_fadd {boundsCheck = true} %arg1 -> %arg0[%c4_i32] : f32 -> memref<4xf32>, i32131 func.return132}133 134// -----135 136// CHECK-LABEL: func @fold_gather_to_lds_of_cast137func.func @fold_gather_to_lds_of_cast(%global: memref<128x72xf32, 1>, %lds: memref<64x64xf32, 3>) {138// CHECK-SAME: %[[GLOBAL:[A-Za-z0-9]+]]: memref<128x72xf32, 1>139 %c0 = arith.constant 0 : index140 %0 = memref.cast %global : memref<128x72xf32, 1> to memref<?x?xf32, 1>141 // CHECK: amdgpu.gather_to_lds %[[GLOBAL]]142 // CHECK-SAME: : f32, memref<128x72xf32, 1>143 amdgpu.gather_to_lds %0[%c0, %c0], %lds[%c0, %c0]144 : f32, memref<?x?xf32, 1>, memref<64x64xf32, 3>145 func.return146}147 148// -----149 150// CHECK-LABEL: func @fold_gather_to_lds_of_cast_dest151func.func @fold_gather_to_lds_of_cast_dest(%global: memref<128x72xf32, 1>, %lds: memref<64x64xf32, 3>) {152// CHECK-SAME: %[[GLOBAL:[A-Za-z0-9]+]]: memref<128x72xf32, 1>153// CHECK-SAME: %[[LDS:[A-Za-z0-9]+]]: memref<64x64xf32, 3>154 %c0 = arith.constant 0 : index155 %0 = memref.cast %lds : memref<64x64xf32, 3> to memref<?x?xf32, 3>156 // CHECK: amdgpu.gather_to_lds %[[GLOBAL]][{{.*}}], %[[LDS]]157 // CHECK-SAME: : f32, memref<128x72xf32, 1>, memref<64x64xf32, 3>158 amdgpu.gather_to_lds %global[%c0, %c0], %0[%c0, %c0]159 : f32, memref<128x72xf32, 1>, memref<?x?xf32, 3>160 func.return161}162 163// -----164 165// CHECK-LABEL: func @scaled_mfma166// CHECK: %[[SCALE_1:.*]] = vector.extract_strided_slice %0 {offsets = [0], sizes = [4], strides = [1]} : vector<16xf8E8M0FNU> to vector<4xf8E8M0FNU>167// CHECK: %[[SCALE_2:.*]] = vector.extract_strided_slice %2 {offsets = [4], sizes = [4], strides = [1]} : vector<16xf8E8M0FNU> to vector<4xf8E8M0FNU>168// CHECK: amdgpu.scaled_mfma 16x16x128 (%[[SCALE_1]][3] * %{{.*}}) * (%[[SCALE_2]][2] * %{{.*}}) {{.*}}169// CHECK: %[[SCALE_3:.*]] = vector.extract_strided_slice %5 {offsets = [8], sizes = [4], strides = [1]} : vector<16xf8E8M0FNU> to vector<4xf8E8M0FNU>170// CHECK: %[[SCALE_4:.*]] = vector.extract_strided_slice %7 {offsets = [12], sizes = [4], strides = [1]} : vector<16xf8E8M0FNU> to vector<4xf8E8M0FNU>171// CHECK: amdgpu.scaled_mfma 16x16x128 (%[[SCALE_3]][1] * %{{.*}}) * (%[[SCALE_4]][0] * %{{.*}}) {{.*}}172func.func @scaled_mfma(%opA: vector<32xf4E2M1FN>, %opB: vector<32xf4E2M1FN>, %scalesA: vector<2x1x8x1xf8E8M0FNU>, %scalesB: vector<2x1x8x1xf8E8M0FNU>) -> (vector<4xf32>, vector<4xf32>) {173 %cst_0 = arith.constant dense<0.000000e+00> : vector<4xf32>174 %cst_1 = arith.constant dense<5.877470e-39> : vector<4xf8E8M0FNU>175 %scaleA = vector.extract %scalesA[0, 0, 3, 0] : f8E8M0FNU from vector<2x1x8x1xf8E8M0FNU>176 %sA = vector.insert %scaleA, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>177 %scaleB = vector.extract %scalesB[0, 0, 6, 0] : f8E8M0FNU from vector<2x1x8x1xf8E8M0FNU>178 %sB = vector.insert %scaleB, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>179 %res_0 = amdgpu.scaled_mfma 16x16x128 (%sA[0] * %opA) * (%sB[0] * %opB) + %cst_0 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>180 %scaleC = vector.extract %scalesA[1, 0, 1, 0] : f8E8M0FNU from vector<2x1x8x1xf8E8M0FNU>181 %sC = vector.insert %scaleC, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>182 %scaleD = vector.extract %scalesB[1, 0, 4, 0] : f8E8M0FNU from vector<2x1x8x1xf8E8M0FNU>183 %sD = vector.insert %scaleD, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>184 %res_1 = amdgpu.scaled_mfma 16x16x128 (%sC[0] * %opA) * (%sD[0] * %opB) + %cst_0 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>185 return %res_0, %res_1 : vector<4xf32>, vector<4xf32>186}187 188// -----189 190// CHECK-LABEL: func @scaled_mfma_less_than_4191// CHECK: vector.extract {{.*}} : f8E8M0FNU from vector<2xf8E8M0FNU>192// CHECK: vector.insert {{.*}} : f8E8M0FNU into vector<4xf8E8M0FNU>193// CHECK: vector.extract {{.*}} : f8E8M0FNU from vector<2xf8E8M0FNU>194// CHECK: vector.insert {{.*}} : f8E8M0FNU into vector<4xf8E8M0FNU>195// CHECK: amdgpu.scaled_mfma 16x16x128 ({{.*}}[0] * {{.*}}) * ({{.*}}[0] * {{.*}}196func.func @scaled_mfma_less_than_4(%opA: vector<32xf4E2M1FN>, %opB: vector<32xf4E2M1FN>, %scalesA: vector<2xf8E8M0FNU>, %scalesB: vector<2xf8E8M0FNU>) -> vector<4xf32> {197 %cst_0 = arith.constant dense<0.000000e+00> : vector<4xf32>198 %cst_1 = arith.constant dense<5.877470e-39> : vector<4xf8E8M0FNU>199 %scaleA = vector.extract %scalesA[0] : f8E8M0FNU from vector<2xf8E8M0FNU>200 %sA = vector.insert %scaleA, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>201 %scaleB = vector.extract %scalesB[1] : f8E8M0FNU from vector<2xf8E8M0FNU>202 %sB = vector.insert %scaleB, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>203 %res_0 = amdgpu.scaled_mfma 16x16x128 (%sA[0] * %opA) * (%sB[0] * %opB) + %cst_0 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>204 return %res_0 : vector<4xf32>205}206 207// -----208 209// CHECK-LABEL: func @scaled_mfma_ugly_shapes210// CHECK: amdgpu.scaled_mfma 16x16x128 (%{{.*}}[0] * %{{.*}}) * (%{{.*}}[3] * %arg1) + %cst : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>211// CHECK: amdgpu.scaled_mfma 16x16x128 (%{{.*}}[1] * %{{.*}}) * (%{{.*}}[3] * %arg1) + %cst : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>212// CHECK: amdgpu.scaled_mfma 16x16x128 (%{{.*}}[2] * %{{.*}}) * (%{{.*}}[2] * %arg1) + %cst : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>213// CHECK: amdgpu.scaled_mfma 16x16x128 (%{{.*}}[3] * %{{.*}}) * (%{{.*}}[1] * %arg1) + %cst : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>214func.func @scaled_mfma_ugly_shapes(%opA: vector<32xf4E2M1FN>, %opB: vector<32xf4E2M1FN>, %scalesA: vector<5x5xf8E8M0FNU>, %scalesB: vector<7x23xf8E8M0FNU>) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) {215 %cst_0 = arith.constant dense<0.000000e+00> : vector<4xf32>216 %cst_1 = arith.constant dense<5.877470e-39> : vector<4xf8E8M0FNU>217 %scaleA_0_4 = vector.extract %scalesA[4, 0] : f8E8M0FNU from vector<5x5xf8E8M0FNU>218 %scaleA_0_5 = vector.extract %scalesA[4, 1] : f8E8M0FNU from vector<5x5xf8E8M0FNU>219 %scaleA_0_6 = vector.extract %scalesA[4, 2] : f8E8M0FNU from vector<5x5xf8E8M0FNU>220 %scaleA_0_7 = vector.extract %scalesA[4, 3] : f8E8M0FNU from vector<5x5xf8E8M0FNU>221 222 // idx = 160 => opsel = 3 (last idx of last 4 bytes)223 %scaleB_6_22 = vector.extract %scalesB[6, 22] : f8E8M0FNU from vector<7x23xf8E8M0FNU>224 // idx = 159 => opsel = 3225 %scaleB_6_21 = vector.extract %scalesB[6, 21] : f8E8M0FNU from vector<7x23xf8E8M0FNU>226 // idx = 158 => opsel = 2227 %scaleB_6_20 = vector.extract %scalesB[6, 20] : f8E8M0FNU from vector<7x23xf8E8M0FNU>228 // idx = 157 => opsel = 1229 %scaleB_6_19 = vector.extract %scalesB[6, 19] : f8E8M0FNU from vector<7x23xf8E8M0FNU>230 231 %sA_0_4 = vector.insert %scaleA_0_4, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>232 %sA_0_5 = vector.insert %scaleA_0_5, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>233 %sA_0_6 = vector.insert %scaleA_0_6, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>234 %sA_0_7 = vector.insert %scaleA_0_7, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>235 236 %sB_6_22 = vector.insert %scaleB_6_22, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>237 %sB_6_21 = vector.insert %scaleB_6_21, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>238 %sB_6_20 = vector.insert %scaleB_6_20, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>239 %sB_6_19 = vector.insert %scaleB_6_19, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>240 241 %res_4 = amdgpu.scaled_mfma 16x16x128 (%sA_0_4[0] * %opA) * (%sB_6_22[0] * %opB) + %cst_0 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>242 %res_5 = amdgpu.scaled_mfma 16x16x128 (%sA_0_5[0] * %opA) * (%sB_6_21[0] * %opB) + %cst_0 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>243 %res_6 = amdgpu.scaled_mfma 16x16x128 (%sA_0_6[0] * %opA) * (%sB_6_20[0] * %opB) + %cst_0 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>244 %res_7 = amdgpu.scaled_mfma 16x16x128 (%sA_0_7[0] * %opA) * (%sB_6_19[0] * %opB) + %cst_0 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>245 return %res_4, %res_5, %res_6, %res_7 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>246}247