Lines Matching full:shape
7 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 @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 %1 = shape.const_shape [4, 5, 6] : tensor<?xindex>
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<?xindex>
59 %1 = shape.const_shape [4, 57, 92] : tensor<?xindex>
60 %2 = shape.broadcast %0, %1 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
65 %0 = shape.const_shape [4, 57, 92] : !shape.shape
66 %1 = shape.const_shape [4, 57, 92] : !shape.shape
67 %2 = "shape.join"(%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.join"(%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.join"(%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<?xindex>
96 %0 = shape.shape_of %arg0 : tensor<?xf32> -> tensor<?xindex>
101 %0 = shape.const_shape [] : !shape.shape
102 %1 = shape.const_shape [1, 2, 3] : !shape.shape
104 %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
105 %w1 = shape.cstr_eq %0, %1
106 %w2 = shape.const_witness true
107 %w3 = shape.const_witness false
108 %w4 = shape.cstr_require %true, "msg"
109 %w_all = shape.assuming_all %w0, %w1, %w2, %w3, %w4
110 shape.assuming %w_all -> !shape.shape {
111 %2 = "shape.any"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
112 shape.assuming_yield %2 : !shape.shape
119 %w0 = shape.cstr_broadcastable %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
123 func @mul(%size_arg : !shape.size, %index_arg : index) {
124 %size_prod = shape.mul %size_arg, %size_arg
125 : !shape.size, !shape.size -> !shape.size
126 %index_prod = shape.mul %index_arg, %index_arg : index, index -> index
127 %mixed_prod = shape.mul %size_arg, %index_arg
128 : !shape.size, index -> !shape.size
132 func @add(%size_arg : !shape.size, %index_arg : index) {
133 %size_sum = shape.add %size_arg, %size_arg
134 : !shape.size, !shape.size -> !shape.size
135 %index_sum = shape.add %index_arg, %index_arg : index, index -> index
136 %mixed_sum = shape.add %size_arg, %index_arg
137 : !shape.size, index -> !shape.size
142 // CHECK: %c1 = shape.const_size 1
143 // CHECK: %c2 = shape.const_size 2
144 // CHECK: %c2_0 = shape.const_size 2
145 %0 = shape.const_size 1
146 %1 = shape.const_size 2
147 %2 = shape.const_size 2
151 func @test_to_extent_tensor(%arg: !shape.shape) -> tensor<3xindex> {
152 %0 = shape.to_extent_tensor %arg : !shape.shape -> tensor<3xindex>
156 func @test_from_extent_tensor(%arg: tensor<?xindex>) -> !shape.shape {
157 %0 = shape.from_extent_tensor %arg : tensor<?xindex>
158 return %0 : !shape.shape
161 func @rank(%shape : !shape.shape) -> !shape.size {
162 %rank = shape.rank %shape : !shape.shape -> !shape.size
163 return %rank : !shape.size
166 func @rank_on_extent_tensor(%shape : tensor<?xindex>) -> index {
167 %rank = shape.rank %shape : tensor<?xindex> -> index
172 func @shape_eq_on_shapes(%a : !shape.shape, %b : !shape.shape) -> i1 {
173 %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
178 %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
182 func @shape_eq_on_mixed(%a : tensor<?xindex>, %b : !shape.shape) -> i1 {
183 %result = shape.shape_eq %a, %b : tensor<?xindex>, !shape.shape
187 func @get_extent_on_shape(%arg : !shape.shape) -> !shape.size {
188 %c0 = shape.const_size 0
189 %result = shape.get_extent %arg, %c0 :
190 !shape.shape, !shape.size -> !shape.size
191 return %result : !shape.size
196 %result = shape.get_extent %arg, %c0 : tensor<?xindex>, index -> index
200 func @get_extent_on_mixed_operands(%arg : tensor<?xindex>) -> !shape.size {
201 %c0 = shape.const_size 0
202 %result = shape.get_extent %arg, %c0 : tensor<?xindex>, !shape.size -> !shape.size
203 return %result : !shape.size
207 %0 = shape.const_shape [1, 2, 3] : !shape.shape
208 %1 = shape.const_shape [4, 5, 6] : !shape.shape
209 %2 = "shape.any"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
210 %3 = shape.const_shape [1, 2, 3] : tensor<?xindex>
211 %4 = shape.const_shape [4, 5, 6] : tensor<?xindex>
212 %5 = "shape.any"(%3, %4) : (tensor<?xindex>, tensor<?xindex>) -> tensor<?xindex>
217 %result = shape.num_elements %arg : tensor<?xindex> -> index
221 func @num_elements_shape(%arg : !shape.shape) -> !shape.size {
222 %result = shape.num_elements %arg : !shape.shape -> !shape.size
223 return %result : !shape.size
226 // Testing invoking shape function from another. shape_equal_shapes is merely
228 func @shape_equal_shapes(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
229 %0 = shape.shape_of %a : !shape.value_shape -> !shape.shape
230 %1 = shape.shape_of %b : !shape.value_shape -> !shape.shape
231 %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
232 return %2 : !shape.shape
234 func @shape_with_shape(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
235 %0 = shape.shape_of %a : !shape.value_shape -> !shape.shape
236 %1 = shape.with_shape %b, %0 : !shape.value_shape, !shape.shape
237 %2 = call @shape_equal_shapes(%a, %1) : (!shape.value_shape, !shape.value_shape) -> !shape.shape
238 return %2 : !shape.shape
241 func @any_on_shape(%a : !shape.shape, %b : !shape.shape, %c : !shape.shape)
242 -> !shape.shape {
243 %result = shape.any %a, %b, %c
244 : !shape.shape, !shape.shape, !shape.shape -> !shape.shape
245 return %result : !shape.shape
250 %c : !shape.shape) -> !shape.shape {
251 %result = shape.any %a, %b, %c
252 : tensor<?xindex>, tensor<?xindex>, !shape.shape -> !shape.shape
253 return %result : !shape.shape
259 %result = shape.any %a, %b, %c
266 %result = shape.is_broadcastable %a, %b
271 func @is_broadcastable_on_shapes(%a : !shape.shape,
272 %b : !shape.shape) -> i1 {
273 %result = shape.is_broadcastable %a, %b
274 : !shape.shape, !shape.shape