1// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints -canonicalize <%s | FileCheck %s --dump-input=fail --check-prefixes=CANON,CHECK-BOTH
2// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints <%s | FileCheck %s --dump-input=fail --check-prefixes=REPLACE,CHECK-BOTH
3
4// -----
5// Check that cstr_broadcastable is removed.
6//
7// CHECK-BOTH: func @f
8func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
9  // REPLACE-NEXT: %[[WITNESS:.+]] = shape.const_witness true
10  // REPLACE-NOT: shape.cstr_eq
11  // REPLACE: shape.assuming %[[WITNESS]]
12  // CANON-NEXT: test.source
13  // CANON-NEXT: return
14  %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
15  %1 = shape.assuming %0 -> index {
16    %2 = "test.source"() : () -> (index)
17    shape.assuming_yield %2 : index
18  }
19  return %1 : index
20}
21
22// -----
23// Check that cstr_eq is removed.
24//
25// CHECK-BOTH: func @f
26func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
27  // REPLACE-NEXT: %[[WITNESS:.+]] = shape.const_witness true
28  // REPLACE-NOT: shape.cstr_eq
29  // REPLACE: shape.assuming %[[WITNESS]]
30  // CANON-NEXT: test.source
31  // CANON-NEXT: return
32  %0 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape
33  %1 = shape.assuming %0 -> index {
34    %2 = "test.source"() : () -> (index)
35    shape.assuming_yield %2 : index
36  }
37  return %1 : index
38}
39
40// -----
41// With a non-const value, we cannot fold away the code, but all constraints
42// should be removed still.
43//
44// CHECK-BOTH: func @f
45func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
46  // CANON-NEXT: test.source
47  // CANON-NEXT: return
48  %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape
49  %1 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape
50  %2 = shape.assuming_all %0, %1
51  %3 = shape.assuming %0 -> index {
52    %4 = "test.source"() : () -> (index)
53    shape.assuming_yield %4 : index
54  }
55  return %3 : index
56}
57