} // namespace
namespace {
+class BufferizeSelectOp : public OpConversionPattern<SelectOp> {
+public:
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(SelectOp op, ArrayRef<Value> operands,
+ ConversionPatternRewriter &rewriter) const override {
+ if (!op.condition().getType().isa<IntegerType>())
+ return rewriter.notifyMatchFailure(op, "requires scalar condition");
+
+ SelectOp::Adaptor adaptor(operands);
+ rewriter.replaceOpWithNewOp<SelectOp>(
+ op, adaptor.condition(), adaptor.true_value(), adaptor.false_value());
+ return success();
+ }
+};
+} // namespace
+
+namespace {
class BufferizeTensorCastOp : public OpConversionPattern<TensorCastOp> {
public:
using OpConversionPattern::OpConversionPattern;
void mlir::populateStdBufferizePatterns(MLIRContext *context,
BufferizeTypeConverter &typeConverter,
OwningRewritePatternList &patterns) {
- patterns
- .insert<BufferizeDynamicTensorFromElementsOp, BufferizeExtractElementOp,
- BufferizeTensorCastOp, BufferizeTensorFromElementsOp>(
- typeConverter, context);
+ patterns.insert<
+ // clang-format off
+ BufferizeDynamicTensorFromElementsOp,
+ BufferizeExtractElementOp,
+ BufferizeSelectOp,
+ BufferizeTensorCastOp,
+ BufferizeTensorFromElementsOp
+ // clang-format on
+ >(typeConverter, context);
}
namespace {
populateStdBufferizePatterns(context, typeConverter, patterns);
target.addIllegalOp<DynamicTensorFromElementsOp, ExtractElementOp,
TensorCastOp, TensorFromElementsOp>();
+ // We only bufferize the case of tensor selected type and scalar condition,
+ // as that boils down to a select over memref descriptors (don't need to
+ // touch the data).
+ target.addDynamicallyLegalOp<SelectOp>([&](SelectOp op) {
+ return typeConverter.isLegal(op.getType()) ||
+ !op.condition().getType().isa<IntegerType>();
+ });
if (failed(
applyPartialConversion(getFunction(), target, std::move(patterns))))
signalPassFailure();
return %0 : f32
}
+// CHECK-LABEL: func @select(
+// CHECK-SAME: %[[PRED:.*]]: i1,
+// CHECK-SAME: %[[TRUE_VAL:.*]]: tensor<f32>,
+// CHECK-SAME: %[[FALSE_VAL:.*]]: tensor<f32>) -> tensor<f32> {
+// CHECK: %[[TRUE_VAL_MEMREF:.*]] = tensor_to_memref %[[TRUE_VAL]] : memref<f32>
+// CHECK: %[[FALSE_VAL_MEMREF:.*]] = tensor_to_memref %[[FALSE_VAL]] : memref<f32>
+// CHECK: %[[RET_MEMREF:.*]] = select %[[PRED]], %[[TRUE_VAL_MEMREF]], %[[FALSE_VAL_MEMREF]] : memref<f32>
+// CHECK: %[[RET:.*]] = tensor_load %[[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>
+}
+
// CHECK-LABEL: func @tensor_cast(
// CHECK-SAME: %[[TENSOR:.*]]: tensor<?xindex>) -> tensor<2xindex> {
// CHECK: %[[MEMREF:.*]] = tensor_to_memref %[[TENSOR]]