57 lines · plain
1// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints -canonicalize <%s | FileCheck %s --dump-input=fail --check-prefixes=CANON,CHECK-BOTH2// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints <%s | FileCheck %s --dump-input=fail --check-prefixes=REPLACE,CHECK-BOTH3 4// -----5// Check that cstr_broadcastable is removed.6//7// CHECK-BOTH: func @f8func.func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {9 // REPLACE-NEXT: %[[WITNESS:.+]] = shape.const_witness true10 // REPLACE-NOT: shape.cstr_eq11 // REPLACE: shape.assuming %[[WITNESS]]12 // CANON-NEXT: test.source13 // CANON-NEXT: return14 %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape15 %1 = shape.assuming %0 -> index {16 %2 = "test.source"() : () -> (index)17 shape.assuming_yield %2 : index18 }19 return %1 : index20}21 22// -----23// Check that cstr_eq is removed.24//25// CHECK-BOTH: func @f26func.func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {27 // REPLACE-NEXT: %[[WITNESS:.+]] = shape.const_witness true28 // REPLACE-NOT: shape.cstr_eq29 // REPLACE: shape.assuming %[[WITNESS]]30 // CANON-NEXT: test.source31 // CANON-NEXT: return32 %0 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape33 %1 = shape.assuming %0 -> index {34 %2 = "test.source"() : () -> (index)35 shape.assuming_yield %2 : index36 }37 return %1 : index38}39 40// -----41// With a non-const value, we cannot fold away the code, but all constraints42// should be removed still.43//44// CHECK-BOTH: func @f45func.func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {46 // CANON-NEXT: test.source47 // CANON-NEXT: return48 %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape49 %1 = shape.cstr_eq %arg0, %arg1 : !shape.shape, !shape.shape50 %2 = shape.assuming_all %0, %151 %3 = shape.assuming %0 -> index {52 %4 = "test.source"() : () -> (index)53 shape.assuming_yield %4 : index54 }55 return %3 : index56}57