Lines Matching full:shape
5 // CHECK: shape.const_shape [2, 3, 4] : tensor<?xindex>
6 %0 = shape.shape_of %arg0 : tensor<2x3x4xf32> -> tensor<?xindex>
14 func @f() -> (!shape.shape, !shape.shape) {
15 // CHECK: shape.const_shape [2, 3] : !shape.shape
16 // CHECK: shape.const_shape [4, 5] : !shape.shape
18 %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape
19 %head, %tail = "shape.split_at"(%0, %c2) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
20 return %head, %tail : !shape.shape, !shape.shape
28 func @f() -> (!shape.shape, !shape.shape) {
29 // CHECK: shape.const_shape [2, 3, 4] : !shape.shape
30 // CHECK: shape.const_shape [5] : !shape.shape
32 %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape
33 %head, %tail = "shape.split_at"(%0, %c-1) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
34 return %head, %tail : !shape.shape, !shape.shape
41 func @f() -> (!shape.shape, !shape.shape) {
42 // CHECK: shape.split_at
44 %0 = shape.const_shape [2, 3, 4, 5] : !shape.shape
45 %head, %tail = "shape.split_at"(%0, %c5) : (!shape.shape, i32) -> (!shape.shape, !shape.shape)
46 return %head, %tail : !shape.shape, !shape.shape
53 func @f() -> !shape.shape {
54 // CHECK: shape.const_shape [7, 2] : !shape.shape
55 %0 = shape.const_shape [1, 2] : !shape.shape
56 %1 = shape.const_shape [7, 1] : !shape.shape
57 %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
58 return %2 : !shape.shape
66 // CHECK: shape.const_shape [7, 2] : tensor<?xindex>
67 %0 = shape.const_shape [1, 2] : tensor<?xindex>
68 %1 = shape.const_shape [7, 1] : tensor<?xindex>
69 %2 = shape.broadcast %0, %1
78 func @broadcast() -> !shape.shape {
79 // CHECK: shape.const_shape [7, 2] : !shape.shape
80 %0 = shape.const_shape [1, 2] : tensor<?xindex>
81 %1 = shape.const_shape [7, 1] : tensor<?xindex>
82 %2 = shape.broadcast %0, %1 : tensor<?xindex>, tensor<?xindex> -> !shape.shape
83 return %2 : !shape.shape
90 func @f(%arg0 : !shape.shape) -> !shape.shape {
92 %0 = shape.const_shape [] : !shape.shape
93 %1 = shape.broadcast %arg0, %0 : !shape.shape, !shape.shape -> !shape.shape
94 return %1 : !shape.shape
101 func @f(%arg0 : !shape.shape) -> !shape.shape {
103 %0 = shape.const_shape [] : !shape.shape
104 %1 = shape.broadcast %0, %arg0 : !shape.shape, !shape.shape -> !shape.shape
105 return %1 : !shape.shape
112 func @f() -> !shape.shape {
113 // CHECK: %[[CST:.*]] = shape.const_shape [1, 2, 3] : !shape.shape
115 %0 = shape.const_shape [] : !shape.shape
116 %1 = shape.const_shape [1, 2, 3] : !shape.shape
117 %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
118 return %2 : !shape.shape
125 func @f() -> !shape.shape {
126 // CHECK: shape.broadcast
127 %0 = shape.const_shape [2] : !shape.shape
128 %1 = shape.const_shape [7] : !shape.shape
129 %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
130 return %2 : !shape.shape
137 func @f() -> !shape.shape {
138 // CHECK: shape.const_shape [0, 1, 2, 3] : !shape.shape
139 %lhs = shape.const_shape [0, 1] : !shape.shape
140 %rhs = shape.const_shape [2, 3] : !shape.shape
141 %0 = shape.concat %lhs, %rhs
142 return %0 : !shape.shape
151 %cs = shape.const_shape [0, 1] : !shape.shape
152 %0 = shape.to_extent_tensor %cs : !shape.shape -> tensor<2xindex>
160 func @f() -> !shape.shape {
161 // CHECK: shape.const_shape [3, 5, 11] : !shape.shape
165 %ret = shape.from_extents %e0, %e1, %e2
166 return %ret : !shape.shape
172 func @no_fold(%arg0: index) -> !shape.shape {
173 // CHECK-NOT: shape.const_shape
175 %ret = shape.from_extents %e0, %arg0
176 return %ret : !shape.shape
184 // CHECK-NOT: shape.index_cast
185 %cs = shape.const_size 123
187 %ci = shape.size_to_index %cs : !shape.size
195 func @const_index_to_size() -> !shape.size {
198 // CHECK: shape.const_size 123
199 %cs = shape.index_to_size %ci
200 return %cs : !shape.size
208 // CHECK-NOT: shape.index_cast
210 %cs0 = shape.index_to_size %ci0
213 %ci1 = shape.size_to_index %cs0 : !shape.size
221 func @nonfoldable_size_to_index(%cs : !shape.size) -> index {
222 // CHECK: shape.size_to_index
223 %ci = shape.size_to_index %cs : !shape.size
231 func @nonfoldable_index_to_size(%ci : index) -> !shape.size {
232 // CHECK: shape.index_to_size
233 %cs = shape.index_to_size %ci
234 return %cs : !shape.size
241 func @num_elements() -> !shape.size {
242 // CHECK-NOT: shape.const_shape
243 %shape = shape.const_shape [4, 5, 6] : !shape.shape
244 // CHECK-NOT: shape.num_elements
245 %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size
246 // CHECK: %[[NUM:.*]] = shape.const_size 120
247 // CHECK-NEXT: return %[[NUM]] : !shape.size
248 return %num_elements : !shape.size
255 func @nonfoldable_num_elements(%shape : !shape.shape) -> !shape.size {
256 // CHECK-NOT: shape.const_{{.*}}
257 %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size
258 return %num_elements : !shape.size
267 %0 = shape.const_shape [0, 1, 2] : tensor<?xindex>
269 %1 = shape.get_extent %0, %c2 : tensor<?xindex>, index -> index
278 // CHECK: shape.const_shape
279 // CHECK: shape.get_extent
280 %0 = shape.const_shape [0, 1, 2] : tensor<?xindex>
282 %1 = shape.get_extent %0, %c3 : tensor<?xindex>, index -> index
291 // CHECK: shape.get_extent
293 %0 = shape.get_extent %arg0, %c3 : tensor<?xindex>, index -> index
301 func @basic() -> !shape.size {
302 // CHECK: shape.const_size 2
303 %0 = shape.const_shape [0, 1, 2] : !shape.shape
304 %c2 = shape.const_size 2
305 %1 = shape.get_extent %0, %c2 : !shape.shape, !shape.size -> !shape.size
306 return %1 : !shape.size
313 func @out_of_bounds() -> !shape.size {
314 // CHECK: shape.const_shape
315 // CHECK: shape.get_extent
316 %0 = shape.const_shape [0, 1, 2] : !shape.shape
317 %c3 = shape.const_size 3
318 %1 = shape.get_extent %0, %c3 : !shape.shape, !shape.size -> !shape.size
319 return %1 : !shape.size
326 func @not_const(%arg0 : !shape.shape) -> !shape.size {
327 // CHECK: shape.get_extent
328 %c3 = shape.const_size 3
329 %0 = shape.get_extent %arg0, %c3 : !shape.shape, !shape.size -> !shape.size
330 return %0 : !shape.size
336 func @f(%arg0 : !shape.shape) {
337 // CHECK-NEXT: shape.const_witness true
340 %0 = shape.cstr_eq %arg0, %arg0, %arg0
341 "consume.witness"(%0) : (!shape.witness) -> ()
349 // CHECK-NEXT: shape.const_witness true
352 %cs0 = shape.const_shape [0, 1] : !shape.shape
353 %cs1 = shape.const_shape [0, 1] : !shape.shape
354 %cs2 = shape.const_shape [0, 1] : !shape.shape
355 %0 = shape.cstr_eq %cs0, %cs1, %cs2
356 "consume.witness"(%0) : (!shape.witness) -> ()
364 // CHECK-NEXT: shape.const_shape
365 // CHECK-NEXT: shape.const_shape
366 // CHECK-NEXT: shape.cstr_eq
369 %cs0 = shape.const_shape [0, 1] : !shape.shape
370 %cs1 = shape.const_shape [3, 1] : !shape.shape
371 %0 = shape.cstr_eq %cs0, %cs1
372 "consume.witness"(%0) : (!shape.witness) -> ()
379 func @f(%arg0: !shape.shape, %arg1: !shape.shape) {
380 // CHECK-NEXT: shape.cstr_eq
383 %0 = shape.cstr_eq %arg0, %arg1
384 "consume.witness"(%0) : (!shape.witness) -> ()
392 // CHECK-NEXT: shape.const_witness true
396 %0 = shape.cstr_require %true, "msg"
397 "consume.witness"(%0) : (!shape.witness) -> ()
405 // CHECK-NEXT: shape.cstr_require
408 %0 = shape.cstr_require %arg0, "msg"
409 "consume.witness"(%0) : (!shape.witness) -> ()
417 // CHECK-NEXT: shape.const_witness true
420 %0 = shape.const_witness true
421 %1 = shape.const_witness true
422 %2 = shape.const_witness true
423 %3 = shape.assuming_all %0, %1, %2
424 "consume.witness"(%3) : (!shape.witness) -> ()
439 // CHECK-NEXT: shape.assuming_all %[[UNKNOWN1]], %[[UNKNOWN2]]
442 %0 = shape.const_witness true
443 %1 = "test.source"() : () -> !shape.witness
444 %2 = "test.source"() : () -> !shape.witness
445 %3 = shape.assuming_all %0, %1, %2
446 "consume.witness"(%3) : (!shape.witness) -> ()
454 func @f(%arg : !shape.shape) -> !shape.shape {
455 // CHECK-NEXT: %[[CS:.*]] = shape.const_shape
457 %0 = shape.const_shape [2, 3, 4] : !shape.shape
458 %1 = shape.any %0, %arg : !shape.shape, !shape.shape -> !shape.shape
459 return %1 : !shape.shape
467 // CHECK-NEXT: %[[CS:.*]] = shape.const_shape [2, 3, 4] : tensor<?xindex>
469 %0 = shape.const_shape [2, 3, 4] : tensor<?xindex>
470 %1 = shape.any %0, %arg : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
478 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape {
479 // CHECK-NEXT: %[[CS:.*]] = shape.any
481 %1 = shape.any %arg0, %arg1 : !shape.shape, !shape.shape -> !shape.shape
482 return %1 : !shape.shape
493 %0 = shape.const_witness true
494 %1 = shape.assuming %0 -> index {
496 shape.assuming_yield %2 : index
508 // CHECK-NEXT: shape.assuming
510 // CHECK-NEXT: shape.assuming_yield
514 %0 = "test.source"() : () -> (!shape.witness)
515 %1 = shape.assuming %0 -> index {
517 shape.assuming_yield %2 : index
527 // CHECK-NEXT: shape.const_witness true
530 %cs0 = shape.const_shape [3, 1] : !shape.shape
531 %cs1 = shape.const_shape [1, 5] : !shape.shape
532 %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
533 "consume.witness"(%0) : (!shape.witness) -> ()
541 // CHECK-NEXT: shape.const_shape
542 // CHECK-NEXT: shape.const_shape
543 // CHECK-NEXT: shape.cstr_broadcastable
546 %cs0 = shape.const_shape [1, 3] : !shape.shape
547 %cs1 = shape.const_shape [1, 5] : !shape.shape
548 %0 = shape.cstr_broadcastable %cs0, %cs1 : !shape.shape, !shape.shape
549 "consume.witness"(%0) : (!shape.witness) -> ()
556 func @f(%arg0 : !shape.shape) {
557 // CHECK-NEXT: shape.const_shape
558 // CHECK-NEXT: shape.cstr_broadcastable
561 %cs0 = shape.const_shape [1, 3] : !shape.shape
562 %0 = shape.cstr_broadcastable %arg0, %cs0 : !shape.shape, !shape.shape
563 "consume.witness"(%0) : (!shape.witness) -> ()
570 func @f(%arg0 : !shape.shape) {
571 // CHECK-NEXT: shape.const_witness true
574 %0 = shape.cstr_broadcastable %arg0, %arg0 : !shape.shape, !shape.shape
575 "consume.witness"(%0) : (!shape.witness) -> ()
584 // CHECK-NEXT: shape.const_witness true
587 %0 = shape.cstr_broadcastable %arg, %arg : tensor<?xindex>, tensor<?xindex>
588 "consume.witness"(%0) : (!shape.witness) -> ()
594 // Fold `rank` based on constant shape.
596 func @fold_rank() -> !shape.size {
597 // CHECK: %[[RESULT:.*]] = shape.const_size 5
598 // CHECK: return %[[RESULT]] : !shape.size
599 %shape = shape.const_shape [3, 4, 5, 6, 7] : !shape.shape
600 %rank = shape.rank %shape : !shape.shape -> !shape.size
601 return %rank : !shape.size
606 // Do not fold `rank` if shape is dynamic.
608 // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) -> !shape.size
609 func @dont_fold_rank(%shape : !shape.shape) -> !shape.size {
610 // CHECK: %[[RESULT:.*]] = shape.rank %[[SHAPE]]
611 // CHECK: return %[[RESULT]] : !shape.size
612 %rank = shape.rank %shape : !shape.shape -> !shape.size
613 return %rank : !shape.size
623 %shape = shape.const_shape [3, 4, 5, 6, 7] : tensor<?xindex>
624 %rank = shape.rank %shape : tensor<?xindex> -> index
632 // CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index
633 func @dont_fold_rank(%shape : tensor<?xindex>) -> index {
634 // CHECK: %[[RESULT:.*]] = shape.rank %[[SHAPE]] : tensor<?xindex> -> index
636 %rank = shape.rank %shape : tensor<?xindex> -> index
642 // Canonicalize `rank` when shape is derived from ranked tensor.
647 %shape = shape.shape_of %arg : tensor<1x2x?xf32> -> tensor<?xindex>
648 %rank = shape.rank %shape : tensor<?xindex> -> index
654 // Canonicalize `rank` when shape is derived from ranked tensor.
656 func @canonicalize_rank_size(%arg : tensor<1x2x?xf32>) -> !shape.size {
657 // CHECK: %[[RESULT:.*]] = shape.const_size 3
658 // CHECK: return %[[RESULT]] : !shape.size
659 %shape = shape.shape_of %arg : tensor<1x2x?xf32> -> !shape.shape
660 %rank = shape.rank %shape : !shape.shape -> !shape.size
661 return %rank : !shape.size
666 // Do not canonicalize `rank` when shape is derived from unranked tensor.
670 // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32> -> tensor<?xindex>
671 // CHECK: %[[SIZE:.*]] = shape.rank %[[SHAPE]]
673 %shape = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>
674 %rank = shape.rank %shape : tensor<?xindex> -> index
685 %size = shape.index_to_size %index
686 %result = shape.size_to_index %size : !shape.size
694 // CHECK-SAME: (%[[SIZE:.*]]: !shape.size) -> !shape.size
695 func @size_to_index_to_size(%size : !shape.size) -> !shape.size {
696 // CHECK: return %[[SIZE]] : !shape.size
697 %idx = shape.size_to_index %size : !shape.size
698 %result = shape.index_to_size %idx
699 return %result : !shape.size
707 // CHECK-NEXT: shape.const_witness true
710 %0 = shape.const_shape [] : !shape.shape
711 %1 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
712 %2 = shape.cstr_broadcastable %0, %1 : !shape.shape, tensor<?xindex>
713 "consume.witness"(%2) : (!shape.witness) -> ()
722 // CHECK-NEXT: shape.shape_of %arg0
723 // CHECK-NEXT: shape.shape_of %arg1
724 // CHECK-NEXT: shape.cstr_broadcastable
727 %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
728 %1 = shape.shape_of %arg1 : tensor<?xf32> -> tensor<?xindex>
729 %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>
730 "consume.witness"(%2) : (!shape.witness) -> ()
739 // CHECK-NEXT: shape.const_witness true
742 %0 = shape.shape_of %arg1 : tensor<index> -> tensor<?xindex>
743 %1 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex>
744 %2 = shape.cstr_broadcastable %0, %1 : tensor<?xindex>, tensor<?xindex>
745 "consume.witness"(%2) : (!shape.witness) -> ()
756 %a = shape.const_shape [1, 2, 3] : !shape.shape
757 %b = shape.const_shape [1, 2, 3] : tensor<?xindex>
758 %result = shape.shape_eq %a, %b : !shape.shape, tensor<?xindex>
769 %a = shape.const_shape [1, 2, 3] : tensor<?xindex>
770 %b = shape.const_shape [4, 5, 6] : tensor<?xindex>
771 %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
782 %a = shape.const_shape [1, 2, 3, 4, 5, 6] : !shape.shape
783 %b = shape.const_shape [1, 2, 3] : !shape.shape
784 %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
792 // CHECK-SAME: (%[[A:.*]]: !shape.shape) -> i1
793 func @shape_eq_do_not_fold(%a : !shape.shape) -> i1 {
794 // CHECK: %[[B:.*]] = shape.const_shape [4, 5, 6]
795 // CHECK: %[[RESULT:.*]] = shape.shape_eq %[[A]], %[[B]] : !shape.shape, !shape.shape
797 %b = shape.const_shape [4, 5, 6] : !shape.shape
798 %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
806 func @fold_mul_size() -> !shape.size {
807 // CHECK: %[[RESULT:.*]] = shape.const_size 6
808 // CHECK: return %[[RESULT]] : !shape.size
809 %c2 = shape.const_size 2
810 %c3 = shape.const_size 3
811 %result = shape.mul %c2, %c3 : !shape.size, !shape.size -> !shape.size
812 return %result : !shape.size
824 %result = shape.mul %c2, %c3 : index, index -> index
832 func @fold_mul_mixed() -> !shape.size {
833 // CHECK: %[[RESULT:.*]] = shape.const_size 6
834 // CHECK: return %[[RESULT]] : !shape.size
835 %c2 = shape.const_size 2
837 %result = shape.mul %c2, %c3 : !shape.size, index -> !shape.size
838 return %result : !shape.size
847 %casted = shape.size_to_index %arg : index
857 %casted = shape.to_extent_tensor %arg : tensor<?xindex> -> tensor<?xindex>
867 %0 = "test.source"() : () -> (!shape.witness)
868 %1 = shape.assuming_all %0
869 "consume.witness"(%1) : (!shape.witness) -> ()
879 %0 = shape.const_shape [2] : tensor<?xindex>
881 %2 = shape.cstr_broadcastable %1, %0 : tensor<1xindex>, tensor<?xindex>
882 "consume.witness"(%2) : (!shape.witness) -> ()