forked from OSchip/llvm-project
16 lines
892 B
MLIR
16 lines
892 B
MLIR
// RUN: mlir-opt %s -std-bufferize | FileCheck %s
|
|
|
|
// CHECK-LABEL: func @select(
|
|
// CHECK-SAME: %[[PRED:.*]]: i1,
|
|
// CHECK-SAME: %[[TRUE_VAL:.*]]: tensor<f32>,
|
|
// CHECK-SAME: %[[FALSE_VAL:.*]]: tensor<f32>) -> tensor<f32> {
|
|
// CHECK-DAG: %[[TRUE_VAL_MEMREF:.*]] = bufferization.to_memref %[[TRUE_VAL]] : memref<f32>
|
|
// CHECK-DAG: %[[FALSE_VAL_MEMREF:.*]] = bufferization.to_memref %[[FALSE_VAL]] : memref<f32>
|
|
// CHECK: %[[RET_MEMREF:.*]] = select %[[PRED]], %[[TRUE_VAL_MEMREF]], %[[FALSE_VAL_MEMREF]] : memref<f32>
|
|
// CHECK: %[[RET:.*]] = bufferization.to_tensor %[[RET_MEMREF]] : memref<f32>
|
|
// CHECK: return %[[RET]] : tensor<f32>
|
|
func @select(%arg0: i1, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
|
|
%0 = select %arg0, %arg1, %arg2 : tensor<f32>
|
|
return %0 : tensor<f32>
|
|
}
|