277 lines · plain
1// RUN: mlir-opt %s -test-linalg-transform-patterns=test-swap-subtensor-padtensor -canonicalize -split-input-file | FileCheck %s2 3// CHECK-LABEL: @static_data_only(4// CHECK-SAME: %[[ARG0:.*]]: tensor<4x5xf32>5// CHECK: %[[RESULT:.*]] = tensor.extract_slice %[[ARG0]][1, 2] [2, 1] [1, 1] : tensor<4x5xf32> to tensor<2x1xf32>6// CHECK: return %[[RESULT]]7func.func @static_data_only(%arg0 : tensor<4x5xf32>, %pad : f32)8 -> tensor<2x1xf32> {9 %0 = tensor.pad %arg0 low[0, 0] high[7, 8] {10 ^bb0(%arg1: index, %arg2: index):11 tensor.yield %pad : f3212 } : tensor<4x5xf32> to tensor<11x13xf32>13 %1 = tensor.extract_slice %0[1, 2] [2, 1] [1, 1] : tensor<11x13xf32> to tensor<2x1xf32>14 return %1 : tensor<2x1xf32>15}16 17// -----18 19// CHECK-LABEL: @static_high_pad_only20// CHECK-SAME: %[[ARG0:.*]]: tensor<4x5xf32>, %[[PAD:.*]]: f3221// CHECK-NOT: tensor.pad22// CHECK-NOT: tensor.extract_slice23// CHECK: %[[RESULT:.*]] = tensor.generate24// CHECK: tensor.yield %[[PAD]]25// CHECK: return %[[RESULT]] : tensor<2x4xf32>26func.func @static_high_pad_only(%arg0 : tensor<4x5xf32>, %pad : f32)27 -> tensor<2x4xf32> {28 %0 = tensor.pad %arg0 low[0, 0] high[7, 8] {29 ^bb0(%arg1: index, %arg2: index):30 tensor.yield %pad : f3231 } : tensor<4x5xf32> to tensor<11x13xf32>32 %1 = tensor.extract_slice %0[4, 5] [2, 4] [1, 1] : tensor<11x13xf32> to tensor<2x4xf32>33 return %1 : tensor<2x4xf32>34}35 36// -----37 38// CHECK-LABEL: @static_low_pad_only39// CHECK-SAME: %[[ARG0:.*]]: tensor<4x5xf32>, %[[PAD:.*]]: f3240// CHECK-NOT: tensor.pad41// CHECK-NOT: tensor.extract_slice42// CHECK: %[[RESULT:.*]] = tensor.generate43// CHECK: tensor.yield %[[PAD]]44// CHECK: return %[[RESULT]] : tensor<2x3xf32>45func.func @static_low_pad_only(%arg0 : tensor<4x5xf32>, %pad : f32)46 -> tensor<2x3xf32> {47 %0 = tensor.pad %arg0 low[3, 7] high[7, 8] {48 ^bb0(%arg1: index, %arg2: index):49 tensor.yield %pad : f3250 } : tensor<4x5xf32> to tensor<14x20xf32>51 %1 = tensor.extract_slice %0[1, 3] [2, 3] [1, 1] : tensor<14x20xf32> to tensor<2x3xf32>52 return %1 : tensor<2x3xf32>53}54 55// -----56 57// CHECK-LABEL: @static_low_pad_only_258// CHECK-SAME: %[[ARG0:.*]]: tensor<4x5xf32>, %[[PAD:.*]]: f3259// CHECK-NOT: tensor.pad60// CHECK-NOT: tensor.extract_slice61// CHECK: %[[RESULT:.*]] = tensor.generate62// CHECK: tensor.yield %[[PAD]]63// CHECK: return %[[RESULT]] : tensor<1x3xf32>64func.func @static_low_pad_only_2(%arg0 : tensor<4x5xf32>, %pad : f32)65 -> tensor<1x3xf32> {66 %0 = tensor.pad %arg0 low[3, 7] high[7, 8] {67 ^bb0(%arg1: index, %arg2: index):68 tensor.yield %pad : f3269 } : tensor<4x5xf32> to tensor<14x20xf32>70 %1 = tensor.extract_slice %0[1, 3] [1, 3] [1, 1] : tensor<14x20xf32> to tensor<1x3xf32>71 return %1 : tensor<1x3xf32>72}73 74// -----75 76// CHECK-LABEL: @static_mixed_data_high_pad77// CHECK-SAME: %[[ARG0:.*]]: tensor<4x5xf32>, %[[PAD:.*]]: f3278// CHECK-NOT: tensor.pad79// CHECK: %[[SUBTENSOR:.*]] = tensor.extract_slice %[[ARG0]][2, 4] [2, 1] [1, 1] : tensor<4x5xf32> to tensor<2x1xf32>80// CHECK: %[[RESULT:.*]] = tensor.pad %[[SUBTENSOR]] low[0, 0] high[1, 3]81// CHECK: tensor.yield %[[PAD]]82// CHECK: return %[[RESULT]] : tensor<3x4xf32>83func.func @static_mixed_data_high_pad(%arg0 : tensor<4x5xf32>, %pad : f32)84 -> tensor<3x4xf32> {85 %0 = tensor.pad %arg0 low[0, 0] high[7, 8] {86 ^bb0(%arg1: index, %arg2: index):87 tensor.yield %pad : f3288 } : tensor<4x5xf32> to tensor<11x13xf32>89 %1 = tensor.extract_slice %0[2, 4] [3, 4] [1, 1] : tensor<11x13xf32> to tensor<3x4xf32>90 return %1 : tensor<3x4xf32>91}92 93// -----94 95// CHECK-LABEL: @static_mixed_data_low_pad96// CHECK-SAME: %[[ARG0:.*]]: tensor<4x5xf32>, %[[PAD:.*]]: f3297// CHECK-NOT: tensor.pad98// CHECK: %[[SUBTENSOR:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [2, 1] [1, 1] : tensor<4x5xf32> to tensor<2x1xf32>99// CHECK: %[[RESULT:.*]] = tensor.pad %[[SUBTENSOR]] low[1, 3] high[0, 0]100// CHECK: tensor.yield %[[PAD]]101// CHECK: return %[[RESULT]] : tensor<3x4xf32>102func.func @static_mixed_data_low_pad(%arg0 : tensor<4x5xf32>, %pad : f32)103 -> tensor<3x4xf32> {104 %0 = tensor.pad %arg0 low[3, 7] high[7, 8] {105 ^bb0(%arg1: index, %arg2: index):106 tensor.yield %pad : f32107 } : tensor<4x5xf32> to tensor<14x20xf32>108 %1 = tensor.extract_slice %0[2, 4] [3, 4] [1, 1] : tensor<14x20xf32> to tensor<3x4xf32>109 return %1 : tensor<3x4xf32>110}111 112// -----113 114// CHECK-LABEL: @static_mixed_data_low_high_pad115// CHECK-SAME: %[[ARG0:.*]]: tensor<4x5xf32>, %[[PAD:.*]]: f32116// CHECK-NOT: tensor.pad117// CHECK: %[[RESULT:.*]] = tensor.pad %[[ARG0]] low[1, 1] high[2, 3]118// CHECK: tensor.yield %[[PAD]]119// CHECK: return %[[RESULT]] : tensor<7x9xf32>120func.func @static_mixed_data_low_high_pad(%arg0 : tensor<4x5xf32>, %pad : f32)121 -> tensor<7x9xf32> {122 %0 = tensor.pad %arg0 low[2, 3] high[7, 8] {123 ^bb0(%arg1: index, %arg2: index):124 tensor.yield %pad : f32125 } : tensor<4x5xf32> to tensor<13x16xf32>126 %1 = tensor.extract_slice %0[1, 2] [7, 9] [1, 1] : tensor<13x16xf32> to tensor<7x9xf32>127 return %1 : tensor<7x9xf32>128}129 130// -----131 132// CHECK-LABEL: @static_rank_reduce133// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16x4xf32>, %[[PADVAL:.*]]: f32134// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[ARG0]][0, 0, 0] [1, 14, 4] [1, 1, 1] : tensor<8x16x4xf32> to tensor<1x14x4xf32>135// CHECK: %[[PADDED:.*]] = tensor.pad %[[SLICE]] low[0, 2, 0] high[0, 0, 0] {136// CHECK: } : tensor<1x14x4xf32> to tensor<1x16x4xf32>137// CHECK: %[[RESULT:.*]] = tensor.extract_slice %[[PADDED]][0, 0, 0] [1, 16, 4] [1, 1, 1] : tensor<1x16x4xf32> to tensor<16x4xf32>138// CHECK: return %[[RESULT]]139func.func @static_rank_reduce(%arg0: tensor<8x16x4xf32>, %pad: f32)140 -> tensor<16x4xf32> {141 %0 = tensor.pad %arg0 low[0, 2, 0] high[0, 0, 0] {142 ^bb0(%i: index, %j: index, %k: index):143 tensor.yield %pad : f32144 } : tensor<8x16x4xf32> to tensor<8x18x4xf32>145 %1 = tensor.extract_slice %0[0, 0, 0] [1, 16, 4] [1, 1, 1]146 : tensor<8x18x4xf32> to tensor<16x4xf32>147 return %1 : tensor<16x4xf32>148}149 150// -----151 152// CHECK-LABEL: @dynamic_high_pad153// CHECK-SAME: %[[ARG0:.*]]: tensor<?x5xf32>154// CHECK-NOT: tensor.pad155// CHECK: %[[C0:.*]] = arith.constant 0 : index156// CHECK: tensor.dim %[[ARG0]], %[[C0]]157// CHECK: %[[RESULT:.*]] = scf.if %{{.*}} -> (tensor<3x4xf32>) {158// CHECK: %[[GEN:.*]] = tensor.generate159// CHECK: scf.yield %[[GEN]]160// CHECK: } else {161// CHECK: %[[SUBTENSOR:.*]] = tensor.extract_slice %[[ARG0]][%{{.*}}, 4] [%{{.*}}, 1] [1, 1] : tensor<?x5xf32> to tensor<?x1xf32>162// CHECK: %[[PADTENSOR:.*]] = tensor.pad %[[SUBTENSOR]] low[0, 0] high[%{{.*}}, 3]163// CHECK: scf.yield %[[PADTENSOR]]164// CHECK: }165// CHECK: return %[[RESULT]]166func.func @dynamic_high_pad(%arg0 : tensor<?x5xf32>, %h1: index, %pad : f32) -> tensor<3x4xf32> {167 %0 = tensor.pad %arg0 low[0, 0] high[%h1, 8] {168 ^bb0(%arg1: index, %arg2: index):169 tensor.yield %pad : f32170 } : tensor<?x5xf32> to tensor<?x13xf32>171 %1 = tensor.extract_slice %0[2, 4] [3, 4] [1, 1] : tensor<?x13xf32> to tensor<3x4xf32>172 return %1 : tensor<3x4xf32>173}174 175// -----176 177// CHECK-LABEL: @dynamic_extract_size178// CHECK-SAME: %[[ARG0:.*]]: tensor<?x5xf32>, %[[ARG1:.*]]: index179// CHECK-NOT: tensor.pad180// CHECK: %[[C0:.*]] = arith.constant 0 : index181// CHECK: tensor.dim %[[ARG0]], %[[C0]]182// CHECK: %[[RESULT:.*]] = scf.if %{{.*}} -> (tensor<?x4xf32>) {183// CHECK: %[[GEN:.*]] = tensor.generate %[[ARG1]]184// CHECK: scf.yield %[[GEN]]185// CHECK: } else {186// CHECK: %[[SUBTENSOR:.*]] = tensor.extract_slice %[[ARG0]][%{{.*}}, 4] [%{{.*}}, 1] [1, 1] : tensor<?x5xf32> to tensor<?x1xf32>187// CHECK: %[[PADTENSOR:.*]] = tensor.pad %[[SUBTENSOR]] low[0, 0] high[%{{.*}}, 3]188// CHECK: scf.yield %[[PADTENSOR]]189// CHECK: }190// CHECK: return %[[RESULT]]191func.func @dynamic_extract_size(%arg0 : tensor<?x5xf32>, %s1: index, %pad : f32) -> tensor<?x4xf32> {192 %0 = tensor.pad %arg0 low[0, 0] high[7, 8] {193 ^bb0(%arg1: index, %arg2: index):194 tensor.yield %pad : f32195 } : tensor<?x5xf32> to tensor<?x13xf32>196 %1 = tensor.extract_slice %0[2, 4] [%s1, 4] [1, 1] : tensor<?x13xf32> to tensor<?x4xf32>197 return %1 : tensor<?x4xf32>198}199 200// -----201 202// CHECK-LABEL: @dynamic_zero_low_padding203// CHECK: scf.if204// CHECK: tensor.generate205// CHECK: else206// CHECK: %[[SLICE:.*]] = tensor.extract_slice207// CHECK: tensor.pad %[[SLICE]] low[0, 0]208func.func @dynamic_zero_low_padding(%arg0 : tensor<?x?xf32>, %pad : f32,209 %o1 : index, %o2 : index,210 %s1 : index, %s2 : index)211 -> tensor<?x?xf32> {212 %0 = tensor.pad %arg0 low[0, 0] high[7, 8] {213 ^bb0(%arg1: index, %arg2: index):214 tensor.yield %pad : f32215 } : tensor<?x?xf32> to tensor<?x?xf32>216 %1 = tensor.extract_slice %0[%o1, %o2] [%s1, %s2] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>217 return %1 : tensor<?x?xf32>218}219 220// -----221 222// CHECK-LABEL: @dynamic_zero_high_padding223// CHECK: scf.if224// CHECK: tensor.generate225// CHECK: else226// CHECK: %[[SLICE:.*]] = tensor.extract_slice227// CHECK: tensor.pad %[[SLICE]] low[%{{.*}}, %{{.*}}] high[0, 0]228func.func @dynamic_zero_high_padding(%arg0 : tensor<?x?xf32>, %pad : f32,229 %o1 : index, %o2 : index,230 %s1 : index, %s2 : index)231 -> tensor<?x?xf32> {232 %0 = tensor.pad %arg0 low[7, 8] high[0, 0] {233 ^bb0(%arg1: index, %arg2: index):234 tensor.yield %pad : f32235 } : tensor<?x?xf32> to tensor<?x?xf32>236 %1 = tensor.extract_slice %0[%o1, %o2] [%s1, %s2] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>237 return %1 : tensor<?x?xf32>238}239 240// -----241 242// CHECK-LABEL: @dynamic_rank_reduce243// CHECK: %[[TEMP:.*]] = scf.if %{{.*}} -> (tensor<1x4xf32>) {244// CHECK: tensor.generate245// CHECK: } else {246// CHECK: %[[SLICE:.*]] = tensor.extract_slice %{{.*}} : tensor<?x5xf32> to tensor<?x1xf32>247// CHECK: tensor.pad %[[SLICE]] low[0, 0] high[%{{.*}}, 3] {248// CHECK: } : tensor<?x1xf32> to tensor<1x4xf32>249// CHECK: }250// CHECK: %[[RESULT:.*]] = tensor.extract_slice %[[TEMP]]{{.*}} : tensor<1x4xf32> to tensor<4xf32>251// CHECK: return %[[RESULT]]252func.func @dynamic_rank_reduce(%arg0 : tensor<?x5xf32>, %s1: index, %pad : f32) -> tensor<4xf32> {253 %0 = tensor.pad %arg0 low[0, 0] high[7, 8] {254 ^bb0(%arg1: index, %arg2: index):255 tensor.yield %pad : f32256 } : tensor<?x5xf32> to tensor<?x13xf32>257 %1 = tensor.extract_slice %0[2, 4] [1, 4] [1, 1] : tensor<?x13xf32> to tensor<4xf32>258 return %1 : tensor<4xf32>259}260 261// -----262// CHECK-LABEL: @nopaddim_with_dynamic_extract(263// CHECK-SAME: %[[ARG0:.*]]: tensor<3x4x5xf32>264// CHECK-SAME: %[[ARG1:.*]]: f32265// CHECK-SAME: %[[ARG2:.*]]: index266// CHECK: %[[RESULT:.*]] = tensor.extract_slice %[[ARG0]][%[[ARG2]], 1, 2] [%[[ARG2]], 2, 1] [1, 1, 1] : tensor<3x4x5xf32> to tensor<?x2x1xf32>267// CHECK: return %[[RESULT]]268func.func @nopaddim_with_dynamic_extract(%arg0 : tensor<3x4x5xf32>, %pad : f32, %index : index)269 -> tensor<?x2x1xf32> {270 %0 = tensor.pad %arg0 low[0, 0, 0] high[0, 7, 8] {271 ^bb0(%arg1: index, %arg2: index, %arg3: index):272 tensor.yield %pad : f32273 } : tensor<3x4x5xf32> to tensor<3x11x13xf32>274 %1 = tensor.extract_slice %0[%index, 1, 2] [%index, 2, 1] [1, 1, 1] : tensor<3x11x13xf32> to tensor<?x2x1xf32>275 return %1 : tensor<?x2x1xf32>276}277