187 lines · plain
1// RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s2 3// CHECK-LABEL: func @broadcast_vec1d_from_scalar4// CHECK-SAME: %[[A:.*0]]: f325// CHECK: %[[T0:.*]] = vector.broadcast %[[A]] : f32 to vector<2xf32>6// CHECK: return %[[T0]] : vector<2xf32>7 8func.func @broadcast_vec1d_from_scalar(%arg0: f32) -> vector<2xf32> {9 %0 = vector.broadcast %arg0 : f32 to vector<2xf32>10 return %0 : vector<2xf32>11}12 13// CHECK-LABEL: func @broadcast_vec2d_from_scalar14// CHECK-SAME: %[[A:.*0]]: f3215// CHECK: %[[T0:.*]] = vector.broadcast %[[A]] : f32 to vector<2x3xf32>16// CHECK: return %[[T0]] : vector<2x3xf32>17 18func.func @broadcast_vec2d_from_scalar(%arg0: f32) -> vector<2x3xf32> {19 %0 = vector.broadcast %arg0 : f32 to vector<2x3xf32>20 return %0 : vector<2x3xf32>21}22 23// CHECK-LABEL: func @broadcast_vec3d_from_scalar24// CHECK-SAME: %[[A:.*0]]: f3225// CHECK: %[[T0:.*]] = vector.broadcast %[[A]] : f32 to vector<2x3x4xf32>26// CHECK: return %[[T0]] : vector<2x3x4xf32>27 28func.func @broadcast_vec3d_from_scalar(%arg0: f32) -> vector<2x3x4xf32> {29 %0 = vector.broadcast %arg0 : f32 to vector<2x3x4xf32>30 return %0 : vector<2x3x4xf32>31}32 33// CHECK-LABEL: func @broadcast_vec1d_from_vec1d34// CHECK-SAME: %[[A:.*0]]: vector<2xf32>35// CHECK: return %[[A]] : vector<2xf32>36 37func.func @broadcast_vec1d_from_vec1d(%arg0: vector<2xf32>) -> vector<2xf32> {38 %0 = vector.broadcast %arg0 : vector<2xf32> to vector<2xf32>39 return %0 : vector<2xf32>40}41 42// CHECK-LABEL: func @broadcast_vec2d_from_vec1d43// CHECK-SAME: %[[A:.*0]]: vector<2xf32>44// CHECK: %[[U0:.*]] = ub.poison : vector<3x2xf32>45// CHECK: %[[T0:.*]] = vector.insert %[[A]], %[[U0]] [0] : vector<2xf32> into vector<3x2xf32>46// CHECK: %[[T1:.*]] = vector.insert %[[A]], %[[T0]] [1] : vector<2xf32> into vector<3x2xf32>47// CHECK: %[[T2:.*]] = vector.insert %[[A]], %[[T1]] [2] : vector<2xf32> into vector<3x2xf32>48// CHECK: return %[[T2]] : vector<3x2xf32>49 50func.func @broadcast_vec2d_from_vec1d(%arg0: vector<2xf32>) -> vector<3x2xf32> {51 %0 = vector.broadcast %arg0 : vector<2xf32> to vector<3x2xf32>52 return %0 : vector<3x2xf32>53}54 55// CHECK-LABEL: func @broadcast_vec3d_from_vec1d56// CHECK-SAME: %[[A:.*0]]: vector<2xf32>57// CHECK-DAG: %[[U0:.*]] = ub.poison : vector<3x2xf32>58// CHECK-DAG: %[[U1:.*]] = ub.poison : vector<4x3x2xf32>59// CHECK: %[[T0:.*]] = vector.insert %[[A]], %[[U0]] [0] : vector<2xf32> into vector<3x2xf32>60// CHECK: %[[T1:.*]] = vector.insert %[[A]], %[[T0]] [1] : vector<2xf32> into vector<3x2xf32>61// CHECK: %[[T2:.*]] = vector.insert %[[A]], %[[T1]] [2] : vector<2xf32> into vector<3x2xf32>62// CHECK: %[[T3:.*]] = vector.insert %[[T2]], %[[U1]] [0] : vector<3x2xf32> into vector<4x3x2xf32>63// CHECK: %[[T4:.*]] = vector.insert %[[T2]], %[[T3]] [1] : vector<3x2xf32> into vector<4x3x2xf32>64// CHECK: %[[T5:.*]] = vector.insert %[[T2]], %[[T4]] [2] : vector<3x2xf32> into vector<4x3x2xf32>65// CHECK: %[[T6:.*]] = vector.insert %[[T2]], %[[T5]] [3] : vector<3x2xf32> into vector<4x3x2xf32>66// CHECK: return %[[T6]] : vector<4x3x2xf32>67 68func.func @broadcast_vec3d_from_vec1d(%arg0: vector<2xf32>) -> vector<4x3x2xf32> {69 %0 = vector.broadcast %arg0 : vector<2xf32> to vector<4x3x2xf32>70 return %0 : vector<4x3x2xf32>71}72 73// CHECK-LABEL: func @broadcast_vec3d_from_vec2d74// CHECK-SAME: %[[A:.*0]]: vector<3x2xf32>75// CHECK: %[[U0:.*]] = ub.poison : vector<4x3x2xf32>76// CHECK: %[[T0:.*]] = vector.insert %[[A]], %[[U0]] [0] : vector<3x2xf32> into vector<4x3x2xf32>77// CHECK: %[[T1:.*]] = vector.insert %[[A]], %[[T0]] [1] : vector<3x2xf32> into vector<4x3x2xf32>78// CHECK: %[[T2:.*]] = vector.insert %[[A]], %[[T1]] [2] : vector<3x2xf32> into vector<4x3x2xf32>79// CHECK: %[[T3:.*]] = vector.insert %[[A]], %[[T2]] [3] : vector<3x2xf32> into vector<4x3x2xf32>80// CHECK: return %[[T3]] : vector<4x3x2xf32>81 82func.func @broadcast_vec3d_from_vec2d(%arg0: vector<3x2xf32>) -> vector<4x3x2xf32> {83 %0 = vector.broadcast %arg0 : vector<3x2xf32> to vector<4x3x2xf32>84 return %0 : vector<4x3x2xf32>85}86 87// CHECK-LABEL: func @broadcast_stretch88// CHECK-SAME: %[[A:.*0]]: vector<1xf32>89// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : f32 from vector<1xf32>90// CHECK: %[[T1:.*]] = vector.broadcast %[[T0]] : f32 to vector<4xf32>91// CHECK: return %[[T1]] : vector<4xf32>92 93func.func @broadcast_stretch(%arg0: vector<1xf32>) -> vector<4xf32> {94 %0 = vector.broadcast %arg0 : vector<1xf32> to vector<4xf32>95 return %0 : vector<4xf32>96}97 98// CHECK-LABEL: func @broadcast_stretch_at_start99// CHECK-SAME: %[[A:.*0]]: vector<1x4xf32>100// CHECK: %[[U0:.*]] = ub.poison : vector<3x4xf32>101// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : vector<4xf32> from vector<1x4xf32>102// CHECK: %[[T1:.*]] = vector.insert %[[T0]], %[[U0]] [0] : vector<4xf32> into vector<3x4xf32>103// CHECK: %[[T2:.*]] = vector.insert %[[T0]], %[[T1]] [1] : vector<4xf32> into vector<3x4xf32>104// CHECK: %[[T3:.*]] = vector.insert %[[T0]], %[[T2]] [2] : vector<4xf32> into vector<3x4xf32>105// CHECK: return %[[T3]] : vector<3x4xf32>106 107func.func @broadcast_stretch_at_start(%arg0: vector<1x4xf32>) -> vector<3x4xf32> {108 %0 = vector.broadcast %arg0 : vector<1x4xf32> to vector<3x4xf32>109 return %0 : vector<3x4xf32>110}111 112// CHECK-LABEL: func @broadcast_stretch_at_end113// CHECK-SAME: %[[A:.*0]]: vector<4x1xf32>114// CHECK: %[[U0:.*]] = ub.poison : vector<4x3xf32>115// CHECK: %[[T0:.*]] = vector.extract %[[A]][0, 0] : f32 from vector<4x1xf32>116// CHECK: %[[T2:.*]] = vector.broadcast %[[T0]] : f32 to vector<3xf32>117// CHECK: %[[T3:.*]] = vector.insert %[[T2]], %[[U0]] [0] : vector<3xf32> into vector<4x3xf32>118// CHECK: %[[T4:.*]] = vector.extract %[[A]][1, 0] : f32 from vector<4x1xf32>119// CHECK: %[[T6:.*]] = vector.broadcast %[[T4]] : f32 to vector<3xf32>120// CHECK: %[[T7:.*]] = vector.insert %[[T6]], %[[T3]] [1] : vector<3xf32> into vector<4x3xf32>121// CHECK: %[[T8:.*]] = vector.extract %[[A]][2, 0] : f32 from vector<4x1xf32>122// CHECK: %[[T10:.*]] = vector.broadcast %[[T8]] : f32 to vector<3xf32>123// CHECK: %[[T11:.*]] = vector.insert %[[T10]], %[[T7]] [2] : vector<3xf32> into vector<4x3xf32>124// CHECK: %[[T12:.*]] = vector.extract %[[A]][3, 0] : f32 from vector<4x1xf32>125// CHECK: %[[T14:.*]] = vector.broadcast %[[T12]] : f32 to vector<3xf32>126// CHECK: %[[T15:.*]] = vector.insert %[[T14]], %[[T11]] [3] : vector<3xf32> into vector<4x3xf32>127// CHECK: return %[[T15]] : vector<4x3xf32>128 129func.func @broadcast_stretch_at_end(%arg0: vector<4x1xf32>) -> vector<4x3xf32> {130 %0 = vector.broadcast %arg0 : vector<4x1xf32> to vector<4x3xf32>131 return %0 : vector<4x3xf32>132}133 134// CHECK-LABEL: func @broadcast_stretch_in_middle135// CHECK-SAME: %[[A:.*0]]: vector<4x1x2xf32>136// CHECK: %[[U0:.*]] = ub.poison : vector<4x3x2xf32>137// CHECK: %[[U1:.*]] = ub.poison : vector<3x2xf32>138// CHECK: %[[T0:.*]] = vector.extract %[[A]][0, 0] : vector<2xf32> from vector<4x1x2xf32>139// CHECK: %[[T2:.*]] = vector.insert %[[T0]], %[[U1]] [0] : vector<2xf32> into vector<3x2xf32>140// CHECK: %[[T3:.*]] = vector.insert %[[T0]], %[[T2]] [1] : vector<2xf32> into vector<3x2xf32>141// CHECK: %[[T4:.*]] = vector.insert %[[T0]], %[[T3]] [2] : vector<2xf32> into vector<3x2xf32>142// CHECK: %[[T5:.*]] = vector.insert %[[T4]], %[[U0]] [0] : vector<3x2xf32> into vector<4x3x2xf32>143// CHECK: %[[T6:.*]] = vector.extract %[[A]][1, 0] : vector<2xf32> from vector<4x1x2xf32>144// CHECK: %[[T8:.*]] = vector.insert %[[T6]], %[[U1]] [0] : vector<2xf32> into vector<3x2xf32>145// CHECK: %[[T9:.*]] = vector.insert %[[T6]], %[[T8]] [1] : vector<2xf32> into vector<3x2xf32>146// CHECK: %[[T10:.*]] = vector.insert %[[T6]], %[[T9]] [2] : vector<2xf32> into vector<3x2xf32>147// CHECK: %[[T11:.*]] = vector.insert %[[T10]], %[[T5]] [1] : vector<3x2xf32> into vector<4x3x2xf32>148// CHECK: %[[T12:.*]] = vector.extract %[[A]][2, 0] : vector<2xf32> from vector<4x1x2xf32>149// CHECK: %[[T14:.*]] = vector.insert %[[T12]], %[[U1]] [0] : vector<2xf32> into vector<3x2xf32>150// CHECK: %[[T15:.*]] = vector.insert %[[T12]], %[[T14]] [1] : vector<2xf32> into vector<3x2xf32>151// CHECK: %[[T16:.*]] = vector.insert %[[T12]], %[[T15]] [2] : vector<2xf32> into vector<3x2xf32>152// CHECK: %[[T17:.*]] = vector.insert %[[T16]], %[[T11]] [2] : vector<3x2xf32> into vector<4x3x2xf32>153// CHECK: %[[T18:.*]] = vector.extract %[[A]][3, 0] : vector<2xf32> from vector<4x1x2xf32>154// CHECK: %[[T20:.*]] = vector.insert %[[T18]], %[[U1]] [0] : vector<2xf32> into vector<3x2xf32>155// CHECK: %[[T21:.*]] = vector.insert %[[T18]], %[[T20]] [1] : vector<2xf32> into vector<3x2xf32>156// CHECK: %[[T22:.*]] = vector.insert %[[T18]], %[[T21]] [2] : vector<2xf32> into vector<3x2xf32>157// CHECK: %[[T23:.*]] = vector.insert %[[T22]], %[[T17]] [3] : vector<3x2xf32> into vector<4x3x2xf32>158// CHECK: return %[[T23]] : vector<4x3x2xf32>159 160func.func @broadcast_stretch_in_middle(%arg0: vector<4x1x2xf32>) -> vector<4x3x2xf32> {161 %0 = vector.broadcast %arg0 : vector<4x1x2xf32> to vector<4x3x2xf32>162 return %0 : vector<4x3x2xf32>163}164 165// CHECK-LABEL: func.func @broadcast_scalable_duplication166// CHECK-SAME: %[[ARG0:.*]]: vector<[32]xf32>)167// CHECK: %[[INIT:.*]] = ub.poison : vector<1x[32]xf32>168// CHECK: %[[RES:.*]] = vector.insert %[[ARG0]], %[[INIT]] [0] : vector<[32]xf32> into vector<1x[32]xf32>169// CHECK: return %[[RES]] : vector<1x[32]xf32>170 171func.func @broadcast_scalable_duplication(%arg0: vector<[32]xf32>) -> vector<1x[32]xf32> {172 %res = vector.broadcast %arg0 : vector<[32]xf32> to vector<1x[32]xf32>173 return %res : vector<1x[32]xf32>174}175 176module attributes {transform.with_named_sequence} {177 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {178 %f = transform.structured.match ops{["func.func"]} in %module_op179 : (!transform.any_op) -> !transform.any_op180 181 transform.apply_patterns to %f {182 transform.apply_patterns.vector.lower_broadcast183 } : !transform.any_op184 transform.yield185 }186}187