1// RUN: mlir-opt %s -split-input-file -verify-diagnostics
2
3func @reduce_op_args_num_mismatch(%shape : !shape.shape, %init : !shape.size) {
4  // expected-error@+1 {{ReduceOp body is expected to have 3 arguments}}
5  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
6    ^bb0(%index: index, %dim: !shape.size):
7      shape.yield %dim : !shape.size
8  }
9  return
10}
11
12// -----
13
14func @reduce_op_arg0_wrong_type(%shape : !shape.shape, %init : !shape.size) {
15  // expected-error@+1 {{argument 0 of ReduceOp body is expected to be of IndexType}}
16  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
17    ^bb0(%index: f32, %dim: !shape.size, %acc: !shape.size):
18      %new_acc = "shape.add"(%acc, %dim)
19          : (!shape.size, !shape.size) -> !shape.size
20      shape.yield %new_acc : !shape.size
21  }
22  return
23}
24
25// -----
26
27func @reduce_op_arg1_wrong_type(%shape : !shape.shape, %init : !shape.size) {
28  // expected-error@+1 {{argument 1 of ReduceOp body is expected to be of SizeType if the ReduceOp operates on a ShapeType}}
29  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
30    ^bb0(%index: index, %dim: f32, %lci: !shape.size):
31      shape.yield
32  }
33  return
34}
35
36// -----
37
38func @reduce_op_arg1_wrong_type(%shape : tensor<?xindex>, %init : index) {
39  // expected-error@+1 {{argument 1 of ReduceOp body is expected to be of IndexType if the ReduceOp operates on an extent tensor}}
40  %num_elements = shape.reduce(%shape, %init) : tensor<?xindex> -> index {
41    ^bb0(%index: index, %dim: f32, %lci: index):
42      shape.yield
43  }
44  return
45}
46
47// -----
48
49func @reduce_op_init_type_mismatch(%shape : !shape.shape, %init : f32) {
50  // expected-error@+1 {{type mismatch between argument 2 of ReduceOp body and initial value 0}}
51  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> f32 {
52    ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
53      shape.yield
54  }
55  return
56}
57
58// -----
59
60func @yield_op_args_num_mismatch(%shape : !shape.shape, %init : !shape.size) {
61  // expected-error@+3 {{number of operands does not match number of results of its parent}}
62  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
63    ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
64      shape.yield %dim, %dim : !shape.size, !shape.size
65  }
66  return
67}
68
69// -----
70
71func @yield_op_type_mismatch(%shape : !shape.shape, %init : !shape.size) {
72  // expected-error@+4 {{types mismatch between yield op and its parent}}
73  %num_elements = shape.reduce(%shape, %init) : !shape.shape -> !shape.size {
74    ^bb0(%index: index, %dim: !shape.size, %lci: !shape.size):
75      %c0 = constant 1 : index
76      shape.yield %c0 : index
77  }
78  return
79}
80
81// -----
82
83func @assuming_all_op_too_few_operands() {
84  // expected-error@+1 {{no operands specified}}
85  %w0 = shape.assuming_all
86  return
87}
88
89// -----
90
91func @shape_of(%value_arg : !shape.value_shape,
92               %shaped_arg : tensor<?x3x4xf32>) {
93  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `shape` to propagate them}}
94  %0 = shape.shape_of %value_arg : !shape.value_shape -> tensor<?xindex>
95  return
96}
97
98// -----
99
100func @rank(%arg : !shape.shape) {
101  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
102  %0 = shape.rank %arg : !shape.shape -> index
103  return
104}
105
106// -----
107
108func @get_extent(%arg : tensor<?xindex>) -> index {
109  %c0 = shape.const_size 0
110  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
111  %result = shape.get_extent %arg, %c0 : tensor<?xindex>, !shape.size -> index
112  return %result : index
113}
114
115// -----
116
117func @mul(%lhs : !shape.size, %rhs : index) -> index {
118  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
119  %result = shape.mul %lhs, %rhs : !shape.size, index -> index
120  return %result : index
121}
122
123// -----
124
125func @num_elements(%arg : !shape.shape) -> index {
126  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
127  %result = shape.num_elements %arg : !shape.shape -> index
128  return %result : index
129}
130
131// -----
132
133func @add(%lhs : !shape.size, %rhs : index) -> index {
134  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `size` to propagate them}}
135  %result = shape.add %lhs, %rhs : !shape.size, index -> index
136  return %result : index
137}
138
139// -----
140
141func @broadcast(%arg0 : !shape.shape, %arg1 : !shape.shape) -> tensor<?xindex> {
142  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `shape` to propagate them}}
143  %result = shape.broadcast %arg0, %arg1
144      : !shape.shape, !shape.shape -> tensor<?xindex>
145  return %result : tensor<?xindex>
146}
147
148
149// -----
150
151func @broadcast(%arg0 : !shape.shape, %arg1 : tensor<?xindex>) -> tensor<?xindex> {
152  // expected-error@+1 {{if at least one of the operands can hold error values then the result must be of type `shape` to propagate them}}
153  %result = shape.broadcast %arg0, %arg1
154      : !shape.shape, tensor<?xindex> -> tensor<?xindex>
155  return %result : tensor<?xindex>
156}
157