270 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s2 3///----------------------------------------------------------------------------------------4/// Tests for vectorization patterns for tensor.pad, i.e.5/// * transform.apply_patterns.linalg.pad_vectorization6///7/// These tests are meant os a lower granule than tests in8/// * pad-with-patterns.mlir.9/// The goal is to test specific patterns. To this end, some inputs already10/// contain Vector ops (on top of `tensor.pad`).11///----------------------------------------------------------------------------------------12 13///----------------------------------------------------------------------------------------14/// [Pattern: PadOpVectorizationWithTransferReadPattern]15///----------------------------------------------------------------------------------------16// CHECK-LABEL: func @pad_and_transfer_read17// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32>18// CHECK-NOT: tensor.pad19// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index20// CHECK-DAG: %[[C5:.*]] = arith.constant 5.021// CHECK: %[[RESULT:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], %[[C5]] : tensor<5x6xf32>, vector<7x9xf32>22// CHECK: return %[[RESULT]]23func.func @pad_and_transfer_read(%arg0: tensor<5x6xf32>) -> vector<7x9xf32> {24 %c0 = arith.constant 0 : index25 %c5 = arith.constant 5.0 : f3226 %c6 = arith.constant 6.0 : f3227 %0 = tensor.pad %arg0 low[0, 0] high[5, 7] {28 ^bb0(%arg1: index, %arg2: index):29 tensor.yield %c5 : f3230 } : tensor<5x6xf32> to tensor<10x13xf32>31 %1 = vector.transfer_read %0[%c0, %c0], %c632 : tensor<10x13xf32>, vector<7x9xf32>33 return %1 : vector<7x9xf32>34}35 36module attributes {transform.with_named_sequence} {37 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {38 %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">39 40 transform.apply_patterns to %func_op {41 transform.apply_patterns.linalg.pad_vectorization42 } : !transform.op<"func.func">43 transform.yield44 }45}46 47// -----48 49///----------------------------------------------------------------------------------------50/// [Pattern: PadOpVectorizationWithTransferWritePattern]51///----------------------------------------------------------------------------------------52func.func private @make_vector() -> vector<7x9xf32>53 54// CHECK-LABEL: func @pad_and_transfer_write_static_low_and_high55// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32>56// CHECK-NOT: tensor.pad57// CHECK: %[[C0:.*]] = arith.constant 0 : index58// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32>59// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[ARG0]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor<5x6xf32>60// CHECK: return %[[RESULT]]61func.func @pad_and_transfer_write_static_low_and_high(62 %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> {63 %c0 = arith.constant 0 : index64 %c5 = arith.constant 5.0 : f3265 %0 = tensor.pad %arg0 low[0, 0] high[5, 7] {66 ^bb0(%arg2: index, %arg3: index):67 tensor.yield %c5 : f3268 } : tensor<5x6xf32> to tensor<10x13xf32>69 %1 = call @make_vector() : () -> vector<7x9xf32>70 %2 = vector.transfer_write %1, %0[%c0, %c0]71 : vector<7x9xf32>, tensor<10x13xf32>72 %3 = tensor.extract_slice %2[0, 0] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32>73 return %3 : tensor<5x6xf32>74}75 76module attributes {transform.with_named_sequence} {77 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {78 %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">79 80 transform.apply_patterns to %func_op {81 transform.apply_patterns.linalg.pad_vectorization82 } : !transform.op<"func.func">83 transform.yield84 }85}86 87// -----88 89func.func private @make_vector() -> vector<7x9xf32>90 91// CHECK-LABEL: func @pad_and_transfer_write_static_low_dynamic_high92// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?xf32>, %[[SIZE:.*]]: index, %[[PADDING:.*]]: index93// CHECK-NOT: tensor.pad94// CHECK: %[[C0:.*]] = arith.constant 0 : index95// CHECK: %[[SUB:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [%[[SIZE]], 6] [1, 1] : tensor<?x?xf32> to tensor<?x6xf32>96// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32>97// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[SUB]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor<?x6xf32>98// CHECK: return %[[RESULT]]99func.func @pad_and_transfer_write_static_low_dynamic_high(100 %arg0: tensor<?x?xf32>, %size: index, %padding: index) -> tensor<?x6xf32> {101 %c0 = arith.constant 0 : index102 %c5 = arith.constant 5.0 : f32103 %s = tensor.extract_slice %arg0[0, 0] [%size, 6] [1, 1]104 : tensor<?x?xf32> to tensor<?x6xf32>105 %0 = tensor.pad %s low[0, 0] high[%padding, 7] {106 ^bb0(%arg2: index, %arg3: index):107 tensor.yield %c5 : f32108 } : tensor<?x6xf32> to tensor<?x13xf32>109 %1 = call @make_vector() : () -> vector<7x9xf32>110 %2 = vector.transfer_write %1, %0[%c0, %c0]111 : vector<7x9xf32>, tensor<?x13xf32>112 %3 = tensor.extract_slice %2[0, 0] [%size, 6] [1, 1] : tensor<?x13xf32> to tensor<?x6xf32>113 return %3 : tensor<?x6xf32>114}115 116module attributes {transform.with_named_sequence} {117 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {118 %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">119 120 transform.apply_patterns to %func_op {121 transform.apply_patterns.linalg.pad_vectorization122 } : !transform.op<"func.func">123 transform.yield124 }125}126 127// -----128 129func.func private @make_vector() -> vector<7x9xf32>130 131// Negative test - low pad is non-zero132 133// CHECK-LABEL: func @pad_and_transfer_write_static_non_zero_low_pad134// CHECK: tensor.pad135func.func @pad_and_transfer_write_static_non_zero_low_pad(136 %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> {137 %c0 = arith.constant 0 : index138 %c5 = arith.constant 5.0 : f32139 %0 = tensor.pad %arg0 low[0, 1] high[5, 6] {140 ^bb0(%arg2: index, %arg3: index):141 tensor.yield %c5 : f32142 } : tensor<5x6xf32> to tensor<10x13xf32>143 %1 = call @make_vector() : () -> vector<7x9xf32>144 %2 = vector.transfer_write %1, %0[%c0, %c0]145 : vector<7x9xf32>, tensor<10x13xf32>146 %3 = tensor.extract_slice %2[0, 0] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32>147 return %3 : tensor<5x6xf32>148}149 150module attributes {transform.with_named_sequence} {151 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {152 %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">153 154 transform.apply_patterns to %func_op {155 transform.apply_patterns.linalg.pad_vectorization156 } : !transform.op<"func.func">157 transform.yield158 }159}160 161// -----162 163// Negative test - TransferWriteOp result is not _directly_ consumed by an164// ExtractSliceOp (noet the non-zero offset).165 166func.func private @make_vector() -> vector<7x9xf32>167 168// CHECK-LABEL: func @pad_and_transfer_write_static_non_zero_offset169// CHECK: tensor.pad170func.func @pad_and_transfer_write_static_non_zero_offset(171 %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> {172 %c0 = arith.constant 0 : index173 %c5 = arith.constant 5.0 : f32174 %0 = tensor.pad %arg0 low[0, 0] high[5, 7] {175 ^bb0(%arg2: index, %arg3: index):176 tensor.yield %c5 : f32177 } : tensor<5x6xf32> to tensor<10x13xf32>178 %1 = call @make_vector() : () -> vector<7x9xf32>179 %2 = vector.transfer_write %1, %0[%c0, %c0]180 : vector<7x9xf32>, tensor<10x13xf32>181 %3 = tensor.extract_slice %2[0, 1] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32>182 return %3 : tensor<5x6xf32>183}184 185module attributes {transform.with_named_sequence} {186 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {187 %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">188 189 transform.apply_patterns to %func_op {190 transform.apply_patterns.linalg.pad_vectorization191 } : !transform.op<"func.func">192 transform.yield193 }194}195 196// -----197 198///----------------------------------------------------------------------------------------199/// [Pattern: PadOpVectorizationWithInsertSlicePattern]200///----------------------------------------------------------------------------------------201 202func.func private @make_vector() -> tensor<12x13xf32>203 204// CHECK-LABEL: func @pad_and_insert_slice_source205// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32>206// CHECK-NOT: tensor.pad207// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index208// CHECK-DAG: %[[C5:.*]] = arith.constant 5.0209// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> tensor<12x13xf32>210// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], %[[C5]] : tensor<5x6xf32>, vector<7x9xf32>211// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[VEC0]][%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<7x9xf32>, tensor<12x13xf32>212// CHECK: return %[[WRITE]]213func.func @pad_and_insert_slice_source(214 %arg0: tensor<5x6xf32>) -> tensor<12x13xf32> {215 %c0 = arith.constant 0 : index216 %c5 = arith.constant 5.0 : f32217 %0 = tensor.pad %arg0 low[0, 0] high[2, 3] {218 ^bb0(%arg2: index, %arg3: index):219 tensor.yield %c5 : f32220 } : tensor<5x6xf32> to tensor<7x9xf32>221 %1 = call @make_vector() : () -> tensor<12x13xf32>222 %r = tensor.insert_slice %0 into %1[0, 0][7, 9][1, 1] : tensor<7x9xf32> into tensor<12x13xf32>223 return %r : tensor<12x13xf32>224}225 226module attributes {transform.with_named_sequence} {227 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {228 %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">229 230 transform.apply_patterns to %func_op {231 transform.apply_patterns.linalg.pad_vectorization232 } : !transform.op<"func.func">233 transform.yield234 }235}236 237// -----238 239func.func private @make_vector() -> tensor<12x13xf32>240 241// The destination of tensor.insert_slice matches the result of tensor.pad -242// not supported.243 244// CHECK-LABEL: func.func @pad_and_insert_slice_dest(245// CHECK-NOT: vector.transfer_read246// CHECK-NOT: vector.transfer_write247 248func.func @pad_and_insert_slice_dest(249 %arg0: tensor<1x5x6xf32>) -> tensor<1x12x13xf32> {250 %c5 = arith.constant 5.0 : f32251 %0 = tensor.pad %arg0 low[0, 0, 0] high[0, 7, 7] {252 ^bb0(%arg2: index, %arg3: index, %arg4: index):253 tensor.yield %c5 : f32254 } : tensor<1x5x6xf32> to tensor<1x12x13xf32>255 %1 = call @make_vector() : () -> tensor<12x13xf32>256 %r = tensor.insert_slice %1 into %0[0, 0, 0][1, 12, 13][1, 1, 1] : tensor<12x13xf32> into tensor<1x12x13xf32>257 return %r : tensor<1x12x13xf32>258}259 260module attributes {transform.with_named_sequence} {261 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {262 %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">263 264 transform.apply_patterns to %func_op {265 transform.apply_patterns.linalg.pad_vectorization266 } : !transform.op<"func.func">267 transform.yield268 }269}270