485 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -split-input-file -verify-diagnostics | FileCheck %s2 3///----------------------------------------------------------------------------------------4/// Tests for vectorizing operations implementing contraction op interface.5/// Ops implementing the contraction interface are vectorized directly to their6/// vector dialect named counterparts.7///----------------------------------------------------------------------------------------8 9func.func @matmul(%A: tensor<8x4xf32>, %B: tensor<4x16xf32>,10 %C: tensor<8x16xf32>) -> tensor<8x16xf32> {11 %0 = linalg.matmul12 ins(%A, %B : tensor<8x4xf32>, tensor<4x16xf32>)13 outs(%C: tensor<8x16xf32>) -> tensor<8x16xf32>14 return %0 : tensor<8x16xf32>15}16 17// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>18// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>19// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>20// CHECK-LABEL: func.func @matmul(21// CHECK-SAME: %[[A:.*]]: tensor<8x4xf32>, %[[B:.*]]: tensor<4x16xf32>,22// CHECK-SAME: %[[C:.*]]: tensor<8x16xf32>)23// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %[[A]]{{.*}}: tensor<8x4xf32>, vector<8x4xf32>24// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %[[B]]{{.*}}: tensor<4x16xf32>, vector<4x16xf32>25// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %[[C]]{{.*}}: tensor<8x16xf32>, vector<8x16xf32>26// CHECK: %[[CONTRACT:.*]] = vector.contract27// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]28// CHECK-SAME: kind = #vector.kind<add>29// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]30// CHECK: vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<8x16xf32>, tensor<8x16xf32>31 32module attributes {transform.with_named_sequence} {33 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {34 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op35 transform.structured.vectorize %0 {create_named_contraction} : !transform.any_op36 transform.yield37 }38}39 40// -----41 42func.func @matmul_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,43 %C: tensor<?x?xf32>) -> tensor<?x?xf32> {44 %0 = linalg.matmul45 ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)46 outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>47 return %0 : tensor<?x?xf32>48}49 50// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>51// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>52// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>53// CHECK-LABEL: func.func @matmul_dynamic(54// CHECK-SAME: %[[A:.*]]: tensor<?x?xf32>, %[[B:.*]]: tensor<?x?xf32>,55// CHECK-SAME: %[[C:.*]]: tensor<?x?xf32>)56 57/// Get the contraction dimensions58// CHECK: %[[MATMUL_DIM_M_IDX:.*]] = arith.constant 0 : index59// CHECK: %[[MATMUL_DIM_M:.*]] = tensor.dim %[[A]], %[[MATMUL_DIM_M_IDX]] : tensor<?x?xf32>60// CHECK: %[[MATMUL_DIM_N_IDX:.*]] = arith.constant 1 : index61// CHECK: %[[MATMUL_DIM_N:.*]] = tensor.dim %[[B]], %[[MATMUL_DIM_N_IDX]] : tensor<?x?xf32>62// CHECK: %[[MATMUL_DIM_K_IDX:.*]] = arith.constant 1 : index63// CHECK: %[[MATMUL_DIM_K:.*]] = tensor.dim %[[A]], %[[MATMUL_DIM_K_IDX]] : tensor<?x?xf32>64 65/// Create a mask for the A matrix66// CHECK: %[[A_OFFSET:.*]] = arith.constant 0 : index67// CHECK: %[[A_DIM_M_IDX:.*]] = arith.constant 0 : index68// CHECK: %[[A_DIM_M:.*]] = tensor.dim %[[A]], %[[A_DIM_M_IDX]] : tensor<?x?xf32>69// CHECK: %[[A_DIM_K_IDX:.*]] = arith.constant 1 : index70// CHECK: %[[A_DIM_K:.*]] = tensor.dim %[[A]], %[[A_DIM_K_IDX]] : tensor<?x?xf32>71// CHECK: %[[LOAD_A_MASK:.*]] = vector.create_mask72// CHECK-SAME: %[[A_DIM_M]], %[[A_DIM_K]] : vector<8x4xi1>73/// Read the A matrix74// CHECK: %[[LOAD_A:.*]] = vector.mask %[[LOAD_A_MASK]]75// CHECK-SAME: { vector.transfer_read %[[A]]{{\[}}%[[A_OFFSET]], %[[A_OFFSET]]{{\]}}76// CHECK-SAME: : tensor<?x?xf32>, vector<8x4xf32> }77// CHECK-SAME: : vector<8x4xi1> -> vector<8x4xf32>78 79/// Create a mask for the B matrix80// CHECK: %[[B_OFFSET:.*]] = arith.constant 0 : index81// CHECK: %[[B_DIM_K_IDX:.*]] = arith.constant 0 : index82// CHECK: %[[B_DIM_K:.*]] = tensor.dim %[[B]], %[[B_DIM_K_IDX]] : tensor<?x?xf32>83// CHECK: %[[B_DIM_N_IDX:.*]] = arith.constant 1 : index84// CHECK: %[[B_DIM_N:.*]] = tensor.dim %[[B]], %[[B_DIM_N_IDX]] : tensor<?x?xf32>85// CHECK: %[[LOAD_B_MASK:.*]] = vector.create_mask86// CHECK-SAME: %[[B_DIM_K]], %[[B_DIM_N]] : vector<4x16xi1>87/// Read the B matrix88// CHECK: %[[LOAD_B:.*]] = vector.mask %[[LOAD_B_MASK]]89// CHECK-SAME: { vector.transfer_read %[[B]]{{\[}}%[[B_OFFSET]], %[[B_OFFSET]]{{\]}}90// CHECK-SAME: : tensor<?x?xf32>, vector<4x16xf32> }91// CHECK-SAME: : vector<4x16xi1> -> vector<4x16xf32>92 93/// Create a mask for the C matrix94// CHECK: %[[C_OFFSET:.*]] = arith.constant 0 : index95// CHECK: %[[C_DIM_M_IDX:.*]] = arith.constant 0 : index96// CHECK: %[[C_DIM_M:.*]] = tensor.dim %[[C]], %[[C_DIM_M_IDX]] : tensor<?x?xf32>97// CHECK: %[[C_DIM_N_IDX:.*]] = arith.constant 1 : index98// CHECK: %[[C_DIM_N:.*]] = tensor.dim %[[C]], %[[C_DIM_N_IDX]] : tensor<?x?xf32>99// CHECK: %[[LOAD_C_MASK:.*]] = vector.create_mask100// CHECK-SAME: %[[C_DIM_M]], %[[C_DIM_N]] : vector<8x16xi1>101/// Read the C matrix102// CHECK: %[[LOAD_C:.*]] = vector.mask %[[LOAD_C_MASK]]103// CHECK-SAME: { vector.transfer_read %[[C]]{{\[}}%[[C_OFFSET]], %[[C_OFFSET]]{{\]}}104// CHECK-SAME: : tensor<?x?xf32>, vector<8x16xf32> }105// CHECK-SAME: : vector<8x16xi1> -> vector<8x16xf32>106 107/// Create a mask for the contraction108// CHECK: %[[CONTRACTION_MASK:.*]] = vector.create_mask109// CHECK-SAME: %[[MATMUL_DIM_M]], %[[MATMUL_DIM_N]], %[[MATMUL_DIM_K]]110// CHECK-SAME: : vector<8x16x4xi1>111/// Perform the contraction112// CHECK: %[[D:.*]] = vector.mask %[[CONTRACTION_MASK]]113// CHECK-SAME: { vector.contract114// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]115// CHECK-SAME: kind = #vector.kind<add>116// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]117// CHECK-SAME: } : vector<8x16x4xi1> -> vector<8x16xf32>118 119/// Create a mask for the result120// CHECK: %[[D_OFFSET:.*]] = arith.constant 0 : index121// CHECK: %[[D_DIM_M_IDX:.*]] = arith.constant 0 : index122// CHECK: %[[D_DIM_M:.*]] = tensor.dim %[[C]], %[[D_DIM_M_IDX]] : tensor<?x?xf32>123// CHECK: %[[D_DIM_N_IDX:.*]] = arith.constant 1 : index124// CHECK: %[[D_DIM_N:.*]] = tensor.dim %[[C]], %[[D_DIM_N_IDX]] : tensor<?x?xf32>125// CHECK: %[[LOAD_D_MASK:.*]] = vector.create_mask126// CHECK-SAME: %[[D_DIM_M]], %[[D_DIM_N]] : vector<8x16xi1>127/// Write the result128// CHECK: vector.mask %[[LOAD_D_MASK]]129// CHECK-SAME: { vector.transfer_write %[[D]], %[[C]]{{\[}}%[[D_OFFSET]], %[[D_OFFSET]]{{\]}}130// CHECK-SAME: : vector<8x16xf32>, tensor<?x?xf32> }131// CHECK-SAME: : vector<8x16xi1> -> tensor<?x?xf32>132 133module attributes {transform.with_named_sequence} {134 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {135 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op136 transform.structured.vectorize %0 vector_sizes [8, 16, 4]137 {create_named_contraction} : !transform.any_op138 transform.yield139 }140}141 142// -----143 144func.func @matmul_dynamic_memref(%A: memref<?x?xf32>, %B: memref<?x?xf32>,145 %C: memref<?x?xf32>) {146 linalg.matmul147 ins(%A, %B : memref<?x?xf32>, memref<?x?xf32>)148 outs(%C: memref<?x?xf32>)149 return150}151 152// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>153// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>154// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>155// CHECK-LABEL: func.func @matmul_dynamic_memref(156// CHECK-SAME: %[[A:.*]]: memref<?x?xf32>, %[[B:.*]]: memref<?x?xf32>,157// CHECK-SAME: %[[C:.*]]: memref<?x?xf32>)158// CHECK: %[[LOAD_A:.*]] = vector.mask{{.*}}{ vector.transfer_read %[[A]]{{.*}}: memref<?x?xf32>, vector<8x4xf32>159// CHECK: %[[LOAD_B:.*]] = vector.mask{{.*}}{ vector.transfer_read %[[B]]{{.*}}: memref<?x?xf32>, vector<4x16xf32>160// CHECK: %[[LOAD_C:.*]] = vector.mask{{.*}}{ vector.transfer_read %[[C]]{{.*}}: memref<?x?xf32>, vector<8x16xf32>161// CHECK: %[[CONTRACT:.*]] = vector.mask{{.*}}{ vector.contract162// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]163// CHECK-SAME: kind = #vector.kind<add>164// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]165// CHECK: vector.mask{{.*}}{ vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<8x16xf32>, memref<?x?xf32>166 167module attributes {transform.with_named_sequence} {168 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {169 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op170 transform.structured.vectorize %0 vector_sizes [8, 16, 4]171 {create_named_contraction} : !transform.any_op172 transform.yield173 }174}175 176// -----177 178func.func @matmul_dynamic_scalable(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,179 %C: tensor<?x?xf32>) -> tensor<?x?xf32> {180 %0 = linalg.matmul181 ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)182 outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>183 return %0 : tensor<?x?xf32>184}185 186// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>187// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>188// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>189// CHECK-LABEL: func.func @matmul_dynamic_scalable(190// CHECK-SAME: %[[A:.*]]: tensor<?x?xf32>, %[[B:.*]]: tensor<?x?xf32>,191// CHECK-SAME: %[[C:.*]]: tensor<?x?xf32>)192// CHECK: %[[LOAD_A:.*]] = vector.mask{{.*}}{ vector.transfer_read %[[A]]{{.*}}: tensor<?x?xf32>, vector<8x4xf32> }193// CHECK-SAME: : vector<8x4xi1> -> vector<8x4xf32>194// CHECK: %[[LOAD_B:.*]] = vector.mask{{.*}}{ vector.transfer_read %[[B]]{{.*}}: tensor<?x?xf32>, vector<4x[16]xf32> }195// CHECK-SAME: : vector<4x[16]xi1> -> vector<4x[16]xf32>196// CHECK: %[[LOAD_C:.*]] = vector.mask{{.*}}{ vector.transfer_read %[[C]]{{.*}}: tensor<?x?xf32>, vector<8x[16]xf32> }197// CHECK-SAME: : vector<8x[16]xi1> -> vector<8x[16]xf32>198// CHECK: %[[CONTRACT:.*]] = vector.mask{{.*}}{ vector.contract199// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]200// CHECK-SAME: kind = #vector.kind<add>201// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]202// CHECK-SAME: } : vector<8x[16]x4xi1> -> vector<8x[16]xf32>203// CHECK: vector.mask{{.*}}{ vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<8x[16]xf32>, tensor<?x?xf32> }204// CHECK-SAME: : vector<8x[16]xi1> -> tensor<?x?xf32>205 206module attributes {transform.with_named_sequence} {207 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {208 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op209 transform.structured.vectorize %0 vector_sizes [8, [16], 4]210 {create_named_contraction} : !transform.any_op211 transform.yield212 }213}214 215// -----216 217func.func @matmul_transpose(%A: tensor<4x8xf32>, %B: tensor<16x4xf32>,218 %C: tensor<8x16xf32>) -> tensor<8x16xf32> {219 %0 = linalg.matmul220 indexing_maps = [affine_map<(m, n, k) -> (k, m)>, // transpose A221 affine_map<(m, n, k) -> (n, k)>, // transpose B222 affine_map<(m, n, k) -> (m, n)>]223 ins(%A, %B : tensor<4x8xf32>, tensor<16x4xf32>)224 outs(%C: tensor<8x16xf32>) -> tensor<8x16xf32>225 return %0 : tensor<8x16xf32>226}227 228// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>229// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>230// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>231// CHECK-LABEL: func.func @matmul_transpose(232// CHECK-SAME: %[[A:.*]]: tensor<4x8xf32>, %[[B:.*]]: tensor<16x4xf32>,233// CHECK-SAME: %[[C:.*]]: tensor<8x16xf32>)234// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %[[A]]{{.*}}: tensor<4x8xf32>, vector<4x8xf32>235// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %[[B]]{{.*}}: tensor<16x4xf32>, vector<16x4xf32>236// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %[[C]]{{.*}}: tensor<8x16xf32>, vector<8x16xf32>237// CHECK: %[[CONTRACT:.*]] = vector.contract238// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]239// CHECK-SAME: kind = #vector.kind<add>240// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]241// CHECK: vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<8x16xf32>, tensor<8x16xf32>242 243module attributes {transform.with_named_sequence} {244 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {245 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op246 transform.structured.vectorize %0 {create_named_contraction} : !transform.any_op247 transform.yield248 }249}250 251// -----252 253func.func @matmul_dynamic_transpose(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,254 %C: tensor<?x?xf32>) -> tensor<?x?xf32> {255 %0 = linalg.matmul256 indexing_maps = [affine_map<(m, n, k) -> (k, m)>, // transpose A257 affine_map<(m, n, k) -> (n, k)>, // transpose B258 affine_map<(m, n, k) -> (m, n)>]259 ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)260 outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>261 return %0 : tensor<?x?xf32>262}263 264// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>265// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>266// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>267// CHECK-LABEL: func.func @matmul_dynamic_transpose(268// CHECK-SAME: %[[A:.*]]: tensor<?x?xf32>, %[[B:.*]]: tensor<?x?xf32>,269// CHECK-SAME: %[[C:.*]]: tensor<?x?xf32>)270// CHECK: %[[LOAD_A:.*]] = vector.mask{{.*}}{ vector.transfer_read %[[A]]{{.*}}: tensor<?x?xf32>, vector<4x8xf32>271// CHECK: %[[LOAD_B:.*]] = vector.mask{{.*}}{ vector.transfer_read %[[B]]{{.*}}: tensor<?x?xf32>, vector<16x4xf32>272// CHECK: %[[LOAD_C:.*]] = vector.mask{{.*}}{ vector.transfer_read %[[C]]{{.*}}: tensor<?x?xf32>, vector<8x16xf32>273// CHECK: %[[CONTRACT:.*]] = vector.mask{{.*}}{ vector.contract274// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]275// CHECK-SAME: kind = #vector.kind<add>276// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]277// CHECK: vector.mask{{.*}}{ vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<8x16xf32>, tensor<?x?xf32>278 279module attributes {transform.with_named_sequence} {280 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {281 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op282 transform.structured.vectorize %0 vector_sizes [8, 16, 4]283 {create_named_contraction} : !transform.any_op284 transform.yield285 }286}287 288// -----289 290/// Contractions with arbitrarty broadcasts are not supported in contraction interface291/// vectorization.292/// Dimension broadcasts are expected to be decomposed first which removes ambiguity293/// caused by possible variants of dimensions materialization.294/// For example, whether the below target LHS input layout is (m, k) or (k, m).295 296func.func @negative_matmul_broadcast(%A: tensor<4xf32>, %B: tensor<4x16xf32>,297 %C: tensor<8x16xf32>) -> tensor<8x16xf32> {298 // expected-error @+1 {{Attempted to vectorize, but failed}}299 %0 = linalg.matmul300 indexing_maps = [affine_map<(m, n, k) -> (k)>, // broadcast301 affine_map<(m, n, k) -> (k, n)>,302 affine_map<(m, n, k) -> (m, n)>]303 ins(%A, %B : tensor<4xf32>, tensor<4x16xf32>)304 outs(%C: tensor<8x16xf32>) -> tensor<8x16xf32>305 return %0 : tensor<8x16xf32>306}307 308module attributes {transform.with_named_sequence} {309 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {310 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op311 transform.structured.vectorize %0 {create_named_contraction} : !transform.any_op312 transform.yield313 }314}315 316// -----317 318func.func @matmul_mixed_precision(%A: tensor<8x4xf16>, %B: tensor<4x16xf16>,319 %C: tensor<8x16xf32>) -> tensor<8x16xf32> {320 %0 = linalg.matmul321 ins(%A, %B : tensor<8x4xf16>, tensor<4x16xf16>)322 outs(%C: tensor<8x16xf32>) -> tensor<8x16xf32>323 return %0 : tensor<8x16xf32>324}325 326// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>327// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>328// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>329// CHECK-LABEL: func.func @matmul_mixed_precision(330// CHECK-SAME: %[[A:.*]]: tensor<8x4xf16>, %[[B:.*]]: tensor<4x16xf16>,331// CHECK-SAME: %[[C:.*]]: tensor<8x16xf32>)332// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %[[A]]{{.*}}: tensor<8x4xf16>, vector<8x4xf16>333// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %[[B]]{{.*}}: tensor<4x16xf16>, vector<4x16xf16>334// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %[[C]]{{.*}}: tensor<8x16xf32>, vector<8x16xf32>335// CHECK: %[[CONTRACT:.*]] = vector.contract336// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]337// CHECK-SAME: kind = #vector.kind<add>338// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]339// CHECK: vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<8x16xf32>, tensor<8x16xf32>340 341module attributes {transform.with_named_sequence} {342 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {343 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op344 transform.structured.vectorize %0 {create_named_contraction} : !transform.any_op345 transform.yield346 }347}348 349// -----350 351func.func @batch_matmul(%A: tensor<3x8x4xf32>, %B: tensor<3x4x16xf32>,352 %C: tensor<3x8x16xf32>) -> tensor<3x8x16xf32> {353 %0 = linalg.batch_matmul354 ins(%A, %B : tensor<3x8x4xf32>, tensor<3x4x16xf32>)355 outs(%C: tensor<3x8x16xf32>) -> tensor<3x8x16xf32>356 return %0 : tensor<3x8x16xf32>357}358 359// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>360// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>361// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>362// CHECK-LABEL: func.func @batch_matmul(363// CHECK-SAME: %[[A:.*]]: tensor<3x8x4xf32>, %[[B:.*]]: tensor<3x4x16xf32>,364// CHECK-SAME: %[[C:.*]]: tensor<3x8x16xf32>)365// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %[[A]]{{.*}}: tensor<3x8x4xf32>, vector<3x8x4xf32>366// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %[[B]]{{.*}}: tensor<3x4x16xf32>, vector<3x4x16xf32>367// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %[[C]]{{.*}}: tensor<3x8x16xf32>, vector<3x8x16xf32>368// CHECK: %[[CONTRACT:.*]] = vector.contract369// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]370// CHECK-SAME: kind = #vector.kind<add>371// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]372// CHECK: vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<3x8x16xf32>, tensor<3x8x16xf32>373 374module attributes {transform.with_named_sequence} {375 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {376 %0 = transform.structured.match ops{["linalg.batch_matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op377 transform.structured.vectorize %0 {create_named_contraction} : !transform.any_op378 transform.yield379 }380}381 382// -----383 384func.func @batch_reduce_matmul(%A: tensor<3x8x4xf32>, %B: tensor<3x4x16xf32>,385 %C: tensor<8x16xf32>) -> tensor<8x16xf32> {386 %0 = linalg.batch_reduce_matmul387 ins(%A, %B : tensor<3x8x4xf32>, tensor<3x4x16xf32>)388 outs(%C: tensor<8x16xf32>) -> tensor<8x16xf32>389 return %0 : tensor<8x16xf32>390}391 392// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>393// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>394// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>395// CHECK-LABEL: func.func @batch_reduce_matmul(396// CHECK-SAME: %[[A:.*]]: tensor<3x8x4xf32>, %[[B:.*]]: tensor<3x4x16xf32>,397// CHECK-SAME: %[[C:.*]]: tensor<8x16xf32>)398// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %[[A]]{{.*}}: tensor<3x8x4xf32>, vector<3x8x4xf32>399// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %[[B]]{{.*}}: tensor<3x4x16xf32>, vector<3x4x16xf32>400// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %[[C]]{{.*}}: tensor<8x16xf32>, vector<8x16xf32>401// CHECK: %[[CONTRACT:.*]] = vector.contract402// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]403// CHECK-SAME: kind = #vector.kind<add>404// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]405// CHECK: vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<8x16xf32>, tensor<8x16xf32>406 407module attributes {transform.with_named_sequence} {408 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {409 %0 = transform.structured.match ops{["linalg.batch_reduce_matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op410 transform.structured.vectorize %0 {create_named_contraction} : !transform.any_op411 transform.yield412 }413}414 415// -----416 417func.func @contract(%A: tensor<4x8x2xf32>, %B: tensor<8x16x2xf32>,418 %C: tensor<4x16xf32>) -> tensor<4x16xf32> {419 %0 = linalg.contract420 indexing_maps = [affine_map<(m, n, k, kk) -> (m, k, kk)>,421 affine_map<(m, n, k, kk) -> (k, n, kk)>,422 affine_map<(m, n, k, kk) -> (m, n)>]423 ins(%A, %B : tensor<4x8x2xf32>, tensor<8x16x2xf32>)424 outs(%C : tensor<4x16xf32>) -> tensor<4x16xf32>425 return %0 : tensor<4x16xf32>426}427 428// CHECK: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>429// CHECK: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3)>430// CHECK: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1)>431// CHECK-LABEL: func.func @contract(432// CHECK-SAME: %[[A:.*]]: tensor<4x8x2xf32>, %[[B:.*]]: tensor<8x16x2xf32>,433// CHECK-SAME: %[[C:.*]]: tensor<4x16xf32>)434// CHECK: %[[LOAD_A:.*]] = vector.transfer_read %[[A]]{{.*}}: tensor<4x8x2xf32>, vector<4x8x2xf32>435// CHECK: %[[LOAD_B:.*]] = vector.transfer_read %[[B]]{{.*}}: tensor<8x16x2xf32>, vector<8x16x2xf32>436// CHECK: %[[LOAD_C:.*]] = vector.transfer_read %[[C]]{{.*}}: tensor<4x16xf32>, vector<4x16xf32>437// CHECK: %[[CONTRACT:.*]] = vector.contract438// CHECK-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]439// CHECK-SAME: kind = #vector.kind<add>440// CHECK-SAME: %[[LOAD_A]], %[[LOAD_B]], %[[LOAD_C]]441// CHECK: vector.transfer_write %[[CONTRACT]], %[[C]]{{.*}}: vector<4x16xf32>, tensor<4x16xf32>442 443module attributes {transform.with_named_sequence} {444 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {445 %0 = transform.structured.match ops{["linalg.contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op446 transform.structured.vectorize %0 {create_named_contraction} : !transform.any_op447 transform.yield448 }449}450 451// -----452 453/// Generic can represent contractions but it does not implement contraction interface.454/// Thus, direct lowering to vector.contract is not supported.455/// Vectorization still works and applies generic rewrite logic.456 457func.func @negative_generic(%A: tensor<8x4xf32>, %B: tensor<4x16xf32>,458 %C: tensor<8x16xf32>) -> tensor<8x16xf32> {459 %0 = linalg.generic {460 indexing_maps = [affine_map<(m, n, k) -> (m, k)>,461 affine_map<(m, n, k) -> (k, n)>,462 affine_map<(m, n, k) -> (m, n)>],463 iterator_types = ["parallel", "parallel", "reduction"]}464 ins(%A, %B : tensor<8x4xf32>, tensor<4x16xf32>)465 outs(%C : tensor<8x16xf32>) {466 ^bb0(%in: f32, %in_0: f32, %out: f32):467 %1 = arith.mulf %in, %in_0 : f32468 %2 = arith.addf %out, %1 : f32469 linalg.yield %2 : f32470 } -> tensor<8x16xf32>471 return %0 : tensor<8x16xf32>472}473 474// CHECK-LABEL: func.func @negative_generic(475// CHECK-NOT: vector.contract476// CHECK: vector.multi_reduction477 478module attributes {transform.with_named_sequence} {479 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {480 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op481 transform.structured.vectorize %0 {create_named_contraction} : !transform.any_op482 transform.yield483 }484}485