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