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