1728 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s2 3///----------------------------------------------------------------------------------------4/// Tests for linalg.generic5///----------------------------------------------------------------------------------------6 7func.func @vectorize_dynamic_identity(%arg0: tensor<?xf32>,8 %arg1: tensor<?xf32>,9 %arg2: tensor<?xf32>) -> tensor<?xf32> {10 %0 = linalg.generic { indexing_maps = [affine_map<(d0) -> (d0)>,11 affine_map<(d0) -> (d0)>,12 affine_map<(d0) -> (d0)>],13 iterator_types = ["parallel"] }14 ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)15 outs(%arg2 : tensor<?xf32>) {16 ^bb(%in0: f32, %in1: f32, %out: f32) :17 %0 = arith.addf %in0, %in1 : f3218 linalg.yield %0 : f3219 } -> tensor<?xf32>20 return %0 : tensor<?xf32>21}22 23// CHECK-LABEL: @vectorize_dynamic_identity24// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index25// CHECK: %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?xf32>26// CHECK: %[[VAL_7:.*]] = vector.create_mask %[[VAL_4]] : vector<4xi1>27// CHECK: %[[VAL_8:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>28// CHECK: %[[VAL_10:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>29// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>30// CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_8]], %[[VAL_10]] : vector<4xf32>31// CHECK: %[[VAL_14:.*]] = vector.mask %[[VAL_7]] { vector.transfer_write %{{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<?xf32> } : vector<4xi1> -> tensor<?xf32>32 33module attributes {transform.with_named_sequence} {34 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {35 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op36 transform.structured.vectorize %0 vector_sizes [4] : !transform.any_op37 transform.yield38 }39}40 41// -----42 43func.func @vectorize_dynamic_identity_scalable(%arg0: tensor<?xf32>,44 %arg1: tensor<?xf32>,45 %arg2: tensor<?xf32>) -> tensor<?xf32> {46 %0 = linalg.generic { indexing_maps = [affine_map<(d0) -> (d0)>,47 affine_map<(d0) -> (d0)>,48 affine_map<(d0) -> (d0)>],49 iterator_types = ["parallel"] }50 ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)51 outs(%arg2 : tensor<?xf32>) {52 ^bb(%in0: f32, %in1: f32, %out: f32) :53 %0 = arith.addf %in0, %in1 : f3254 linalg.yield %0 : f3255 } -> tensor<?xf32>56 return %0 : tensor<?xf32>57}58 59// CHECK-LABEL: @vectorize_dynamic_identity_scalable60// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index61// CHECK: %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?xf32>62// CHECK: %[[VAL_7:.*]] = vector.create_mask %[[VAL_4]] : vector<[4]xi1>63// CHECK: %[[VAL_8:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<[4]xf32> } : vector<[4]xi1> -> vector<[4]xf32>64// CHECK: %[[VAL_10:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<[4]xf32> } : vector<[4]xi1> -> vector<[4]xf32>65// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<[4]xf32> } : vector<[4]xi1> -> vector<[4]xf32>66// CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_8]], %[[VAL_10]] : vector<[4]xf32>67// CHECK: %[[VAL_14:.*]] = vector.mask %[[VAL_7]] { vector.transfer_write %{{.*}} {in_bounds = [true]} : vector<[4]xf32>, tensor<?xf32> } : vector<[4]xi1> -> tensor<?xf32>68 69module attributes {transform.with_named_sequence} {70 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {71 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op72 transform.structured.vectorize %0 vector_sizes [[4]] : !transform.any_op73 transform.yield74 }75}76 77// -----78 79func.func @vectorize_dynamic_identity_with_constant(%arg0: tensor<?xf32>,80 %arg1: tensor<?xf32>,81 %arg2: tensor<?xf32>) -> tensor<?xf32> {82 %c4 = arith.constant 4 : index83 %0 = linalg.generic { indexing_maps = [affine_map<(d0) -> (d0)>,84 affine_map<(d0) -> (d0)>,85 affine_map<(d0) -> (d0)>],86 iterator_types = ["parallel"] }87 ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)88 outs(%arg2 : tensor<?xf32>) {89 ^bb(%in0: f32, %in1: f32, %out: f32) :90 %0 = arith.addf %in0, %in1 : f3291 linalg.yield %0 : f3292 } -> tensor<?xf32>93 return %0 : tensor<?xf32>94}95 96// CHECK-LABEL: @vectorize_dynamic_identity_with_constant97// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index98// CHECK: %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?xf32>99// CHECK: %[[VAL_7:.*]] = vector.create_mask %[[VAL_4]] : vector<4xi1>100// CHECK: %[[VAL_8:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>101// CHECK: %[[VAL_10:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>102// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>103// CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_8]], %[[VAL_10]] : vector<4xf32>104// CHECK: %[[VAL_14:.*]] = vector.mask %[[VAL_7]] { vector.transfer_write %{{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<?xf32> } : vector<4xi1> -> tensor<?xf32>105 106module attributes {transform.with_named_sequence} {107 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {108 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op109 %size = transform.structured.match ops{["arith.constant"]} in %arg1 : (!transform.any_op) -> !transform.any_op110 transform.structured.vectorize %0 vector_sizes [%size] : !transform.any_op, !transform.any_op111 transform.yield112 }113}114 115// -----116 117func.func @vectorize_dynamic_identity_with_param(%arg0: tensor<?xf32>,118 %arg1: tensor<?xf32>,119 %arg2: tensor<?xf32>) -> tensor<?xf32> {120 %0 = linalg.generic { indexing_maps = [affine_map<(d0) -> (d0)>,121 affine_map<(d0) -> (d0)>,122 affine_map<(d0) -> (d0)>],123 iterator_types = ["parallel"] }124 ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)125 outs(%arg2 : tensor<?xf32>) {126 ^bb(%in0: f32, %in1: f32, %out: f32) :127 %0 = arith.addf %in0, %in1 : f32128 linalg.yield %0 : f32129 } -> tensor<?xf32>130 return %0 : tensor<?xf32>131}132 133// CHECK-LABEL: @vectorize_dynamic_identity_with_param134// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index135// CHECK: %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?xf32>136// CHECK: %[[VAL_7:.*]] = vector.create_mask %[[VAL_4]] : vector<4xi1>137// CHECK: %[[VAL_8:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>138// CHECK: %[[VAL_10:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>139// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>140// CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_8]], %[[VAL_10]] : vector<4xf32>141// CHECK: %[[VAL_14:.*]] = vector.mask %[[VAL_7]] { vector.transfer_write %{{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<?xf32> } : vector<4xi1> -> tensor<?xf32>142 143module attributes {transform.with_named_sequence} {144 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {145 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op146 %vector_size = transform.param.constant 4 : i64 -> !transform.param<i64>147 transform.structured.vectorize %0 vector_sizes [%vector_size] : !transform.any_op, !transform.param<i64>148 transform.yield149 }150}151 152// -----153 154func.func @vectorize_dynamic_1d_broadcast(%arg0: tensor<?xf32>,155 %arg1: tensor<?xf32>,156 %arg2: tensor<?xf32>) -> tensor<?xf32> {157 %0 = linalg.generic { indexing_maps = [affine_map<(d0) -> (0)>,158 affine_map<(d0) -> (d0)>,159 affine_map<(d0) -> (d0)>],160 iterator_types = ["parallel"] }161 ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)162 outs(%arg2 : tensor<?xf32>) {163 ^bb(%in0: f32, %in1: f32, %out: f32) :164 %0 = arith.addf %in0, %in1 : f32165 linalg.yield %0 : f32166 } -> tensor<?xf32>167 return %0 : tensor<?xf32>168}169 170// CHECK-LABEL: @vectorize_dynamic_1d_broadcast171// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index172// CHECK: %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?xf32>173// CHECK: %[[VAL_7:.*]] = vector.transfer_read %{{.*}} {permutation_map = #{{.*}}} : tensor<?xf32>, vector<4xf32>174// CHECK: %[[VAL_9:.*]] = vector.create_mask %[[VAL_4]] : vector<4xi1>175// CHECK: %[[VAL_10:.*]] = vector.mask %[[VAL_9]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>176// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_9]] { vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>177// CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_7]], %[[VAL_10]] : vector<4xf32>178// CHECK: %[[VAL_14:.*]] = vector.mask %{{.*}} { vector.transfer_write %[[VAL_13]], {{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<?xf32> } : vector<4xi1> -> tensor<?xf32>179 180module attributes {transform.with_named_sequence} {181 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {182 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op183 transform.structured.vectorize %0 vector_sizes [4] : !transform.any_op184 transform.yield185 }186}187 188// -----189 190#map = affine_map<(d0, d1) -> (d0, d1)>191#map1 = affine_map<(d0, d1) -> (d0, 0)>192 193func.func @dynamic_generic_with_reduction_and_broadcast(%arg0: tensor<?x?xf32>, %init: tensor<?x?xf32>) -> (tensor<?x?xf32>) {194 %0 = linalg.generic { indexing_maps = [#map, #map1],195 iterator_types = ["parallel", "reduction"]}196 ins(%arg0 : tensor<?x?xf32>)197 outs(%init : tensor<?x?xf32>) {198 ^bb0(%in: f32, %out: f32):199 %1 = arith.addf %in, %out : f32200 linalg.yield %1 : f32201 } -> tensor<?x?xf32>202 return %0 : tensor<?x?xf32>203}204// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1) -> (d0)>205 206// CHECK-LABEL: func.func @dynamic_generic_with_reduction_and_broadcast(207// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?xf32>,208// CHECK-SAME: %[[VAL_1:.*]]: tensor<?x?xf32>) -> tensor<?x?xf32> {209// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index210// CHECK: %[[VAL_3:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?xf32>211// CHECK: %[[VAL_4:.*]] = arith.constant 1 : index212// CHECK: %[[VAL_5:.*]] = tensor.dim %[[VAL_0]], %[[VAL_4]] : tensor<?x?xf32>213// CHECK: %[[VAL_6:.*]] = arith.constant 0 : index214// CHECK: %[[VAL_7:.*]] = ub.poison : f32215// CHECK: %[[VAL_8:.*]] = vector.create_mask %[[VAL_3]], %[[VAL_5]] : vector<4x4xi1>216// CHECK: %[[VAL_9:.*]] = vector.mask %[[VAL_8]] { vector.transfer_read %[[VAL_0]]{{\[}}%[[VAL_6]], %[[VAL_6]]], %[[VAL_7]] {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x4xf32> } : vector<4x4xi1> -> vector<4x4xf32>217// CHECK: %[[VAL_10:.*]] = ub.poison : f32218// CHECK: %[[VAL_11:.*]] = vector.create_mask %[[VAL_3]] : vector<4xi1>219// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_11]] { vector.transfer_read %[[VAL_1]]{{\[}}%[[VAL_6]], %[[VAL_6]]], %[[VAL_10]] {in_bounds = [true], permutation_map = #[[$MAP]]} : tensor<?x?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>220// CHECK: %[[VAL_13:.*]] = vector.mask %[[VAL_8]] { vector.multi_reduction <add>, %[[VAL_9]], %[[VAL_12]] [1] : vector<4x4xf32> to vector<4xf32> } : vector<4x4xi1> -> vector<4xf32>221// CHECK: %[[VAL_14:.*]] = arith.constant 0 : index222// CHECK: %[[VAL_15:.*]] = vector.mask %[[VAL_11]] { vector.transfer_write %[[VAL_13]], %[[VAL_1]]{{\[}}%[[VAL_14]], %[[VAL_14]]] {in_bounds = [true], permutation_map = #[[$MAP]]} : vector<4xf32>, tensor<?x?xf32> } : vector<4xi1> -> tensor<?x?xf32>223// CHECK: return %[[VAL_15]] : tensor<?x?xf32>224 225module attributes {transform.with_named_sequence} {226 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {227 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op228 transform.structured.vectorize %0 vector_sizes [4, 4] : !transform.any_op229 transform.yield230 }231}232 233// -----234 235func.func @vectorize_dynamic_2d_transpose(%arg0: tensor<?x?xf32>,236 %arg1: tensor<?x?xf32>,237 %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {238 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>,239 affine_map<(d0, d1) -> (d0, d1)>,240 affine_map<(d0, d1) -> (d0, d1)>],241 iterator_types = ["parallel", "parallel"] }242 ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)243 outs(%arg2 : tensor<?x?xf32>) {244 ^bb(%in0: f32, %in1: f32, %out: f32) :245 %0 = arith.addf %in0, %in1 : f32246 linalg.yield %0 : f32247 } -> tensor<?x?xf32>248 return %0 : tensor<?x?xf32>249}250 251// CHECK-LABEL: @vectorize_dynamic_2d_transpose252// CHECK: %[[VAL_3:.*]] = arith.constant 1 : index253// CHECK: %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?x?xf32>254// CHECK: %[[VAL_5:.*]] = arith.constant 0 : index255// CHECK: %[[VAL_6:.*]] = tensor.dim %{{.*}}, %[[VAL_5]] : tensor<?x?xf32>256// CHECK: %[[VAL_9:.*]] = vector.create_mask %[[VAL_6]], %[[VAL_4]] : vector<8x4xi1>257// CHECK: %[[VAL_10:.*]] = vector.mask %[[VAL_9]] { vector.transfer_read %{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : tensor<?x?xf32>, vector<4x8xf32> } : vector<8x4xi1> -> vector<4x8xf32>258// CHECK: %[[VAL_12:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_6]] : vector<4x8xi1>259// CHECK: %[[VAL_13:.*]] = vector.mask %[[VAL_12]] { vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>260// CHECK: %[[VAL_14:.*]] = ub.poison : f32261// CHECK: %[[VAL_15:.*]] = vector.mask %[[VAL_12]] { vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>262// CHECK: %[[VAL_16:.*]] = arith.addf %[[VAL_10]], %[[VAL_13]] : vector<4x8xf32>263// CHECK: %[[VAL_17:.*]] = vector.mask %[[VAL_12]] { vector.transfer_write %[[VAL_16]], %{{.*}} {in_bounds = [true, true]} : vector<4x8xf32>, tensor<?x?xf32> } : vector<4x8xi1> -> tensor<?x?xf32>264 265module attributes {transform.with_named_sequence} {266 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {267 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op268 transform.structured.vectorize %0 vector_sizes [4, 8] : !transform.any_op269 transform.yield270 }271}272 273// -----274 275func.func @vectorize_dynamic_generic_2d_broadcast(%arg0: tensor<?x?xf32>,276 %arg1: tensor<?x?xf32>,277 %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {278 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (0, d1)>,279 affine_map<(d0, d1) -> (d0, d1)>,280 affine_map<(d0, d1) -> (d0, d1)>],281 iterator_types = ["parallel", "parallel"] }282 ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)283 outs(%arg2 : tensor<?x?xf32>) {284 ^bb(%in0: f32, %in1: f32, %out: f32) :285 %0 = arith.addf %in0, %in1 : f32286 linalg.yield %0 : f32287 } -> tensor<?x?xf32>288 return %0 : tensor<?x?xf32>289}290 291// CHECK-LABEL: @vectorize_dynamic_generic_2d_broadcast292// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index293// CHECK: %[[VAL_4:.*]] = tensor.dim %{{.*}}, %[[VAL_3]] : tensor<?x?xf32>294// CHECK: %[[VAL_5:.*]] = arith.constant 1 : index295// CHECK: %[[VAL_6:.*]] = tensor.dim %{{.*}}, %[[VAL_5]] : tensor<?x?xf32>296// CHECK: %[[VAL_9:.*]] = vector.create_mask %[[VAL_6]] : vector<8xi1>297// CHECK: %[[VAL_10:.*]] = vector.mask %[[VAL_9]] { vector.transfer_read %{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : tensor<?x?xf32>, vector<4x8xf32> } : vector<8xi1> -> vector<4x8xf32>298// CHECK: %[[VAL_12:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_6]] : vector<4x8xi1>299// CHECK: %[[VAL_13:.*]] = vector.mask %[[VAL_12]] { vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>300// CHECK: %[[VAL_15:.*]] = vector.mask %[[VAL_12]] { vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>301// CHECK: %[[VAL_16:.*]] = arith.addf %[[VAL_10]], %[[VAL_13]] : vector<4x8xf32>302// CHECK: %[[VAL_18:.*]] = vector.mask %[[VAL_12]] { vector.transfer_write %{{.*}} {in_bounds = [true, true]} : vector<4x8xf32>, tensor<?x?xf32> } : vector<4x8xi1> -> tensor<?x?xf32>303 304module attributes {transform.with_named_sequence} {305 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {306 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op307 transform.structured.vectorize %0 vector_sizes [4, 8] : !transform.any_op308 transform.yield309 }310}311 312// -----313 314func.func @vectorize_dynamic_reduction_2d(%arg0: tensor<?x?xf32>,315 %arg1: tensor<?xf32>) -> tensor<?xf32> {316 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,317 affine_map<(d0, d1) -> (d0)>],318 iterator_types = ["parallel", "reduction"] }319 ins(%arg0 : tensor<?x?xf32>)320 outs(%arg1 : tensor<?xf32>) {321 ^bb(%in: f32, %out: f32) :322 %0 = arith.addf %in, %out : f32323 linalg.yield %0 : f32324 } -> tensor<?xf32>325 return %0 : tensor<?xf32>326}327 328module attributes {transform.with_named_sequence} {329 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {330 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op331 transform.structured.vectorize %0 vector_sizes [4, 8] : !transform.any_op332 transform.yield333 }334}335 336// CHECK-LABEL: @vectorize_dynamic_reduction_2d(337// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?xf32>,338// CHECK-SAME: %[[VAL_1:.*]]: tensor<?xf32>) -> tensor<?xf32> {339// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index340// CHECK: %[[VAL_3:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?xf32>341// CHECK: %[[VAL_4:.*]] = arith.constant 1 : index342// CHECK: %[[VAL_5:.*]] = tensor.dim %[[VAL_0]], %[[VAL_4]] : tensor<?x?xf32>343// CHECK: %[[VAL_8:.*]] = vector.create_mask %[[VAL_3]], %[[VAL_5]] : vector<4x8xi1>344// CHECK: %[[VAL_9:.*]] = vector.mask %[[VAL_8]] { vector.transfer_read %[[VAL_0]]{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x8xf32> } : vector<4x8xi1> -> vector<4x8xf32>345// CHECK: %[[VAL_11:.*]] = vector.create_mask %[[VAL_3]] : vector<4xi1>346// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_11]] { vector.transfer_read %[[VAL_1]]{{.*}} {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>347// CHECK: %[[VAL_13:.*]] = vector.mask %[[VAL_8]] { vector.multi_reduction <add>, %[[VAL_9]], %[[VAL_12]] [1] : vector<4x8xf32> to vector<4xf32> } : vector<4x8xi1> -> vector<4xf32>348// CHECK: %[[VAL_15:.*]] = vector.mask %[[VAL_11]] { vector.transfer_write %[[VAL_13]], %[[VAL_1]]{{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<?xf32> } : vector<4xi1> -> tensor<?xf32>349// CHECK: return %[[VAL_15]] : tensor<?xf32>350// CHECK: }351 352// -----353 354func.func @vectorize_dynamic_reduction_2d_scalable(%arg0: tensor<?x?xf32>,355 %arg1: tensor<?xf32>) -> tensor<?xf32> {356 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,357 affine_map<(d0, d1) -> (d0)>],358 iterator_types = ["parallel", "reduction"] }359 ins(%arg0 : tensor<?x?xf32>)360 outs(%arg1 : tensor<?xf32>) {361 ^bb(%in: f32, %out: f32) :362 %0 = arith.addf %in, %out : f32363 linalg.yield %0 : f32364 } -> tensor<?xf32>365 return %0 : tensor<?xf32>366}367 368// CHECK-LABEL: func.func @vectorize_dynamic_reduction_2d_scalable(369// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x?xf32>, %[[ARG_1:.*]]: tensor<?xf32>) -> tensor<?xf32> {370// CHECK: %[[C0_IDX:.*]] = arith.constant 0 : index371// CHECK: %[[DIM_A0_0:.*]] = tensor.dim %[[ARG_0]], %[[C0_IDX]] : tensor<?x?xf32>372// CHECK: %[[C1_IDX:.*]] = arith.constant 1 : index373// CHECK: %[[DIM_A0_1:.*]] = tensor.dim %[[ARG_0]], %[[C1_IDX]] : tensor<?x?xf32>374// CHECK: %[[C0_IDX:.*]] = arith.constant 0 : index375// CHECK: %[[PV:.*]] = ub.poison : f32376// CHECK: %[[MASK_2D:.*]] = vector.create_mask %[[DIM_A0_0]], %[[DIM_A0_1]] : vector<4x[8]xi1>377// CHECK: %[[VEC_RD_0:.*]] = vector.mask %[[MASK_2D]] { vector.transfer_read %[[ARG_0]][%[[C0_IDX]], %[[C0_IDX]]], %[[PV]] {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x[8]xf32> } : vector<4x[8]xi1> -> vector<4x[8]xf32>378// CHECK: %[[PV:.*]] = ub.poison : f32379// CHECK: %[[MASK_1D:.*]] = vector.create_mask %[[DIM_A0_0]] : vector<4xi1>380// CHECK: %[[VEC_RD_1:.*]] = vector.mask %[[MASK_1D]] { vector.transfer_read %[[ARG_1]][%[[C0_IDX]]], %[[PV]] {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>381// CHECK: %[[REDUCE:.*]] = vector.mask %[[MASK_2D]] { vector.multi_reduction <add>, %[[VEC_RD_0]], %[[VEC_RD_1]] [1] : vector<4x[8]xf32> to vector<4xf32> } : vector<4x[8]xi1> -> vector<4xf32>382// CHECK: %[[C0_IDX:.*]] = arith.constant 0 : index383// CHECK: %{{.*}} = vector.mask %[[MASK_1D]] { vector.transfer_write %[[REDUCE]], %[[ARG_1]][%[[C0_IDX]]] {in_bounds = [true]} : vector<4xf32>, tensor<?xf32> } : vector<4xi1> -> tensor<?xf32>384 385module attributes {transform.with_named_sequence} {386 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {387 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op388 transform.structured.vectorize %0 vector_sizes [4, [8]] : !transform.any_op389 transform.yield390 }391}392 393// -----394 395func.func @vectorize_dynamic_reduction_scalable_1d(%arg0: tensor<?xf32>,396 %arg1: tensor<f32>) -> tensor<f32> {397 398 %0 = linalg.reduce ins(%arg0 : tensor<?xf32>) outs(%arg1 : tensor<f32>) dimensions = [0]399 (%in: f32, %init: f32) {400 %0 = arith.addf %in, %init : f32401 linalg.yield %0 : f32402 }403 return %0 : tensor<f32>404}405 406// CHECK-LABEL: func.func @vectorize_dynamic_reduction_scalable_1d(407// CHECK-SAME: %[[ARG_0:.*]]: tensor<?xf32>, %[[ARG_1:.*]]: tensor<f32>) -> tensor<f32> {408// CHECK: %[[C0_IDX:.*]] = arith.constant 0 : index409// CHECK: %[[DIM_A0_0:.*]] = tensor.dim %[[ARG_0]], %[[C0_IDX]] : tensor<?xf32>410// CHECK: %[[C0_IDX:.*]] = arith.constant 0 : index411// CHECK: %[[PV:.*]] = ub.poison : f32412// CHECK: %[[MASK:.*]] = vector.create_mask %[[DIM_A0_0]] : vector<[4]xi1>413// CHECK: %[[VEC_RD_0:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[ARG_0]][%[[C0_IDX]]], %[[PV]] {in_bounds = [true]} : tensor<?xf32>, vector<[4]xf32> } : vector<[4]xi1> -> vector<[4]xf32>414// CHECK: %[[PV:.*]] = ub.poison : f32415// CHECK: %[[VEC_RD_1:.*]] = vector.transfer_read %[[ARG_1]][], %[[PV]] : tensor<f32>, vector<f32>416// CHECK: %[[ACC_f32:.*]] = vector.extract %[[VEC_RD_1]][] : f32 from vector<f32>417// CHECK: %[[REDUCE:.*]] = vector.mask %[[MASK]] { vector.multi_reduction <add>, %[[VEC_RD_0]], %[[ACC_f32]] [0] : vector<[4]xf32> to f32 } : vector<[4]xi1> -> f32418// CHECK: %[[VEC_f32:.*]] = vector.broadcast %[[REDUCE]] : f32 to vector<f32>419// CHECK: %{{.*}} = vector.transfer_write %[[VEC_f32]], %[[ARG_1]][] : vector<f32>, tensor<f32>420 421module attributes {transform.with_named_sequence} {422 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {423 %0 = transform.structured.match ops{["linalg.reduce"]} in %arg1 : (!transform.any_op) -> !transform.any_op424 transform.structured.vectorize %0 vector_sizes [[4]] : !transform.any_op425 transform.yield426 }427}428 429// -----430 431func.func @vectorize_dynamic_transpose_reduction(%arg0: tensor<?x?x?xf32>,432 %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {433 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,434 affine_map<(d0, d1, d2) -> (d2, d1)>],435 iterator_types = ["reduction", "parallel", "parallel"] }436 ins(%arg0 : tensor<?x?x?xf32>)437 outs(%arg1 : tensor<?x?xf32>) {438 ^bb(%in: f32, %out: f32) :439 %0 = arith.addf %in, %out : f32440 linalg.yield %0 : f32441 } -> tensor<?x?xf32>442 return %0 : tensor<?x?xf32>443}444 445module attributes {transform.with_named_sequence} {446 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {447 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op448 transform.structured.vectorize %0 vector_sizes [4, 8, 16] : !transform.any_op449 transform.yield450 }451}452 453// CHECK-LABEL: @vectorize_dynamic_transpose_reduction(454// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?x?xf32>,455// CHECK-SAME: %[[VAL_1:.*]]: tensor<?x?xf32>) -> tensor<?x?xf32> {456// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index457// CHECK: %[[VAL_3:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?x?xf32>458// CHECK: %[[VAL_4:.*]] = arith.constant 1 : index459// CHECK: %[[VAL_5:.*]] = tensor.dim %[[VAL_0]], %[[VAL_4]] : tensor<?x?x?xf32>460// CHECK: %[[VAL_6:.*]] = arith.constant 2 : index461// CHECK: %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_6]] : tensor<?x?x?xf32>462// CHECK: %[[VAL_10:.*]] = vector.create_mask %[[VAL_3]], %[[VAL_5]], %[[VAL_7]] : vector<4x8x16xi1>463// CHECK: %[[VAL_11:.*]] = vector.mask %[[VAL_10]] { vector.transfer_read %[[VAL_0]]{{.*}} {in_bounds = [true, true, true]} : tensor<?x?x?xf32>, vector<4x8x16xf32> } : vector<4x8x16xi1> -> vector<4x8x16xf32>464// CHECK: %[[VAL_13:.*]] = vector.create_mask %[[VAL_7]], %[[VAL_5]] : vector<16x8xi1>465// CHECK: %[[VAL_14:.*]] = vector.mask %[[VAL_13]] { vector.transfer_read %[[VAL_1]]{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : tensor<?x?xf32>, vector<8x16xf32> } : vector<16x8xi1> -> vector<8x16xf32>466// CHECK: %[[VAL_15:.*]] = vector.mask %[[VAL_10]] { vector.multi_reduction <add>, %[[VAL_11]], %[[VAL_14]] [0] : vector<4x8x16xf32> to vector<8x16xf32> } : vector<4x8x16xi1> -> vector<8x16xf32>467// CHECK: %[[VAL_17:.*]] = vector.mask %[[VAL_13]] { vector.transfer_write %[[VAL_15]], %{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : vector<8x16xf32>, tensor<?x?xf32> } : vector<16x8xi1> -> tensor<?x?xf32>468 469// -----470 471func.func @vectorize_dynamic_transpose_reduction_with_params(%arg0: tensor<?x?x?xf32>,472 %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {473 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,474 affine_map<(d0, d1, d2) -> (d2, d1)>],475 iterator_types = ["reduction", "parallel", "parallel"] }476 ins(%arg0 : tensor<?x?x?xf32>)477 outs(%arg1 : tensor<?x?xf32>) {478 ^bb(%in: f32, %out: f32) :479 %0 = arith.addf %in, %out : f32480 linalg.yield %0 : f32481 } -> tensor<?x?xf32>482 return %0 : tensor<?x?xf32>483}484 485module attributes {transform.with_named_sequence} {486 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {487 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op488 %vector_size_0 = transform.param.constant 4 : i64 -> !transform.param<i64>489 %vector_size_2 = transform.param.constant 16 : i64 -> !transform.param<i64>490 transform.structured.vectorize %0 vector_sizes491 [%vector_size_0, 8, %vector_size_2] : !transform.any_op, !transform.param<i64>, !transform.param<i64>492 transform.yield493 }494}495 496// CHECK-LABEL: @vectorize_dynamic_transpose_reduction_with_params(497// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?x?xf32>,498// CHECK-SAME: %[[VAL_1:.*]]: tensor<?x?xf32>) -> tensor<?x?xf32> {499// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index500// CHECK: %[[VAL_3:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?x?xf32>501// CHECK: %[[VAL_4:.*]] = arith.constant 1 : index502// CHECK: %[[VAL_5:.*]] = tensor.dim %[[VAL_0]], %[[VAL_4]] : tensor<?x?x?xf32>503// CHECK: %[[VAL_6:.*]] = arith.constant 2 : index504// CHECK: %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_6]] : tensor<?x?x?xf32>505// CHECK: %[[VAL_10:.*]] = vector.create_mask %[[VAL_3]], %[[VAL_5]], %[[VAL_7]] : vector<4x8x16xi1>506// CHECK: %[[VAL_11:.*]] = vector.mask %[[VAL_10]] { vector.transfer_read %[[VAL_0]]{{.*}} {in_bounds = [true, true, true]} : tensor<?x?x?xf32>, vector<4x8x16xf32> } : vector<4x8x16xi1> -> vector<4x8x16xf32>507// CHECK: %[[VAL_13:.*]] = vector.create_mask %[[VAL_7]], %[[VAL_5]] : vector<16x8xi1>508// CHECK: %[[VAL_14:.*]] = vector.mask %[[VAL_13]] { vector.transfer_read %[[VAL_1]]{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : tensor<?x?xf32>, vector<8x16xf32> } : vector<16x8xi1> -> vector<8x16xf32>509// CHECK: %[[VAL_15:.*]] = vector.mask %[[VAL_10]] { vector.multi_reduction <add>, %[[VAL_11]], %[[VAL_14]] [0] : vector<4x8x16xf32> to vector<8x16xf32> } : vector<4x8x16xi1> -> vector<8x16xf32>510// CHECK: %[[VAL_17:.*]] = vector.mask %[[VAL_13]] { vector.transfer_write %[[VAL_15]], %{{.*}} {in_bounds = [true, true], permutation_map = #{{.*}}} : vector<8x16xf32>, tensor<?x?xf32> } : vector<16x8xi1> -> tensor<?x?xf32>511 512// -----513 514func.func @vectorize_partial_dynamic_identity(%arg0: tensor<8x?xf32>,515 %arg1: tensor<8x?xf32>,516 %arg2: tensor<8x?xf32>) -> tensor<8x?xf32> {517 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,518 affine_map<(d0, d1) -> (d0, d1)>,519 affine_map<(d0, d1) -> (d0, d1)>],520 iterator_types = ["parallel", "parallel"] }521 ins(%arg0, %arg1 : tensor<8x?xf32>, tensor<8x?xf32>)522 outs(%arg2 : tensor<8x?xf32>) {523 ^bb(%in0: f32, %in1: f32, %out: f32) :524 %0 = arith.addf %in0, %in1 : f32525 linalg.yield %0 : f32526 } -> tensor<8x?xf32>527 return %0 : tensor<8x?xf32>528}529 530// CHECK-LABEL: func.func @vectorize_partial_dynamic_identity(531// CHECK-SAME: %[[VAL_0:.*]]: tensor<8x?xf32>, %[[VAL_1:.*]]: tensor<8x?xf32>, %[[VAL_2:.*]]: tensor<8x?xf32>) -> tensor<8x?xf32> {532// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 1 : index533// CHECK-DAG: %[[VAL_4:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<8x?xf32>534// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index535// CHECK-DAG: %[[VAL_6:.*]] = ub.poison : f32536// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 8 : index537// CHECK: %[[VAL_8:.*]] = vector.create_mask %[[VAL_7]], %[[VAL_4]] : vector<8x32xi1>538// CHECK: %[[VAL_9:.*]] = vector.mask %[[VAL_8]] { vector.transfer_read %[[VAL_0]][%[[VAL_5]], %[[VAL_5]]], %[[VAL_6]] {in_bounds = [true, true]} : tensor<8x?xf32>, vector<8x32xf32> } : vector<8x32xi1> -> vector<8x32xf32>539// CHECK: %[[VAL_10:.*]] = ub.poison : f32540// CHECK: %[[VAL_11:.*]] = vector.mask %[[VAL_8]] { vector.transfer_read %[[VAL_1]][%[[VAL_5]], %[[VAL_5]]], %[[VAL_10]] {in_bounds = [true, true]} : tensor<8x?xf32>, vector<8x32xf32> } : vector<8x32xi1> -> vector<8x32xf32>541// CHECK: %[[VAL_12:.*]] = ub.poison : f32542// CHECK: %[[VAL_13:.*]] = vector.mask %[[VAL_8]] { vector.transfer_read %[[VAL_2]][%[[VAL_5]], %[[VAL_5]]], %[[VAL_12]] {in_bounds = [true, true]} : tensor<8x?xf32>, vector<8x32xf32> } : vector<8x32xi1> -> vector<8x32xf32>543// CHECK: %[[VAL_14:.*]] = arith.addf %[[VAL_9]], %[[VAL_11]] : vector<8x32xf32>544// CHECK: %[[VAL_15:.*]] = arith.constant 0 : index545// CHECK: %[[VAL_16:.*]] = vector.mask %[[VAL_8]] { vector.transfer_write %[[VAL_14]], %[[VAL_2]][%[[VAL_15]], %[[VAL_15]]] {in_bounds = [true, true]} : vector<8x32xf32>, tensor<8x?xf32> } : vector<8x32xi1> -> tensor<8x?xf32>546 547 548module attributes {transform.with_named_sequence} {549 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {550 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op551 transform.structured.vectorize %0 vector_sizes [8, 32] : !transform.any_op552 transform.yield553 }554}555 556// -----557 558func.func @vectorize_partial_dynamic_identity_scalable(%arg0: tensor<8x?xf32>,559 %arg1: tensor<8x?xf32>,560 %arg2: tensor<8x?xf32>) -> tensor<8x?xf32> {561 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,562 affine_map<(d0, d1) -> (d0, d1)>,563 affine_map<(d0, d1) -> (d0, d1)>],564 iterator_types = ["parallel", "parallel"] }565 ins(%arg0, %arg1 : tensor<8x?xf32>, tensor<8x?xf32>)566 outs(%arg2 : tensor<8x?xf32>) {567 ^bb(%in0: f32, %in1: f32, %out: f32) :568 %0 = arith.addf %in0, %in1 : f32569 linalg.yield %0 : f32570 } -> tensor<8x?xf32>571 return %0 : tensor<8x?xf32>572}573 574// CHECK-LABEL: func.func @vectorize_partial_dynamic_identity_scalable575// CHECK-SAME: %[[VAL_0:.*]]: tensor<8x?xf32>, %[[VAL_1:.*]]: tensor<8x?xf32>, %[[VAL_2:.*]]: tensor<8x?xf32>) -> tensor<8x?xf32> {576// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 1 : index577// CHECK-DAG: %[[VAL_4:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<8x?xf32>578// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index579// CHECK-DAG: %[[VAL_6:.*]] = ub.poison : f32580// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 8 : index581// CHECK: %[[VAL_8:.*]] = vector.create_mask %[[VAL_7]], %[[VAL_4]] : vector<8x[32]xi1>582// CHECK: %[[VAL_9:.*]] = vector.mask %[[VAL_8]] { vector.transfer_read %[[VAL_0]][%[[VAL_5]], %[[VAL_5]]], %[[VAL_6]] {in_bounds = [true, true]} : tensor<8x?xf32>, vector<8x[32]xf32> } : vector<8x[32]xi1> -> vector<8x[32]xf32>583// CHECK: %[[VAL_10:.*]] = ub.poison : f32584// CHECK: %[[VAL_11:.*]] = vector.mask %[[VAL_8]] { vector.transfer_read %[[VAL_1]][%[[VAL_5]], %[[VAL_5]]], %[[VAL_10]] {in_bounds = [true, true]} : tensor<8x?xf32>, vector<8x[32]xf32> } : vector<8x[32]xi1> -> vector<8x[32]xf32>585// CHECK: %[[VAL_12:.*]] = ub.poison : f32586// CHECK: %[[VAL_13:.*]] = vector.mask %[[VAL_8]] { vector.transfer_read %[[VAL_2]][%[[VAL_5]], %[[VAL_5]]], %[[VAL_12]] {in_bounds = [true, true]} : tensor<8x?xf32>, vector<8x[32]xf32> } : vector<8x[32]xi1> -> vector<8x[32]xf32>587// CHECK: %[[VAL_14:.*]] = arith.addf %[[VAL_9]], %[[VAL_11]] : vector<8x[32]xf32>588// CHECK: %[[VAL_15:.*]] = arith.constant 0 : index589// CHECK: %[[VAL_16:.*]] = vector.mask %[[VAL_8]] { vector.transfer_write %[[VAL_14]], %[[VAL_2]][%[[VAL_15]], %[[VAL_15]]] {in_bounds = [true, true]} : vector<8x[32]xf32>, tensor<8x?xf32> } : vector<8x[32]xi1> -> tensor<8x?xf32>590 591 592module attributes {transform.with_named_sequence} {593 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {594 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op595 transform.structured.vectorize %0 vector_sizes [8, [32]] : !transform.any_op596 transform.yield597 }598}599 600// -----601 602func.func @do_not_generate_masks(%arg0: tensor<8x32xf32>,603 %arg1: tensor<8x32xf32>,604 %arg2: tensor<8x32xf32>) -> tensor<8x32xf32> {605 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,606 affine_map<(d0, d1) -> (d0, d1)>,607 affine_map<(d0, d1) -> (d0, d1)>],608 iterator_types = ["parallel", "parallel"] }609 ins(%arg0, %arg1 : tensor<8x32xf32>, tensor<8x32xf32>)610 outs(%arg2 : tensor<8x32xf32>) {611 ^bb(%in0: f32, %in1: f32, %out: f32) :612 %0 = arith.addf %in0, %in1 : f32613 linalg.yield %0 : f32614 } -> tensor<8x32xf32>615 return %0 : tensor<8x32xf32>616}617 618// CHECK-LABEL: func.func @do_not_generate_masks619// CHECK-NOT: vector.mask620 621module attributes {transform.with_named_sequence} {622 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {623 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op624 transform.structured.vectorize %0 vector_sizes [8, 32] : !transform.any_op625 transform.yield626 }627}628 629// -----630 631func.func @vectorize_static_shape_with_mask(%arg0: tensor<8x30xf32>,632 %arg1: tensor<8x30xf32>,633 %arg2: tensor<8x30xf32>) -> tensor<8x30xf32> {634 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,635 affine_map<(d0, d1) -> (d0, d1)>,636 affine_map<(d0, d1) -> (d0, d1)>],637 iterator_types = ["parallel", "parallel"] }638 ins(%arg0, %arg1 : tensor<8x30xf32>, tensor<8x30xf32>)639 outs(%arg2 : tensor<8x30xf32>) {640 ^bb(%in0: f32, %in1: f32, %out: f32) :641 %0 = arith.addf %in0, %in1 : f32642 linalg.yield %0 : f32643 } -> tensor<8x30xf32>644 return %0 : tensor<8x30xf32>645}646 647// CHECK-LABEL: func.func @vectorize_static_shape_with_mask(648// CHECK-SAME: %[[VAL_0:.*]]: tensor<8x30xf32>, %[[VAL_1:.*]]: tensor<8x30xf32>, %[[VAL_2:.*]]: tensor<8x30xf32>) -> tensor<8x30xf32> {649// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index650// CHECK-DAG: %[[VAL_4:.*]] = ub.poison : f32651// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index652// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 30 : index653// CHECK: %[[VAL_7:.*]] = vector.create_mask %[[VAL_5]], %[[VAL_6]] : vector<8x32xi1>654// CHECK: %[[VAL_8:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %[[VAL_0]][%[[VAL_3]], %[[VAL_3]]], %[[VAL_4]] {in_bounds = [true, true]} : tensor<8x30xf32>, vector<8x32xf32> } : vector<8x32xi1> -> vector<8x32xf32>655// CHECK: %[[VAL_9:.*]] = ub.poison : f32656// CHECK: %[[VAL_10:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %[[VAL_1]][%[[VAL_3]], %[[VAL_3]]], %[[VAL_9]] {in_bounds = [true, true]} : tensor<8x30xf32>, vector<8x32xf32> } : vector<8x32xi1> -> vector<8x32xf32>657// CHECK: %[[VAL_11:.*]] = ub.poison : f32658// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %[[VAL_2]][%[[VAL_3]], %[[VAL_3]]], %[[VAL_11]] {in_bounds = [true, true]} : tensor<8x30xf32>, vector<8x32xf32> } : vector<8x32xi1> -> vector<8x32xf32>659// CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_8]], %[[VAL_10]] : vector<8x32xf32>660// CHECK: %[[VAL_14:.*]] = arith.constant 0 : index661// CHECK: %[[VAL_15:.*]] = vector.mask %[[VAL_7]] { vector.transfer_write %[[VAL_13]], %[[VAL_2]][%[[VAL_14]], %[[VAL_14]]] {in_bounds = [true, true]} : vector<8x32xf32>, tensor<8x30xf32> } : vector<8x32xi1> -> tensor<8x30xf32>662 663module attributes {transform.with_named_sequence} {664 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {665 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op666 transform.structured.vectorize %0 vector_sizes [8, 32] : !transform.any_op667 transform.yield668 }669}670 671// -----672 673func.func @vectorize_static_shape_with_mask_scalable(%arg0: tensor<8x30xf32>,674 %arg1: tensor<8x30xf32>,675 %arg2: tensor<8x30xf32>) -> tensor<8x30xf32> {676 %0 = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,677 affine_map<(d0, d1) -> (d0, d1)>,678 affine_map<(d0, d1) -> (d0, d1)>],679 iterator_types = ["parallel", "parallel"] }680 ins(%arg0, %arg1 : tensor<8x30xf32>, tensor<8x30xf32>)681 outs(%arg2 : tensor<8x30xf32>) {682 ^bb(%in0: f32, %in1: f32, %out: f32) :683 %0 = arith.addf %in0, %in1 : f32684 linalg.yield %0 : f32685 } -> tensor<8x30xf32>686 return %0 : tensor<8x30xf32>687}688 689// CHECK-LABEL: func.func @vectorize_static_shape_with_mask_scalable(690// CHECK-SAME: %[[VAL_0:.*]]: tensor<8x30xf32>, %[[VAL_1:.*]]: tensor<8x30xf32>, %[[VAL_2:.*]]: tensor<8x30xf32>) -> tensor<8x30xf32> {691// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index692// CHECK-DAG: %[[VAL_4:.*]] = ub.poison : f32693// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index694// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 30 : index695// CHECK: %[[VAL_7:.*]] = vector.create_mask %[[VAL_5]], %[[VAL_6]] : vector<8x[32]xi1>696// CHECK: %[[VAL_8:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %[[VAL_0]][%[[VAL_3]], %[[VAL_3]]], %[[VAL_4]] {in_bounds = [true, true]} : tensor<8x30xf32>, vector<8x[32]xf32> } : vector<8x[32]xi1> -> vector<8x[32]xf32>697// CHECK: %[[VAL_9:.*]] = ub.poison : f32698// CHECK: %[[VAL_10:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %[[VAL_1]][%[[VAL_3]], %[[VAL_3]]], %[[VAL_9]] {in_bounds = [true, true]} : tensor<8x30xf32>, vector<8x[32]xf32> } : vector<8x[32]xi1> -> vector<8x[32]xf32>699// CHECK: %[[VAL_11:.*]] = ub.poison : f32700// CHECK: %[[VAL_12:.*]] = vector.mask %[[VAL_7]] { vector.transfer_read %[[VAL_2]][%[[VAL_3]], %[[VAL_3]]], %[[VAL_11]] {in_bounds = [true, true]} : tensor<8x30xf32>, vector<8x[32]xf32> } : vector<8x[32]xi1> -> vector<8x[32]xf32>701// CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_8]], %[[VAL_10]] : vector<8x[32]xf32>702// CHECK: %[[VAL_14:.*]] = arith.constant 0 : index703// CHECK: %[[VAL_15:.*]] = vector.mask %[[VAL_7]] { vector.transfer_write %[[VAL_13]], %[[VAL_2]][%[[VAL_14]], %[[VAL_14]]] {in_bounds = [true, true]} : vector<8x[32]xf32>, tensor<8x30xf32> } : vector<8x[32]xi1> -> tensor<8x30xf32>704 705module attributes {transform.with_named_sequence} {706 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {707 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op708 transform.structured.vectorize %0 vector_sizes [8, [32]] : !transform.any_op709 transform.yield710 }711}712 713// -----714 715///----------------------------------------------------------------------------------------716/// Tests for linalg.matvec717///----------------------------------------------------------------------------------------718 719// Scalable _reduction_ dimension.720 721func.func @vectorize_dynamic_matvec_trailing_reduction_dim(%arg0: tensor<?x?xf32>,722 %arg1: tensor<?xf32>,723 %arg2: tensor<?xf32>) {724 linalg.matvec ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)725 outs(%arg2 : tensor<?xf32>) -> tensor<?xf32>726 return727}728 729// CHECK-LABEL: func.func @vectorize_dynamic_matvec_trailing_reduction_dim(730// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x?xf32>, %[[ARG_1:.*]]: tensor<?xf32>, %[[ARG_2:.*]]: tensor<?xf32>) {731// CHECK: %[[C0_idx:.*]] = arith.constant 0 : index732// CHECK: %[[DIM_A0_0:.*]] = tensor.dim %[[ARG_0]], %[[C0_idx]] : tensor<?x?xf32>733// CHECK: %[[C1_idx:.*]] = arith.constant 1 : index734// CHECK: %[[DIM_A0_1:.*]] = tensor.dim %[[ARG_0]], %[[C1_idx]] : tensor<?x?xf32>735// CHECK: %[[C0_idx:.*]] = arith.constant 0 : index736// CHECK: %[[PV:.*]] = ub.poison : f32737// CHECK: %[[MASK_2d:.*]] = vector.create_mask %[[DIM_A0_0]], %[[DIM_A0_1]] : vector<4x[4]xi1>738// CHECK: %[[VEC_RD_0:.*]] = vector.mask %[[MASK_2d]] { vector.transfer_read %[[ARG_0]][%[[C0_idx]], %[[C0_idx]]], %[[PV]] {in_bounds = [true, true]} : tensor<?x?xf32>, vector<4x[4]xf32> } : vector<4x[4]xi1> -> vector<4x[4]xf32>739// CHECK: %[[PV:.*]] = ub.poison : f32740// CHECK: %[[MASK_d1:.*]] = vector.create_mask %[[DIM_A0_1]] : vector<[4]xi1>741// CHECK: %[[VEC_RD_1:.*]] = vector.mask %[[MASK_d1]] { vector.transfer_read %[[ARG_1]][%[[C0_idx]]], %[[PV]] {in_bounds = [true, true], permutation_map = #map} : tensor<?xf32>, vector<4x[4]xf32> } : vector<[4]xi1> -> vector<4x[4]xf32>742// CHECK: %[[PV:.*]] = ub.poison : f32743// CHECK: %[[MASK_d2:.*]] = vector.create_mask %[[DIM_A0_0]] : vector<4xi1>744// CHECK: %[[VEC_RD_2:.*]] = vector.mask %[[MASK_d2]] { vector.transfer_read %[[ARG_2]][%[[C0_idx]]], %[[PV]] {in_bounds = [true]} : tensor<?xf32>, vector<4xf32> } : vector<4xi1> -> vector<4xf32>745// CHECK: %[[MUL:.*]] = arith.mulf %[[VEC_RD_0:.*]], %[[VEC_RD_1:.*]] : vector<4x[4]xf32>746// CHECK: %[[REDUCE:.*]] = vector.mask %[[MASK_2d]] { vector.multi_reduction <add>, %[[MUL]], %[[VEC_RD_2]] [1] : vector<4x[4]xf32> to vector<4xf32> } : vector<4x[4]xi1> -> vector<4xf32>747// CHECK: %[[C0_idx:.*]] = arith.constant 0 : index748// CHECK: %{{.*}} = vector.mask %[[MASK_d2]] { vector.transfer_write %[[REDUCE]], %[[ARG_2]][%[[C0_idx]]] {in_bounds = [true]} : vector<4xf32>, tensor<?xf32> } : vector<4xi1> -> tensor<?xf32>749 750module attributes {transform.with_named_sequence} {751 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {752 %0 = transform.structured.match ops{["linalg.matvec"]} in %arg1 : (!transform.any_op) -> !transform.any_op753 transform.structured.vectorize %0 vector_sizes [4, [4]] : !transform.any_op754 transform.yield755 }756}757 758// -----759 760// Scalable _parallel_ dimension.761 762func.func @vectorize_dynamic_matvec_trailing_reduction_dim(%arg0: tensor<?x?xf32>,763 %arg1: tensor<?xf32>,764 %arg2:765 tensor<?xf32>) ->766 tensor<?xf32>{767 %0 = linalg.matvec ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)768 outs(%arg2 : tensor<?xf32>) -> tensor<?xf32>769 return %0 : tensor<?xf32>770}771 772// CHECK-LABEL: func.func @vectorize_dynamic_matvec_trailing_reduction_dim(773// CHECK-SAME: %[[ARG_0:.*]]: tensor<?x?xf32>, %[[ARG_1:.*]]: tensor<?xf32>, %[[ARG_2:.*]]: tensor<?xf32>) -> tensor<?xf32> {774// CHECK: %[[C0_idx:.*]] = arith.constant 0 : index775// CHECK: %[[DIM_A0_0:.*]] = tensor.dim %[[ARG_0]], %[[C0_idx]] : tensor<?x?xf32>776// CHECK: %[[C1_idx:.*]] = arith.constant 1 : index777// CHECK: %[[DIM_A0_1:.*]] = tensor.dim %[[ARG_0]], %[[C1_idx]] : tensor<?x?xf32>778// CHECK: %[[C0_idx:.*]] = arith.constant 0 : index779// CHECK: %[[PV:.*]] = ub.poison : f32780// CHECK: %[[MASK_2d:.*]] = vector.create_mask %[[DIM_A0_0]], %[[DIM_A0_1]] : vector<[4]x4xi1>781// CHECK: %[[VEC_RD_0:.*]] = vector.mask %[[MASK_2d]] { vector.transfer_read %[[ARG_0]][%[[C0_idx]], %[[C0_idx]]], %[[PV]] {in_bounds = [true, true]} : tensor<?x?xf32>, vector<[4]x4xf32> } : vector<[4]x4xi1> -> vector<[4]x4xf32>782// CHECK: %[[PV:.*]] = ub.poison : f32783// CHECK: %[[MASK_d1:.*]] = vector.create_mask %[[DIM_A0_1]] : vector<4xi1>784// CHECK: %[[VEC_RD_1:.*]] = vector.mask %[[MASK_d1]] { vector.transfer_read %[[ARG_1]][%[[C0_idx]]], %[[PV]] {in_bounds = [true, true], permutation_map = #map} : tensor<?xf32>, vector<[4]x4xf32> } : vector<4xi1> -> vector<[4]x4xf32>785// CHECK: %[[PV:.*]] = ub.poison : f32786// CHECK: %[[MASK_d2:.*]] = vector.create_mask %[[DIM_A0_0]] : vector<[4]xi1>787// CHECK: %[[VEC_RD_2:.*]] = vector.mask %[[MASK_d2]] { vector.transfer_read %[[ARG_2]][%[[C0_idx]]], %[[PV]] {in_bounds = [true]} : tensor<?xf32>, vector<[4]xf32> } : vector<[4]xi1> -> vector<[4]xf32>788// CHECK: %[[MUL:.*]] = arith.mulf %[[VEC_RD_0:.*]], %[[VEC_RD_1:.*]] : vector<[4]x4xf32>789// CHECK: %[[REDUCE:.*]] = vector.mask %[[MASK_2d]] { vector.multi_reduction <add>, %[[MUL]], %[[VEC_RD_2]] [1] : vector<[4]x4xf32> to vector<[4]xf32> } : vector<[4]x4xi1> -> vector<[4]xf32>790// CHECK: %[[C0_idx:.*]] = arith.constant 0 : index791// CHECK: %{{.*}} = vector.mask %[[MASK_d2]] { vector.transfer_write %[[REDUCE]], %[[ARG_2]][%[[C0_idx]]] {in_bounds = [true]} : vector<[4]xf32>, tensor<?xf32> } : vector<[4]xi1> -> tensor<?xf32>792 793module attributes {transform.with_named_sequence} {794 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {795 %0 = transform.structured.match ops{["linalg.matvec"]} in %arg1 : (!transform.any_op) -> !transform.any_op796 transform.structured.vectorize %0 vector_sizes [[4], 4] : !transform.any_op797 transform.yield798 }799}800 801// -----802 803///----------------------------------------------------------------------------------------804/// Tests for linalg.index805///----------------------------------------------------------------------------------------806 807#map = affine_map<(d0) -> (d0)>808func.func @vectorize_linalg_index_scalable(%dest: tensor<?xindex>) -> tensor<?xindex> {809 %0 = linalg.generic {810 indexing_maps = [#map],811 iterator_types = ["parallel"]812 } outs(%dest : tensor<?xindex>) {813 ^bb0(%in: index):814 %1 = linalg.index 0 : index815 linalg.yield %1: index816 } -> tensor<?xindex>817 return %0 : tensor<?xindex>818}819 820// CHECK-LABEL: func.func @vectorize_linalg_index_scalable(821// CHECK-SAME: %[[DEST:.*]]: tensor<?xindex>) -> tensor<?xindex> {822// CHECK: %[[C0:.*]] = arith.constant 0 : index823// CHECK: %[[D0:.*]] = tensor.dim %[[DEST]], %[[C0]] : tensor<?xindex>824// CHECK: %[[C0_1:.*]] = arith.constant 0 : index825// CHECK: %[[PV:.*]] = ub.poison : index826// CHECK: %[[MASK:.*]] = vector.create_mask %[[D0]] : vector<[4]xi1>827// TODO: This xfer_read is not used - avoid creating it.828// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[DEST]]{{\[}}%[[C0_1]]], %[[PV]] {in_bounds = [true]} : tensor<?xindex>, vector<[4]xindex> } : vector<[4]xi1> -> vector<[4]xindex>829// CHECK: %[[STEP:.*]] = vector.step : vector<[4]xindex>830// CHECK: %[[C0_3:.*]] = arith.constant 0 : index831// CHECK: %[[WRITE:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[STEP]], %[[DEST]]{{\[}}%[[C0_3]]] {in_bounds = [true]} : vector<[4]xindex>, tensor<?xindex> } : vector<[4]xi1> -> tensor<?xindex>832// CHECK: return %[[WRITE]] : tensor<?xindex>833 834module attributes {transform.with_named_sequence} {835 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {836 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op837 transform.structured.vectorize %0 vector_sizes [[4]] : !transform.any_op838 839 transform.yield840 }841}842 843// -----844 845///----------------------------------------------------------------------------------------846/// Tests for linalg.mmt4d847///----------------------------------------------------------------------------------------848 849func.func @mmt4d(%A: memref<16x16x8x1xf32>, %B: memref<16x16x8x1xf32>, %C_in: memref<16x16x8x8xf32>) {850 linalg.mmt4d ins(%A, %B: memref<16x16x8x1xf32>, memref<16x16x8x1xf32>)851 outs(%C_in: memref<16x16x8x8xf32>)852 return853}854 855// CHECK-LABEL: func.func @mmt4d(856// CHECK-SAME: %[[A:.*]]: memref<16x16x8x1xf32>, %[[B:.*]]: memref<16x16x8x1xf32>, %[[C:.*]]: memref<16x16x8x8xf32>) {857// CHECK: %[[VEC_A:.*]] = vector.transfer_read %[[A]]{{.*}} : memref<16x16x8x1xf32>, vector<16x16x16x8x8x1xf32>858// CHECK: %[[VEC_B:.*]] = vector.transfer_read %[[B]]{{.*}} : memref<16x16x8x1xf32>, vector<16x16x16x8x8x1xf32>859// CHECK: %[[VEC_C:.*]] = vector.transfer_read %[[C]]{{.*}} : memref<16x16x8x8xf32>, vector<16x16x8x8xf32>860// CHECK: %[[MUL:.*]] = arith.mulf %[[VEC_A]], %[[VEC_B]] : vector<16x16x16x8x8x1xf32>861// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[VEC_C]] [2, 5] : vector<16x16x16x8x8x1xf32> to vector<16x16x8x8xf32>862// CHECK: vector.transfer_write %[[RED]], %[[C]]{{.*}} : vector<16x16x8x8xf32>, memref<16x16x8x8xf32>863 864module attributes {transform.with_named_sequence} {865 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {866 %mmt4d = transform.structured.match ops{["linalg.mmt4d"]} in %arg1 : (!transform.any_op) -> !transform.any_op867 transform.structured.vectorize %mmt4d : !transform.any_op868 transform.yield869 }870}871 872// -----873 874func.func @mmt4d_scalable(%A: memref<16x16x8x1xf32>, %B: memref<16x16x?x1xf32>, %C_in: memref<16x16x8x?xf32>) {875 linalg.mmt4d ins(%A, %B: memref<16x16x8x1xf32>, memref<16x16x?x1xf32>)876 outs(%C_in: memref<16x16x8x?xf32>)877 return878}879// CHECK-LABEL: func.func @mmt4d_scalable(880// CHECK-SAME: %[[A:.*]]: memref<16x16x8x1xf32>,881// CHECK-SAME: %[[B:.*]]: memref<16x16x?x1xf32>,882// CHECK-SAME: %[[C_IN:.*]]: memref<16x16x8x?xf32>) {883// CHECK: %[[C16_M:.*]] = arith.constant 16 : index884// CHECK: %[[C16_N:.*]] = arith.constant 16 : index885// CHECK: %[[C16_K:.*]] = arith.constant 16 : index886// CHECK: %[[C8:.*]] = arith.constant 8 : index887// CHECK: %[[C2:.*]] = arith.constant 2 : index888// CHECK: %[[DIM_2:.*]] = memref.dim %[[B]], %[[C2]] : memref<16x16x?x1xf32>889// CHECK: %[[C1:.*]] = arith.constant 1 : index890// CHECK: %[[VEC_A:.*]] = vector.transfer_read %[[A]]{{.*}} : memref<16x16x8x1xf32>, vector<16x16x16x8x[4]x1xf32>891// CHECK: %[[MASK_1:.*]] = vector.create_mask %[[C16_N]], %[[C16_K]], %[[DIM_2]], %[[C1]] : vector<16x16x[4]x1xi1>892// CHECK: %[[VEC_B:.*]] = vector.mask %[[MASK_1]] { vector.transfer_read %[[B]]{{.*}} : memref<16x16x?x1xf32>, vector<16x16x16x8x[4]x1xf32> } : vector<16x16x[4]x1xi1> -> vector<16x16x16x8x[4]x1xf32>893// CHECK: %[[MASK_2:.*]] = vector.create_mask %[[C16_M]], %[[C16_N]], %[[C8]], %[[DIM_2]] : vector<16x16x8x[4]xi1>894// CHECK: %[[VEC_C:.*]] = vector.mask %[[MASK_2]] { vector.transfer_read %[[C_IN]]{{.*}} : memref<16x16x8x?xf32>, vector<16x16x8x[4]xf32> } : vector<16x16x8x[4]xi1> -> vector<16x16x8x[4]xf32>895// CHECK: %[[MUL:.*]] = arith.mulf %[[VEC_A]], %[[VEC_B]] : vector<16x16x16x8x[4]x1xf32>896// CHECK: %[[MASK_3:.*]] = vector.create_mask %[[C16_M]], %[[C16_N]], %[[C16_K]], %[[C8]], %[[DIM_2]], %[[C1]] : vector<16x16x16x8x[4]x1xi1>897// CHECK: %[[RED:.*]] = vector.mask %[[MASK_3]] { vector.multi_reduction <add>, %[[MUL]], %[[VEC_C]] [2, 5] : vector<16x16x16x8x[4]x1xf32> to vector<16x16x8x[4]xf32> } : vector<16x16x16x8x[4]x1xi1> -> vector<16x16x8x[4]xf32>898// CHECK: vector.mask %[[MASK_2]] { vector.transfer_write %[[RED]], %[[C_IN]]{{.*}} : vector<16x16x8x[4]xf32>, memref<16x16x8x?xf32> } : vector<16x16x8x[4]xi1>899 900 901module attributes {transform.with_named_sequence} {902 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {903 %mmt4d = transform.structured.match ops{["linalg.mmt4d"]} in %arg1 : (!transform.any_op) -> !transform.any_op904 transform.structured.vectorize %mmt4d vector_sizes [16, 16, 16, 8, [4], 1] : !transform.any_op905 transform.yield906 }907}908 909// -----910 911func.func @mmt4d_scalable_with_assume(%A: memref<16x16x8x1xf32>, %B: memref<16x16x?x1xf32>, %C_in: memref<16x16x8x?xf32>) {912 linalg.mmt4d ins(%A, %B: memref<16x16x8x1xf32>, memref<16x16x?x1xf32>)913 outs(%C_in: memref<16x16x8x?xf32>)914 return915}916// CHECK-LABEL: func.func @mmt4d_scalable_with_assume(917// CHECK-SAME: %[[A:.*]]: memref<16x16x8x1xf32>,918// CHECK-SAME: %[[B:.*]]: memref<16x16x?x1xf32>,919// CHECK-SAME: %[[C_IN:.*]]: memref<16x16x8x?xf32>) {920// CHECK-NOT: mask921// CHECK: %[[VEC_A:.*]] = vector.transfer_read %[[A]]922// CHECK-SAME: memref<16x16x8x1xf32>, vector<16x16x16x8x[4]x1xf32>923// CHECK: %[[VEC_B:.*]] = vector.transfer_read %[[B]]924// `in-bounds` are set to true for dynamic dims with assume, static sizes will be inferred elsewhere.925// CHECK-SAME: in_bounds = [false, false, false, false, true, false]{{.*}} : memref<16x16x?x1xf32>, vector<16x16x16x8x[4]x1xf32>926// CHECK: %[[VEC_C:.*]] = vector.transfer_read %[[C_IN]]927// CHECK-SAME: in_bounds = [false, false, false, true]{{.*}} : memref<16x16x8x?xf32>, vector<16x16x8x[4]xf32>928// CHECK: %[[MUL:.*]] = arith.mulf %[[VEC_A]], %[[VEC_B]] : vector<16x16x16x8x[4]x1xf32>929// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[VEC_C]] [2, 5] : vector<16x16x16x8x[4]x1xf32> to vector<16x16x8x[4]xf32>930// CHECK: vector.transfer_write %[[RED]], %[[C_IN]]931// CHECK-SAME: in_bounds = [false, false, false, true]{{.*}} : vector<16x16x8x[4]xf32>, memref<16x16x8x?xf32>932 933module attributes {transform.with_named_sequence} {934 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {935 %mmt4d = transform.structured.match ops{["linalg.mmt4d"]} in %arg1 : (!transform.any_op) -> !transform.any_op936 transform.structured.vectorize %mmt4d vector_sizes [16, 16, 16, 8, [4], 1] {assume_dynamic_dims_match_vec_sizes} : !transform.any_op937 transform.yield938 }939}940 941// -----942 943///----------------------------------------------------------------------------------------944/// Tests for linalg.batch_mmt4d945///----------------------------------------------------------------------------------------946 947func.func @batch_mmt4d(%A: memref<2x16x16x8x1xf32>, %B: memref<2x16x16x8x1xf32>, %C_in: memref<2x16x16x8x8xf32>) {948 linalg.batch_mmt4d ins(%A, %B: memref<2x16x16x8x1xf32>, memref<2x16x16x8x1xf32>)949 outs(%C_in: memref<2x16x16x8x8xf32>)950 return951}952 953// CHECK-LABEL: func.func @batch_mmt4d(954// CHECK-SAME: %[[A:.*]]: memref<2x16x16x8x1xf32>, %[[B:.*]]: memref<2x16x16x8x1xf32>, %[[C:.*]]: memref<2x16x16x8x8xf32>) {955// CHECK: %[[VEC_A:.*]] = vector.transfer_read %[[A]]{{.*}} : memref<2x16x16x8x1xf32>, vector<2x16x16x16x8x8x1xf32>956// CHECK: %[[VEC_B:.*]] = vector.transfer_read %[[B]]{{.*}} : memref<2x16x16x8x1xf32>, vector<2x16x16x16x8x8x1xf32>957// CHECK: %[[VEC_C:.*]] = vector.transfer_read %[[C]]{{.*}} : memref<2x16x16x8x8xf32>, vector<2x16x16x8x8xf32>958// CHECK: %[[MUL:.*]] = arith.mulf %[[VEC_A]], %[[VEC_B]] : vector<2x16x16x16x8x8x1xf32>959// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[VEC_C]] [3, 6] : vector<2x16x16x16x8x8x1xf32> to vector<2x16x16x8x8xf32>960// CHECK: vector.transfer_write %[[RED]], %[[C]]{{.*}} : vector<2x16x16x8x8xf32>, memref<2x16x16x8x8xf32>961 962module attributes {transform.with_named_sequence} {963 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {964 %batch_mmt4d = transform.structured.match ops{["linalg.batch_mmt4d"]} in %arg1 : (!transform.any_op) -> !transform.any_op965 transform.structured.vectorize %batch_mmt4d : !transform.any_op966 transform.yield967 }968}969 970// -----971 972func.func @batch_mmt4d_scalable(%A: memref<2x16x16x8x1xf32>, %B: memref<2x16x16x?x1xf32>, %C_in: memref<2x16x16x8x?xf32>) {973 linalg.batch_mmt4d ins(%A, %B: memref<2x16x16x8x1xf32>, memref<2x16x16x?x1xf32>)974 outs(%C_in: memref<2x16x16x8x?xf32>)975 return976}977// CHECK-LABEL: func.func @batch_mmt4d_scalable(978// CHECK-SAME: %[[A:.*]]: memref<2x16x16x8x1xf32>,979// CHECK-SAME: %[[B:.*]]: memref<2x16x16x?x1xf32>,980// CHECK-SAME: %[[C_IN:.*]]: memref<2x16x16x8x?xf32>) {981// CHECK: %[[C2:.*]] = arith.constant 2 : index982// CHECK: %[[C16_M:.*]] = arith.constant 16 : index983// CHECK: %[[C16_N:.*]] = arith.constant 16 : index984// CHECK: %[[C16_K:.*]] = arith.constant 16 : index985// CHECK: %[[C8:.*]] = arith.constant 8 : index986// CHECK: %[[C3:.*]] = arith.constant 3 : index987// CHECK: %[[DIM_N_IN:.*]] = memref.dim %[[B]], %[[C3]] : memref<2x16x16x?x1xf32>988// CHECK: %[[C1:.*]] = arith.constant 1 : index989// CHECK: %[[VEC_A:.*]] = vector.transfer_read %[[A]]{{.*}} : memref<2x16x16x8x1xf32>, vector<2x16x16x16x8x[4]x1xf32>990// CHECK: %[[MASK_1:.*]] = vector.create_mask %[[C2]], %[[C16_N]], %[[C16_K]], %[[DIM_N_IN]], %[[C1]] : vector<2x16x16x[4]x1xi1>991// CHECK: %[[VEC_B:.*]] = vector.mask %[[MASK_1]] { vector.transfer_read %[[B]]{{.*}} : memref<2x16x16x?x1xf32>, vector<2x16x16x16x8x[4]x1xf32> } : vector<2x16x16x[4]x1xi1> -> vector<2x16x16x16x8x[4]x1xf32>992// CHECK: %[[MASK_2:.*]] = vector.create_mask %[[C2]], %[[C16_M]], %[[C16_N]], %[[C8]], %[[DIM_N_IN]] : vector<2x16x16x8x[4]xi1>993// CHECK: %[[VEC_C:.*]] = vector.mask %[[MASK_2]] { vector.transfer_read %[[C_IN]]{{.*}} : memref<2x16x16x8x?xf32>, vector<2x16x16x8x[4]xf32> } : vector<2x16x16x8x[4]xi1> -> vector<2x16x16x8x[4]xf32>994// CHECK: %[[MUL:.*]] = arith.mulf %[[VEC_A]], %[[VEC_B]] : vector<2x16x16x16x8x[4]x1xf32>995// CHECK: %[[MASK_3:.*]] = vector.create_mask %[[C2]], %[[C16_M]], %[[C16_N]], %[[C16_K]], %[[C8]], %[[DIM_N_IN]], %[[C1]] : vector<2x16x16x16x8x[4]x1xi1>996// CHECK: %[[RED:.*]] = vector.mask %[[MASK_3]] { vector.multi_reduction <add>, %[[MUL]], %[[VEC_C]] [3, 6] : vector<2x16x16x16x8x[4]x1xf32> to vector<2x16x16x8x[4]xf32> } : vector<2x16x16x16x8x[4]x1xi1> -> vector<2x16x16x8x[4]xf32>997// CHECK: vector.mask %[[MASK_2]] { vector.transfer_write %[[RED]], %[[C_IN]]{{.*}} : vector<2x16x16x8x[4]xf32>, memref<2x16x16x8x?xf32> } : vector<2x16x16x8x[4]xi1>998 999module attributes {transform.with_named_sequence} {1000 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {1001 %batch_mmt4d = transform.structured.match ops{["linalg.batch_mmt4d"]} in %arg1 : (!transform.any_op) -> !transform.any_op1002 transform.structured.vectorize %batch_mmt4d vector_sizes [2, 16, 16, 16, 8, [4], 1] : !transform.any_op1003 transform.yield1004 }1005}1006 1007// -----1008 1009func.func @batch_mmt4d_scalable_with_assume(%A: memref<2x16x16x8x1xf32>, %B: memref<2x16x16x?x1xf32>, %C_in: memref<2x16x16x8x?xf32>) {1010 linalg.batch_mmt4d ins(%A, %B: memref<2x16x16x8x1xf32>, memref<2x16x16x?x1xf32>)1011 outs(%C_in: memref<2x16x16x8x?xf32>)1012 return1013}1014// CHECK-LABEL: func.func @batch_mmt4d_scalable_with_assume(1015// CHECK-SAME: %[[A:.*]]: memref<2x16x16x8x1xf32>,1016// CHECK-SAME: %[[B:.*]]: memref<2x16x16x?x1xf32>,1017// CHECK-SAME: %[[C_IN:.*]]: memref<2x16x16x8x?xf32>) {1018// CHECK-NOT: mask1019// CHECK: %[[VEC_A:.*]] = vector.transfer_read %[[A]]1020// CHECK-SAME: memref<2x16x16x8x1xf32>, vector<2x16x16x16x8x[4]x1xf32>1021// CHECK: %[[VEC_B:.*]] = vector.transfer_read %[[B]]1022// `in-bounds` are set to true for dynamic dims with assume, static sizes will be inferred elsewhere.1023// CHECK-SAME: in_bounds = [false, false, false, false, false, true, false]{{.*}} : memref<2x16x16x?x1xf32>, vector<2x16x16x16x8x[4]x1xf32>1024// CHECK: %[[VEC_C:.*]] = vector.transfer_read %[[C_IN]]1025// CHECK-SAME: in_bounds = [false, false, false, false, true]{{.*}} : memref<2x16x16x8x?xf32>, vector<2x16x16x8x[4]xf32>1026// CHECK: %[[MUL:.*]] = arith.mulf %[[VEC_A]], %[[VEC_B]] : vector<2x16x16x16x8x[4]x1xf32>1027// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[VEC_C]] [3, 6] : vector<2x16x16x16x8x[4]x1xf32> to vector<2x16x16x8x[4]xf32>1028// CHECK: vector.transfer_write %[[RED]], %[[C_IN]]1029// CHECK-SAME: in_bounds = [false, false, false, false, true]{{.*}} : vector<2x16x16x8x[4]xf32>, memref<2x16x16x8x?xf32>1030 1031module attributes {transform.with_named_sequence} {1032 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {1033 %batch_mmt4d = transform.structured.match ops{["linalg.batch_mmt4d"]} in %arg1 : (!transform.any_op) -> !transform.any_op1034 transform.structured.vectorize %batch_mmt4d vector_sizes [2, 16, 16, 16, 8, [4], 1] {assume_dynamic_dims_match_vec_sizes} : !transform.any_op1035 transform.yield1036 }1037}1038 1039 1040// -----1041 1042///----------------------------------------------------------------------------------------1043/// Tests for linalg.unpack1044///----------------------------------------------------------------------------------------1045 1046// CHECK-LABEL: func @test_vectorize_dynamic_shapes_unpack1047// CHECK-SAME: %[[DEST:.*]]: tensor<?x?xf32>,1048// CHECK-SAME: %[[SRC:.*]]: tensor<?x?x16x2xf32>1049func.func @test_vectorize_dynamic_shapes_unpack(%dest: tensor<?x?xf32>, %src: tensor<?x?x16x2xf32>) -> tensor<?x?xf32> {1050 // CHECK: %[[C0:.*]] = arith.constant 0 : index1051 // CHECK: %[[C0_1:.*]] = arith.constant 0 : index1052 // CHECK: %[[DIM_0:.*]] = tensor.dim %[[SRC]], %[[C0_1]] : tensor<?x?x16x2xf32>1053 // CHECK: %[[C1:.*]] = arith.constant 11054 // CHECK: %[[DIM6:.*]] = tensor.dim %[[SRC]], %[[C1]] : tensor<?x?x16x2xf32>1055 // CHECK: %[[CNST16:.*]] = arith.constant 16 : index1056 // CHECK: %[[CNST2:.*]] = arith.constant 2 : index1057 // CHECK: %[[MASK_READ:.*]] = vector.create_mask %[[DIM_0]], %[[DIM6]], %[[CNST16]], %[[CNST2]] : vector<2x1x16x2xi1>1058 // CHECK: %[[READ:.*]] = vector.mask %[[MASK_READ]] {{.*}} vector.transfer_read %{{.*}} : tensor<?x?x16x2xf32>, vector<2x1x16x2xf32> } : vector<2x1x16x2xi1> -> vector<2x1x16x2xf32>1059 // CHECK: %[[TR:.*]] = vector.transpose %[[READ]], [0, 3, 1, 2] : vector<2x1x16x2xf32> to vector<2x2x1x16xf32>1060 // CHECK: %[[SC:.*]] = vector.shape_cast %[[TR]] : vector<2x2x1x16xf32> to vector<4x16xf32>1061 // CHECK: %[[MASK_WRITE:.*]] = vector.create_mask {{.*}} : vector<4x16xi1>1062 // CHECK: %[[WRITE:.*]] = vector.mask %[[MASK_WRITE:.*]] {{.*}} vector.transfer_write %[[SC]], %[[DEST]]1063 // CHECK: return %[[WRITE]]1064 %ret = linalg.unpack %src inner_dims_pos = [1, 0] inner_tiles = [16, 2] into %dest : tensor<?x?x16x2xf32> -> tensor<?x?xf32>1065 return %ret : tensor<?x?xf32>1066}1067module attributes {transform.with_named_sequence} {1068 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1069 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1070 transform.structured.vectorize %0 vector_sizes [2, 1, 16, 2] : !transform.any_op1071 transform.yield1072 }1073}1074 1075// -----1076 1077// CHECK-LABEL: func @test_vectorize_dynamic_shapes_unpack_scalable_vec1078// CHECK-SAME: %[[DEST:.*]]: tensor<?x?xf32>,1079// CHECK-SAME: %[[SRC:.*]]: tensor<?x?x16x2xf32>1080func.func @test_vectorize_dynamic_shapes_unpack_scalable_vec(%dest: tensor<?x?xf32>, %src: tensor<?x?x16x2xf32>) -> tensor<?x?xf32> {1081 // CHECK-DAG: %[[PAD:.*]] = ub.poison : f321082 // CHECK-DAG: %[[C01:.*]] = arith.constant 01083 // CHECK-DAG: %[[C02:.*]] = arith.constant 01084 // CHECK: %[[DIM4:.*]] = tensor.dim %[[SRC]], %[[C02]] : tensor<?x?x16x2xf32>1085 // CHECK: %[[CNST14:.*]] = arith.constant 11086 // CHECK: %[[DIM6:.*]] = tensor.dim %[[SRC]], %[[CNST14]] : tensor<?x?x16x2xf32>1087 // CHECK: %[[CNST16:.*]] = arith.constant 16 : index1088 // CHECK: %[[CNST2:.*]] = arith.constant 2 : index1089 // CHECK: %[[MASK_READ:.*]] = vector.create_mask %[[DIM4]], %[[DIM6]], %[[CNST16]], %[[CNST2]] : vector<2x1x[16]x2xi1>1090 // CHECK: %[[READ:.*]] = vector.mask %[[MASK_READ]] {{.*}} vector.transfer_read %{{.*}} %[[PAD]] {{.*}}: tensor<?x?x16x2xf32>, vector<2x1x[16]x2xf32> } : vector<2x1x[16]x2xi1> -> vector<2x1x[16]x2xf32>1091 // CHECK: %[[TR:.*]] = vector.transpose %[[READ]], [0, 3, 1, 2] : vector<2x1x[16]x2xf32> to vector<2x2x1x[16]xf32>1092 // CHECK: %[[SC:.*]] = vector.shape_cast %[[TR]] : vector<2x2x1x[16]xf32> to vector<4x[16]xf32>1093 // CHECK: %[[MASK_WRITE:.*]] = vector.create_mask {{.*}} : vector<4x[16]xi1>1094 // CHECK: %[[WRITE:.*]] = vector.mask %[[MASK_WRITE:.*]] {{.*}} vector.transfer_write %[[SC]], %[[DEST]]1095 // CHECK: return %[[WRITE]]1096 %ret = linalg.unpack %src inner_dims_pos = [1, 0] inner_tiles = [16, 2] into %dest : tensor<?x?x16x2xf32> -> tensor<?x?xf32>1097 return %ret : tensor<?x?xf32>1098}1099module attributes {transform.with_named_sequence} {1100 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1101 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1102 transform.structured.vectorize %0 vector_sizes [2, 1, [16], 2] : !transform.any_op1103 transform.yield1104 }1105}1106 1107// -----1108 1109// CHECK-LABEL: func @test_vectorize_dynamic_shapes_unpack_scalable_vec_and_tile_size1110// CHECK-SAME: %[[DEST:.*]]: tensor<?x?xf32>,1111// CHECK-SAME: %[[SRC:.*]]: tensor<?x?x?x2xf32>1112func.func @test_vectorize_dynamic_shapes_unpack_scalable_vec_and_tile_size(%dest: tensor<?x?xf32>, %src: tensor<?x?x?x2xf32>) -> tensor<?x?xf32> {1113 // CHECK-DAG: %[[PAD:.*]] = ub.poison : f321114 // CHECK-DAG: %[[C01:.*]] = arith.constant 01115 // CHECK-DAG: %[[C02:.*]] = arith.constant 01116 // CHECK: %[[DIM4:.*]] = tensor.dim %[[SRC]], %[[C02]] : tensor<?x?x?x2xf32>1117 // CHECK: %[[C1_2:.*]] = arith.constant 11118 // CHECK: %[[DIM6:.*]] = tensor.dim %[[SRC]], %[[C1_2]] : tensor<?x?x?x2xf32>1119 // CHECK: %[[C2:.*]] = arith.constant 2 : index1120 // CHECK: %[[DIM_2:.*]] = tensor.dim %[[SRC]], %[[C2]] : tensor<?x?x?x2xf32>1121 // CHECK: %[[C2_1:.*]] = arith.constant 2 : index1122 // CHECK: %[[MASK_READ:.*]] = vector.create_mask %[[DIM4]], %[[DIM6]], %[[DIM_2]], %[[C2_1]] : vector<2x1x[16]x2xi1>1123 // CHECK: %[[READ:.*]] = vector.mask %[[MASK_READ]] {{.*}} vector.transfer_read %{{.*}} %[[PAD]] {{.*}}: tensor<?x?x?x2xf32>, vector<2x1x[16]x2xf32> } : vector<2x1x[16]x2xi1> -> vector<2x1x[16]x2xf32>1124 // CHECK: %[[TR:.*]] = vector.transpose %[[READ]], [0, 3, 1, 2] : vector<2x1x[16]x2xf32> to vector<2x2x1x[16]xf32>1125 // CHECK: %[[SC:.*]] = vector.shape_cast %[[TR]] : vector<2x2x1x[16]xf32> to vector<4x[16]xf32>1126 // CHECK: %[[MASK_WRITE:.*]] = vector.create_mask {{.*}} : vector<4x[16]xi1>1127 // CHECK: %[[WRITE:.*]] = vector.mask %[[MASK_WRITE:.*]] {{.*}} vector.transfer_write %[[SC]], %[[DEST]]1128 // CHECK: return %[[WRITE]]1129 1130 %vs = vector.vscale1131 %c16 = arith.constant 16 : index1132 %tile_size = arith.muli %vs, %c16 : index1133 1134 %ret = linalg.unpack %src inner_dims_pos = [1, 0] inner_tiles = [%tile_size, 2] into %dest : tensor<?x?x?x2xf32> -> tensor<?x?xf32>1135 return %ret : tensor<?x?xf32>1136}1137module attributes {transform.with_named_sequence} {1138 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1139 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1140 transform.structured.vectorize %0 vector_sizes [2, 1, [16], 2] : !transform.any_op1141 transform.yield1142 }1143}1144 1145// -----1146 1147// CHECK-LABEL: func @test_vectorize_unpack1148// CHECK-SAME: %[[SRC:.*]]: tensor<8x8x32x16xf32>1149// CHECK-SAME: %[[DEST:.*]]: tensor<256x128xf32>1150func.func @test_vectorize_unpack(%source: tensor<8x8x32x16xf32>, %dest: tensor<256x128xf32>) -> tensor<256x128xf32> {1151 // CHECK-DAG: %[[PAD:.*]] = ub.poison : f321152 // CHECK-DAG: %[[C0:.*]]= arith.constant 0 : index1153 // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index1154 // CHECK-DAG: %[[C80:.*]] = arith.constant 8 : index1155 // CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index1156 // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index1157 // CHECK: %[[MSK0:.*]] = vector.create_mask %[[C8]], %[[C80]], %[[C32]], %[[C16]] : vector<16x8x32x16xi1>1158 // CHECK: %[[READ0:.*]] = vector.mask %[[MSK0]] { vector.transfer_read %[[SRC]]{{.*}} %[[PAD]] {{.*}} : vector<16x8x32x16xi1> -> vector<16x8x32x16xf32>1159 // CHECK: %[[TRANSP0:.*]] = vector.transpose %[[READ0]], [0, 2, 1, 3] : vector<16x8x32x16xf32> to vector<16x32x8x16xf32>1160 // CHECK: %[[SHAPC:.*]] = vector.shape_cast %[[TRANSP0]] : vector<16x32x8x16xf32> to vector<512x128xf32>1161 // CHECK: %[[C01:.*]] = arith.constant 0 : index1162 // CHECK: %[[C256:.*]] = arith.constant 256 : index1163 // CHECK: %[[C128:.*]] = arith.constant 128 : index1164 // CHECK: %[[WRITEMSK:.*]] = vector.create_mask %[[C256]], %[[C128]] : vector<512x128xi1>1165 // CHECK: %[[WRIT:.*]] = vector.mask %[[WRITEMSK]] { vector.transfer_write %[[SHAPC]], %[[DEST]]{{.*}}} : vector<512x128xi1> -> tensor<256x128xf32>1166 // CHECK: return %[[WRIT]] : tensor<256x128xf32>1167 %0 = linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %dest : tensor<8x8x32x16xf32> -> tensor<256x128xf32>1168 return %0 : tensor<256x128xf32>1169 }1170 module attributes {transform.with_named_sequence} {1171 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1172 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1173 transform.structured.vectorize %0 vector_sizes [16, 8, 32, 16] : !transform.any_op1174 transform.yield1175 }1176}1177 1178// -----1179 1180// CHECK-LABEL: func @test_vectorize_unpack_no_masks1181// CHECK-SAME: %[[SRC:.*]]: tensor<8x8x32x16xf32>1182// CHECK-SAME: %[[DEST:.*]]: tensor<256x128xf32>1183func.func @test_vectorize_unpack_no_masks(%source: tensor<8x8x32x16xf32>, %dest: tensor<256x128xf32>) -> tensor<256x128xf32> {1184 // CHECK-DAG: %[[PAD:.*]] = ub.poison : f321185 // CHECK-AD: %[[C0:.*]] = arith.constant 0 : index1186 // CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]]{{.*}} %[[PAD]] {{.*}} : tensor<8x8x32x16xf32>, vector<8x8x32x16xf32> 1187 // CHECK: %[[TRANSP:.*]] = vector.transpose %[[READ]], [0, 2, 1, 3] : vector<8x8x32x16xf32> to vector<8x32x8x16xf32>1188 // CHECK: %[[SHAPC:.*]] = vector.shape_cast %[[TRANSP]] : vector<8x32x8x16xf32> to vector<256x128xf32>1189 // CHECK: %[[C00:.*]] = arith.constant 0 : index1190 // CHECK: %[[WRIT:.*]] = vector.transfer_write %[[SHAPC]], %[[DEST]]{{.*}}} : vector<256x128xf32>, tensor<256x128xf32> 1191 // CHECK: return %[[WRIT]] : tensor<256x128xf32>1192 %0 = linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %dest : tensor<8x8x32x16xf32> -> tensor<256x128xf32>1193 return %0 : tensor<256x128xf32>1194 }1195 module attributes {transform.with_named_sequence} {1196 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1197 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1198 transform.structured.vectorize %0 vector_sizes [8, 8, 32, 16] : !transform.any_op1199 transform.yield1200 }1201 }1202 1203// -----1204 1205// CHECK-LABEL: test_vectorize_unpack_with_outer_perm1206// CHECK-SAME: %[[SRC:.*]]: tensor<8x8x32x16xf32>1207// CHECK-SAME: %[[DEST:.*]]: tensor<256x128xf32>1208 func.func @test_vectorize_unpack_with_outer_perm(%source: tensor<8x8x32x16xf32>, %dest: tensor<256x128xf32>) -> tensor<256x128xf32> {1209 // CHECK-DAG: %[[PAD:.*]] = ub.poison : f321210 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index1211 // CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]]{{.*}} %[[PAD]] {{.*}} : tensor<8x8x32x16xf32>, vector<8x8x32x16xf32> 1212 // CHECK: %[[TRANSP:.*]] = vector.transpose %[[READ]], [1, 2, 0, 3] : vector<8x8x32x16xf32> to vector<8x32x8x16xf32>1213 // CHECK: %[[SHAPC:.*]] = vector.shape_cast %[[TRANSP]] : vector<8x32x8x16xf32> to vector<256x128xf32>1214 // CHECK: %[[C00:.*]] = arith.constant 0 : index1215 // CHECK: %[[WRIT:.*]] = vector.transfer_write %[[SHAPC]], %[[DEST]]{{.*}}} : vector<256x128xf32>, tensor<256x128xf32> 1216 // CHECK: return %[[WRIT]] : tensor<256x128xf32>1217 %0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %dest : tensor<8x8x32x16xf32> -> tensor<256x128xf32>1218 return %0 : tensor<256x128xf32>1219 }1220 module attributes {transform.with_named_sequence} {1221 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1222 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1223 transform.structured.vectorize %0 vector_sizes [8, 8, 32, 16] : !transform.any_op1224 transform.yield1225 }1226}1227 1228// -----1229 1230// CHECK-LABEL: test_vectorize_unpack_no_vector_sizes1231// CHECK-SAME: %[[SRC:.*]]: tensor<8x8x32x16xf32>1232// CHECK-SAME: %[[DEST:.*]]: tensor<256x128xf32>1233func.func @test_vectorize_unpack_no_vector_sizes(%source: tensor<8x8x32x16xf32>, %dest: tensor<256x128xf32>) -> tensor<256x128xf32> {1234 // CHECK-DAG: %[[PAD:.*]] = ub.poison : f321235 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index1236 // CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]]{{.*}} %[[PAD]] {{.*}} : tensor<8x8x32x16xf32>, vector<8x8x32x16xf32> 1237 // CHECK: %[[TRANSP:.*]] = vector.transpose %[[READ]], [0, 2, 1, 3] : vector<8x8x32x16xf32> to vector<8x32x8x16xf32>1238 // CHECK: %[[SHAPC:.*]] = vector.shape_cast %[[TRANSP]] : vector<8x32x8x16xf32> to vector<256x128xf32>1239 // CHECK: %[[C00:.*]] = arith.constant 0 : index1240 // CHECK: %[[WRIT:.*]] = vector.transfer_write %[[SHAPC]], %[[DEST]]{{.*}}} : vector<256x128xf32>, tensor<256x128xf32> 1241 // CHECK: return %[[WRIT]] : tensor<256x128xf32>1242 %0 = linalg.unpack %source inner_dims_pos = [0, 1] inner_tiles = [32, 16] into %dest : tensor<8x8x32x16xf32> -> tensor<256x128xf32>1243 return %0 : tensor<256x128xf32>1244 }1245 module attributes {transform.with_named_sequence} {1246 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1247 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1248 transform.structured.vectorize %0 : !transform.any_op1249 transform.yield1250 }1251 }1252 1253// -----1254 1255// CHECK-LABEL: test_vectorize_unpack_no_vector_sizes_slice_output1256// CHECK-SAME: %[[SRC:.*]]: tensor<8x4x16x16xf32>1257// CHECK-SAME: %[[DEST:.*]]: tensor<64x127xf32>1258func.func @test_vectorize_unpack_no_vector_sizes_slice_output(%source: tensor<8x4x16x16xf32>, %dest: tensor<64x127xf32>) -> tensor<64x127xf32> {1259 // CHECK-DAG: %[[PAD:.*]] = ub.poison : f321260 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index1261 // CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]]{{.*}} %[[PAD]] {{.*}} : tensor<8x4x16x16xf32>, vector<8x4x16x16xf32>1262 // CHECK: %[[TRANSP:.*]] = vector.transpose %[[READ]], [1, 2, 0, 3] : vector<8x4x16x16xf32> to vector<4x16x8x16xf32>1263 // CHECK: %[[SHAPC:.*]] = vector.shape_cast %[[TRANSP]] : vector<4x16x8x16xf32> to vector<64x128xf32>1264 // CHECK: %[[C00:.*]] = arith.constant 0 : index1265 // CHECK: %[[WRIT:.*]] = vector.transfer_write %[[SHAPC]], %[[DEST]]1266 // CHECK-SAME: {in_bounds = [true, false]} : vector<64x128xf32>, tensor<64x127xf32>1267 // CHECK: return %[[WRIT]] : tensor<64x127xf32>1268 %0 = linalg.unpack %source outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [16, 16] into %dest : tensor<8x4x16x16xf32> -> tensor<64x127xf32>1269 return %0 : tensor<64x127xf32>1270 }1271 module attributes {transform.with_named_sequence} {1272 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1273 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1274 transform.structured.vectorize %0 : !transform.any_op1275 transform.yield1276 }1277 }1278 1279// -----1280 1281// CHECK-LABEL: test_vectorize_unpack_no_vector_sizes_permute1282// CHECK-SAME: %[[SRC:.*]]: tensor<4x7x4xf32>1283// CHECK-SAME: %[[DEST:.*]]: tensor<7x16xf32>1284func.func @test_vectorize_unpack_no_vector_sizes_permute(%source: tensor<4x7x4xf32>, %dest: tensor<7x16xf32>) -> tensor<7x16xf32> {1285 %0 = linalg.unpack %source outer_dims_perm=[1, 0] inner_dims_pos = [1] inner_tiles = [4] into %dest : tensor<4x7x4xf32> -> tensor<7x16xf32>1286 return %0 : tensor<7x16xf32>1287 }1288 // CHECK-DAG: %[[PAD:.*]] = ub.poison : f321289 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index1290 // CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]]{{.*}} %[[PAD]] {{.*}} : tensor<4x7x4xf32>, vector<4x7x4xf32>1291 // CHECK: %[[TRANSP:.*]] = vector.transpose %[[READ]], [1, 0, 2] : vector<4x7x4xf32> to vector<7x4x4xf32>1292 // CHECK: %[[SHAPC:.*]] = vector.shape_cast %[[TRANSP]] : vector<7x4x4xf32> to vector<7x16xf32>1293 // CHECK: %[[C00:.*]] = arith.constant 0 : index1294 // CHECK: %[[WRIT:.*]] = vector.transfer_write %[[SHAPC]], %[[DEST]]{{.*}}} : vector<7x16xf32>, tensor<7x16xf32> 1295 // CHECK: return %[[WRIT]] : tensor<7x16xf32>1296 module attributes {transform.with_named_sequence} {1297 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1298 %0 = transform.structured.match ops{["linalg.unpack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1299 transform.structured.vectorize %0 : !transform.any_op1300 transform.yield1301 }1302 }1303 1304// -----1305 1306///----------------------------------------------------------------------------------------1307/// Tests for linalg.pack1308///----------------------------------------------------------------------------------------1309 1310// This packing requires no padding, so no out-of-bounds read/write vector Ops.1311 1312// Note, see a similar test in:1313// * vectorization-with-patterns.mlir1314// The output is identical (the input vector sizes == the inferred vector1315// sizes, i.e. the tensor sizes).1316 1317// CHECK-LABEL: func @pack_no_padding1318// CHECK-SAME: %[[SRC:.*]]: tensor<32x8x16xf32>,1319// CHECK-SAME: %[[DEST:.*]]: tensor<4x1x32x16x2xf32>1320func.func @pack_no_padding(%src: tensor<32x8x16xf32>, %dest: tensor<4x1x32x16x2xf32>) -> tensor<4x1x32x16x2xf32> {1321 %pack = linalg.pack %src outer_dims_perm = [1, 2, 0] inner_dims_pos = [2, 1] inner_tiles = [16, 2] into %dest : tensor<32x8x16xf32> -> tensor<4x1x32x16x2xf32>1322 return %pack : tensor<4x1x32x16x2xf32>1323}1324// CHECK-DAG: %[[CST:.*]] = ub.poison : f321325// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index1326// CHECK: %[[READ:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %[[CST]]1327// CHECK-SAME: {in_bounds = [true, true, true]} : tensor<32x8x16xf32>, vector<32x8x16xf32>1328// CHECK: %[[SC:.*]] = vector.shape_cast %[[READ]] : vector<32x8x16xf32> to vector<32x4x2x1x16xf32>1329// CHECK: %[[TR:.*]] = vector.transpose %[[SC]], [1, 3, 0, 4, 2] : vector<32x4x2x1x16xf32> to vector<4x1x32x16x2xf32>1330// CHECK-DAG: %[[C0_1:.*]] = arith.constant 0 : index1331// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[TR]], %[[DEST]][%[[C0_1]], %[[C0_1]], %[[C0_1]], %[[C0_1]], %[[C0_1]]]1332// CHECK-SAME: {in_bounds = [true, true, true, true, true]} : vector<4x1x32x16x2xf32>, tensor<4x1x32x16x2xf32>1333// CHECK: return %[[WRITE]] : tensor<4x1x32x16x2xf32>1334 1335module attributes {transform.with_named_sequence} {1336 transform.named_sequence @__transform_main(%src: !transform.any_op {transform.readonly}) {1337 %0 = transform.structured.match ops{["linalg.pack"]} in %src : (!transform.any_op) -> !transform.any_op1338 transform.structured.vectorize %0 vector_sizes [4, 1, 32, 16, 2] : !transform.any_op1339 transform.yield1340 }1341}1342 1343// -----1344 1345// This packing does require padding, so there are out-of-bounds read/write1346// vector Ops.1347 1348// Note, see a similar test in:1349// * vectorization-with-patterns.mlir.1350// The output is different (the input vector sizes != inferred vector sizes,1351// i.e. the tensor sizes).1352 1353// CHECK-LABEL: func @pack_with_padding1354// CHECK-SAME: %[[SRC:.*]]: tensor<32x7x15xf32>,1355// CHECK-SAME: %[[DEST:.*]]: tensor<32x4x1x16x2xf32>1356func.func @pack_with_padding(%src: tensor<32x7x15xf32>, %dest: tensor<32x4x1x16x2xf32>) -> tensor<32x4x1x16x2xf32> {1357 %pad = arith.constant 0.000000e+00 : f321358 %pack = linalg.pack %src padding_value(%pad : f32) inner_dims_pos = [2, 1] inner_tiles = [16, 2] into %dest : tensor<32x7x15xf32> -> tensor<32x4x1x16x2xf32>1359 return %pack : tensor<32x4x1x16x2xf32>1360}1361// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f321362// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index1363// CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index1364// CHECK-DAG: %[[C7:.*]] = arith.constant 7 : index1365// CHECK-DAG: %[[C15:.*]] = arith.constant 15 : index1366// CHECK: %[[MASK:.*]] = vector.create_mask %[[C32]], %[[C7]], %[[C15]] : vector<32x8x16xi1>1367// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] {1368// CHECK-SAME: vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %[[CST]]1369// CHECK-SAME: {in_bounds = [true, true, true]} : tensor<32x7x15xf32>, vector<32x8x16xf32>1370// CHECK-SAME: } : vector<32x8x16xi1> -> vector<32x8x16xf32>1371// CHECK: %[[SC:.*]] = vector.shape_cast %[[READ]] : vector<32x8x16xf32> to vector<32x4x2x1x16xf32>1372// CHECK: %[[TR:.*]] = vector.transpose %[[SC]], [0, 1, 3, 4, 2] : vector<32x4x2x1x16xf32> to vector<32x4x1x16x2xf32>1373// CHECK-DAG: %[[C0_1:.*]] = arith.constant 0 : index1374// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[TR]], %[[DEST]][%[[C0_1]], %[[C0_1]], %[[C0_1]], %[[C0_1]], %[[C0_1]]]1375// CHECK-SAME: {in_bounds = [true, true, true, true, true]} : vector<32x4x1x16x2xf32>, tensor<32x4x1x16x2xf32>1376// CHECK: return %[[WRITE]] : tensor<32x4x1x16x2xf32>1377 1378module attributes {transform.with_named_sequence} {1379 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1380 %0 = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1381 transform.structured.vectorize %0 vector_sizes [32, 4, 1, 16, 2] : !transform.any_op1382 transform.yield1383 }1384}1385 1386// -----1387 1388// This packing does require padding, so there are out-of-bounds read/write1389// vector Ops.1390 1391// Note, see a similar test in:1392// * vectorization-with-patterns.mlir.1393// The output is identical (in both cases the vector sizes are inferred).1394 1395// CHECK-LABEL: func @pack_with_padding_no_vector_sizes1396// CHECK-SAME: %[[SRC:.*]]: tensor<32x7x15xf32>,1397// CHECK-SAME: %[[DEST:.*]]: tensor<32x4x1x16x2xf32>1398func.func @pack_with_padding_no_vector_sizes(%src: tensor<32x7x15xf32>, %dest: tensor<32x4x1x16x2xf32>) -> tensor<32x4x1x16x2xf32> {1399 %pad = arith.constant 0.000000e+00 : f321400 %pack = linalg.pack %src padding_value(%pad : f32) inner_dims_pos = [2, 1] inner_tiles = [16, 2] into %dest : tensor<32x7x15xf32> -> tensor<32x4x1x16x2xf32>1401 return %pack : tensor<32x4x1x16x2xf32>1402}1403// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f321404// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index1405// CHECK: %[[READ:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %[[CST]]1406// CHECK-SAME: {in_bounds = [true, false, false]} : tensor<32x7x15xf32>, vector<32x8x16xf32>1407// CHECK: %[[SC:.*]] = vector.shape_cast %[[READ]] : vector<32x8x16xf32> to vector<32x4x2x1x16xf32>1408// CHECK: %[[TR:.*]] = vector.transpose %[[SC]], [0, 1, 3, 4, 2] : vector<32x4x2x1x16xf32> to vector<32x4x1x16x2xf32>1409// CHECK-DAG: %[[C0_1:.*]] = arith.constant 0 : index1410// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[TR]], %[[DEST]][%[[C0_1]], %[[C0_1]], %[[C0_1]], %[[C0_1]], %[[C0_1]]]1411// CHECK-SAME: {in_bounds = [true, true, true, true, true]} : vector<32x4x1x16x2xf32>, tensor<32x4x1x16x2xf32>1412// CHECK: return %[[WRITE]] : tensor<32x4x1x16x2xf32>1413 1414module attributes {transform.with_named_sequence} {1415 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1416 %0 = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1417 transform.structured.vectorize %0 : !transform.any_op1418 transform.yield1419 }1420}1421 1422// -----1423 1424// CHECK-LABEL: func @pack_with_dynamic_dims1425// CHECK-SAME: %[[SRC:.*]]: tensor<?x?xf32>,1426// CHECK-SAME: %[[DEST:.*]]: tensor<?x?x16x2xf32>1427func.func @pack_with_dynamic_dims(1428 %src: tensor<?x?xf32>, 1429 %dest: tensor<?x?x16x2xf32>) -> tensor<?x?x16x2xf32> {1430 %pack = linalg.pack %src 1431 inner_dims_pos = [1, 0]1432 inner_tiles = [16, 2]1433 into %dest : tensor<?x?xf32> -> tensor<?x?x16x2xf32>1434 return %pack : tensor<?x?x16x2xf32>1435}1436 1437// CHECK-DAG: %[[CST:.*]] = ub.poison : f321438// CHECK-DAG: %[[C0_1:.*]] = arith.constant 0 : index1439// CHECK-DAG: %[[C0_0:.*]] = arith.constant 0 : index1440// CHECK-DAG: %[[C1_0:.*]] = arith.constant 1 : index1441 1442/// Compute mask for xfer_read1443// CHECK-DAG: %[[D0_0:.*]] = tensor.dim {{.*}} %[[C0_0]] : tensor<?x?xf32>1444// CHECK-DAG: %[[D1_0:.*]] = tensor.dim {{.*}} %[[C1_0]] : tensor<?x?xf32>1445// CHECK: %[[MASK:.*]] = vector.create_mask %[[D0_0]], %[[D1_0]] : vector<8x16xi1>1446 1447/// --= read =---1448// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] {1449// CHECK-SAME: vector.transfer_read %{{.*}}[%[[C0_1]], %[[C0_1]]], %[[CST]]1450// CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<8x16xf32>1451// CHECK-SAME: } : vector<8x16xi1> -> vector<8x16xf32>1452 1453/// --= shape_cast =---1454// CHECK: %[[SC:.*]] = vector.shape_cast %[[READ]] : vector<8x16xf32> to vector<4x2x1x16xf32>1455 1456/// --= transpose =---1457// CHECK: %[[TR:.*]] = vector.transpose %[[SC]], [0, 2, 3, 1] : vector<4x2x1x16xf32> to vector<4x1x16x2xf32>1458 1459/// Compute mask for xfer_write1460// CHECK-DAG: %[[C0_2:.*]] = arith.constant 0 : index1461// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index1462// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index1463// CHECK-DAG: %[[D2:.*]] = tensor.dim %[[DEST]], {{.*}} : tensor<?x?x16x2xf32>1464// CHECK-DAG: %[[D3:.*]] = tensor.dim %[[DEST]], {{.*}} : tensor<?x?x16x2xf32>1465// CHECK: %[[MASK_0:.*]] = vector.create_mask %[[D2]], %[[D3]], %[[C16]], %[[C2]] : vector<4x1x16x2xi1>1466 1467/// --= write =---1468// CHECK: %[[WRITE:.*]] = vector.mask %[[MASK_0]] {1469// CHECK-SAME: vector.transfer_write %[[TR]], %[[DEST]][%[[C0_2]], %[[C0_2]], %[[C0_2]], %[[C0_2]]]1470// CHECK-SAME: {in_bounds = [true, true, true, true]} : vector<4x1x16x2xf32>, tensor<?x?x16x2xf32>1471 1472// CHECK: return %[[WRITE]] : tensor<?x?x16x2xf32>1473 1474module attributes {transform.with_named_sequence} {1475 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1476 %0 = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1477 transform.structured.vectorize %0 vector_sizes [4, 1, 16, 2] : !transform.any_op1478 transform.yield1479 }1480}1481 1482// -----1483 1484/// Similar to the test above, but one of the inner tile sizes is dynamic. As a1485/// result, more output dims are dynamic (and, e.g., output mask calcuation is a bit different).1486 1487// CHECK-LABEL: func @pack_with_dynamic_dims_and_dynamic_inner_tile1488// CHECK-SAME: %[[SRC:.*]]: tensor<?x?xf32>,1489// CHECK-SAME: %[[DEST:.*]]: tensor<?x?x?x2xf32>1490func.func @pack_with_dynamic_dims_and_dynamic_inner_tile(1491 %src: tensor<?x?xf32>,1492 %dest: tensor<?x?x?x2xf32>) -> tensor<?x?x?x2xf32> {1493 %c16 = arith.constant 16 : index1494 %pack = linalg.pack %src1495 inner_dims_pos = [1, 0]1496 inner_tiles = [%c16, 2]1497 into %dest : tensor<?x?xf32> -> tensor<?x?x?x2xf32>1498 return %pack : tensor<?x?x?x2xf32>1499}1500 1501// CHECK-DAG: %[[CST:.*]] = ub.poison : f321502// CHECK-DAG: %[[C0_1:.*]] = arith.constant 0 : index1503// CHECK-DAG: %[[C0_0:.*]] = arith.constant 0 : index1504// CHECK-DAG: %[[C1_0:.*]] = arith.constant 1 : index1505 1506/// Compute mask for xfer_read1507// CHECK-DAG: %[[D0_0:.*]] = tensor.dim {{.*}} %[[C0_0]] : tensor<?x?xf32>1508// CHECK-DAG: %[[D1_0:.*]] = tensor.dim {{.*}} %[[C1_0]] : tensor<?x?xf32>1509// CHECK: %[[MASK:.*]] = vector.create_mask %[[D0_0]], %[[D1_0]] : vector<8x16xi1>1510 1511/// --= read =---1512// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] {1513// CHECK-SAME: vector.transfer_read %{{.*}}[%[[C0_1]], %[[C0_1]]], %[[CST]]1514// CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<8x16xf32>1515// CHECK-SAME: } : vector<8x16xi1> -> vector<8x16xf32>1516 1517/// --= shape_cast =---1518// CHECK: %[[SC:.*]] = vector.shape_cast %[[READ]] : vector<8x16xf32> to vector<4x2x1x16xf32>1519 1520/// --= transpose =---1521// CHECK: %[[TR:.*]] = vector.transpose %[[SC]], [0, 2, 3, 1] : vector<4x2x1x16xf32> to vector<4x1x16x2xf32>1522 1523/// Compute mask for xfer_write1524// CHECK-DAG: %[[C0_2:.*]] = arith.constant 0 : index1525// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index1526// CHECK-DAG: %[[C2_2:.*]] = arith.constant 2 : index1527// CHECK-DAG: %[[D2:.*]] = tensor.dim %[[DEST]], {{.*}} : tensor<?x?x?x2xf32>1528// CHECK-DAG: %[[D3:.*]] = tensor.dim %[[DEST]], {{.*}} : tensor<?x?x?x2xf32>1529// CHECK-DAG: %[[D4:.*]] = tensor.dim %[[DEST]], {{.*}} : tensor<?x?x?x2xf32>1530// CHECK: %[[MASK_0:.*]] = vector.create_mask %[[D2]], %[[D3]], %[[D4]], %[[C2_2]] : vector<4x1x16x2xi1>1531 1532/// --= write =---1533// CHECK: %[[WRITE:.*]] = vector.mask %[[MASK_0]] {1534// CHECK-SAME: vector.transfer_write %[[TR]], %[[DEST]][%[[C0_2]], %[[C0_2]], %[[C0_2]], %[[C0_2]]]1535// CHECK-SAME: {in_bounds = [true, true, true, true]} : vector<4x1x16x2xf32>, tensor<?x?x?x2xf32>1536 1537// CHECK: return %[[WRITE]] : tensor<?x?x?x2xf32>1538 1539module attributes {transform.with_named_sequence} {1540 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {1541 %0 = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.any_op1542 transform.structured.vectorize %0 vector_sizes [4, 1, 16, 2] : !transform.any_op1543 transform.yield1544 }1545}1546 1547///----------------------------------------------------------------------------------------1548/// Tests for other Ops1549///----------------------------------------------------------------------------------------1550 1551// -----1552 1553func.func @vectorize_dynamic_fill(%A : tensor<?x?xf32>, %arg0 : f32) -> tensor<?x?xf32> {1554 %0 = linalg.fill ins(%arg0 : f32) outs(%A : tensor<?x?xf32>) -> tensor<?x?xf32>1555 return %0 : tensor<?x?xf32>1556}1557 1558// CHECK-LABEL: func.func @vectorize_dynamic_fill1559// CHECK: %[[DIM0:.*]] = tensor.dim1560// CHECK: %[[DIM1:.*]] = tensor.dim1561// CHECK: %[[MASK:.*]] = vector.create_mask %[[DIM0]], %[[DIM1]] : vector<8x16xi1>1562// CHECK: %[[BCAST:.*]] = vector.broadcast %{{.*}} : f32 to vector<8x16xf32>1563// CHECK: vector.mask %[[MASK]] { vector.transfer_write %[[BCAST]], {{.*}} {in_bounds = [true, true]} : vector<8x16xf32>, tensor<?x?xf32> } : vector<8x16xi1>1564 1565module attributes {transform.with_named_sequence} {1566 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {1567 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.any_op1568 transform.structured.vectorize %0 vector_sizes [8, 16] : !transform.any_op1569 transform.yield1570 }1571}1572 1573// -----1574 1575// NOTE: Often, non-trailing scalable sizes are problematic - there are no1576// "scalable" arrays of vectors at the LLVM level (multi-dim vectors are1577// decomposed into arrays of aggregates). However, the trailing dim in this1578// case is 1 and that can be folded away later.1579 1580// NOTE: This is similar to the example above, but the trailing dim was set to1581// 1 to make it foldable + vectorizable.1582 1583func.func @vectorize_dynamic_fill_scalable(%A : tensor<?x?xf32>, %arg0 : f32) -> tensor<?x?xf32> {1584 %0 = linalg.fill ins(%arg0 : f32) outs(%A : tensor<?x?xf32>) -> tensor<?x?xf32>1585 return %0 : tensor<?x?xf32>1586}1587 1588// CHECK-LABEL: func.func @vectorize_dynamic_fill_scalable1589// CHECK: %[[DIM0:.*]] = tensor.dim1590// CHECK: %[[DIM1:.*]] = tensor.dim1591// CHECK: %[[MASK:.*]] = vector.create_mask %[[DIM0]], %[[DIM1]] : vector<[8]x1xi1>1592// CHECK: %[[BCAST:.*]] = vector.broadcast %{{.*}} : f32 to vector<[8]x1xf32>1593// CHECK: vector.mask %[[MASK]] { vector.transfer_write %[[BCAST]], {{.*}} {in_bounds = [true, true]} : vector<[8]x1xf32>, tensor<?x?xf32> } : vector<[8]x1xi1>1594 1595module attributes {transform.with_named_sequence} {1596 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {1597 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.any_op1598 transform.structured.vectorize %0 vector_sizes [[8], 1] : !transform.any_op1599 transform.yield1600 }1601}1602 1603// -----1604 1605// CHECK: #[[MAP:.*]] = affine_map<(d0, d1) -> (d1, d0)>1606// CHECK: func @test_masked_vectorize_linalg_transpose1607func.func @test_masked_vectorize_linalg_transpose(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {1608 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index1609 // CHECK-DAG: %[[D0:.*]] = tensor.dim %arg0, %[[C0]]1610 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index1611 // CHECK-DAG: %[[D1:.*]] = tensor.dim %arg0, %[[C1]]1612 // CHECK: %[[MASK0:.*]] = vector.create_mask %[[D0]], %[[D1]]1613 // CHECK: %[[LOAD:.*]] = vector.mask %[[MASK0]] { vector.transfer_read %arg0{{.+}} permutation_map = #[[MAP]]{{.+}} }1614 // CHECK-SAME: vector<4x2xi1> -> vector<2x4xf32>1615 // CHECK: %[[MASK1:.*]] = vector.create_mask %[[D1]], %[[D0]]1616 // CHECK: %[[WRITE:.*]] = vector.mask %[[MASK1]] { vector.transfer_write %[[LOAD]], %arg1{{.+}} }1617 // CHECK-SAME: vector<2x4xi1> -> tensor<?x?xf32>1618 // CHECK: return %[[WRITE]]1619 %0 = linalg.transpose ins(%arg0 : tensor<?x?xf32>) outs(%arg1 : tensor<?x?xf32>) permutation = [1, 0]1620 return %0 : tensor<?x?xf32>1621}1622 1623module attributes {transform.with_named_sequence} {1624 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {1625 %0 = transform.structured.match ops{["linalg.transpose"]} in %arg1 : (!transform.any_op) -> !transform.any_op1626 transform.structured.vectorize %0 vector_sizes [2, 4] : !transform.any_op1627 transform.yield1628 }1629}1630 1631// -----1632 1633// CHECK-LABEL: func @test_masked_vectorize_linalg_copy1634func.func @test_masked_vectorize_linalg_copy(%A : memref<?x?xf32>, %B : memref<?x?xf32>) {1635 // CHECK: %[[c0:.*]] = arith.constant 0 : index1636 // CHECK: %[[d0:.*]] = memref.dim %{{.*}}, %[[c0]] : memref<?x?xf32>1637 // CHECK: %[[c1:.*]] = arith.constant 1 : index1638 // CHECK: %[[d1:.*]] = memref.dim %{{.*}}, %[[c1]] : memref<?x?xf32>1639 // CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1>1640 // CHECK: vector.mask %[[mask]] {{.*}} vector.transfer_read %{{.*}} {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x4xf32> } : vector<2x4xi1> -> vector<2x4xf32>1641 // CHECK: vector.mask %[[mask]] {{.*}} vector.transfer_write %{{.*}} {in_bounds = [true, true]} : vector<2x4xf32>, memref<?x?xf32> } : vector<2x4xi1>1642 linalg.copy ins(%A : memref<?x?xf32>) outs(%B : memref<?x?xf32>)1643 return1644}1645 1646module attributes {transform.with_named_sequence} {1647 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {1648 %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 : (!transform.any_op) -> !transform.any_op1649 transform.structured.vectorize %0 vector_sizes [2, 4] : !transform.any_op1650 transform.yield1651 }1652}1653 1654 1655 1656// -----1657 1658func.func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {1659 linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)1660 outs(%C: memref<?x?xf32>)1661 return1662}1663 1664// CHECK-LABEL: func.func @matmul(1665// CHECK-SAME: %[[A:.*]]: memref<?x?xf32>, %[[B:.*]]: memref<?x?xf32>, %[[C:.*]]: memref<?x?xf32>) {1666// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index1667// CHECK-DAG: %[[VAL_4:.*]] = memref.dim %[[A]], %[[VAL_3]] : memref<?x?xf32>1668// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index1669// CHECK-DAG: %[[VAL_6:.*]] = memref.dim %[[B]], %[[VAL_5]] : memref<?x?xf32>1670// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index1671// CHECK-DAG: %[[VAL_8:.*]] = memref.dim %[[A]], %[[VAL_7]] : memref<?x?xf32>1672// CHECK: %[[MASK_A:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_8]] : vector<8x4xi1>1673// CHECK: %[[LOAD_A:.*]] = vector.mask %[[MASK_A]] { vector.transfer_read %[[A]]{{\[}}%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true, true], permutation_map = #{{.*}}} : memref<?x?xf32>, vector<8x16x4xf32> } : vector<8x4xi1> -> vector<8x16x4xf32>1674// CHECK: %[[MASK_B:.*]] = vector.create_mask %[[VAL_8]], %[[VAL_6]] : vector<4x16xi1>1675// CHECK: %[[LOAD_B:.*]] = vector.mask %[[MASK_B]] { vector.transfer_read %[[B]]{{\[}}%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true, true], permutation_map = #{{.*}}} : memref<?x?xf32>, vector<8x16x4xf32> } : vector<4x16xi1> -> vector<8x16x4xf32>1676// CHECK: %[[MASK_C:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_6]] : vector<8x16xi1>1677// CHECK: %[[LOAD_C:.*]] = vector.mask %[[MASK_C]] { vector.transfer_read %[[C]]{{\[}}%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true]} : memref<?x?xf32>, vector<8x16xf32> } : vector<8x16xi1> -> vector<8x16xf32>1678// CHECK: %[[MULF:.*]] = arith.mulf %[[LOAD_A]], %[[LOAD_B]] : vector<8x16x4xf32>1679// CHECK: %[[MASK_MULIT_RED:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_6]], %[[VAL_8]] : vector<8x16x4xi1>1680// CHECK: %[[MULTI_RED:.*]] = vector.mask %[[MASK_MULIT_RED]] { vector.multi_reduction <add>, %[[MULF]], %[[LOAD_C]] [2] : vector<8x16x4xf32> to vector<8x16xf32> } : vector<8x16x4xi1> -> vector<8x16xf32>1681// CHECK: %[[C2:.*]] = arith.constant 0 : index1682// CHECK: vector.mask %[[MASK_C]] { vector.transfer_write %[[MULTI_RED]], %[[C]]{{\[}}%[[C2]], %[[C2]]] {in_bounds = [true, true]} : vector<8x16xf32>, memref<?x?xf32> } : vector<8x16xi1>1683 1684module attributes {transform.with_named_sequence} {1685 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {1686 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op1687 transform.structured.vectorize %matmul vector_sizes [8, 16, 4] : !transform.any_op1688 transform.yield1689 }1690}1691 1692 1693// -----1694 1695func.func @matmul_scalable(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {1696 linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)1697 outs(%C: memref<?x?xf32>)1698 return1699}1700 1701// CHECK-LABEL: func.func @matmul_scalable(1702// CHECK-SAME: %[[A:.*]]: memref<?x?xf32>, %[[B:.*]]: memref<?x?xf32>, %[[C:.*]]: memref<?x?xf32>) {1703// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index1704// CHECK-DAG: %[[VAL_4:.*]] = memref.dim %[[A]], %[[VAL_3]] : memref<?x?xf32>1705// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index1706// CHECK-DAG: %[[VAL_6:.*]] = memref.dim %[[B]], %[[VAL_5]] : memref<?x?xf32>1707// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index1708// CHECK-DAG: %[[VAL_8:.*]] = memref.dim %[[A]], %[[VAL_7]] : memref<?x?xf32>1709// CHECK: %[[MASK_A:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_8]] : vector<8x4xi1>1710// CHECK: %[[LOAD_A:.*]] = vector.mask %[[MASK_A]] { vector.transfer_read %[[A]]{{\[}}%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true, true], permutation_map = #{{.*}}} : memref<?x?xf32>, vector<8x[16]x4xf32> } : vector<8x4xi1> -> vector<8x[16]x4xf32>1711// CHECK: %[[MASK_B:.*]] = vector.create_mask %[[VAL_8]], %[[VAL_6]] : vector<4x[16]xi1>1712// CHECK: %[[LOAD_B:.*]] = vector.mask %[[MASK_B]] { vector.transfer_read %[[B]]{{\[}}%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true, true], permutation_map = #{{.*}}} : memref<?x?xf32>, vector<8x[16]x4xf32> } : vector<4x[16]xi1> -> vector<8x[16]x4xf32>1713// CHECK: %[[MASK_C:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_6]] : vector<8x[16]xi1>1714// CHECK: %[[LOAD_C:.*]] = vector.mask %[[MASK_C]] { vector.transfer_read %[[C]]{{\[}}%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true]} : memref<?x?xf32>, vector<8x[16]xf32> } : vector<8x[16]xi1> -> vector<8x[16]xf32>1715// CHECK: %[[MULF:.*]] = arith.mulf %[[LOAD_A]], %[[LOAD_B]] : vector<8x[16]x4xf32>1716// CHECK: %[[MASK_MULIT_RED:.*]] = vector.create_mask %[[VAL_4]], %[[VAL_6]], %[[VAL_8]] : vector<8x[16]x4xi1>1717// CHECK: %[[MULTI_RED:.*]] = vector.mask %[[MASK_MULIT_RED]] { vector.multi_reduction <add>, %[[MULF]], %[[LOAD_C]] [2] : vector<8x[16]x4xf32> to vector<8x[16]xf32> } : vector<8x[16]x4xi1> -> vector<8x[16]xf32>1718// CHECK: %[[C2:.*]] = arith.constant 0 : index1719// CHECK: vector.mask %[[MASK_C]] { vector.transfer_write %[[MULTI_RED]], %[[C]]{{\[}}%[[C2]], %[[C2]]] {in_bounds = [true, true]} : vector<8x[16]xf32>, memref<?x?xf32> } : vector<8x[16]xi1>1720 1721module attributes {transform.with_named_sequence} {1722 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {1723 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op1724 transform.structured.vectorize %matmul vector_sizes [8, [16], 4] : !transform.any_op1725 transform.yield1726 }1727}1728