269 lines · plain
1// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file -allow-unregistered-dialect | FileCheck %s2 3// This file contains some tests of folding/canonicalizing vector.from_elements4 5///===----------------------------------------------===//6/// Tests of `rewriteFromElementsAsSplat`7///===----------------------------------------------===//8 9// CHECK-LABEL: func @extract_scalar_from_from_elements(10// CHECK-SAME: %[[A:.*]]: f32, %[[B:.*]]: f32)11func.func @extract_scalar_from_from_elements(%a: f32, %b: f32) -> (f32, f32, f32, f32, f32, f32, f32) {12 // Extract from 0D.13 %0 = vector.from_elements %a : vector<f32>14 %1 = vector.extract %0[] : f32 from vector<f32>15 16 // Extract from 1D.17 %2 = vector.from_elements %a : vector<1xf32>18 %3 = vector.extract %2[0] : f32 from vector<1xf32>19 %4 = vector.from_elements %a, %b, %a, %a, %b : vector<5xf32>20 %5 = vector.extract %4[4] : f32 from vector<5xf32>21 22 // Extract from 2D.23 %6 = vector.from_elements %a, %a, %a, %b, %b, %b : vector<2x3xf32>24 %7 = vector.extract %6[0, 0] : f32 from vector<2x3xf32>25 %8 = vector.extract %6[0, 1] : f32 from vector<2x3xf32>26 %9 = vector.extract %6[1, 1] : f32 from vector<2x3xf32>27 %10 = vector.extract %6[1, 2] : f32 from vector<2x3xf32>28 29 // CHECK: return %[[A]], %[[A]], %[[B]], %[[A]], %[[A]], %[[B]], %[[B]]30 return %1, %3, %5, %7, %8, %9, %10 : f32, f32, f32, f32, f32, f32, f3231}32 33// -----34 35// CHECK-LABEL: func @extract_1d_from_from_elements(36// CHECK-SAME: %[[A:.*]]: f32, %[[B:.*]]: f32)37func.func @extract_1d_from_from_elements(%a: f32, %b: f32) -> (vector<3xf32>, vector<3xf32>) {38 %0 = vector.from_elements %a, %a, %a, %b, %b, %b : vector<2x3xf32>39 // CHECK: %[[SPLAT1:.*]] = vector.broadcast %[[A]] : f32 to vector<3xf32>40 %1 = vector.extract %0[0] : vector<3xf32> from vector<2x3xf32>41 // CHECK: %[[SPLAT2:.*]] = vector.broadcast %[[B]] : f32 to vector<3xf32>42 %2 = vector.extract %0[1] : vector<3xf32> from vector<2x3xf32>43 // CHECK: return %[[SPLAT1]], %[[SPLAT2]]44 return %1, %2 : vector<3xf32>, vector<3xf32>45}46 47// -----48 49// CHECK-LABEL: func @extract_2d_from_from_elements(50// CHECK-SAME: %[[A:.*]]: f32, %[[B:.*]]: f32)51func.func @extract_2d_from_from_elements(%a: f32, %b: f32) -> (vector<2x2xf32>, vector<2x2xf32>) {52 %0 = vector.from_elements %a, %a, %a, %b, %b, %b, %b, %a, %b, %a, %a, %b : vector<3x2x2xf32>53 // CHECK: %[[SPLAT1:.*]] = vector.from_elements %[[A]], %[[A]], %[[A]], %[[B]] : vector<2x2xf32>54 %1 = vector.extract %0[0] : vector<2x2xf32> from vector<3x2x2xf32>55 // CHECK: %[[SPLAT2:.*]] = vector.from_elements %[[B]], %[[B]], %[[B]], %[[A]] : vector<2x2xf32>56 %2 = vector.extract %0[1] : vector<2x2xf32> from vector<3x2x2xf32>57 // CHECK: return %[[SPLAT1]], %[[SPLAT2]]58 return %1, %2 : vector<2x2xf32>, vector<2x2xf32>59}60 61// -----62 63// CHECK-LABEL: func @from_elements_to_splat(64// CHECK-SAME: %[[A:.*]]: f32, %[[B:.*]]: f32)65func.func @from_elements_to_splat(%a: f32, %b: f32) -> (vector<2x3xf32>, vector<2x3xf32>, vector<f32>) {66 // CHECK: %[[SPLAT:.*]] = vector.broadcast %[[A]] : f32 to vector<2x3xf32>67 %0 = vector.from_elements %a, %a, %a, %a, %a, %a : vector<2x3xf32>68 // CHECK: %[[FROM_EL:.*]] = vector.from_elements {{.*}} : vector<2x3xf32>69 %1 = vector.from_elements %a, %a, %a, %a, %b, %a : vector<2x3xf32>70 // CHECK: %[[SPLAT2:.*]] = vector.broadcast %[[A]] : f32 to vector<f32>71 %2 = vector.from_elements %a : vector<f32>72 // CHECK: return %[[SPLAT]], %[[FROM_EL]], %[[SPLAT2]]73 return %0, %1, %2 : vector<2x3xf32>, vector<2x3xf32>, vector<f32>74}75 76// -----77 78///===----------------------------------------------===//79/// Tests of `FromElementsToShapeCast`80///===----------------------------------------------===//81 82// CHECK-LABEL: func @to_shape_cast_rank2_to_rank1(83// CHECK-SAME: %[[A:.*]]: vector<1x2xi8>)84// CHECK: %[[EXTRACT:.*]] = vector.extract %[[A]][0] : vector<2xi8> from vector<1x2xi8>85// CHECK: return %[[EXTRACT]] : vector<2xi8>86func.func @to_shape_cast_rank2_to_rank1(%arg0: vector<1x2xi8>) -> vector<2xi8> {87 %0 = vector.extract %arg0[0, 0] : i8 from vector<1x2xi8>88 %1 = vector.extract %arg0[0, 1] : i8 from vector<1x2xi8>89 %4 = vector.from_elements %0, %1 : vector<2xi8>90 return %4 : vector<2xi8>91}92 93// -----94 95// CHECK-LABEL: func @to_shape_cast_rank1_to_rank3(96// CHECK-SAME: %[[A:.*]]: vector<8xi8>)97// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[A]] : vector<8xi8> to vector<2x2x2xi8>98// CHECK: return %[[SHAPE_CAST]] : vector<2x2x2xi8>99func.func @to_shape_cast_rank1_to_rank3(%arg0: vector<8xi8>) -> vector<2x2x2xi8> {100 %0 = vector.extract %arg0[0] : i8 from vector<8xi8>101 %1 = vector.extract %arg0[1] : i8 from vector<8xi8>102 %2 = vector.extract %arg0[2] : i8 from vector<8xi8>103 %3 = vector.extract %arg0[3] : i8 from vector<8xi8>104 %4 = vector.extract %arg0[4] : i8 from vector<8xi8>105 %5 = vector.extract %arg0[5] : i8 from vector<8xi8>106 %6 = vector.extract %arg0[6] : i8 from vector<8xi8>107 %7 = vector.extract %arg0[7] : i8 from vector<8xi8>108 %8 = vector.from_elements %0, %1, %2, %3, %4, %5, %6, %7 : vector<2x2x2xi8>109 return %8 : vector<2x2x2xi8>110}111 112// -----113 114// CHECK-LABEL: func @source_larger_than_out(115// CHECK-SAME: %[[A:.*]]: vector<2x3x4xi8>)116// CHECK: %[[EXTRACT:.*]] = vector.extract %[[A]][1] : vector<3x4xi8> from vector<2x3x4xi8>117// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[EXTRACT]] : vector<3x4xi8> to vector<12xi8>118// CHECK: return %[[SHAPE_CAST]] : vector<12xi8>119func.func @source_larger_than_out(%arg0: vector<2x3x4xi8>) -> vector<12xi8> {120 %0 = vector.extract %arg0[1, 0, 0] : i8 from vector<2x3x4xi8>121 %1 = vector.extract %arg0[1, 0, 1] : i8 from vector<2x3x4xi8>122 %2 = vector.extract %arg0[1, 0, 2] : i8 from vector<2x3x4xi8>123 %3 = vector.extract %arg0[1, 0, 3] : i8 from vector<2x3x4xi8>124 %4 = vector.extract %arg0[1, 1, 0] : i8 from vector<2x3x4xi8>125 %5 = vector.extract %arg0[1, 1, 1] : i8 from vector<2x3x4xi8>126 %6 = vector.extract %arg0[1, 1, 2] : i8 from vector<2x3x4xi8>127 %7 = vector.extract %arg0[1, 1, 3] : i8 from vector<2x3x4xi8>128 %8 = vector.extract %arg0[1, 2, 0] : i8 from vector<2x3x4xi8>129 %9 = vector.extract %arg0[1, 2, 1] : i8 from vector<2x3x4xi8>130 %10 = vector.extract %arg0[1, 2, 2] : i8 from vector<2x3x4xi8>131 %11 = vector.extract %arg0[1, 2, 3] : i8 from vector<2x3x4xi8>132 %12 = vector.from_elements %0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11 : vector<12xi8>133 return %12 : vector<12xi8>134}135 136// -----137 138// This test is similar to `source_larger_than_out` except here the number of elements139// extracted contigously starting from the first position [0,0] could be 6 instead of 3140// and the pattern would still match.141// CHECK-LABEL: func @suffix_with_excess_zeros(142// CHECK: %[[EXT:.*]] = vector.extract {{.*}}[0] : vector<3xi8> from vector<2x3xi8>143// CHECK: return %[[EXT]] : vector<3xi8>144func.func @suffix_with_excess_zeros(%arg0: vector<2x3xi8>) -> vector<3xi8> {145 %0 = vector.extract %arg0[0, 0] : i8 from vector<2x3xi8>146 %1 = vector.extract %arg0[0, 1] : i8 from vector<2x3xi8>147 %2 = vector.extract %arg0[0, 2] : i8 from vector<2x3xi8>148 %3 = vector.from_elements %0, %1, %2 : vector<3xi8>149 return %3 : vector<3xi8>150}151 152// -----153 154// CHECK-LABEL: func @large_source_with_shape_cast_required(155// CHECK-SAME: %[[A:.*]]: vector<2x2x2x2xi8>)156// CHECK: %[[EXTRACT:.*]] = vector.extract %[[A]][0, 1] : vector<2x2xi8> from vector<2x2x2x2xi8>157// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[EXTRACT]] : vector<2x2xi8> to vector<1x4x1xi8>158// CHECK: return %[[SHAPE_CAST]] : vector<1x4x1xi8>159func.func @large_source_with_shape_cast_required(%arg0: vector<2x2x2x2xi8>) -> vector<1x4x1xi8> {160 %0 = vector.extract %arg0[0, 1, 0, 0] : i8 from vector<2x2x2x2xi8>161 %1 = vector.extract %arg0[0, 1, 0, 1] : i8 from vector<2x2x2x2xi8>162 %2 = vector.extract %arg0[0, 1, 1, 0] : i8 from vector<2x2x2x2xi8>163 %3 = vector.extract %arg0[0, 1, 1, 1] : i8 from vector<2x2x2x2xi8>164 %4 = vector.from_elements %0, %1, %2, %3 : vector<1x4x1xi8>165 return %4 : vector<1x4x1xi8>166}167 168// -----169 170// Could match, but handled by `rewriteFromElementsAsSplat`.171// CHECK-LABEL: func @extract_single_elm(172// CHECK-NEXT: vector.extract173// CHECK-NEXT: vector.broadcast174// CHECK-NEXT: return175func.func @extract_single_elm(%arg0 : vector<2x3xi8>) -> vector<1xi8> {176 %0 = vector.extract %arg0[0, 0] : i8 from vector<2x3xi8>177 %1 = vector.from_elements %0 : vector<1xi8>178 return %1 : vector<1xi8>179}180 181// -----182 183// CHECK-LABEL: func @negative_source_contiguous_but_not_suffix(184// CHECK-NOT: shape_cast185// CHECK: from_elements186func.func @negative_source_contiguous_but_not_suffix(%arg0: vector<2x3xi8>) -> vector<3xi8> {187 %0 = vector.extract %arg0[0, 1] : i8 from vector<2x3xi8>188 %1 = vector.extract %arg0[0, 2] : i8 from vector<2x3xi8>189 %2 = vector.extract %arg0[1, 0] : i8 from vector<2x3xi8>190 %3 = vector.from_elements %0, %1, %2 : vector<3xi8>191 return %3 : vector<3xi8>192}193 194// -----195 196// The extracted elements are recombined into a single vector, but in a new order.197// CHECK-LABEL: func @negative_nonascending_order(198// CHECK-NOT: shape_cast199// CHECK: from_elements200func.func @negative_nonascending_order(%arg0: vector<1x2xi8>) -> vector<2xi8> {201 %0 = vector.extract %arg0[0, 1] : i8 from vector<1x2xi8>202 %1 = vector.extract %arg0[0, 0] : i8 from vector<1x2xi8>203 %2 = vector.from_elements %0, %1 : vector<2xi8>204 return %2 : vector<2xi8>205}206 207// -----208 209// CHECK-LABEL: func @negative_nonstatic_extract(210// CHECK-NOT: shape_cast211// CHECK: from_elements212func.func @negative_nonstatic_extract(%arg0: vector<1x2xi8>, %i0 : index, %i1 : index) -> vector<2xi8> {213 %0 = vector.extract %arg0[0, %i0] : i8 from vector<1x2xi8>214 %1 = vector.extract %arg0[0, %i1] : i8 from vector<1x2xi8>215 %2 = vector.from_elements %0, %1 : vector<2xi8>216 return %2 : vector<2xi8>217}218 219// -----220 221// CHECK-LABEL: func @negative_different_sources(222// CHECK-NOT: shape_cast223// CHECK: from_elements224func.func @negative_different_sources(%arg0: vector<1x2xi8>, %arg1: vector<1x2xi8>) -> vector<2xi8> {225 %0 = vector.extract %arg0[0, 0] : i8 from vector<1x2xi8>226 %1 = vector.extract %arg1[0, 1] : i8 from vector<1x2xi8>227 %2 = vector.from_elements %0, %1 : vector<2xi8>228 return %2 : vector<2xi8>229}230 231// -----232 233// CHECK-LABEL: func @negative_source_not_suffix(234// CHECK-NOT: shape_cast235// CHECK: from_elements236func.func @negative_source_not_suffix(%arg0: vector<1x3xi8>) -> vector<2xi8> {237 %0 = vector.extract %arg0[0, 0] : i8 from vector<1x3xi8>238 %1 = vector.extract %arg0[0, 1] : i8 from vector<1x3xi8>239 %2 = vector.from_elements %0, %1 : vector<2xi8>240 return %2 : vector<2xi8>241}242 243// -----244 245// The inserted elements are a subset of the extracted elements.246// [0, 1, 2] -> [1, 1, 2]247// CHECK-LABEL: func @negative_nobijection_order(248// CHECK-NOT: shape_cast249// CHECK: from_elements250func.func @negative_nobijection_order(%arg0: vector<1x3xi8>) -> vector<3xi8> {251 %0 = vector.extract %arg0[0, 1] : i8 from vector<1x3xi8>252 %1 = vector.extract %arg0[0, 2] : i8 from vector<1x3xi8>253 %2 = vector.from_elements %0, %0, %1 : vector<3xi8>254 return %2 : vector<3xi8>255}256 257// -----258 259// CHECK-LABEL: func @negative_source_too_small(260// CHECK-NOT: shape_cast261// CHECK: from_elements262func.func @negative_source_too_small(%arg0: vector<2xi8>) -> vector<4xi8> {263 %0 = vector.extract %arg0[0] : i8 from vector<2xi8>264 %1 = vector.extract %arg0[1] : i8 from vector<2xi8>265 %2 = vector.from_elements %0, %1, %1, %1 : vector<4xi8>266 return %2 : vector<4xi8>267}268 269