1105 lines · plain
1// RUN: mlir-opt -fold-memref-alias-ops -split-input-file %s | FileCheck %s2 3func.func @fold_static_stride_subview_with_load(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) -> f32 {4 %0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] : memref<12x32xf32> to memref<4x4xf32, strided<[64, 3], offset: ?>>5 %1 = memref.load %0[%arg3, %arg4] : memref<4x4xf32, strided<[64, 3], offset: ?>>6 return %1 : f327}8// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>9// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 3)>10// CHECK: func @fold_static_stride_subview_with_load11// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<12x32xf32>12// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index13// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index14// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: index15// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: index16// CHECK-DAG: %[[I1:.+]] = affine.apply #[[MAP0]]()[%[[ARG1]], %[[ARG3]]]17// CHECK-DAG: %[[I2:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[ARG4]]]18// CHECK: memref.load %[[ARG0]][%[[I1]], %[[I2]]]19 20// -----21 22func.func @fold_dynamic_stride_subview_with_load(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index) -> f32 {23 %0 = memref.subview %arg0[%arg1, %arg2][4, 4][%arg5, %arg6] :24 memref<12x32xf32> to memref<4x4xf32, strided<[?, ?], offset: ?>>25 %1 = memref.load %0[%arg3, %arg4] : memref<4x4xf32, strided<[?, ?], offset: ?>>26 return %1 : f3227}28// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 + s1 * s2)>29// CHECK: func @fold_dynamic_stride_subview_with_load30// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<12x32xf32>31// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index32// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index33// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: index34// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: index35// CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: index36// CHECK-SAME: %[[ARG6:[a-zA-Z0-9_]+]]: index37// CHECK-DAG: %[[I1:.+]] = affine.apply #[[MAP]]()[%[[ARG1]], %[[ARG3]], %[[ARG5]]]38// CHECK-DAG: %[[I2:.+]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[ARG4]], %[[ARG6]]]39// CHECK: memref.load %[[ARG0]][%[[I1]], %[[I2]]]40 41// -----42 43func.func @fold_static_stride_subview_with_store(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : f32) {44 %0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] :45 memref<12x32xf32> to memref<4x4xf32, strided<[64, 3], offset: ?>>46 memref.store %arg5, %0[%arg3, %arg4] : memref<4x4xf32, strided<[64, 3], offset: ?>>47 return48}49// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>50// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 3)>51// CHECK: func @fold_static_stride_subview_with_store52// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<12x32xf32>53// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index54// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index55// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: index56// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: index57// CHECK-DAG: %[[I1:.+]] = affine.apply #[[MAP0]]()[%[[ARG1]], %[[ARG3]]]58// CHECK-DAG: %[[I2:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[ARG4]]]59// CHECK: memref.store %{{.+}}, %[[ARG0]][%[[I1]], %[[I2]]]60 61// -----62 63func.func @fold_dynamic_stride_subview_with_store(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index, %arg7 : f32) {64 %0 = memref.subview %arg0[%arg1, %arg2][4, 4][%arg5, %arg6] :65 memref<12x32xf32> to memref<4x4xf32, strided<[?, ?], offset: ?>>66 memref.store %arg7, %0[%arg3, %arg4] : memref<4x4xf32, strided<[?, ?], offset: ?>>67 return68}69// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 + s1 * s2)>70// CHECK: func @fold_dynamic_stride_subview_with_store71// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<12x32xf32>72// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index73// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index74// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: index75// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: index76// CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: index77// CHECK-SAME: %[[ARG6:[a-zA-Z0-9_]+]]: index78// CHECK-DAG: %[[I1:.+]] = affine.apply #[[MAP]]()[%[[ARG1]], %[[ARG3]], %[[ARG5]]]79// CHECK-DAG: %[[I2:.+]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[ARG4]], %[[ARG6]]]80// CHECK: memref.store %{{.+}}, %[[ARG0]][%[[I1]], %[[I2]]]81 82// -----83 84func.func @fold_subview_with_transfer_read_0d(85 %arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index)86 -> vector<f32> {87 %f1 = arith.constant 1.0 : f3288 %0 = memref.subview %arg0[%arg1, %arg2][1, 1][1, 1] : memref<12x32xf32> to memref<f32, strided<[], offset: ?>>89 %1 = vector.transfer_read %0[], %f1 : memref<f32, strided<[], offset: ?>>, vector<f32>90 return %1 : vector<f32>91}92// CHECK: func @fold_subview_with_transfer_read_0d93// CHECK-SAME: %[[MEM:[a-zA-Z0-9_]+]]: memref<12x32xf32>94// CHECK-SAME: %[[SZ0:[a-zA-Z0-9_]+]]: index95// CHECK-SAME: %[[SZ1:[a-zA-Z0-9_]+]]: index96// CHECK-SAME: %[[ST1:[a-zA-Z0-9_]+]]: index97// CHECK: vector.transfer_read %[[MEM]][%[[SZ0]], %[[SZ1]]]98 99// -----100 101func.func @fold_subview_with_transfer_read(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index) -> vector<4xf32> {102 %f1 = arith.constant 1.0 : f32103 104 %0 = memref.subview %arg0[%arg1, %arg2][4, 4][%arg5, %arg6] : memref<12x32xf32> to memref<4x4xf32, strided<[?, ?], offset: ?>>105 %1 = vector.transfer_read %0[%arg3, %arg4], %f1 {in_bounds = [true]} : memref<4x4xf32, strided<[?, ?], offset: ?>>, vector<4xf32>106 return %1 : vector<4xf32>107}108// CHECK: func @fold_subview_with_transfer_read109// Can't fold this atm since we don't emit the proper vector.extract_strided_slice.110// CHECK: memref.subview111 112// -----113 114func.func @fold_static_stride_subview_with_transfer_write_0d(115 %arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index,116 %v : vector<f32>) {117 %f1 = arith.constant 1.0 : f32118 %0 = memref.subview %arg0[%arg1, %arg2][1, 1][1, 1] : memref<12x32xf32> to memref<f32, strided<[], offset: ?>>119 vector.transfer_write %v, %0[] {in_bounds = []} : vector<f32>, memref<f32, strided<[], offset: ?>>120 return121}122// CHECK: func @fold_static_stride_subview_with_transfer_write_0d123// CHECK-SAME: %[[MEM:[a-zA-Z0-9_]+]]: memref<12x32xf32>124// CHECK-SAME: %[[SZ0:[a-zA-Z0-9_]+]]: index125// CHECK-SAME: %[[SZ1:[a-zA-Z0-9_]+]]: index126// CHECK-SAME: %[[ST1:[a-zA-Z0-9_]+]]: index127// CHECK-SAME: %[[V:[a-zA-Z0-9_]+]]: vector<f32>128// CHECK: vector.transfer_write %[[V]], %[[MEM]][%[[SZ0]], %[[SZ1]]]129 130// -----131 132func.func @fold_static_stride_subview_with_transfer_write(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5: index, %arg6 : index, %arg7 : vector<4xf32>) {133 %0 = memref.subview %arg0[%arg1, %arg2][4, 4][%arg5, %arg6] :134 memref<12x32xf32> to memref<4x4xf32, strided<[?, ?], offset: ?>>135 vector.transfer_write %arg7, %0[%arg3, %arg4] {in_bounds = [true]} : vector<4xf32>, memref<4x4xf32, strided<[?, ?], offset: ?>>136 return137}138// CHECK: func @fold_static_stride_subview_with_transfer_write139// Can't fold this atm since we don't emit the proper vector.extract_strided_slice.140// CHECK: memref.subview141 142// -----143 144func.func @fold_rank_reducing_subview_with_load145 (%arg0 : memref<?x?x?x?x?x?xf32>, %arg1 : index, %arg2 : index,146 %arg3 : index, %arg4 : index, %arg5 : index, %arg6 : index,147 %arg7 : index, %arg8 : index, %arg9 : index, %arg10: index,148 %arg11 : index, %arg12 : index, %arg13 : index, %arg14: index,149 %arg15 : index, %arg16 : index) -> f32 {150 %0 = memref.subview %arg0[%arg1, %arg2, %arg3, %arg4, %arg5, %arg6][4, 1, 1, 4, 1, 1][%arg7, %arg8, %arg9, %arg10, %arg11, %arg12] : memref<?x?x?x?x?x?xf32> to memref<4x1x4x1xf32, strided<[?, ?, ?, ?], offset: ?>>151 %1 = memref.load %0[%arg13, %arg14, %arg15, %arg16] : memref<4x1x4x1xf32, strided<[?, ?, ?, ?], offset: ?>>152 return %1 : f32153}154// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 + s1 * s2)>155// CHECK: func @fold_rank_reducing_subview_with_load156// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?x?x?x?x?xf32>157// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index158// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index159// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: index160// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: index161// CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: index162// CHECK-SAME: %[[ARG6:[a-zA-Z0-9_]+]]: index163// CHECK-SAME: %[[ARG7:[a-zA-Z0-9_]+]]: index164// CHECK-SAME: %[[ARG8:[a-zA-Z0-9_]+]]: index165// CHECK-SAME: %[[ARG9:[a-zA-Z0-9_]+]]: index166// CHECK-SAME: %[[ARG10:[a-zA-Z0-9_]+]]: index167// CHECK-SAME: %[[ARG11:[a-zA-Z0-9_]+]]: index168// CHECK-SAME: %[[ARG12:[a-zA-Z0-9_]+]]: index169// CHECK-SAME: %[[ARG13:[a-zA-Z0-9_]+]]: index170// CHECK-SAME: %[[ARG14:[a-zA-Z0-9_]+]]: index171// CHECK-SAME: %[[ARG15:[a-zA-Z0-9_]+]]: index172// CHECK-SAME: %[[ARG16:[a-zA-Z0-9_]+]]: index173// CHECK-DAG: %[[I0:.+]] = affine.apply #[[MAP]]()[%[[ARG1]], %[[ARG13]], %[[ARG7]]]174// CHECK-DAG: %[[I2:.+]] = affine.apply #[[MAP]]()[%[[ARG3]], %[[ARG14]], %[[ARG9]]]175// CHECK-DAG: %[[I3:.+]] = affine.apply #[[MAP]]()[%[[ARG4]], %[[ARG15]], %[[ARG10]]]176// CHECK-DAG: %[[I4:.+]] = affine.apply #[[MAP]]()[%[[ARG5]], %[[ARG16]], %[[ARG11]]]177// CHECK: memref.load %[[ARG0]][%[[I0]], %[[ARG2]], %[[I2]], %[[I3]], %[[I4]], %[[ARG6]]]178 179// -----180 181func.func @fold_vector_transfer_read_with_rank_reduced_subview(182 %arg0 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>,183 %arg1: index, %arg2 : index, %arg3 : index, %arg4: index, %arg5 : index,184 %arg6 : index) -> vector<4xf32> {185 %cst = arith.constant 0.0 : f32186 %0 = memref.subview %arg0[0, %arg1, %arg2] [1, %arg3, %arg4] [1, 1, 1]187 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to188 memref<?x?xf32, strided<[?, ?], offset: ?>>189 %1 = vector.transfer_read %0[%arg5, %arg6], %cst {in_bounds = [true]}190 : memref<?x?xf32, strided<[?, ?], offset: ?>>, vector<4xf32>191 return %1 : vector<4xf32>192}193// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>194// CHECK: func @fold_vector_transfer_read_with_rank_reduced_subview195// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>196// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index197// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index198// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index199// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index200// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: index201// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: index202// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index203// CHECK-DAG: %[[IDX0:.+]] = affine.apply #[[MAP1]]()[%[[ARG1]], %[[ARG5]]]204// CHECK-DAG: %[[IDX1:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[ARG6]]]205// CHECK: vector.transfer_read %[[ARG0]][%[[C0]], %[[IDX0]], %[[IDX1]]], %{{.*}} : memref<?x?x?xf32206 207// -----208 209func.func @fold_vector_transfer_write_with_rank_reduced_subview(210 %arg0 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>,211 %arg1 : vector<4xf32>, %arg2: index, %arg3 : index, %arg4 : index,212 %arg5: index, %arg6 : index, %arg7 : index) {213 %cst = arith.constant 0.0 : f32214 %0 = memref.subview %arg0[0, %arg2, %arg3] [1, %arg4, %arg5] [1, 1, 1]215 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to216 memref<?x?xf32, strided<[?, ?], offset: ?>>217 vector.transfer_write %arg1, %0[%arg6, %arg7] {in_bounds = [true]}218 : vector<4xf32>, memref<?x?xf32, strided<[?, ?], offset: ?>>219 return220}221// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>222// CHECK: func @fold_vector_transfer_write_with_rank_reduced_subview223// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>224// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: vector<4xf32>225// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index226// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index227// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index228// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: index229// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: index230// CHECK-SAME: %[[ARG7:[a-zA-Z0-9]+]]: index231// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index232// CHECK-DAG: %[[IDX0:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[ARG6]]]233// CHECK-DAG: %[[IDX1:.+]] = affine.apply #[[MAP1]]()[%[[ARG3]], %[[ARG7]]]234// CHECK-DAG: vector.transfer_write %[[ARG1]], %[[ARG0]][%[[C0]], %[[IDX0]], %[[IDX1]]] {in_bounds = [true]} : vector<4xf32>, memref<?x?x?xf32235 236// -----237 238func.func @fold_vector_transfer_write_with_inner_rank_reduced_subview(239 %arg0 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>,240 %arg1 : vector<4xf32>, %arg2: index, %arg3 : index, %arg4 : index,241 %arg5: index, %arg6 : index, %arg7 : index) {242 %cst = arith.constant 0.0 : f32243 %0 = memref.subview %arg0[%arg2, %arg3, 0] [%arg4, %arg5, 1] [1, 1, 1]244 : memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>> to245 memref<?x?xf32, strided<[?, ?], offset: ?>>246 vector.transfer_write %arg1, %0[%arg6, %arg7] {in_bounds = [true]}247 : vector<4xf32>, memref<?x?xf32, strided<[?, ?], offset: ?>>248 return249}250// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>251// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1)>252// CHECK: func @fold_vector_transfer_write_with_inner_rank_reduced_subview253// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>254// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: vector<4xf32>255// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index256// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index257// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index258// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: index259// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: index260// CHECK-SAME: %[[ARG7:[a-zA-Z0-9]+]]: index261// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index262// CHECK-DAG: %[[IDX0:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[ARG6]]]263// CHECK-DAG: %[[IDX1:.+]] = affine.apply #[[MAP1]]()[%[[ARG3]], %[[ARG7]]]264// CHECK-DAG: vector.transfer_write %[[ARG1]], %[[ARG0]][%[[IDX0]], %[[IDX1]], %[[C0]]]265// CHECK-SAME: {in_bounds = [true], permutation_map = #[[MAP2]]} : vector<4xf32>, memref<?x?x?xf32266 267// -----268 269func.func @fold_masked_vector_transfer_read_with_subview(270 %arg0 : memref<?x?xf32, strided<[?, ?], offset: ?>>,271 %arg1: index, %arg2 : index, %arg3 : index, %arg4: index, %arg5 : index,272 %arg6 : index, %mask : vector<4xi1>) -> vector<4xf32> {273 %cst = arith.constant 0.0 : f32274 %0 = memref.subview %arg0[%arg1, %arg2] [%arg3, %arg4] [1, 1]275 : memref<?x?xf32, strided<[?, ?], offset: ?>> to276 memref<?x?xf32, strided<[?, ?], offset: ?>>277 %1 = vector.transfer_read %0[%arg5, %arg6], %cst, %mask {in_bounds = [true]}278 : memref<?x?xf32, strided<[?, ?], offset: ?>>, vector<4xf32>279 return %1 : vector<4xf32>280}281// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>282// CHECK: func @fold_masked_vector_transfer_read_with_subview283// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?xf32, strided<[?, ?], offset: ?>>284// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index285// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index286// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index287// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index288// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: index289// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: index290// CHECK-SAME: %[[MASK:[a-zA-Z0-9]+]]: vector<4xi1>291// CHECK-DAG: %[[IDX0:.+]] = affine.apply #[[MAP1]]()[%[[ARG1]], %[[ARG5]]]292// CHECK-DAG: %[[IDX1:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[ARG6]]]293// CHECK: vector.transfer_read %[[ARG0]][%[[IDX0]], %[[IDX1]]], %{{.*}}, %[[MASK]] {{.*}} : memref<?x?xf32294 295// -----296 297func.func @fold_masked_vector_transfer_read_with_rank_reducing_subview(298 %arg0 : memref<?x?x?x?xf32, strided<[?, ?, ?, ?], offset: ?>>,299 %arg1: index, %arg2 : index, %arg3 : index, %arg4: index, %arg5 : index,300 %arg6 : index, %mask : vector<4x3xi1>) -> vector<3x4xf32> {301 %cst = arith.constant 0.0 : f32302 %0 = memref.subview %arg0[0, %arg1, 0, %arg2] [1, %arg3, 1, %arg4] [1, 1, 1, 1]303 : memref<?x?x?x?xf32, strided<[?, ?, ?, ?], offset: ?>> to304 memref<?x?xf32, strided<[?, ?], offset: ?>>305 %1 = vector.transfer_read %0[%arg5, %arg6], %cst, %mask {306 permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds = [true, true]}307 : memref<?x?xf32, strided<[?, ?], offset: ?>>, vector<3x4xf32>308 return %1 : vector<3x4xf32>309}310// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>311// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d1)>312// CHECK: func @fold_masked_vector_transfer_read_with_rank_reducing_subview313// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32, strided<[?, ?, ?, ?], offset: ?>>314// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index315// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index316// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index317// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index318// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: index319// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: index320// CHECK-SAME: %[[MASK:[a-zA-Z0-9]+]]: vector<4x3xi1>321// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index322// CHECK-DAG: %[[PAD:.+]] = arith.constant 0.000000e+00 : f32323// CHECK-DAG: %[[IDX0:.+]] = affine.apply #[[MAP0]]()[%[[ARG1]], %[[ARG5]]]324// CHECK-DAG: %[[IDX1:.+]] = affine.apply #[[MAP0]]()[%[[ARG2]], %[[ARG6]]]325// CHECK: vector.transfer_read %[[ARG0]][%[[C0]], %[[IDX0]], %[[C0]], %[[IDX1]]], %[[PAD]], %[[MASK]] {{.*}} permutation_map = #[[MAP1]]} : memref<?x?x?x?xf32326 327// -----328 329func.func @fold_masked_vector_transfer_write_with_subview(330 %arg0 : memref<?x?xf32, strided<[?, ?], offset: ?>>,331 %arg1 : vector<4xf32>, %arg2: index, %arg3 : index, %arg4 : index,332 %arg5: index, %arg6 : index, %arg7 : index, %mask : vector<4xi1>) {333 %cst = arith.constant 0.0 : f32334 %0 = memref.subview %arg0[%arg2, %arg3] [%arg4, %arg5] [1, 1]335 : memref<?x?xf32, strided<[?, ?], offset: ?>> to336 memref<?x?xf32, strided<[?, ?], offset: ?>>337 vector.transfer_write %arg1, %0[%arg6, %arg7], %mask {in_bounds = [true]}338 : vector<4xf32>, memref<?x?xf32, strided<[?, ?], offset: ?>>339 return340}341// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>342// CHECK: func @fold_masked_vector_transfer_write_with_subview343// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?xf32, strided<[?, ?], offset: ?>>344// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: vector<4xf32>345// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index346// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index347// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index348// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: index349// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: index350// CHECK-SAME: %[[ARG7:[a-zA-Z0-9]+]]: index351// CHECK-SAME: %[[MASK:[a-zA-Z0-9]+]]: vector<4xi1>352// CHECK-DAG: %[[IDX0:.+]] = affine.apply #[[MAP1]]()[%[[ARG2]], %[[ARG6]]]353// CHECK-DAG: %[[IDX1:.+]] = affine.apply #[[MAP1]]()[%[[ARG3]], %[[ARG7]]]354// CHECK-DAG: vector.transfer_write %[[ARG1]], %[[ARG0]][%[[IDX0]], %[[IDX1]]], %[[MASK]] {in_bounds = [true]} : vector<4xf32>, memref<?x?xf32355 356// -----357 358func.func @fold_masked_vector_transfer_write_with_rank_reducing_subview(359 %arg0 : memref<?x?x?x?xf32, strided<[?, ?, ?, ?], offset: ?>>,360 %arg1 : vector<3x4xf32>, %arg2: index, %arg3 : index, %arg4 : index,361 %arg5: index, %arg6 : index, %arg7 : index, %mask : vector<4x3xi1>) {362 %cst = arith.constant 0.0 : f32363 %0 = memref.subview %arg0[0, %arg2, 0, %arg3] [1, %arg4, 1, %arg5] [1, 1, 1, 1]364 : memref<?x?x?x?xf32, strided<[?, ?, ?, ?], offset: ?>> to365 memref<?x?xf32, strided<[?, ?], offset: ?>>366 vector.transfer_write %arg1, %0[%arg6, %arg7], %mask {367 permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds = [true, true]}368 : vector<3x4xf32>, memref<?x?xf32, strided<[?, ?], offset: ?>>369 return370}371// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>372// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d1)>373// CHECK: func @fold_masked_vector_transfer_write_with_rank_reducing_subview374// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32, strided<[?, ?, ?, ?], offset: ?>>375// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: vector<3x4xf32>376// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index377// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index378// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index379// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: index380// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: index381// CHECK-SAME: %[[ARG7:[a-zA-Z0-9]+]]: index382// CHECK-SAME: %[[MASK:[a-zA-Z0-9]+]]: vector<4x3xi1>383// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index384// CHECK-DAG: %[[IDX0:.+]] = affine.apply #[[MAP0]]()[%[[ARG2]], %[[ARG6]]]385// CHECK-DAG: %[[IDX1:.+]] = affine.apply #[[MAP0]]()[%[[ARG3]], %[[ARG7]]]386// CHECK-DAG: vector.transfer_write %[[ARG1]], %[[ARG0]][%[[C0]], %[[IDX0]], %[[C0]], %[[IDX1]]], %[[ARG8]] {in_bounds = [true, true], permutation_map = #[[MAP1]]} : vector<3x4xf32>, memref<?x?x?x?xf32387 388// -----389 390// Test with affine.load/store ops. We only do a basic test here since the391// logic is identical to that with memref.load/store ops. The same affine.apply392// ops would be generated.393 394// CHECK-LABEL: func @fold_static_stride_subview_with_affine_load_store395func.func @fold_static_stride_subview_with_affine_load_store(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) -> f32 {396 %0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] : memref<12x32xf32> to memref<4x4xf32, strided<[64, 3], offset: ?>>397 %1 = affine.load %0[%arg3, %arg4] : memref<4x4xf32, strided<[64, 3], offset: ?>>398 // CHECK-NEXT: affine.apply399 // CHECK-NEXT: affine.apply400 // CHECK-NEXT: affine.load401 affine.store %1, %0[%arg3, %arg4] : memref<4x4xf32, strided<[64, 3], offset: ?>>402 // CHECK-NEXT: affine.apply403 // CHECK-NEXT: affine.apply404 // CHECK-NEXT: affine.store405 // CHECK-NEXT: return406 return %1 : f32407}408 409// -----410 411// CHECK-LABEL: fold_static_stride_subview_with_affine_load_store_expand_shape412// CHECK-SAME: (%[[ARG0:.*]]: memref<12x32xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index) -> f32 {413func.func @fold_static_stride_subview_with_affine_load_store_expand_shape(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index) -> f32 {414 %0 = memref.expand_shape %arg0 [[0, 1], [2]] output_shape [2, 6, 32] : memref<12x32xf32> into memref<2x6x32xf32>415 %1 = affine.load %0[%arg1, %arg2, %arg3] : memref<2x6x32xf32>416 return %1 : f32417}418// CHECK: %[[INDEX:.*]] = affine.linearize_index disjoint [%[[ARG1]], %[[ARG2]]] by (2, 6)419// CHECK-NEXT: %[[RESULT:.*]] = affine.load %[[ARG0]][%[[INDEX]], %[[ARG3]]] : memref<12x32xf32>420// CHECK-NEXT: return %[[RESULT]] : f32421 422// -----423 424// CHECK-LABEL: @fold_static_stride_subview_with_affine_load_store_collapse_shape425// CHECK-SAME: (%[[ARG0:.*]]: memref<2x6x32xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)426func.func @fold_static_stride_subview_with_affine_load_store_collapse_shape(%arg0 : memref<2x6x32xf32>, %arg1 : index, %arg2 : index) -> f32 {427 %0 = memref.collapse_shape %arg0 [[0, 1], [2]] : memref<2x6x32xf32> into memref<12x32xf32>428 %1 = affine.load %0[%arg1, %arg2] : memref<12x32xf32>429 return %1 : f32430}431// CHECK-NEXT: %[[MODIFIED_INDEXES:.*]]:2 = affine.delinearize_index %[[ARG1]] into (2, 6)432// CHECK-NEXT: %[[RESULT:.*]] = affine.load %[[ARG0]][%[[MODIFIED_INDEXES]]#0, %[[MODIFIED_INDEXES]]#1, %[[ARG2]]] : memref<2x6x32xf32>433// CHECK-NEXT: return %[[RESULT]] : f32434 435// -----436 437// CHECK-LABEL: @fold_dynamic_size_collapse_shape_with_affine_load438// CHECK-SAME: (%[[ARG0:.*]]: memref<?x6x32xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)439func.func @fold_dynamic_size_collapse_shape_with_affine_load(%arg0 : memref<?x6x32xf32>, %arg1 : index, %arg2 : index) -> f32 {440 %0 = memref.collapse_shape %arg0 [[0, 1], [2]] : memref<?x6x32xf32> into memref<?x32xf32>441 %1 = affine.load %0[%arg1, %arg2] : memref<?x32xf32>442 return %1 : f32443}444// CHECK-NEXT: %{{.*}}, %{{.*}}, %[[SIZES:.*]]:3, %{{.*}}:3 = memref.extract_strided_metadata %[[ARG0]]445// CHECK-NEXT: %[[MODIFIED_INDEXES:.*]]:2 = affine.delinearize_index %[[ARG1]] into (%[[SIZES]]#0, 6)446// CHECK-NEXT: %[[RESULT:.*]] = affine.load %[[ARG0]][%[[MODIFIED_INDEXES]]#0, %[[MODIFIED_INDEXES]]#1, %[[ARG2]]] : memref<?x6x32xf32>447// CHECK-NEXT: return %[[RESULT]] : f32448 449// -----450 451// CHECK-LABEL: @fold_fully_dynamic_size_collapse_shape_with_affine_load452// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)453func.func @fold_fully_dynamic_size_collapse_shape_with_affine_load(%arg0 : memref<?x?x?xf32>, %arg1 : index, %arg2 : index) -> f32 {454 %0 = memref.collapse_shape %arg0 [[0, 1], [2]] : memref<?x?x?xf32> into memref<?x?xf32>455 %1 = affine.load %0[%arg1, %arg2] : memref<?x?xf32>456 return %1 : f32457}458// CHECK-NEXT: %{{.*}}, %{{.*}}, %[[SIZES:.*]]:3, %{{.*}}:3 = memref.extract_strided_metadata %[[ARG0]]459// CHECK-NEXT: %[[MODIFIED_INDEXES:.*]]:2 = affine.delinearize_index %[[ARG1]] into (%[[SIZES]]#0, %[[SIZES]]#1)460// CHECK-NEXT: %[[RESULT:.*]] = affine.load %[[ARG0]][%[[MODIFIED_INDEXES]]#0, %[[MODIFIED_INDEXES]]#1, %[[ARG2]]] : memref<?x?x?xf32>461// CHECK-NEXT: return %[[RESULT]] : f32462 463 464// -----465 466// CHECK-LABEL: fold_static_stride_subview_with_affine_load_store_expand_shape_3d467// CHECK-SAME: (%[[ARG0:.*]]: memref<12x32xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index) -> f32 {468func.func @fold_static_stride_subview_with_affine_load_store_expand_shape_3d(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4: index) -> f32 {469 %0 = memref.expand_shape %arg0 [[0, 1, 2], [3]] output_shape [2, 2, 3, 32] : memref<12x32xf32> into memref<2x2x3x32xf32>470 %1 = affine.load %0[%arg1, %arg2, %arg3, %arg4] : memref<2x2x3x32xf32>471 return %1 : f32472}473// CHECK: %[[INDEX:.*]] = affine.linearize_index disjoint [%[[ARG1]], %[[ARG2]], %[[ARG3]]] by (2, 2, 3)474// CHECK-NEXT: %[[RESULT:.*]] = affine.load %[[ARG0]][%[[INDEX]], %[[ARG4]]] : memref<12x32xf32>475// CHECK-NEXT: return %[[RESULT]] : f32476 477// -----478 479// CHECK-LABEL: fold_dynamic_subview_with_memref_load_expand_shape480// CHECK-SAME: (%[[ARG0:.*]]: memref<16x?xf32, strided<[16, 1]>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index) -> f32481func.func @fold_dynamic_subview_with_memref_load_expand_shape(%arg0 : memref<16x?xf32, strided<[16, 1]>>, %arg1 : index, %arg2 : index, %sz0: index) -> f32 {482 %c0 = arith.constant 0 : index483 %expand_shape = memref.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [1, 16, %sz0, 1] : memref<16x?xf32, strided<[16, 1]>> into memref<1x16x?x1xf32, strided<[256, 16, 1, 1]>>484 %0 = memref.load %expand_shape[%c0, %arg1, %arg2, %c0] {nontemporal = true} : memref<1x16x?x1xf32, strided<[256, 16, 1, 1]>>485 return %0 : f32486}487// CHECK-NEXT: %[[C0:.*]] = arith.constant 0488// CHECK-NEXT: %[[INDEX1:.*]] = affine.linearize_index disjoint [%[[C0]], %[[ARG1]]] by (1, 16)489// CHECK-NEXT: %[[INDEX2:.*]] = affine.linearize_index disjoint [%[[ARG2]], %[[C0]]] by (%[[ARG3]], 1)490// CHECK-NEXT: %[[VAL1:.*]] = memref.load %[[ARG0]][%[[INDEX1]], %[[INDEX2]]] {nontemporal = true} : memref<16x?xf32, strided<[16, 1]>>491// CHECK-NEXT: return %[[VAL1]] : f32492 493// -----494 495// CHECK-LABEL: fold_dynamic_subview_with_memref_store_expand_shape496// CHECK-SAME: (%[[ARG0:.*]]: memref<16x?xf32, strided<[16, 1]>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)497func.func @fold_dynamic_subview_with_memref_store_expand_shape(%arg0 : memref<16x?xf32, strided<[16, 1]>>, %arg1 : index, %arg2 : index, %sz0 : index) {498 %c0 = arith.constant 0 : index499 %c1f32 = arith.constant 1.0 : f32500 %expand_shape = memref.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [1, 16, %sz0, 1] : memref<16x?xf32, strided<[16, 1]>> into memref<1x16x?x1xf32, strided<[256, 16, 1, 1]>>501 memref.store %c1f32, %expand_shape[%c0, %arg1, %arg2, %c0] {nontemporal = true} : memref<1x16x?x1xf32, strided<[256, 16, 1, 1]>>502 return503}504// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index505// CHECK-DAG: %[[C1F32:.*]] = arith.constant 1.000000e+00 : f32506// CHECK-NEXT: %[[INDEX1:.*]] = affine.linearize_index disjoint [%[[C0]], %[[ARG1]]] by (1, 16)507// CHECK-NEXT: %[[INDEX2:.*]] = affine.linearize_index disjoint [%[[ARG2]], %[[C0]]] by (%[[ARG3]], 1)508// CHECK-NEXT: memref.store %[[C1F32]], %[[ARG0]][%[[INDEX1]], %[[INDEX2]]] {nontemporal = true} : memref<16x?xf32, strided<[16, 1]>>509// CHECK-NEXT: return510 511// -----512 513// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0)[s0] -> (d0 + s0)>514// CHECK-LABEL: fold_memref_alias_expand_shape_subview_load_store_dynamic_dim515// CHECK-SAME: (%[[ARG0:.*]]: memref<2048x16xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index)516func.func @fold_memref_alias_expand_shape_subview_load_store_dynamic_dim(%alloc: memref<2048x16xf32>, %c10: index, %c5: index, %c0: index, %sz0: index) {517 %subview = memref.subview %alloc[%c5, 0] [%c10, 16] [1, 1] : memref<2048x16xf32> to memref<?x16xf32, strided<[16, 1], offset: ?>>518 %expand_shape = memref.expand_shape %subview [[0], [1, 2, 3]] output_shape [%sz0, 1, 8, 2] : memref<?x16xf32, strided<[16, 1], offset: ?>> into memref<?x1x8x2xf32, strided<[16, 16, 2, 1], offset: ?>>519 %dim = memref.dim %expand_shape, %c0 : memref<?x1x8x2xf32, strided<[16, 16, 2, 1], offset: ?>>520 521 affine.for %arg6 = 0 to %dim step 64 {522 affine.for %arg7 = 0 to 16 step 16 {523 %dummy_load = affine.load %expand_shape[%arg6, 0, %arg7, %arg7] : memref<?x1x8x2xf32, strided<[16, 16, 2, 1], offset: ?>>524 affine.store %dummy_load, %subview[%arg6, %arg7] : memref<?x16xf32, strided<[16, 1], offset: ?>>525 }526 }527 return528}529// CHECK-NEXT: %[[C0:.*]] = arith.constant 0530// CHECK-NEXT: memref.subview531// CHECK-NEXT: %[[EXPAND_SHAPE:.*]] = memref.expand_shape532// CHECK-NEXT: %[[DIM:.*]] = memref.dim %[[EXPAND_SHAPE]], %[[ARG3]] : memref<?x1x8x2xf32, strided<[16, 16, 2, 1], offset: ?>>533// CHECK-NEXT: affine.for %[[ARG5:.*]] = 0 to %[[DIM]] step 64 {534// CHECK-NEXT: affine.for %[[ARG6:.*]] = 0 to 16 step 16 {535// CHECK-NEXT: %[[VAL0:.*]] = affine.linearize_index disjoint [%[[C0]], %[[ARG6]], %[[ARG6]]] by (1, 8, 2)536// CHECK-NEXT: %[[VAL1:.*]] = affine.apply #[[$MAP0]](%[[ARG5]])[%[[ARG2]]]537// CHECK-NEXT: %[[VAL2:.*]] = affine.load %[[ARG0]][%[[VAL1]], %[[VAL0]]] : memref<2048x16xf32>538// CHECK-NEXT: %[[VAL3:.*]] = affine.apply #[[$MAP0]](%[[ARG5]])[%[[ARG2]]]539// CHECK-NEXT: affine.store %[[VAL2]], %[[ARG0]][%[[VAL3]], %[[ARG6]]] : memref<2048x16xf32>540 541// -----542 543// CHECK-LABEL: fold_static_stride_subview_with_affine_load_store_expand_shape544// CHECK-SAME: (%[[ARG0:.*]]: memref<1024x1024xf32>, %[[ARG1:.*]]: memref<1xf32>, %[[ARG2:.*]]: index)545func.func @fold_static_stride_subview_with_affine_load_store_expand_shape(%arg0: memref<1024x1024xf32>, %arg1: memref<1xf32>, %arg2: index) -> f32 {546 %0 = memref.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [1, 1024, 1024, 1] : memref<1024x1024xf32> into memref<1x1024x1024x1xf32>547 affine.for %arg3 = 0 to 1 {548 affine.for %arg4 = 0 to 1024 {549 affine.for %arg5 = 0 to 1020 {550 affine.for %arg6 = 0 to 1 {551 %1 = affine.load %0[%arg3, %arg4, %arg5, %arg6] : memref<1x1024x1024x1xf32>552 affine.store %1, %arg1[%arg2] : memref<1xf32>553 }554 }555 }556 }557 %2 = affine.load %arg1[%arg2] : memref<1xf32>558 return %2 : f32559}560// CHECK-NEXT: affine.for %[[ARG3:.*]] = 0 to 1 {561// CHECK-NEXT: affine.for %[[ARG4:.*]] = 0 to 1024 {562// CHECK-NEXT: affine.for %[[ARG5:.*]] = 0 to 1020 {563// CHECK-NEXT: affine.for %[[ARG6:.*]] = 0 to 1 {564// CHECK-NEXT: %[[IDX1:.*]] = affine.linearize_index disjoint [%[[ARG3]], %[[ARG4]]] by (1, 1024)565// CHECK-NEXT: %[[IDX2:.*]] = affine.linearize_index disjoint [%[[ARG5]], %[[ARG6]]] by (1024, 1)566// CHECK-NEXT: affine.load %[[ARG0]][%[[IDX1]], %[[IDX2]]] : memref<1024x1024xf32>567 568// -----569 570// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0 + d1)>571// CHECK-LABEL: fold_static_stride_subview_with_affine_load_store_expand_shape_when_access_index_is_an_expression572// CHECK-SAME: (%[[ARG0:.*]]: memref<1024x1024xf32>, %[[ARG1:.*]]: memref<1xf32>, %[[ARG2:.*]]: index)573func.func @fold_static_stride_subview_with_affine_load_store_expand_shape_when_access_index_is_an_expression(%arg0: memref<1024x1024xf32>, %arg1: memref<1xf32>, %arg2: index) -> f32 {574 %0 = memref.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [1, 1024, 1024, 1] : memref<1024x1024xf32> into memref<1x1024x1024x1xf32>575 affine.for %arg3 = 0 to 1 {576 affine.for %arg4 = 0 to 1024 {577 affine.for %arg5 = 0 to 1020 {578 affine.for %arg6 = 0 to 1 {579 %1 = affine.load %0[%arg3, %arg4 + %arg3, %arg5, %arg6] : memref<1x1024x1024x1xf32>580 affine.store %1, %arg1[%arg2] : memref<1xf32>581 }582 }583 }584 }585 %2 = affine.load %arg1[%arg2] : memref<1xf32>586 return %2 : f32587}588// CHECK-NEXT: affine.for %[[ARG3:.*]] = 0 to 1 {589// CHECK-NEXT: affine.for %[[ARG4:.*]] = 0 to 1024 {590// CHECK-NEXT: affine.for %[[ARG5:.*]] = 0 to 1020 {591// CHECK-NEXT: affine.for %[[ARG6:.*]] = 0 to 1 {592// CHECK-NEXT: %[[TMP1:.*]] = affine.apply #[[$MAP0]](%[[ARG3]], %[[ARG4]])593// CHECK-NEXT: %[[TMP2:.*]] = affine.linearize_index disjoint [%[[ARG3]], %[[TMP1]]] by (1, 1024)594// CHECK-NEXT: %[[TMP3:.*]] = affine.linearize_index disjoint [%[[ARG5]], %[[ARG6]]] by (1024, 1)595// CHECK-NEXT: affine.load %[[ARG0]][%[[TMP2]], %[[TMP3]]] : memref<1024x1024xf32>596 597// -----598 599// CHECK-LABEL: fold_static_stride_subview_with_affine_load_store_expand_shape_with_constant_access_index600// CHECK-SAME: (%[[ARG0:.*]]: memref<1024x1024xf32>, %[[ARG1:.*]]: memref<1xf32>, %[[ARG2:.*]]: index)601func.func @fold_static_stride_subview_with_affine_load_store_expand_shape_with_constant_access_index(%arg0: memref<1024x1024xf32>, %arg1: memref<1xf32>, %arg2: index) -> f32 {602 %0 = memref.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [1, 1024, 1024, 1] : memref<1024x1024xf32> into memref<1x1024x1024x1xf32>603 %cst = arith.constant 0 : index604 affine.for %arg3 = 0 to 1 {605 affine.for %arg4 = 0 to 1024 {606 affine.for %arg5 = 0 to 1020 {607 affine.for %arg6 = 0 to 1 {608 %1 = memref.load %0[%arg3, %cst, %arg5, %arg6] : memref<1x1024x1024x1xf32>609 memref.store %1, %arg1[%arg2] : memref<1xf32>610 }611 }612 }613 }614 %2 = memref.load %arg1[%arg2] : memref<1xf32>615 return %2 : f32616}617// CHECK-NEXT: %[[C0:.*]] = arith.constant 0618// CHECK-NEXT: affine.for %[[ARG3:.*]] = 0 to 1 {619// CHECK-NEXT: affine.for %[[ARG4:.*]] = 0 to 1024 {620// CHECK-NEXT: affine.for %[[ARG5:.*]] = 0 to 1020 {621// CHECK-NEXT: affine.for %[[ARG6:.*]] = 0 to 1 {622// CHECK-NEXT: %[[TMP1:.*]] = affine.linearize_index disjoint [%[[ARG3]], %[[C0]]] by (1, 1024)623// CHECK-NEXT: %[[TMP2:.*]] = affine.linearize_index disjoint [%[[ARG5]], %[[ARG6]]] by (1024, 1)624// CHECK-NEXT: memref.load %[[ARG0]][%[[TMP1]], %[[TMP2]]] : memref<1024x1024xf32>625 626// -----627 628// CHECK-LABEL: fold_static_stride_subview_with_affine_load_store_collapse_shape_with_0d_result629// CHECK-SAME: (%[[ARG0:.*]]: memref<1xf32>, %[[ARG1:.*]]: memref<1xf32>)630func.func @fold_static_stride_subview_with_affine_load_store_collapse_shape_with_0d_result(%arg0: memref<1xf32>, %arg1: memref<1xf32>) -> memref<1xf32> {631 %0 = memref.collapse_shape %arg0 [] : memref<1xf32> into memref<f32>632 affine.for %arg2 = 0 to 3 {633 %1 = affine.load %0[] : memref<f32>634 affine.store %1, %arg1[0] : memref<1xf32>635 }636 return %arg1 : memref<1xf32>637}638// CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0 : index639// CHECK-NEXT: affine.for %{{.*}} = 0 to 3 {640// CHECK-NEXT: affine.load %[[ARG0]][%[[ZERO]]] : memref<1xf32>641 642// -----643 644// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 + 2)>645// CHECK-LABEL: func @subview_of_subview(646// CHECK-SAME: %[[m:.*]]: memref<8x1024xf32, 3>, %[[pos:.*]]: index647// CHECK: %[[add:.*]] = affine.apply #[[$map]]()[%arg1]648// CHECK: memref.subview %arg0[4, %[[add]]] [1, 1] [1, 1] : memref<8x1024xf32, 3> to memref<f32, strided<[], offset: ?>, 3>649func.func @subview_of_subview(%m: memref<8x1024xf32, 3>, %pos: index)650 -> memref<f32, strided<[], offset: ?>, 3>651{652 %0 = memref.subview %m[3, %pos] [5, 7] [1, 1]653 : memref<8x1024xf32, 3>654 to memref<5x7xf32, strided<[1024, 1], offset: ?>, 3>655 %1 = memref.subview %0[1, 2] [1, 1] [1, 1]656 : memref<5x7xf32, strided<[1024, 1], offset: ?>, 3>657 to memref<f32, strided<[], offset: ?>, 3>658 return %1 : memref<f32, strided<[], offset: ?>, 3>659}660 661// -----662 663// CHECK-LABEL: func @subview_of_subview_rank_reducing(664// CHECK-SAME: %[[m:.*]]: memref<?x?x?xf32>665// CHECK: memref.subview %arg0[3, 7, 8] [1, 1, 1] [1, 1, 1] : memref<?x?x?xf32> to memref<f32, strided<[], offset: ?>>666func.func @subview_of_subview_rank_reducing(%m: memref<?x?x?xf32>,667 %sz: index, %pos: index)668 -> memref<f32, strided<[], offset: ?>>669{670 %0 = memref.subview %m[3, 1, 8] [1, %sz, 1] [1, 1, 1]671 : memref<?x?x?xf32>672 to memref<?xf32, strided<[?], offset: ?>>673 %1 = memref.subview %0[6] [1] [1]674 : memref<?xf32, strided<[?], offset: ?>>675 to memref<f32, strided<[], offset: ?>>676 return %1 : memref<f32, strided<[], offset: ?>>677}678 679// -----680 681// CHECK-LABEL: func @fold_load_keep_nontemporal(682// CHECK: memref.load %{{.+}}[%{{.+}}, %{{.+}}] {nontemporal = true}683func.func @fold_load_keep_nontemporal(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) -> f32 {684 %0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] : memref<12x32xf32> to memref<4x4xf32, strided<[64, 3], offset: ?>>685 %1 = memref.load %0[%arg3, %arg4] {nontemporal = true }: memref<4x4xf32, strided<[64, 3], offset: ?>>686 return %1 : f32687}688 689// -----690 691// CHECK-LABEL: func @fold_store_keep_nontemporal(692// CHECK: memref.store %{{.+}}, %{{.+}}[%{{.+}}, %{{.+}}] {nontemporal = true} : memref<12x32xf32>693func.func @fold_store_keep_nontemporal(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : f32) {694 %0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] :695 memref<12x32xf32> to memref<4x4xf32, strided<[64, 3], offset: ?>>696 memref.store %arg5, %0[%arg3, %arg4] {nontemporal=true}: memref<4x4xf32, strided<[64, 3], offset: ?>>697 return698}699 700// -----701 702func.func @fold_gpu_subgroup_mma_load_matrix_1d(%src: memref<?xvector<4xf32>>, %offset: index, %i: index) -> !gpu.mma_matrix<16x16xf16, "COp"> {703 %subview = memref.subview %src[%offset] [81920] [1] : memref<?xvector<4xf32>> to memref<81920xvector<4xf32>, strided<[1], offset: ?>>704 %matrix = gpu.subgroup_mma_load_matrix %subview[%i] {leadDimension = 160 : index} : memref<81920xvector<4xf32>, strided<[1], offset: ?>> -> !gpu.mma_matrix<16x16xf16, "COp">705 return %matrix: !gpu.mma_matrix<16x16xf16, "COp">706}707 708// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>709// CHECK: func.func @fold_gpu_subgroup_mma_load_matrix_1d710// CHECK-SAME: (%[[SRC:.+]]: memref<?xvector<4xf32>>, %[[OFFSET:.+]]: index, %[[I:.+]]: index)711// CHECK: %[[APPLY:.+]] = affine.apply #[[MAP]]()[%[[OFFSET]], %[[I]]]712// CHECK: %[[LOAD:.+]] = gpu.subgroup_mma_load_matrix %[[SRC]][%[[APPLY]]] {leadDimension = 160 : index} : memref<?xvector<4xf32>> -> !gpu.mma_matrix<16x16xf16, "COp">713// CHECK: return %[[LOAD]]714 715// -----716 717func.func @fold_gpu_subgroup_mma_store_matrix_1d(%dst: memref<?xvector<4xf32>>, %offset: index, %i: index, %matrix: !gpu.mma_matrix<16x16xf16, "COp">) {718 %subview = memref.subview %dst[%offset] [81920] [1] : memref<?xvector<4xf32>> to memref<81920xvector<4xf32>, strided<[1], offset: ?>>719 gpu.subgroup_mma_store_matrix %matrix, %subview[%i] {leadDimension = 160 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<81920xvector<4xf32>, strided<[1], offset: ?>>720 return721}722 723// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>724// CHECK: func.func @fold_gpu_subgroup_mma_store_matrix_1d725// CHECK-SAME: (%[[DST:.+]]: memref<?xvector<4xf32>>, %[[OFFSET:.+]]: index, %[[I0:.+]]: index, %[[VAL:.+]]: !gpu.mma_matrix<16x16xf16, "COp">)726// CHECK: %[[APPLY:.+]] = affine.apply #[[MAP]]()[%[[OFFSET]], %[[I0]]]727// CHECK: gpu.subgroup_mma_store_matrix %[[VAL]], %[[DST]][%[[APPLY]]] {leadDimension = 160 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<?xvector<4xf32>>728 729// -----730 731// CHECK-LABEL: func.func @fold_gpu_subgroup_mma_load_matrix_2d732// CHECK-SAME: %[[SRC:.+]]: memref<128x128xf32>733func.func @fold_gpu_subgroup_mma_load_matrix_2d(%arg0 : memref<128x128xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) -> !gpu.mma_matrix<16x16xf16, "COp"> {734 %subview = memref.subview %arg0[%arg1, %arg2][64, 32][2, 1] : memref<128x128xf32> to memref<64x32xf32, strided<[256, 1], offset: ?>>735 // CHECK: gpu.subgroup_mma_load_matrix %[[SRC]][{{.+}}] {leadDimension = 32 : index} : memref<128x128xf32> -> !gpu.mma_matrix<16x16xf16, "COp">736 %matrix = gpu.subgroup_mma_load_matrix %subview[%arg3, %arg4] {leadDimension = 32 : index} : memref<64x32xf32, strided<[256, 1], offset: ?>> -> !gpu.mma_matrix<16x16xf16, "COp">737 return %matrix : !gpu.mma_matrix<16x16xf16, "COp">738}739 740// -----741 742// CHECK-LABEL: func.func @fold_gpu_subgroup_mma_load_matrix_2d743// CHECK-SAME: %[[DST:.+]]: memref<128x128xf32>744func.func @fold_gpu_subgroup_mma_load_matrix_2d(%arg0 : memref<128x128xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %matrix: !gpu.mma_matrix<16x16xf16, "COp">) {745 %subview = memref.subview %arg0[%arg1, %arg2][64, 32][2, 1] : memref<128x128xf32> to memref<64x32xf32, strided<[256, 1], offset: ?>>746 // CHECK: gpu.subgroup_mma_store_matrix %{{.+}}, %[[DST]][{{.+}}] {leadDimension = 32 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<128x128xf32>747 gpu.subgroup_mma_store_matrix %matrix, %subview[%arg3, %arg4] {leadDimension = 32 : index} : !gpu.mma_matrix<16x16xf16, "COp">, memref<64x32xf32, strided<[256, 1], offset: ?>>748 return749}750 751// -----752 753 754func.func @fold_nvgpu_device_async_copy_zero_sub_idx(%gmem_memref_3d : memref<2x128x768xf16>, %idx_1 : index, %idx_2 : index, %idx_3 : index) {755 756 %c0 = arith.constant 0 : index757 %smem_memref_4d = memref.alloc() : memref<5x1x64x64xf16, #gpu.address_space<workgroup>>758 %gmem_memref_subview_2d = memref.subview %gmem_memref_3d[%idx_1, %idx_2, %idx_3] [1, 1, 8] [1, 1, 1] : memref<2x128x768xf16> to memref<1x8xf16, strided<[98304, 1], offset: ?>>759 %async_token = nvgpu.device_async_copy %gmem_memref_subview_2d[%c0, %c0], %smem_memref_4d[%c0, %c0, %c0, %c0], 8 {bypassL1} : memref<1x8xf16, strided<[98304, 1], offset: ?>> to memref<5x1x64x64xf16, #gpu.address_space<workgroup>>760 return761}762 763// CHECK-LABEL: func.func @fold_nvgpu_device_async_copy_zero_sub_idx764// CHECK-SAME: (%[[GMEM_MEMREF_3d:.+]]: memref<2x128x768xf16>, %[[IDX_1:.+]]: index, %[[IDX_2:.+]]: index, %[[IDX_3:.+]]: index)765// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index766// CHECK-DAG: %[[SMEM_MEMREF_4d:.+]] = memref.alloc() : memref<5x1x64x64xf16, #gpu.address_space<workgroup>>767// CHECK: nvgpu.device_async_copy %[[GMEM_MEMREF_3d]][%[[IDX_1]], %[[IDX_2]], %[[IDX_3]]], %[[SMEM_MEMREF_4d]][%[[c0]], %[[c0]], %[[c0]], %[[c0]]], 8 {bypassL1} : memref<2x128x768xf16> to memref<5x1x64x64xf16, #gpu.address_space<workgroup>>768 769// -----770 771 772func.func @fold_src_nvgpu_device_async_copy(%gmem_memref_3d : memref<2x128x768xf16>, %src_idx_0 : index, %src_idx_1 : index, %src_idx_2 : index, %src_sub_idx_0 : index, %src_sub_idx_1 : index) {773 %c0 = arith.constant 0 : index774 %smem_memref_4d = memref.alloc() : memref<5x1x64x64xf16, #gpu.address_space<workgroup>>775 %gmem_memref_subview_2d = memref.subview %gmem_memref_3d[%src_idx_0, %src_idx_1, %src_idx_2] [1, 1, 8] [1, 1, 1] : memref<2x128x768xf16> to memref<1x8xf16, strided<[98304, 1], offset: ?>>776 %async_token = nvgpu.device_async_copy %gmem_memref_subview_2d[%src_sub_idx_0, %src_sub_idx_1], %smem_memref_4d[%c0, %c0, %c0, %c0], 8 {bypassL1} : memref<1x8xf16, strided<[98304, 1], offset: ?>> to memref<5x1x64x64xf16, #gpu.address_space<workgroup>>777 return778}779 780// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>781// CHECK: func.func @fold_src_nvgpu_device_async_copy782// CHECK-SAME: (%[[GMEM_MEMREF_3d:.+]]: memref<2x128x768xf16>, %[[SRC_IDX_0:.+]]: index, %[[SRC_IDX_1:.+]]: index, %[[SRC_IDX_2:.+]]: index, %[[SRC_SUB_IDX_0:.+]]: index, %[[SRC_SUB_IDX_1:.+]]: index)783// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index784// CHECK-DAG: %[[RESOLVED_SRC_IDX_0:.+]] = affine.apply #[[MAP]]()[%[[SRC_IDX_0]], %[[SRC_SUB_IDX_0]]]785// CHECK-DAG: %[[RESOLVED_SRC_IDX_1:.+]] = affine.apply #[[MAP]]()[%[[SRC_IDX_2]], %[[SRC_SUB_IDX_1]]]786// CHECK-DAG: nvgpu.device_async_copy %[[GMEM_MEMREF_3d]][%[[RESOLVED_SRC_IDX_0]], %[[SRC_IDX_1]], %[[RESOLVED_SRC_IDX_1]]], %[[SMEM_MEMREF_4d]][%[[c0]], %[[c0]], %[[c0]], %[[c0]]], 8 {bypassL1} : memref<2x128x768xf16> to memref<5x1x64x64xf16, #gpu.address_space<workgroup>>787 788// -----789 790 791func.func @fold_src_fold_dest_nvgpu_device_async_copy(%gmem_memref_3d : memref<2x128x768xf16>, %src_idx_0 : index, %src_idx_1 : index, %src_idx_2 : index, %src_sub_idx_0 : index, %src_sub_idx_1 : index, %dest_idx_0 : index, %dest_idx_1 : index, %dest_idx_2 : index, %dest_idx_3 : index, %dest_sub_idx_0 : index, %dest_sub_idx_1 : index) {792 %c0 = arith.constant 0 : index793 %smem_memref_4d = memref.alloc() : memref<5x1x64x64xf16, #gpu.address_space<workgroup>>794 %gmem_memref_subview_2d = memref.subview %gmem_memref_3d[%src_idx_0, %src_idx_1, %src_idx_2] [1, 1, 8] [1, 1, 1] : memref<2x128x768xf16> to memref<1x8xf16, strided<[98304, 1], offset: ?>>795 %smem_memref_2d = memref.subview %smem_memref_4d[%dest_idx_0, %dest_idx_1, %dest_idx_2, %dest_idx_3] [1, 1, 1, 8] [1, 1, 1, 1] : memref<5x1x64x64xf16, #gpu.address_space<workgroup>> to memref<1x8xf16, strided<[4096, 1], offset: ?>, #gpu.address_space<workgroup>>796 %async_token = nvgpu.device_async_copy %gmem_memref_subview_2d[%src_sub_idx_0, %src_sub_idx_1], %smem_memref_2d[%dest_sub_idx_0, %dest_sub_idx_1], 8 {bypassL1} : memref<1x8xf16, strided<[98304, 1], offset: ?>> to memref<1x8xf16, strided<[4096, 1], offset: ?>, #gpu.address_space<workgroup>>797 return798}799 800// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>801// CHECK: func.func @fold_src_fold_dest_nvgpu_device_async_copy802// CHECK-SAME: (%[[GMEM_MEMREF_3d:.+]]: memref<2x128x768xf16>, %[[SRC_IDX_0:.+]]: index, %[[SRC_IDX_1:.+]]: index, %[[SRC_IDX_2:.+]]: index, %[[SRC_SUB_IDX_0:.+]]: index, %[[SRC_SUB_IDX_1:.+]]: index, %[[DEST_IDX_0:.+]]: index, %[[DEST_IDX_1:.+]]: index, %[[DEST_IDX_2:.+]]: index, %[[DEST_IDX_3:.+]]: index, %[[DEST_SUB_IDX_0:.+]]: index, %[[DEST_SUB_IDX_1:.+]]: index)803// CHECK-DAG: %[[RESOLVED_SRC_IDX_0:.+]] = affine.apply #[[MAP]]()[%[[SRC_IDX_0]], %[[SRC_SUB_IDX_0]]]804// CHECK-DAG: %[[RESOLVED_SRC_IDX_1:.+]] = affine.apply #[[MAP]]()[%[[SRC_IDX_2]], %[[SRC_SUB_IDX_1]]]805// CHECK-DAG: %[[RESOLVED_DST_IDX_1:.+]] = affine.apply #[[MAP]]()[%[[DEST_IDX_1]], %[[DEST_SUB_IDX_0]]]806// CHECK-DAG: %[[RESOLVED_DST_IDX_3:.+]] = affine.apply #[[MAP]]()[%[[DEST_IDX_3]], %[[DEST_SUB_IDX_1]]]807// CHECK-DAG: nvgpu.device_async_copy %[[GMEM_MEMREF_3d]][%[[RESOLVED_SRC_IDX_0]], %[[SRC_IDX_1]], %[[RESOLVED_SRC_IDX_1]]], %[[SMEM_MEMREF_4d]][%[[DEST_IDX_0]], %[[RESOLVED_DST_IDX_1]], %[[DEST_IDX_2]], %[[RESOLVED_DST_IDX_3]]], 8 {bypassL1} : memref<2x128x768xf16> to memref<5x1x64x64xf16, #gpu.address_space<workgroup>>808 809// -----810 811#map = affine_map<()[s0] -> (-s0 + 4)>812#map1 = affine_map<()[s0] -> (-s0 + 32)>813 814func.func @test_ldmatrix(%arg0: memref<4x32x32xf16, 3>, %arg1: index, %arg2: index, %arg3: index) -> vector<4x2xf16> {815 %c0 = arith.constant 0 : index816 %0 = affine.apply #map()[%arg1]817 %1 = affine.apply #map1()[%arg2]818 %2 = affine.apply #map1()[%arg3]819 %subview = memref.subview %arg0[%arg1, %arg2, %arg3] [%0, %1, %2] [1, 1, 1] : memref<4x32x32xf16, 3> to memref<?x?x?xf16, strided<[1024, 32, 1], offset: ?>, 3>820 %3 = nvgpu.ldmatrix %subview[%c0, %c0, %c0] {numTiles = 4 : i32, transpose = false} : memref<?x?x?xf16, strided<[1024, 32, 1], offset: ?>, 3> -> vector<4x2xf16>821 return %3 : vector<4x2xf16>822}823 824// CHECK: func @test_ldmatrix825// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x32x32xf16, 3>826// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index827// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index828// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: index829// CHECK: nvgpu.ldmatrix %[[ARG0]][%[[ARG1]], %[[ARG2]], %[[ARG3]]] {numTiles = 4 : i32, transpose = false} : memref<4x32x32xf16, 3> -> vector<4x2xf16>830 831// -----832 833func.func @fold_vector_load_subview(%src : memref<24x64xf32>,834 %off1 : index,835 %off2 : index,836 %dim1 : index,837 %dim2 : index,838 %idx : index) -> vector<12x32xf32> {839 840 %0 = memref.subview %src[%off1, %off2][%dim1, %dim2][1, 1] : memref<24x64xf32> to memref<?x?xf32, strided<[64, 1], offset: ?>>841 %1 = vector.load %0[%idx, %idx] : memref<?x?xf32, strided<[64, 1], offset: ?>>, vector<12x32xf32>842 return %1 : vector<12x32xf32>843}844 845// CHECK: #[[$ATTR_46:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>846// CHECK-LABEL: func.func @fold_vector_load_subview(847// CHECK-SAME: %[[SRC:[a-zA-Z0-9$._-]*]]: memref<24x64xf32>,848// CHECK-SAME: %[[OFF_1:[a-zA-Z0-9$._-]*]]: index,849// CHECK-SAME: %[[OFF_2:[a-zA-Z0-9$._-]*]]: index,850// CHECK-SAME: %[[DIM_1:[a-zA-Z0-9$._-]*]]: index,851// CHECK-SAME: %[[DIM_2:[a-zA-Z0-9$._-]*]]: index,852// CHECK-SAME: %[[IDX:[a-zA-Z0-9$._-]*]]: index) -> vector<12x32xf32> {853// CHECK: %[[VAL_6:.*]] = affine.apply #[[$ATTR_46]](){{\[}}%[[OFF_1]], %[[IDX]]]854// CHECK: %[[VAL_7:.*]] = affine.apply #[[$ATTR_46]](){{\[}}%[[OFF_2]], %[[IDX]]]855// CHECK: %[[VAL_8:.*]] = vector.load %[[SRC]]{{\[}}%[[VAL_6]], %[[VAL_7]]] : memref<24x64xf32>, vector<12x32xf32>856 857// -----858 859func.func @fold_vector_maskedload_subview(860 %arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3: vector<32xi1>, %arg4: vector<32xf32>) -> vector<32xf32> {861 %0 = memref.subview %arg0[%arg1, %arg2][1, 1][1, 1] : memref<12x32xf32> to memref<f32, strided<[], offset: ?>>862 %1 = vector.maskedload %0[], %arg3, %arg4 : memref<f32, strided<[], offset: ?>>, vector<32xi1>, vector<32xf32> into vector<32xf32>863 return %1 : vector<32xf32>864}865 866// CHECK: func @fold_vector_maskedload_subview867// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<12x32xf32>868// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index869// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index870// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: vector<32xi1>871// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: vector<32xf32>872// CHECK: vector.maskedload %[[ARG0]][%[[ARG1]], %[[ARG2]]], %[[ARG3]], %[[ARG4]] : memref<12x32xf32>, vector<32xi1>, vector<32xf32> into vector<32xf32>873 874// -----875 876func.func @fold_vector_store_subview(%src : memref<24x64xf32>,877 %off1 : index,878 %off2 : index,879 %vec: vector<2x32xf32>,880 %idx : index,881 %dim1 : index,882 %dim2 : index) -> () {883 884 %0 = memref.subview %src[%off1, %off2][%dim1, %dim2][1, 1] : memref<24x64xf32> to memref<?x?xf32, strided<[64, 1], offset: ?>>885 vector.store %vec, %0[%idx, %idx] : memref<?x?xf32, strided<[64, 1], offset: ?>> , vector<2x32xf32>886 return887}888 889// CHECK: #[[$ATTR_47:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>890 891// CHECK-LABEL: func.func @fold_vector_store_subview(892// CHECK-SAME: %[[SRC:[a-zA-Z0-9$._-]*]]: memref<24x64xf32>,893// CHECK-SAME: %[[OFF1:[a-zA-Z0-9$._-]*]]: index,894// CHECK-SAME: %[[OFF_2:[a-zA-Z0-9$._-]*]]: index,895// CHECK-SAME: %[[VEC:[a-zA-Z0-9$._-]*]]: vector<2x32xf32>,896// CHECK-SAME: %[[IDX:[a-zA-Z0-9$._-]*]]: index,897// CHECK-SAME: %[[VAL_5:[a-zA-Z0-9$._-]*]]: index,898// CHECK-SAME: %[[VAL_6:[a-zA-Z0-9$._-]*]]: index) {899// CHECK: %[[VAL_7:.*]] = affine.apply #[[$ATTR_47]](){{\[}}%[[OFF1]], %[[IDX]]]900// CHECK: %[[VAL_8:.*]] = affine.apply #[[$ATTR_47]](){{\[}}%[[OFF_2]], %[[IDX]]]901// CHECK: vector.store %[[VEC]], %[[SRC]]{{\[}}%[[VAL_7]], %[[VAL_8]]] : memref<24x64xf32>, vector<2x32xf32>902 903// -----904 905func.func @fold_vector_maskedstore_subview(906 %arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3: vector<32xi1>, %arg4: vector<32xf32>) -> () {907 %0 = memref.subview %arg0[%arg1, %arg2][1, 1][1, 1] : memref<12x32xf32> to memref<f32, strided<[], offset: ?>>908 vector.maskedstore %0[], %arg3, %arg4 : memref<f32, strided<[], offset: ?>>, vector<32xi1>, vector<32xf32>909 return910}911 912// CHECK: func @fold_vector_maskedstore_subview913// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<12x32xf32>914// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index915// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index916// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: vector<32xi1>917// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: vector<32xf32>918// CHECK: vector.maskedstore %[[ARG0]][%[[ARG1]], %[[ARG2]]], %[[ARG3]], %[[ARG4]] : memref<12x32xf32>, vector<32xi1>, vector<32xf32>919// CHECK: return920 921// -----922 923func.func @fold_vector_load_expand_shape(924 %arg0 : memref<32xf32>, %arg1 : index) -> vector<8xf32> {925 %c0 = arith.constant 0 : index926 %0 = memref.expand_shape %arg0 [[0, 1]] output_shape [4, 8] : memref<32xf32> into memref<4x8xf32>927 %1 = vector.load %0[%arg1, %c0] {nontemporal = true} : memref<4x8xf32>, vector<8xf32>928 return %1 : vector<8xf32>929}930 931// CHECK-LABEL: func @fold_vector_load_expand_shape932// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<32xf32>933// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index934// CHECK: %[[C0:.*]] = arith.constant 0935// CHECK: %[[IDX:.*]] = affine.linearize_index [%[[ARG1]], %[[C0]]] by (4, 8)936// CHECK: vector.load %[[ARG0]][%[[IDX]]] {nontemporal = true}937 938// -----939 940func.func @fold_vector_maskedload_expand_shape(941 %arg0 : memref<32xf32>, %arg1 : index, %arg3: vector<8xi1>, %arg4: vector<8xf32>) -> vector<8xf32> {942 %c0 = arith.constant 0 : index943 %0 = memref.expand_shape %arg0 [[0, 1]] output_shape [4, 8] : memref<32xf32> into memref<4x8xf32>944 %1 = vector.maskedload %0[%arg1, %c0], %arg3, %arg4 : memref<4x8xf32>, vector<8xi1>, vector<8xf32> into vector<8xf32>945 return %1 : vector<8xf32>946}947 948// CHECK-LABEL: func @fold_vector_maskedload_expand_shape949// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<32xf32>950// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index951// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: vector<8xi1>952// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: vector<8xf32>953// CHECK: %[[C0:.*]] = arith.constant 0954// CHECK: %[[IDX:.*]] = affine.linearize_index [%[[ARG1]], %[[C0]]] by (4, 8)955// CHECK: vector.maskedload %[[ARG0]][%[[IDX]]], %[[ARG3]], %[[ARG4]]956 957// -----958 959func.func @fold_vector_store_expand_shape(960 %arg0 : memref<32xf32>, %arg1 : index, %val : vector<8xf32>) {961 %c0 = arith.constant 0 : index962 %0 = memref.expand_shape %arg0 [[0, 1]] output_shape [4, 8] : memref<32xf32> into memref<4x8xf32>963 vector.store %val, %0[%arg1, %c0] {nontemporal = true} : memref<4x8xf32>, vector<8xf32>964 return965}966 967// CHECK-LABEL: func @fold_vector_store_expand_shape968// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<32xf32>969// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index970// CHECK: %[[C0:.*]] = arith.constant 0971// CHECK: %[[IDX:.*]] = affine.linearize_index [%[[ARG1]], %[[C0]]] by (4, 8)972// CHECK: vector.store %{{.*}}, %[[ARG0]][%[[IDX]]] {nontemporal = true}973 974// -----975 976func.func @fold_vector_maskedstore_expand_shape(977 %arg0 : memref<32xf32>, %arg1 : index, %arg3: vector<8xi1>, %arg4: vector<8xf32>) {978 %c0 = arith.constant 0 : index979 %0 = memref.expand_shape %arg0 [[0, 1]] output_shape [4, 8] : memref<32xf32> into memref<4x8xf32>980 vector.maskedstore %0[%arg1, %c0], %arg3, %arg4 : memref<4x8xf32>, vector<8xi1>, vector<8xf32>981 return982}983 984// CHECK-LABEL: func @fold_vector_maskedstore_expand_shape985// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<32xf32>986// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index987// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: vector<8xi1>988// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: vector<8xf32>989// CHECK: %[[C0:.*]] = arith.constant 0990// CHECK: %[[IDX:.*]] = affine.linearize_index [%[[ARG1]], %[[C0]]] by (4, 8)991// CHECK: vector.maskedstore %[[ARG0]][%[[IDX]]], %[[ARG3]], %[[ARG4]]992 993// -----994 995func.func @fold_vector_transfer_read_expand_shape(996 %arg0 : memref<32xf32>, %arg1 : index) -> vector<8xf32> {997 %c0 = arith.constant 0 : index998 %pad = ub.poison : f32999 %0 = memref.expand_shape %arg0 [[0, 1]] output_shape [4, 8] : memref<32xf32> into memref<4x8xf32>1000 %1 = vector.transfer_read %0[%arg1, %c0], %pad {in_bounds = [true]} : memref<4x8xf32>, vector<8xf32>1001 return %1 : vector<8xf32>1002}1003 1004// CHECK-LABEL: func @fold_vector_transfer_read_expand_shape1005// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<32xf32>1006// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index1007// CHECK: %[[C0:.*]] = arith.constant 01008// CHECK: %[[PAD:.*]] = ub.poison : f321009// CHECK: %[[IDX:.*]] = affine.linearize_index [%[[ARG1]], %[[C0]]] by (4, 8)1010// CHECK: vector.transfer_read %[[ARG0]][%[[IDX]]], %[[PAD]] {in_bounds = [true]}1011 1012// -----1013 1014func.func @fold_vector_transfer_read_with_perm_map(1015 %arg0 : memref<32xf32>, %arg1 : index) -> vector<4x4xf32> {1016 %c0 = arith.constant 0 : index1017 %pad = ub.poison : f321018 %0 = memref.expand_shape %arg0 [[0, 1]] output_shape [4, 8] : memref<32xf32> into memref<4x8xf32>1019 %1 = vector.transfer_read %0[%arg1, %c0], %pad { permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds = [true, true]} : memref<4x8xf32>, vector<4x4xf32>1020 return %1 : vector<4x4xf32>1021}1022 1023// CHECK-LABEL: func @fold_vector_transfer_read_with_perm_map1024// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<32xf32>1025// CHECK: memref.expand_shape %[[ARG0]] {{\[}}[0, 1]] output_shape [4, 8] : memref<32xf32> into memref<4x8xf32>1026 1027// -----1028 1029func.func @fold_vector_transfer_read_rank_mismatch(1030 %arg0 : memref<32xf32>, %arg1 : index) -> vector<4x4xf32> {1031 %c0 = arith.constant 0 : index1032 %pad = ub.poison : f321033 %0 = memref.expand_shape %arg0 [[0, 1, 2]] output_shape [2, 4, 4] : memref<32xf32> into memref<2x4x4xf32>1034 %1 = vector.transfer_read %0[%arg1, %c0, %c0], %pad {in_bounds = [true, true]} : memref<2x4x4xf32>, vector<4x4xf32>1035 return %1 : vector<4x4xf32>1036}1037 1038// CHECK-LABEL: func @fold_vector_transfer_read_rank_mismatch1039// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<32xf32>1040// CHECK: memref.expand_shape %[[ARG0]] {{\[}}[0, 1, 2]] output_shape [2, 4, 4] : memref<32xf32> into memref<2x4x4xf32>1041 1042// -----1043 1044func.func @fold_vector_load_collapse_shape(1045 %arg0 : memref<4x8xf32>, %arg1 : index) -> vector<8xf32> {1046 %0 = memref.collapse_shape %arg0 [[0, 1]] : memref<4x8xf32> into memref<32xf32>1047 %1 = vector.load %0[%arg1] {nontemporal = true} : memref<32xf32>, vector<8xf32>1048 return %1 : vector<8xf32>1049}1050 1051// CHECK-LABEL: func @fold_vector_load_collapse_shape1052// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x8xf32>1053// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index1054// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)1055// CHECK: vector.load %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1] {nontemporal = true}1056 1057// -----1058 1059func.func @fold_vector_maskedload_collapse_shape(1060 %arg0 : memref<4x8xf32>, %arg1 : index, %arg3: vector<8xi1>, %arg4: vector<8xf32>) -> vector<8xf32> {1061 %0 = memref.collapse_shape %arg0 [[0, 1]] : memref<4x8xf32> into memref<32xf32>1062 %1 = vector.maskedload %0[%arg1], %arg3, %arg4 : memref<32xf32>, vector<8xi1>, vector<8xf32> into vector<8xf32>1063 return %1 : vector<8xf32>1064}1065 1066// CHECK-LABEL: func @fold_vector_maskedload_collapse_shape1067// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x8xf32>1068// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index1069// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: vector<8xi1>1070// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: vector<8xf32>1071// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)1072// CHECK: vector.maskedload %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1], %[[ARG3]], %[[ARG4]]1073 1074// -----1075 1076func.func @fold_vector_store_collapse_shape(1077 %arg0 : memref<4x8xf32>, %arg1 : index, %val : vector<8xf32>) {1078 %0 = memref.collapse_shape %arg0 [[0, 1]] : memref<4x8xf32> into memref<32xf32>1079 vector.store %val, %0[%arg1] {nontemporal = true} : memref<32xf32>, vector<8xf32>1080 return1081}1082 1083// CHECK-LABEL: func @fold_vector_store_collapse_shape1084// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x8xf32>1085// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index1086// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)1087// CHECK: vector.store %{{.*}}, %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1] {nontemporal = true}1088 1089// -----1090 1091func.func @fold_vector_maskedstore_collapse_shape(1092 %arg0 : memref<4x8xf32>, %arg1 : index, %arg3: vector<8xi1>, %arg4: vector<8xf32>) {1093 %0 = memref.collapse_shape %arg0 [[0, 1]] : memref<4x8xf32> into memref<32xf32>1094 vector.maskedstore %0[%arg1], %arg3, %arg4 : memref<32xf32>, vector<8xi1>, vector<8xf32>1095 return1096}1097 1098// CHECK-LABEL: func @fold_vector_maskedstore_collapse_shape1099// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x8xf32>1100// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index1101// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: vector<8xi1>1102// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: vector<8xf32>1103// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 8)1104// CHECK: vector.maskedstore %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1], %[[ARG3]], %[[ARG4]]1105