[mlir] Add bufferization for std.select op.

mlir/lib/Dialect/StandardOps/Transforms/Bufferize.cpp

matchAndRewrite(ExtractElementOp op, ArrayRef<Value> operands,
     rewriter.replaceOpWithNewOp<LoadOp>(op, adaptor.aggregate(),
                                         adaptor.indices());
     return success();
   }
 };
 } // 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;
   LogicalResult
   matchAndRewrite(TensorCastOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto resultType = getTypeConverter()->convertType(op.getType());
     rewriter.replaceOpWithNewOp<MemRefCastOp>(op, resultType, operands[0]);
matchAndRewrite(TensorFromElementsOp op, ArrayRef<Value> operands,
     return success();
   }
 };
 } // namespace
 
 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 {
 struct StdBufferizePass : public StdBufferizeBase<StdBufferizePass> {
   void runOnFunction() override {
     auto *context = &getContext();
     BufferizeTypeConverter typeConverter;
     OwningRewritePatternList patterns;
     ConversionTarget target(*context);
 
     target.addLegalDialect<StandardOpsDialect>();
     target.addLegalDialect<scf::SCFDialect>();
 
     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, patterns)))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createStdBufferizePass() {
   return std::make_unique<StdBufferizePass>();
 }

mlir/test/Dialect/Standard/bufferize.mlir

 // CHECK:           %[[RET:.*]] = load %[[MEMREF]][%[[IDX]]] : memref<?xf32>
 // CHECK:           return %[[RET]] : f32
 // CHECK:         }
 func @extract_element(%arg0: tensor<?xf32>, %arg1: index) -> f32 {
   %0 = extract_element %arg0[%arg1] : tensor<?xf32>
   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]]
 // CHECK:           %[[CASTED:.*]] = memref_cast %[[MEMREF]] : memref<?xindex> to memref<2xindex>
 // CHECK:           %[[RET:.*]] = tensor_load %[[CASTED]]
 // CHECK:           return %[[RET]] : tensor<2xindex>
 func @tensor_cast(%arg0: tensor<?xindex>) -> tensor<2xindex> {
   %0 = tensor_cast %arg0 : tensor<?xindex> to tensor<2xindex>