1// RUN: mlir-opt -shape-to-shape-lowering -split-input-file %s | FileCheck %s
2
3// CHECK-LABEL: func @num_elements_to_reduce
4// CHECK-SAME:  ([[ARG:%.*]]: !shape.shape) -> !shape.size
5func @num_elements_to_reduce(%shape : !shape.shape) -> !shape.size {
6  %num_elements = shape.num_elements %shape : !shape.shape -> !shape.size
7  return %num_elements : !shape.size
8}
9// CHECK: [[C1:%.*]] = shape.const_size 1
10// CHECK: [[NUM_ELEMENTS:%.*]] = shape.reduce([[ARG]], [[C1]]) : !shape.shape -> !shape.size
11// CHECK: ^bb0({{.*}}: index, [[DIM:%.*]]: !shape.size, [[ACC:%.*]]: !shape.size
12// CHECK:   [[NEW_ACC:%.*]] = shape.mul [[DIM]], [[ACC]]
13// CHECK:   shape.yield [[NEW_ACC]] : !shape.size
14// CHECK: }
15// CHECK: return [[NUM_ELEMENTS]] : !shape.size
16
17// -----
18
19// CHECK-LABEL: func @num_elements_to_reduce_on_index
20// CHECK-SAME:  ([[ARG:%.*]]: tensor<?xindex>) -> index
21func @num_elements_to_reduce_on_index(%shape : tensor<?xindex>) -> index {
22  %num_elements = shape.num_elements %shape : tensor<?xindex> -> index
23  return %num_elements : index
24}
25// CHECK: [[C1:%.*]] = constant 1 : index
26// CHECK: [[NUM_ELEMENTS:%.*]] = shape.reduce([[ARG]], [[C1]]) : tensor<?xindex> -> index
27// CHECK: ^bb0({{.*}}: index, [[DIM:%.*]]: index, [[ACC:%.*]]: index
28// CHECK:   [[NEW_ACC:%.*]] = shape.mul [[DIM]], [[ACC]]
29// CHECK:   shape.yield [[NEW_ACC]] : index
30// CHECK: }
31// CHECK: return [[NUM_ELEMENTS]] : index
32