| // RUN: mlir-opt -split-input-file -shape-bufferize <%s | FileCheck %s |
| |
| // ----- |
| |
| // CHECK-LABEL: func @shape_assuming() { |
| // CHECK: %[[WTRUE:.*]] = shape.const_witness true |
| // CHECK: %[[MEMREF:.*]] = shape.assuming %[[WTRUE]] -> (memref<2xf16>) { |
| // CHECK: %[[TENSOR_VAL:.*]] = "test.source"() : () -> tensor<2xf16> |
| // CHECK: %[[YIELDED_MEMREF:.*]] = tensor_to_memref %[[TENSOR_VAL]] : memref<2xf16> |
| // CHECK: shape.assuming_yield %[[YIELDED_MEMREF]] : memref<2xf16> |
| // CHECK: } |
| // CHECK: %[[TENSOR:.*]] = tensor_load %[[MEMREF:.*]] : memref<2xf16> |
| // CHECK: "test.sink"(%[[TENSOR]]) : (tensor<2xf16>) -> () |
| // CHECK: return |
| // CHECK: } |
| func @shape_assuming() { |
| %0 = shape.const_witness true |
| %1 = shape.assuming %0 -> (tensor<2xf16>) { |
| %2 = "test.source"() : () -> (tensor<2xf16>) |
| shape.assuming_yield %2 : tensor<2xf16> |
| } |
| "test.sink"(%1) : (tensor<2xf16>) -> () |
| return |
| } |