Lines Matching full:shape
7 func.func @shape_num_elements(%shape : !shape.shape) -> !shape.size {
8 %init = shape.const_size 1
9 %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
10 ^bb0(%index : index, %extent : !shape.size, %acc : !shape.size):
11 %acc_next = shape.mul %acc, %extent
12 : !shape.size, !shape.size -> !shape.size
13 shape.yield %acc_next : !shape.size
15 return %num_elements : !shape.size
19 func.func @extent_tensor_num_elements(%shape : tensor<?xindex>) -> index {
21 %num_elements = shape.reduce(%shape, %init) : tensor<?xindex> -> index {
23 %acc_next = shape.mul %acc, %extent : index, index -> index
24 shape.yield %acc_next : index
30 %0 = "shape.unknown_shape"() : () -> !shape.shape
31 %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
32 %2 = "shape.print"(%1) : (!shape.size) -> !shape.size
37 %0 = shape.const_shape [1, 2, 3] : !shape.shape
38 %2 = shape.const_shape [4, 5, 6] : tensor<3xindex>
43 %0 = shape.const_shape [1, 57, 92] : !shape.shape
44 %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
45 %3 = "shape.print"(%1) : (!shape.size) -> !shape.size
50 %0 = shape.const_shape [10, 1, 57, 92] : !shape.shape
51 %1 = shape.const_shape [4, 57, 92] : !shape.shape
52 %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
53 %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
58 %0 = shape.const_shape [10, 1, 57, 92] : tensor<4xindex>
59 %1 = shape.const_shape [4, 57, 92] : tensor<3xindex>
60 %2 = shape.broadcast %0, %1 : tensor<4xindex>, tensor<3xindex> -> tensor<4xindex>
65 %0 = shape.const_shape [4, 57, 92] : !shape.shape
66 %1 = shape.const_shape [4, 57, 92] : !shape.shape
67 %2 = "shape.meet"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
68 %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
73 %0 = shape.const_shape [4, -1, 92] : !shape.shape
74 %1 = shape.const_shape [-1, 57, 92] : !shape.shape
75 %2 = "shape.meet"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
76 %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
81 %0 = shape.const_shape [4, 57, 92] : !shape.shape
82 %1 = shape.const_shape [2, 57, 92] : !shape.shape
83 %2 = "shape.meet"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
84 %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
89 %0 = shape.const_shape [] : !shape.shape
90 %1 = shape.const_shape [1, 2, 3] : !shape.shape
91 %2 = shape.const_shape [1, 2, 3] : tensor<3xindex>
96 %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
100 func.func @test_value_of(%arg0: !shape.value_shape) -> tensor<?xf32> {
101 %0 = shape.value_of %arg0 : tensor<?xf32>
106 %0 = shape.const_shape [] : !shape.shape
107 %1 = shape.const_shape [1, 2, 3] : !shape.shape
109 %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
110 %w1 = shape.cstr_eq %0, %1 : !shape.shape, !shape.shape
111 %w2 = shape.const_witness true
112 %w3 = shape.const_witness false
113 %w4 = shape.cstr_require %true, "msg"
114 %w_all = shape.assuming_all %w0, %w1, %w2, %w3, %w4
115 shape.assuming %w_all -> !shape.shape {
116 %2 = "shape.any"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
117 shape.assuming_yield %2 : !shape.shape
124 %w0 = shape.cstr_eq %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
130 %w0 = shape.cstr_broadcastable %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
134 func.func @mul(%size_arg : !shape.size, %index_arg : index) {
135 %size_prod = shape.mul %size_arg, %size_arg
136 : !shape.size, !shape.size -> !shape.size
137 %index_prod = shape.mul %index_arg, %index_arg : index, index -> index
138 %mixed_prod = shape.mul %size_arg, %index_arg
139 : !shape.size, index -> !shape.size
143 func.func @div(%size_arg : !shape.size, %index_arg : index) {
144 %size_div = shape.div %size_arg, %size_arg
145 : !shape.size, !shape.size -> !shape.size
146 %index_div = shape.div %index_arg, %index_arg : index, index -> index
147 %mixed_div = shape.div %size_arg, %index_arg
148 : !shape.size, index -> !shape.size
152 func.func @add(%size_arg : !shape.size, %index_arg : index) {
153 %size_sum = shape.add %size_arg, %size_arg
154 : !shape.size, !shape.size -> !shape.size
155 %index_sum = shape.add %index_arg, %index_arg : index, index -> index
156 %mixed_sum = shape.add %size_arg, %index_arg
157 : !shape.size, index -> !shape.size
162 // CHECK: %c1 = shape.const_size 1
163 // CHECK: %c2 = shape.const_size 2
164 // CHECK: %c2_0 = shape.const_size 2
165 %0 = shape.const_size 1
166 %1 = shape.const_size 2
167 %2 = shape.const_size 2
171 func.func @test_to_extent_tensor(%arg: !shape.shape) -> tensor<3xindex> {
172 %0 = shape.to_extent_tensor %arg : !shape.shape -> tensor<3xindex>
177 %0 = shape.to_extent_tensor %arg : tensor<3xindex> -> tensor<3xindex>
181 func.func @test_from_extent_tensor(%arg: tensor<?xindex>) -> !shape.shape {
182 %0 = shape.from_extent_tensor %arg : tensor<?xindex>
183 return %0 : !shape.shape
186 func.func @rank(%shape : !shape.shape) -> !shape.size {
187 %rank = shape.rank %shape : !shape.shape -> !shape.size
188 return %rank : !shape.size
191 func.func @rank_on_extent_tensor(%shape : tensor<?xindex>) -> index {
192 %rank = shape.rank %shape : tensor<?xindex> -> index
196 func.func @shape_eq_on_shapes(%a : !shape.shape, %b : !shape.shape) -> i1 {
197 %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
202 %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
206 func.func @shape_eq_on_mixed(%a : tensor<?xindex>, %b : !shape.shape) -> i1 {
207 %result = shape.shape_eq %a, %b : tensor<?xindex>, !shape.shape
211 func.func @get_extent_on_shape(%arg : !shape.shape) -> !shape.size {
212 %c0 = shape.const_size 0
213 %result = shape.get_extent %arg, %c0 :
214 !shape.shape, !shape.size -> !shape.size
215 return %result : !shape.size
220 %result = shape.get_extent %arg, %c0 : tensor<?xindex>, index -> index
226 %result = shape.dim %arg, %c0 : memref<?x?xindex>, index -> index
230 func.func @get_extent_on_mixed_operands(%arg : tensor<?xindex>) -> !shape.size {
231 %c0 = shape.const_size 0
232 %result = shape.get_extent %arg, %c0 : tensor<?xindex>, !shape.size -> !shape.size
233 return %result : !shape.size
237 %0 = shape.const_shape [1, 2, 3] : !shape.shape
238 %1 = shape.const_shape [4, 5, 6] : !shape.shape
239 %2 = "shape.any"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
240 %3 = shape.const_shape [1, 2, 3] : tensor<3xindex>
241 %4 = shape.const_shape [4, 5, 6] : tensor<3xindex>
242 %5 = "shape.any"(%3, %4) : (tensor<3xindex>, tensor<3xindex>) -> tensor<3xindex>
247 %result = shape.num_elements %arg : tensor<?xindex> -> index
251 func.func @num_elements_shape(%arg : !shape.shape) -> !shape.size {
252 %result = shape.num_elements %arg : !shape.shape -> !shape.size
253 return %result : !shape.size
256 // Testing invoking shape function from another. shape_equal_shapes is merely
258 func.func @shape_equal_shapes(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
259 %0 = shape.shape_of %a : !shape.value_shape -> !shape.shape
260 %1 = shape.shape_of %b : !shape.value_shape -> !shape.shape
261 %2 = "shape.meet"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
262 return %2 : !shape.shape
264 func.func @shape_with_shape(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
265 %0 = shape.shape_of %a : !shape.value_shape -> !shape.shape
266 %1 = shape.with_shape %b, %0 : !shape.value_shape, !shape.shape
267 %2 = call @shape_equal_shapes(%a, %1) : (!shape.value_shape, !shape.value_shape) -> !shape.shape
268 return %2 : !shape.shape
271 …e_with_shape_extent_tensor_type(%a : tensor<?x?x?xf32>, %b : !shape.value_shape) -> !shape.value_s…
272 %0 = shape.shape_of %a : tensor<?x?x?xf32> -> tensor<3xindex>
273 %1 = shape.with_shape %b, %0 : !shape.value_shape, tensor<3xindex>
274 return %1 : !shape.value_shape
277 func.func @any_on_shape(%a : !shape.shape, %b : !shape.shape, %c : !shape.shape)
278 -> !shape.shape {
279 %result = shape.any %a, %b, %c
280 : !shape.shape, !shape.shape, !shape.shape -> !shape.shape
281 return %result : !shape.shape
286 %c : !shape.shape) -> !shape.shape {
287 %result = shape.any %a, %b, %c
288 : tensor<?xindex>, tensor<?xindex>, !shape.shape -> !shape.shape
289 return %result : !shape.shape
295 %result = shape.any %a, %b, %c
302 %result = shape.is_broadcastable %a, %b
307 func.func @is_broadcastable_on_shapes(%a : !shape.shape,
308 %b : !shape.shape) -> i1 {
309 %result = shape.is_broadcastable %a, %b
310 : !shape.shape, !shape.shape
314 func.func @shape_upper_bounded_by_constant(%a: !shape.shape) -> !shape.shape {
315 %0 = shape.const_shape [4, 57, 92] : !shape.shape
316 %1 = shape.max %a, %0 : !shape.shape, !shape.shape -> !shape.shape
317 %2 = shape.meet %0, %1, error="exceeded element-wise upper bound" :
318 !shape.shape, !shape.shape -> !shape.shape
319 return %2 : !shape.shape
322 func.func @shape_lower_bounded_by_constant(%a: !shape.shape) -> !shape.shape {
323 %0 = shape.const_shape [4, 57, 92] : !shape.shape
324 %1 = shape.min %a, %0 : !shape.shape, !shape.shape -> !shape.shape
325 %2 = shape.meet %0, %1, error="lower bound element-wise exceeded" :
326 !shape.shape, !shape.shape -> !shape.shape
327 return %2 : !shape.shape
330 func.func @size_upper_bounded_by_constant(%a: !shape.size) -> !shape.size {
331 %0 = shape.const_size 5
332 %1 = shape.max %a, %0 : !shape.size, !shape.size -> !shape.size
333 %2 = shape.meet %0, %1, error="exceeded element-wise upper bound" :
334 !shape.size, !shape.size -> !shape.size
335 return %2 : !shape.size
338 func.func @size_lower_bounded_by_constant(%a: !shape.size) -> !shape.size {
339 %0 = shape.const_size 9
340 %1 = shape.min %a, %0 : !shape.size, !shape.size -> !shape.size
341 %2 = shape.meet %0, %1, error="lower bound element-wise exceeded" :
342 !shape.size, !shape.size -> !shape.size
343 return %2 : !shape.size
347 %result = shape.meet %arg0, %arg1 : index, index -> index