322 lines · plain
1// RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s2 3#dotp_accesses = [4 affine_map<(i) -> (i)>,5 affine_map<(i) -> (i)>,6 affine_map<(i) -> ()>7]8#dotp_trait = {9 indexing_maps = #dotp_accesses,10 iterator_types = ["reduction"]11}12 13// CHECK-LABEL: func @extract_contract114// CHECK-SAME: %[[A:.*0]]: vector<4xf32>,15// CHECK-SAME: %[[B:.*1]]: vector<4xf32>,16// CHECK-SAME: %[[C:.*2]]: f3217// CHECK: %[[F:.*]] = arith.mulf %[[A]], %[[B]] : vector<4xf32>18// CHECK: %[[R:.*]] = vector.reduction <add>, %[[F]], %[[C]] : vector<4xf32> into f3219// CHECK: return %[[R]] : f3220 21func.func @extract_contract1(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: f32) -> f32 {22 %0 = vector.contract #dotp_trait %arg0, %arg1, %arg223 : vector<4xf32>, vector<4xf32> into f3224 return %0 : f3225}26 27// CHECK-LABEL: func @masked_extract_contract128// CHECK-SAME: %[[A:.*0]]: vector<4xf32>, %[[B:.*1]]: vector<4xf32>, %[[C:.*2]]: f3229// CHECK-SAME: %[[M:.*]]: vector<4xi1>30// CHECK: %[[F:.*]] = arith.mulf %[[A]], %[[B]] : vector<4xf32>31// CHECK: %[[R:.*]] = vector.mask %[[M]] { vector.reduction <add>, %0, %arg2 : vector<4xf32> into f32 } : vector<4xi1> -> f3232// CHECK: return %[[R]] : f3233 34func.func @masked_extract_contract1(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: f32, %mask: vector<4xi1>) -> f32 {35 %0 = vector.mask %mask { vector.contract #dotp_trait %arg0, %arg1, %arg2 : vector<4xf32>, vector<4xf32> into f32 } : vector<4xi1> -> f3236 return %0 : f3237}38 39// CHECK-LABEL: func @extract_contract1_int40// CHECK-SAME: %[[A:.*0]]: vector<4xi32>,41// CHECK-SAME: %[[B:.*1]]: vector<4xi32>,42// CHECK-SAME: %[[C:.*2]]: i3243// CHECK: %[[F:.*]] = arith.muli %[[A]], %[[B]] : vector<4xi32>44// CHECK: %[[R:.*]] = vector.reduction <add>, %[[F]], %[[C]] : vector<4xi32> into i3245// CHECK: return %[[R]] : i3246 47func.func @extract_contract1_int(%arg0: vector<4xi32>, %arg1: vector<4xi32>, %arg2: i32) -> i32 {48 %0 = vector.contract #dotp_trait %arg0, %arg1, %arg249 : vector<4xi32>, vector<4xi32> into i3250 return %0 : i3251}52 53#matvec_accesses = [54 affine_map<(i, j) -> (i, j)>,55 affine_map<(i, j) -> (j)>,56 affine_map<(i, j) -> (i)>57]58#matvec_trait = {59 indexing_maps = #matvec_accesses,60 iterator_types = ["parallel", "reduction"]61}62 63// CHECK-LABEL: func @extract_contract264// CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,65// CHECK-SAME: %[[B:.*1]]: vector<3xf32>,66// CHECK-SAME: %[[C:.*2]]: vector<2xf32>67// CHECK: %[[R:.*]] = arith.constant dense<0.000000e+00> : vector<2xf32>68// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : vector<3xf32> from vector<2x3xf32>69// CHECK: %[[T2:.*]] = arith.mulf %[[T0]], %[[B]] : vector<3xf32>70// CHECK: %[[T3:.*]] = vector.reduction <add>, %[[T2]] : vector<3xf32> into f3271// CHECK: %[[T4:.*]] = vector.insert %[[T3]], %[[R]] [0] : f32 into vector<2xf32>72// CHECK: %[[T5:.*]] = vector.extract %[[A]][1] : vector<3xf32> from vector<2x3xf32>73// CHECK: %[[T7:.*]] = arith.mulf %[[T5]], %[[B]] : vector<3xf32>74// CHECK: %[[T8:.*]] = vector.reduction <add>, %[[T7]] : vector<3xf32> into f3275// CHECK: %[[T9:.*]] = vector.insert %[[T8]], %[[T4]] [1] : f32 into vector<2xf32>76// CHECK: %[[T10:.*]] = arith.addf %[[T9]], %[[C]] : vector<2xf32>77// CHECK: return %[[T10]] : vector<2xf32>78 79func.func @extract_contract2(%arg0: vector<2x3xf32>,80 %arg1: vector<3xf32>,81 %arg2: vector<2xf32>) -> vector<2xf32> {82 %0 = vector.contract #matvec_trait %arg0, %arg1, %arg283 : vector<2x3xf32>, vector<3xf32> into vector<2xf32>84 return %0 : vector<2xf32>85}86 87// CHECK-LABEL: func @extract_contract2_int88// CHECK-SAME: %[[A:.*0]]: vector<2x3xi32>,89// CHECK-SAME: %[[B:.*1]]: vector<3xi32>,90// CHECK-SAME: %[[C:.*2]]: vector<2xi32>91// CHECK: %[[R:.*]] = arith.constant dense<0> : vector<2xi32>92// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : vector<3xi32> from vector<2x3xi32>93// CHECK: %[[T2:.*]] = arith.muli %[[T0]], %[[B]] : vector<3xi32>94// CHECK: %[[T3:.*]] = vector.reduction <add>, %[[T2]] : vector<3xi32> into i3295// CHECK: %[[T4:.*]] = vector.insert %[[T3]], %[[R]] [0] : i32 into vector<2xi32>96// CHECK: %[[T5:.*]] = vector.extract %[[A]][1] : vector<3xi32> from vector<2x3xi32>97// CHECK: %[[T7:.*]] = arith.muli %[[T5]], %[[B]] : vector<3xi32>98// CHECK: %[[T8:.*]] = vector.reduction <add>, %[[T7]] : vector<3xi32> into i3299// CHECK: %[[T9:.*]] = vector.insert %[[T8]], %[[T4]] [1] : i32 into vector<2xi32>100// CHECK: %[[T10:.*]] = arith.addi %[[T9]], %[[C]] : vector<2xi32>101// CHECK: return %[[T10]] : vector<2xi32>102func.func @extract_contract2_int(%arg0: vector<2x3xi32>,103 %arg1: vector<3xi32>,104 %arg2: vector<2xi32>) -> vector<2xi32> {105 %0 = vector.contract #matvec_trait %arg0, %arg1, %arg2106 : vector<2x3xi32>, vector<3xi32> into vector<2xi32>107 return %0 : vector<2xi32>108}109 110#vecmat_accesses = [111 affine_map<(i, j) -> (j)>,112 affine_map<(i, j) -> (i, j)>,113 affine_map<(i, j) -> (i)>114]115#vecmat_trait = {116 indexing_maps = #vecmat_accesses,117 iterator_types = ["parallel", "reduction"]118}119 120// CHECK-LABEL: func @extract_contract3121// CHECK-SAME: %[[A:.*0]]: vector<3xf32>,122// CHECK-SAME: %[[B:.*1]]: vector<2x3xf32>,123// CHECK-SAME: %[[C:.*2]]: vector<2xf32>124// CHECK: %[[R:.*]] = arith.constant dense<0.000000e+00> : vector<2xf32>125// CHECK: %[[T0:.*]] = vector.extract %[[B]][0] : vector<3xf32> from vector<2x3xf32>126// CHECK: %[[T2:.*]] = arith.mulf %[[T0]], %[[A]] : vector<3xf32>127// CHECK: %[[T3:.*]] = vector.reduction <add>, %[[T2]] : vector<3xf32> into f32128// CHECK: %[[T4:.*]] = vector.insert %[[T3]], %[[R]] [0] : f32 into vector<2xf32>129// CHECK: %[[T5:.*]] = vector.extract %[[B]][1] : vector<3xf32> from vector<2x3xf32>130// CHECK: %[[T7:.*]] = arith.mulf %[[T5]], %[[A]] : vector<3xf32>131// CHECK: %[[T8:.*]] = vector.reduction <add>, %[[T7]] : vector<3xf32> into f32132// CHECK: %[[T9:.*]] = vector.insert %[[T8]], %[[T4]] [1] : f32 into vector<2xf32>133// CHECK: %[[T10:.*]] = arith.addf %[[T9]], %[[C]] : vector<2xf32>134// CHECK: return %[[T10]] : vector<2xf32>135 136func.func @extract_contract3(%arg0: vector<3xf32>,137 %arg1: vector<2x3xf32>,138 %arg2: vector<2xf32>) -> vector<2xf32> {139 %0 = vector.contract #vecmat_trait %arg0, %arg1, %arg2140 : vector<3xf32>, vector<2x3xf32> into vector<2xf32>141 return %0 : vector<2xf32>142}143 144#matmat_accesses = [145 affine_map<(i, j, k) -> (i, k)>,146 affine_map<(i, j, k) -> (k, j)>,147 affine_map<(i, j, k) -> (i, j)>148]149#matmat_trait = {150 indexing_maps = #matmat_accesses,151 iterator_types = ["parallel", "parallel", "reduction"]152}153 154// CHECK-LABEL: func @contract_to_dot_matmat155// CHECK-SAME: %[[LHS:.*0]]: vector<2x2xf32>,156// CHECK-SAME: %[[RHS:.*1]]: vector<2x2xf32>,157// CHECK-SAME: %[[OUT:.*2]]: vector<2x2xf32>158//159// The `vector.contract` to dot lowering will 'unroll' a matrix-matrix160// multiplication into individual dot products betweem rows of the LHS with columns161// of the RHS. In the following test we expect 4 extract-dotproduct-insert sequences of162// ops that correspond to the 4 dot products resulting from unrolling a matmul between163// two matrices of size (2, 2).164//165// CHECK: %[[INIT:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf32>166//167// First, The RHS will be transposed to make it easier to extract individual columns168// using vector.extract.169//170// CHECK: %[[RHS_T:.*]] = vector.transpose %[[RHS]], [1, 0] : vector<2x2xf32> to vector<2x2xf32>171//172// Next, we expect 4 sequences of extracting rows of the RHS, LHS, performing a dot173// product and then inserting it into the result.174//175// CHECK: %[[LHS0:.*]] = vector.extract %[[LHS]][0] : vector<2xf32> from vector<2x2xf32>176// CHECK: %[[RHS_T0:.*]] = vector.extract %[[RHS_T]][0] : vector<2xf32> from vector<2x2xf32>177// CHECK: %[[PROD0:.*]] = arith.mulf %[[LHS0]], %[[RHS_T0]] : vector<2xf32>178// CHECK: %[[SUM0:.*]] = vector.reduction <add>, %[[PROD0]] : vector<2xf32> into f32179// CHECK: %[[RES0:.*]] = vector.insert %[[SUM0]], %[[INIT]] [0, 0] : f32 into vector<2x2xf32>180//181// CHECK: %[[RHS_T1:.*]] = vector.extract %[[RHS_T]][1] : vector<2xf32> from vector<2x2xf32>182// CHECK: %[[PROD1:.*]] = arith.mulf %[[LHS0]], %[[RHS_T1]] : vector<2xf32>183// CHECK: %[[SUM1:.*]] = vector.reduction <add>, %[[PROD1]] : vector<2xf32> into f32184// CHECK: %[[RES1:.*]] = vector.insert %[[SUM1]], %[[RES0]] [0, 1] : f32 into vector<2x2xf32>185//186// CHECK: %[[LHS1:.*]] = vector.extract %[[LHS]][1] : vector<2xf32> from vector<2x2xf32>187// CHECK: %[[PROD2:.*]] = arith.mulf %[[LHS1]], %[[RHS_T0]] : vector<2xf32>188// CHECK: %[[SUM2:.*]] = vector.reduction <add>, %[[PROD2]] : vector<2xf32> into f32189// CHECK: %[[RES2:.*]] = vector.insert %[[SUM2]], %[[RES1]] [1, 0] : f32 into vector<2x2xf32>190//191// CHECK: %[[PROD3:.*]] = arith.mulf %[[LHS1]], %[[RHS_T1]] : vector<2xf32>192// CHECK: %[[SUM3:.*]] = vector.reduction <add>, %[[PROD3]] : vector<2xf32> into f32193// CHECK: %[[RES3:.*]] = vector.insert %[[SUM3]], %[[RES2]] [1, 1] : f32 into vector<2x2xf32>194//195// CHECK: %[[RES:.*]] = arith.addf %[[RES3]], %[[OUT]] : vector<2x2xf32>196// CHECK: return %[[RES]] : vector<2x2xf32>197 198func.func @contract_to_dot_matmat(%lhs: vector<2x2xf32>,199 %rhs: vector<2x2xf32>,200 %init: vector<2x2xf32>) -> vector<2x2xf32> {201 %res = vector.contract #matmat_trait %lhs, %rhs, %init202 : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>203 return %res : vector<2x2xf32>204}205 206 207#contraction2d_accesses = [208 affine_map<(i, j) -> (i, j)>,209 affine_map<(i, j) -> (i, j)>,210 affine_map<(i, j) -> ()>211]212#contraction2d_trait = {213 indexing_maps = #contraction2d_accesses,214 iterator_types = ["reduction", "reduction"]215}216 217// CHECK-LABEL: func @full_contract1218// CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,219// CHECK-SAME: %[[B:.*1]]: vector<2x3xf32>,220// CHECK-SAME: %[[C:.*2]]: f32221// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : vector<3xf32> from vector<2x3xf32>222// CHECK: %[[T1:.*]] = vector.extract %[[B]][0] : vector<3xf32> from vector<2x3xf32>223// CHECK: %[[T2:.*]] = arith.mulf %[[T0]], %[[T1]] : vector<3xf32>224// CHECK: %[[T3:.*]] = vector.reduction <add>, %[[T2]], %[[C]] : vector<3xf32> into f32225// CHECK: %[[T5:.*]] = vector.extract %[[A]][1] : vector<3xf32> from vector<2x3xf32>226// CHECK: %[[T6:.*]] = vector.extract %[[B]][1] : vector<3xf32> from vector<2x3xf32>227// CHECK: %[[T7:.*]] = arith.mulf %[[T5]], %[[T6]] : vector<3xf32>228// CHECK: %[[T8:.*]] = vector.reduction <add>, %[[T7]], %[[T3]] : vector<3xf32> into f32229// CHECK: return %[[T8]] : f32230 231func.func @full_contract1(%arg0: vector<2x3xf32>,232 %arg1: vector<2x3xf32>,233 %arg2: f32) -> f32 {234 %0 = vector.contract #contraction2d_trait %arg0, %arg1, %arg2235 : vector<2x3xf32>, vector<2x3xf32> into f32236 return %0 : f32237}238 239#contraction2d_trans_accesses = [240 affine_map<(i, j) -> (i, j)>,241 affine_map<(i, j) -> (j, i)>,242 affine_map<(i, j) -> ()>243]244#contraction2d_trans_trait = {245 indexing_maps = #contraction2d_trans_accesses,246 iterator_types = ["reduction", "reduction"]247}248 249// CHECK-LABEL: func @full_contract2250// CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,251// CHECK-SAME: %[[B:.*1]]: vector<3x2xf32>,252// CHECK-SAME: %[[C:.*2]]: f32253// CHECK: %[[Z:.*]] = arith.constant dense<0.000000e+00> : vector<3xf32>254// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : vector<3xf32> from vector<2x3xf32>255// CHECK: %[[T1:.*]] = vector.extract %[[B]][0, 0] : f32 from vector<3x2xf32>256// CHECK: %[[T3:.*]] = vector.insert %[[T1]], %[[Z]] [0] : f32 into vector<3xf32>257// CHECK: %[[T4:.*]] = vector.extract %[[B]][1, 0] : f32 from vector<3x2xf32>258// CHECK: %[[T6:.*]] = vector.insert %[[T4]], %[[T3]] [1] : f32 into vector<3xf32>259// CHECK: %[[T7:.*]] = vector.extract %[[B]][2, 0] : f32 from vector<3x2xf32>260// CHECK: %[[T9:.*]] = vector.insert %[[T7]], %[[T6]] [2] : f32 into vector<3xf32>261// CHECK: %[[T10:.*]] = arith.mulf %[[T0]], %[[T9]] : vector<3xf32>262// CHECK: %[[T11:.*]] = vector.reduction <add>, %[[T10]], %[[C]] : vector<3xf32> into f32263//264// CHECK: %[[T12:.*]] = vector.extract %[[A]][1] : vector<3xf32> from vector<2x3xf32>265// CHECK: %[[T13:.*]] = vector.extract %[[B]][0, 1] : f32 from vector<3x2xf32>266// CHECK: %[[T15:.*]] = vector.insert %[[T13]], %[[Z]] [0] : f32 into vector<3xf32>267// CHECK: %[[T16:.*]] = vector.extract %[[B]][1, 1] : f32 from vector<3x2xf32>268// CHECK: %[[T18:.*]] = vector.insert %[[T16]], %[[T15]] [1] : f32 into vector<3xf32>269// CHECK: %[[T19:.*]] = vector.extract %[[B]][2, 1] : f32 from vector<3x2xf32>270// CHECK: %[[T21:.*]] = vector.insert %[[T19]], %[[T18]] [2] : f32 into vector<3xf32>271// CHECK: %[[T22:.*]] = arith.mulf %[[T12]], %[[T21]] : vector<3xf32>272// CHECK: %[[T23:.*]] = vector.reduction <add>, %[[T22]], %[[T11]] : vector<3xf32> into f32273// CHECK: return %[[T23]] : f32274 275func.func @full_contract2(%arg0: vector<2x3xf32>,276 %arg1: vector<3x2xf32>,277 %arg2: f32) -> f32 {278 %0 = vector.contract #contraction2d_trans_trait %arg0, %arg1, %arg2279 : vector<2x3xf32>, vector<3x2xf32> into f32280 return %0 : f32281}282 283// CHECK-LABEL: @contract_one_sided_unit_reduction_dim284// CHECK-SAME: (%[[A0:.+]]: vector<1x2xi32>, %[[A1:.+]]: vector<2x2xi32>, %[[A2:.+]]: vector<2xi32>)285// CHECK-DAG: %[[C:.+]] = arith.constant dense<0> : vector<2xi32>286// CHECK-DAG: %[[E00:.+]] = vector.extract %[[A0]][0] : vector<2xi32> from vector<1x2xi32>287// CHECK-DAG: %[[E10:.+]] = vector.extract %[[A1]][0] : vector<2xi32> from vector<2x2xi32>288// CHECK: %[[M0:.+]] = arith.muli %[[E10]], %[[E00]] : vector<2xi32>289// CHECK: %[[R0:.+]] = vector.reduction <add>, %[[M0]] : vector<2xi32> into i32290// CHECK: %[[I0:.+]] = vector.insert %[[R0]], %[[C]] [0] : i32 into vector<2xi32>291// CHECK: %[[E11:.+]] = vector.extract %[[A1]][1] : vector<2xi32> from vector<2x2xi32>292// CHECK: %[[M1:.+]] = arith.muli %[[E11]], %[[E00]] : vector<2xi32>293// CHECK: %[[R1:.+]] = vector.reduction <add>, %[[M1]] : vector<2xi32> into i32294// CHECK: %[[I1:.+]] = vector.insert %[[R1]], %[[I0]] [1] : i32 into vector<2xi32>295// CHECK: %[[S:.+]] = arith.addi %[[I1]], %[[A2]] : vector<2xi32>296// CHECK: return %[[S]] : vector<2xi32>297 298func.func @contract_one_sided_unit_reduction_dim(%arg0 : vector<1x2xi32>, %arg1 : vector<2x2xi32>, %arg2 : vector<2xi32>) -> vector<2xi32> {299 %res = vector.contract {300 indexing_maps = [301 affine_map<(d0, d1, d2) -> (d0, d2)>,302 affine_map<(d0, d1, d2) -> (d1, d2)>,303 affine_map<(d0, d1, d2) -> (d1)>304 ],305 iterator_types = ["reduction", "parallel", "reduction"],306 kind = #vector.kind<add>307 } %arg0, %arg1, %arg2 : vector<1x2xi32>, vector<2x2xi32>, vector<2xi32> into vector<2xi32>308 return %res : vector<2xi32>309}310 311module attributes {transform.with_named_sequence} {312 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {313 %f = transform.structured.match ops{["func.func"]} in %module_op314 : (!transform.any_op) -> !transform.any_op315 316 transform.apply_patterns to %f {317 transform.apply_patterns.vector.lower_contraction lowering_strategy = "dot"318 } : !transform.any_op319 transform.yield320 }321}322