Lines Matching full:shape
1 // RUN: mlir-opt -shape-to-shape-lowering -split-input-file %s | FileCheck %s
4 // CHECK-SAME: ([[ARG:%.*]]: !shape.shape) -> !shape.size
5 func @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
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
15 // CHECK: return [[NUM_ELEMENTS]] : !shape.size
21 func @num_elements_to_reduce_on_index(%shape : tensor<?xindex>) -> index {
22 %num_elements = shape.num_elements %shape : tensor<?xindex> -> index
26 // CHECK: [[NUM_ELEMENTS:%.*]] = shape.reduce([[ARG]], [[C1]]) : tensor<?xindex> -> index
28 // CHECK: [[NEW_ACC:%.*]] = shape.mul [[DIM]], [[ACC]]
29 // CHECK: shape.yield [[NEW_ACC]] : index