xref: /llvm-project/mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir (revision bb6d5c220004a5d7e466a669324001285a688918)
1// RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s
2
3// CHECK-LABEL: func @broadcast_vec1d_from_scalar
4// CHECK-SAME: %[[A:.*0]]: f32
5// CHECK:      %[[T0:.*]] = vector.splat %[[A]] : 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_scalar
14// CHECK-SAME: %[[A:.*0]]: f32
15// CHECK:      %[[T0:.*]] = vector.splat %[[A]] : 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_scalar
24// CHECK-SAME: %[[A:.*0]]: f32
25// CHECK:      %[[T0:.*]] = vector.splat %[[A]] : 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_vec1d
34// 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_vec1d
43// CHECK-SAME: %[[A:.*0]]: vector<2xf32>
44// CHECK:      %[[C0:.*]] = arith.constant dense<0.000000e+00> : vector<3x2xf32>
45// CHECK:      %[[T0:.*]] = vector.insert %[[A]], %[[C0]] [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_vec1d
56// CHECK-SAME: %[[A:.*0]]: vector<2xf32>
57// CHECK-DAG:  %[[C0:.*]] = arith.constant dense<0.000000e+00> : vector<3x2xf32>
58// CHECK-DAG:  %[[C1:.*]] = arith.constant dense<0.000000e+00> : vector<4x3x2xf32>
59// CHECK:      %[[T0:.*]] = vector.insert %[[A]], %[[C0]] [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]], %[[C1]] [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_vec2d
74// CHECK-SAME: %[[A:.*0]]: vector<3x2xf32>
75// CHECK:      %[[C0:.*]] = arith.constant dense<0.000000e+00> : vector<4x3x2xf32>
76// CHECK:      %[[T0:.*]] = vector.insert %[[A]], %[[C0]] [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_stretch
88// CHECK-SAME: %[[A:.*0]]: vector<1xf32>
89// CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : f32 from vector<1xf32>
90// CHECK:      %[[T1:.*]] = vector.splat %[[T0]] : 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_start
99// CHECK-SAME: %[[A:.*0]]: vector<1x4xf32>
100// CHECK:      %[[C0:.*]] = arith.constant dense<0.000000e+00> : vector<3x4xf32>
101// CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : vector<4xf32> from vector<1x4xf32>
102// CHECK:      %[[T1:.*]] = vector.insert %[[T0]], %[[C0]] [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_end
113// CHECK-SAME: %[[A:.*0]]: vector<4x1xf32>
114// CHECK:      %[[C0:.*]] = arith.constant dense<0.000000e+00> : vector<4x3xf32>
115// CHECK:      %[[T0:.*]] = vector.extract %[[A]][0, 0] : f32 from vector<4x1xf32>
116// CHECK:      %[[T2:.*]] = vector.splat %[[T0]] : vector<3xf32>
117// CHECK:      %[[T3:.*]] = vector.insert %[[T2]], %[[C0]] [0] : vector<3xf32> into vector<4x3xf32>
118// CHECK:      %[[T4:.*]] = vector.extract %[[A]][1, 0] : f32 from vector<4x1xf32>
119// CHECK:      %[[T6:.*]] = vector.splat %[[T4]] : 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.splat %[[T8]] : 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.splat %[[T12]] : 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_middle
135// CHECK-SAME: %[[A:.*0]]: vector<4x1x2xf32>
136// CHECK:      %[[C0:.*]] = arith.constant dense<0.000000e+00> : vector<4x3x2xf32>
137// CHECK:      %[[C1:.*]] = arith.constant dense<0.000000e+00> : vector<3x2xf32>
138// CHECK:      %[[T0:.*]] = vector.extract %[[A]][0, 0] : vector<2xf32> from vector<4x1x2xf32>
139// CHECK:      %[[T2:.*]] = vector.insert %[[T0]], %[[C1]] [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]], %[[C0]] [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]], %[[C1]] [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]], %[[C1]] [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]], %[[C1]] [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_duplication
166// CHECK-SAME:      %[[ARG0:.*]]: vector<[32]xf32>)
167// CHECK:           %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<1x[32]xf32>
168// CHECK:           %[[RES:.*]] = vector.insert %[[ARG0]], %[[CST]] [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_op
179      : (!transform.any_op) -> !transform.any_op
180
181    transform.apply_patterns to %f {
182      transform.apply_patterns.vector.lower_broadcast
183    } : !transform.any_op
184    transform.yield
185  }
186}
187