Lines Matching refs:xindex
31 // CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index
32 func @rank(%shape : tensor<?xindex>) -> index {
36 %rank = shape.rank %shape : tensor<?xindex> -> index
44 func @get_extent(%shape : tensor<?xindex>, %idx : !shape.size) -> !shape.size {
47 : tensor<?xindex>, !shape.size -> !shape.size
70 %shape = shape.shape_of %arg : tensor<2x3xf32> -> tensor<?xindex>
71 %result = shape.get_extent %shape, %idx : tensor<?xindex>, index -> index
79 // CHECK-SAME: (%[[EXTENTS:.*]]: tensor<?xindex>, %[[IDX:.*]]: index) -> index
80 func @get_extent_from_extent_tensor(%extents : tensor<?xindex>, %idx : index)
82 // CHECK: %[[RESULT:.*]] = extract_element %[[EXTENTS]][%[[IDX]]] : tensor<?xindex>
84 %result = shape.get_extent %extents, %idx : tensor<?xindex>, index -> index
92 // CHECK-SAME: () -> tensor<?xindex>
93 func @const_shape() -> tensor<?xindex> {
98 // CHECK: %[[RESULT:.*]] = tensor_cast %[[TENSOR3]] : tensor<3xindex> to tensor<?xindex>
99 // CHECK: return %[[RESULT]] : tensor<?xindex>
100 %shape = shape.const_shape [1, 2, 3] : tensor<?xindex>
101 return %shape : tensor<?xindex>
108 // CHECK-SAME: () -> tensor<?xindex>
109 func @const_shape_zero_elements() -> tensor<?xindex> {
110 // CHECK: %[[TENSOR:.*]] = tensor_from_elements : tensor<0xindex>
111 // CHECK: %[[RESULT:.*]] = tensor_cast %[[TENSOR]] : tensor<0xindex> to tensor<?xindex>
112 // CHECK: return %[[RESULT]] : tensor<?xindex>
113 %shape = shape.const_shape [] : tensor<?xindex>
114 return %shape : tensor<?xindex>
121 …/ CHECK-SAME: (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<?xindex>…
122 func @any_of_three(%a : tensor<?xindex>,
123 %b : tensor<?xindex>,
124 %c : tensor<?xindex>) -> tensor<?xindex> {
125 // CHECK: return %[[A]] : tensor<?xindex>
126 …%result = "shape.any"(%a, %b, %c) : (tensor<?xindex>, tensor<?xindex>, tensor<?xindex>) -> tensor<…
127 return %result : tensor<?xindex>
134 // CHECK-SAME: (%[[A:.*]]: tensor<?xindex>) -> tensor<?xindex>
135 func @any_of_one(%a : tensor<?xindex>) -> tensor<?xindex> {
136 // CHECK: return %[[A]] : tensor<?xindex>
137 %result = "shape.any"(%a) : (tensor<?xindex>) -> tensor<?xindex>
138 return %result : tensor<?xindex>
158 // CHECK-SAME: (%[[ARG:.*]]: tensor<?xindex>
159 func @to_extent_tensor(%arg: tensor<?xindex>) -> tensor<3xindex> {
161 // CHECK: %[[RES:.*]] = tensor_cast %[[ARG]] : tensor<?xindex> to tensor<3xindex
162 %casted = shape.to_extent_tensor %arg : tensor<?xindex> -> tensor<3xindex>
164 return %casted : tensor<3xindex>
168 // CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index
169 func @shape_reduce(%shape : tensor<?xindex>) -> index {
171 %num_elements = shape.reduce(%shape, %init) : tensor<?xindex> -> index {
181 // CHECK-NEXT: %[[RANK:.*]] = dim %[[SHAPE]], %[[C0]] : tensor<?xindex>
211 // CHECK: } : tensor<?xindex>
212 %shape = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>
236 …HECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements %[[C1]], %[[C2]], %[[C3]] : tensor<3xindex>
237 %shape = shape.shape_of %arg : tensor<1x2x3xf32> -> tensor<?xindex>
247 // CHECK-DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements : tensor<0xindex>
248 %shape = shape.shape_of %arg : tensor<f32> -> tensor<?xindex>
262 …DAG: %[[SHAPE_UNCASTED:.*]] = tensor_from_elements %[[C1]], %[[C5]], %[[DYN_DIM]] : tensor<3xindex>
263 %shape = shape.shape_of %arg : tensor<1x5x?xf32> -> tensor<?xindex>
270 // CHECK-SAME: (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>) -> i1
271 func @shape_eq(%a : tensor<?xindex>, %b : tensor<?xindex>) -> i1 {
273 // CHECK: %[[RANK_A:.*]] = dim %[[A]], %[[C0]] : tensor<?xindex>
274 // CHECK: %[[RANK_B:.*]] = dim %[[B]], %[[C0]] : tensor<?xindex>
280 // CHECK: %[[EXTENT_A:.*]] = extract_element %[[A]][%[[I]]] : tensor<?xindex>
281 // CHECK: %[[EXTENT_B:.*]] = extract_element %[[B]][%[[I]]] : tensor<?xindex>
292 %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
300 func @broadcast(%a : tensor<?xindex>, %b : !shape.shape) -> !shape.shape {
302 %c = shape.broadcast %a, %b : tensor<?xindex>, !shape.shape -> !shape.shape
309 // CHECK-SAME: %[[LHS:.*]]: tensor<?xindex>,
310 // CHECK-SAME: %[[RHS:.*]]: tensor<?xindex>) {
311 func @broadcast_unknown_extents(%a : tensor<?xindex>, %b : tensor<?xindex>) {
314 // CHECK: %[[LHS_RANK:.*]] = dim %[[LHS]], %[[C0]] : tensor<?xindex>
315 // CHECK: %[[RHS_RANK:.*]] = dim %[[RHS]], %[[C0]] : tensor<?xindex>
319 // CHECK: %[[ERASED_LHS:.*]] = tensor_cast %[[LHS]] : tensor<?xindex> to tensor<?xindex>
320 // CHECK: %[[ERASED_RHS:.*]] = tensor_cast %[[RHS]] : tensor<?xindex> to tensor<?xindex>
321 …ER_RANK_OPERAND:.*]] = select %[[LHS_RANK_ULE]], %[[ERASED_LHS]], %[[ERASED_RHS]] : tensor<?xindex>
322 …ER_RANK_OPERAND:.*]] = select %[[LHS_RANK_ULE]], %[[ERASED_RHS]], %[[ERASED_LHS]] : tensor<?xindex>
327 …ND_EXTENT:.*]] = extract_element %[[GREATER_RANK_OPERAND]][%[[OUTPUT_DIMENSION]]] : tensor<?xindex>
332 …]] = extract_element %[[LESSER_RANK_OPERAND]][%[[LESSER_RANK_OPERAND_DIMENSION]]] : tensor<?xindex>
338 // CHECK: } : tensor<?xindex>
342 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
349 // CHECK-SAME: %[[LHS:.*]]: tensor<2xindex>,
350 // CHECK-SAME: %[[RHS:.*]]: tensor<3xindex>) {
351 func @broadcast_known_different_extents(%a : tensor<2xindex>, %b : tensor<3xindex>) {
354 // CHECK: %[[LHS_RANK:.*]] = dim %[[LHS]], %[[C0]] : tensor<2xindex>
355 // CHECK: %[[RHS_RANK:.*]] = dim %[[RHS]], %[[C0]] : tensor<3xindex>
359 // CHECK: %[[ERASED_LHS:.*]] = tensor_cast %[[LHS]] : tensor<2xindex> to tensor<?xindex>
360 // CHECK: %[[ERASED_RHS:.*]] = tensor_cast %[[RHS]] : tensor<3xindex> to tensor<?xindex>
361 …ER_RANK_OPERAND:.*]] = select %[[LHS_RANK_ULE]], %[[ERASED_LHS]], %[[ERASED_RHS]] : tensor<?xindex>
362 …ER_RANK_OPERAND:.*]] = select %[[LHS_RANK_ULE]], %[[ERASED_RHS]], %[[ERASED_LHS]] : tensor<?xindex>
367 …ND_EXTENT:.*]] = extract_element %[[GREATER_RANK_OPERAND]][%[[OUTPUT_DIMENSION]]] : tensor<?xindex>
372 …]] = extract_element %[[LESSER_RANK_OPERAND]][%[[LESSER_RANK_OPERAND_DIMENSION]]] : tensor<?xindex>
378 // CHECK: } : tensor<?xindex>
382 : tensor<2xindex>, tensor<3xindex> -> tensor<?xindex>
388 func @try_is_broadcastable(%a : tensor<3xindex>, %b : tensor<?xindex>) -> i1 {
389 %0 = shape.is_broadcastable %a, %b : tensor<3xindex>, tensor<?xindex>
394 // CHECK-SAME: %[[LHS:.*]]: tensor<3xindex>,
395 // CHECK-SAME: %[[RHS:.*]]: tensor<?xindex>) -> i1 {
398 // CHECK: %[[LHS_RANK:.*]] = dim %[[LHS]], %[[C0]] : tensor<3xindex>
399 // CHECK: %[[RHS_RANK:.*]] = dim %[[RHS]], %[[C0]] : tensor<?xindex>
403 …HECK: %[[RANK_ERASED_LHS:.*]] = tensor_cast %[[LHS]] : tensor<3xindex> to tensor<?xindex>
404 …HECK: %[[RANK_ERASED_RHS:.*]] = tensor_cast %[[RHS]] : tensor<?xindex> to tensor<?xindex>
405 …_SHAPE:.*]] = select %[[LHS_SMALLER]], %[[RANK_ERASED_LHS]], %[[RANK_ERASED_RHS]] : tensor<?xindex>
406 …_SHAPE:.*]] = select %[[LHS_SMALLER]], %[[RANK_ERASED_RHS]], %[[RANK_ERASED_LHS]] : tensor<?xindex>
410 … %[[LARGER_EXTENT:.*]] = extract_element %[[LARGER_SHAPE]]{{\[}}%[[I]]] : tensor<?xindex>
413 …_EXTENT:.*]] = extract_element %[[SMALLER_SHAPE]]{{\[}}%[[SMALLER_EXTENT_INDEX]]] : tensor<?xindex>
426 func @broadcast(%a : tensor<?xindex>, %b : tensor<?xindex>) -> !shape.witness {
427 %0 = shape.cstr_broadcastable %a, %b : tensor<?xindex>, tensor<?xindex>
432 // CHECK-SAME: %[[LHS:.*]]: tensor<?xindex>,
433 // CHECK-SAME: %[[RHS:.*]]: tensor<?xindex>) -> !shape.witness {
436 // CHECK: %[[LHS_RANK:.*]] = dim %[[LHS]], %[[C0]] : tensor<?xindex>
437 // CHECK: %[[RHS_RANK:.*]] = dim %[[RHS]], %[[C0]] : tensor<?xindex>
441 …HECK: %[[RANK_ERASED_LHS:.*]] = tensor_cast %[[LHS]] : tensor<?xindex> to tensor<?xindex>
442 …HECK: %[[RANK_ERASED_RHS:.*]] = tensor_cast %[[RHS]] : tensor<?xindex> to tensor<?xindex>
443 …_SHAPE:.*]] = select %[[LHS_SMALLER]], %[[RANK_ERASED_LHS]], %[[RANK_ERASED_RHS]] : tensor<?xindex>
444 …_SHAPE:.*]] = select %[[LHS_SMALLER]], %[[RANK_ERASED_RHS]], %[[RANK_ERASED_LHS]] : tensor<?xindex>
448 … %[[LARGER_EXTENT:.*]] = extract_element %[[LARGER_SHAPE]]{{\[}}%[[VAL_16]]] : tensor<?xindex>
451 …LLER_EXTENT:.*]] = extract_element %[[SMALLER_SHAPE]]{{\[}}%[[LHS_EXTENT_INDEX]]] : tensor<?xindex>