[mlir][linalg][bufferize] Support std.select bufferization

[springerm]

Summary

This op is an example for how to deal with ops who's OpResult may aliasing with one of multiple OpOperands.

[springerm]

fix test

[nicolasvasilache]

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

1226

TODO: we also need a test case where std.select needs to bufferize out of place.

[springerm]

add test cases

[nicolasvasilache]

Please leave a TODO somewhere that you can improve this to making 0 or 1 copies (e.g. dynamically if you put your copy operations under the same conditional as the switch; or something else as we discussed offline).
Atm you make 0, 1 or 2 copies quite greedily.

[springerm]

address comments

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h

-//===- LinalgInterfaceImpl.h - Linalg Impl. of BufferizableOpInterface ----===//
+//===- AffineInterfaceImpl.h - Affine Impl. of BufferizableOpInterface ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_AFFINE_INTERFACE_IMPL_H

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

public:
   ///   |          |          |          |
   ///   +----------+----------+----------+
   ///              |
   ///              1
   ///
   /// In the above example, Values with a star satisfy the condition. When
   /// starting the traversal from Value 1, the resulting SetVector is:
   /// { 2, 7, 8, 5 }
-  llvm::SetVector<Value> findValueInReverseUseDefChain(
+  SetVector<Value> findValueInReverseUseDefChain(
       Value value, llvm::function_ref<bool(Value)> condition) const;
 
-  /// Find the Value of the last preceding write of a given Value.
+  /// Find the Values of the last preceding write of a given Value.
   ///
   /// Note: Unknown ops are handled conservatively and assumed to be writes.
   /// Furthermore, BlockArguments are also assumed to be writes. There is no
   /// analysis across block boundaries.
   ///
   /// Note: When reaching an end of the reverse SSA use-def chain, that value
   /// is returned regardless of whether it is a memory write or not.
-  Value findLastPrecedingWrite(Value value) const;
+  SetVector<Value> findLastPrecedingWrite(Value value) const;
 
   /// Creates a memref allocation.
   FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
                                ArrayRef<Value> dynShape) const;
 
   /// Creates a memref allocation for the given shaped value. This function may
   /// perform additional optimizations such as buffer allocation hoisting. If
   /// `createDealloc`, a deallocation op is inserted at the point where the

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h

 /// call analysis to determine which function arguments are inplaceable. Then
 /// analyzes and bufferizes FuncOps one-by-one with Comprehensive Bufferization.
 LogicalResult
 runComprehensiveBufferize(ModuleOp moduleOp,
                           std::unique_ptr<BufferizationOptions> options);
 
 namespace std_ext {
 
-void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
+void registerModuleBufferizationExternalModels(DialectRegistry &registry);
 
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_MODULE_BUFFERIZATION_H

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h

-//===- LinalgInterfaceImpl.h - Linalg Impl. of BufferizableOpInterface ----===//
+//===- StdInterfaceImpl.h - Standard Impl. of BufferizableOpInterface- ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_AFFINE_INTERFACE_IMPL_H
-#define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_AFFINE_INTERFACE_IMPL_H
+#ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_STD_INTERFACE_IMPL_H
+#define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_STD_INTERFACE_IMPL_H
 
 namespace mlir {
 
 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
-namespace affine_ext {
+namespace std_ext {
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
-} // namespace affine_ext
+} // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
-#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_AFFINE_INTERFACE_IMPL_H
+#endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_STD_INTERFACE_IMPL_H

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp

while (!workingSet.empty()) {
 
     for (OpOperand *o : opOperands)
       workingSet.insert(o->get());
   }
 
   return result;
 }
 
-// Find the Value of the last preceding write of a given Value.
-Value mlir::linalg::comprehensive_bufferize::BufferizationState::
-    findLastPrecedingWrite(Value value) const {
-  SetVector<Value> result =
-      findValueInReverseUseDefChain(value, [&](Value value) {
-        Operation *op = value.getDefiningOp();
-        if (!op)
-          return true;
-        auto bufferizableOp = options.dynCastBufferizableOp(op);
-        if (!bufferizableOp)
-          return true;
-        return bufferizableOp.isMemoryWrite(value.cast<OpResult>(), *this);
-      });
-
-  // To simplify the analysis, `scf.if` ops are considered memory writes. There
-  // are currently no other ops where one OpResult may alias with multiple
-  // OpOperands. Therefore, this function should return exactly one result at
-  // the moment.
-  assert(result.size() == 1 && "expected exactly one result");
-  return result.front();
+// Find the Values of the last preceding write of a given Value.
+llvm::SetVector<Value> mlir::linalg::comprehensive_bufferize::
+    BufferizationState::findLastPrecedingWrite(Value value) const {
+  return findValueInReverseUseDefChain(value, [&](Value value) {
 }
 
 mlir::linalg::comprehensive_bufferize::BufferizationState::BufferizationState(
     Operation *op, const BufferizationOptions &options)
     : aliasInfo(op), options(options) {
   // Set up alias sets for OpResults that must bufferize in-place. This should
   // be done before making any other bufferization decisions.
   op->walk([&](BufferizableOpInterface bufferizableOp) {
mlir::linalg::comprehensive_bufferize::BufferizationState::getBuffer(
   // Move insertion point right after `operandBuffer`. That is where the
   // allocation should be inserted (in the absence of allocation hoisting).
   setInsertionPointAfter(rewriter, operandBuffer);
   // Allocate the result buffer.
   FailureOr<Value> resultBuffer =
       createAlloc(rewriter, loc, operandBuffer, options.createDeallocs);
   if (failed(resultBuffer))
     return failure();
-  // Do not copy if the last preceding write of `operand` is an op that does
+  // Do not copy if the last preceding writes of `operand` are ops that do
   // not write (skipping ops that merely create aliases). E.g., InitTensorOp.
   // Note: If `findLastPrecedingWrite` reaches the end of the reverse SSA
   // use-def chain, it returns that value, regardless of whether it is a
   // memory write or not.
-  Value lastWrite = findLastPrecedingWrite(operand);
+  SetVector<Value> lastWrites = findLastPrecedingWrite(operand);
+  if (llvm::none_of(lastWrites, [&](Value lastWrite) {
         if (auto bufferizableOp = options.dynCastBufferizableOp(lastWrite))
-    if (!bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>(), *this))
-      return resultBuffer;
+          return bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>(),
+                                              *this);
+        return true;
+      }))
   // Do not copy if the copied data is never read.
   OpResult aliasingOpResult = getAliasingOpResult(opOperand);
   if (aliasingOpResult && !bufferizesToMemoryRead(opOperand) &&
       !isValueRead(aliasingOpResult))
     return resultBuffer;
   // Do not copy if this op does not read the data, but writes it.
   if (bufferizesToMemoryWrite(opOperand) && !bufferizesToMemoryRead(opOperand))
     return resultBuffer;

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AffineInterfaceImpl.cpp
   ArithInterfaceImpl.cpp
   BufferizableOpInterface.cpp
   BufferizationInterfaceImpl.cpp
   ComprehensiveBufferize.cpp
   LinalgInterfaceImpl.cpp
   ModuleBufferization.cpp
   SCFInterfaceImpl.cpp
+  StdInterfaceImpl.cpp
   TensorInterfaceImpl.cpp
   VectorInterfaceImpl.cpp
 )
 
 add_mlir_dialect_library(MLIRBufferizableOpInterface
   BufferizableOpInterface.cpp
 
   DEPENDS
add_mlir_dialect_library(MLIRSCFBufferizableOpInterfaceImpl
   SCFInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRBufferizableOpInterface
   MLIRIR
   MLIRSCF
 )
 
+add_mlir_dialect_library(MLIRStdBufferizableOpInterfaceImpl
+  StdInterfaceImpl.cpp
+
+  LINK_LIBS PUBLIC
+  MLIRBufferizableOpInterface
+  MLIRStandard
+)
+
 add_mlir_dialect_library(MLIRTensorBufferizableOpInterfaceImpl
   TensorInterfaceImpl.cpp
 
   LINK_LIBS PUBLIC
   MLIRBufferizableOpInterface
   MLIRIR
   MLIRMemRef
   MLIRTensor

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

static bool hasReadAfterWriteInterference(
     const DenseSet<OpOperand *> &usesRead,
     const DenseSet<OpOperand *> &usesWrite, const DominanceInfo &domInfo,
     BufferizationState &state, const BufferizationAliasInfo &aliasInfo) {
   const BufferizationOptions &options = state.getOptions();
 
   for (OpOperand *uRead : usesRead) {
     Operation *readingOp = uRead->getOwner();
 
-    // Find most recent write of uRead by following the SSA use-def chain. E.g.:
+    // Find most recent writes of uRead by following the SSA use-def chain.
+    // E.g.:
     //
     // %0 = "writing_op"(%t) : tensor<?x32> -> tensor<?xf32>
     // %1 = "aliasing_op"(%0) : tensor<?x32> -> tensor<?xf32>
     // %2 = "reading_op"(%1) : : tensor<?x32> -> not_a_tensor_type
     //
     // In the above example, if uRead is the OpOperand of reading_op, lastWrite
     // is %0. Note that operations that create an alias but do not write (such
     // as ExtractSliceOp) are skipped.
-    Value lastWrite = state.findLastPrecedingWrite(uRead->get());
+    SetVector<Value> lastWrites = state.findLastPrecedingWrite(uRead->get());
 
     // Look for conflicting memory writes. Potential conflicts are writes to an
     // alias that have been decided to bufferize inplace.
     for (OpOperand *uConflictingWrite : usesWrite) {
       // Throughout this loop, check for multiple requirements that have to be
       // met for uConflictingWrite to be an actual conflict.
       Operation *conflictingWritingOp = uConflictingWrite->getOwner();
 
for (OpOperand *uConflictingWrite : usesWrite) {
           if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state,
                                               aliasInfo))
             continue;
 
       // Ops are not conflicting if they are in mutually exclusive regions.
       if (insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp))
         continue;
 
+      // Check all possible last writes.
+      for (Value lastWrite : lastWrites) {
         // No conflict if the conflicting write happens before the last
         // write.
         if (Operation *writingOp = lastWrite.getDefiningOp()) {
           if (happensBefore(conflictingWritingOp, writingOp, domInfo))
             // conflictingWritingOp happens before writingOp. No conflict.
             continue;
           // No conflict if conflictingWritingOp is contained in writingOp.
           if (writingOp->isProperAncestor(conflictingWritingOp))
             continue;
         } else {
           auto bbArg = lastWrite.cast<BlockArgument>();
           Block *block = bbArg.getOwner();
           if (!block->findAncestorOpInBlock(*conflictingWritingOp))
             // conflictingWritingOp happens outside of the block. No
             // conflict.
             continue;
         }
 
         // No conflict if the conflicting write and the last write are the same
         // use.
         if (state.getAliasingOpResult(*uConflictingWrite) == lastWrite)
           continue;
 
         // All requirements are met. Conflict found!
 
         if (options.printConflicts)
           annotateConflict(uRead, uConflictingWrite, lastWrite);
 
         return true;
       }
     }
+  }
 
   return false;
 }
 
 /// Return true if bufferizing `operand` inplace would create a conflict. A read
 /// R and a write W of the same alias set is a conflict if inplace bufferization
 /// of W changes the value read by R to a value different from the one that
 /// would be expected by tracing back R's origin through SSA use-def chains.

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

 };
 
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 void mlir::linalg::comprehensive_bufferize::std_ext::
-    registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
+    registerModuleBufferizationExternalModels(DialectRegistry &registry) {
   registry.addOpInterface<CallOp, std_ext::CallOpInterface>();
   registry.addOpInterface<ReturnOp, std_ext::ReturnOpInterface>();
   registry.addOpInterface<FuncOp, std_ext::FuncOpInterface>();
 }
 
 /// Set the attribute that triggers inplace bufferization on a FuncOp argument
 /// `bbArg`.
 static void setInPlaceFuncArgument(BlockArgument bbArg, bool inPlace) {

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.cpp

+//===- StdInterfaceImpl.cpp - Standard Impl. of BufferizableOpInterface ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h"
+
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Operation.h"
+
+namespace mlir {
+namespace linalg {
+namespace comprehensive_bufferize {
+namespace std_ext {
+
+/// Bufferization of std.select. Just replace the operands.
+struct SelectOpInterface
+    : public BufferizableOpInterface::ExternalModel<SelectOpInterface,
+                                                    SelectOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const BufferizationState &state) const {
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const BufferizationState &state) const {
+    return false;
+  }
+
+  OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                               const BufferizationState &state) const {
+    return op->getOpResult(0) /*result*/;
+  }
+
+  SmallVector<OpOperand *>
+  getAliasingOpOperand(Operation *op, OpResult opResult,
+                       const BufferizationState &state) const {
+    return {&op->getOpOperand(1) /*true_value*/,
+            &op->getOpOperand(2) /*false_value*/};
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationState &state) const {
+    auto selectOp = cast<SelectOp>(op);
+    // `getBuffer` introduces copies if an OpOperand bufferizes out-of-place.
+    // TODO: It would be more efficient to copy the result of the `select` op
+    // instead of its OpOperands. In the worst case, 2 copies are inserted at
+    // the moment (one for each tensor). When copying the op result, only one
+    // copy would be needed.
+    Value trueBuffer =
+        *state.getBuffer(rewriter, selectOp->getOpOperand(1) /*true_value*/);
+    Value falseBuffer =
+        *state.getBuffer(rewriter, selectOp->getOpOperand(2) /*false_value*/);
+    replaceOpWithNewBufferizedOp<SelectOp>(
+        rewriter, op, selectOp.getCondition(), trueBuffer, falseBuffer);
+    return success();
+  }
+
+  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+                                const BufferizationAliasInfo &aliasInfo,
+                                const BufferizationState &state) const {
+    return BufferRelation::None;
+  }
+};
+
+} // namespace std_ext
+} // namespace comprehensive_bufferize
+} // namespace linalg
+} // namespace mlir
+
+void mlir::linalg::comprehensive_bufferize::std_ext::
+    registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
+  registry.addOpInterface<SelectOp, std_ext::SelectOpInterface>();
+}

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgTransforms
   MLIRMemRef
   MLIRLinalg
   MLIRLinalgAnalysis
   MLIRLinalgBufferizableOpInterfaceImpl
   MLIRLinalgUtils
   MLIRSCF
   MLIRSCFBufferizableOpInterfaceImpl
   MLIRSCFTransforms
+  MLIRStdBufferizableOpInterfaceImpl
   MLIRPass
   MLIRStandard
   MLIRStandardOpsTransforms
   MLIRStandardToLLVM
   MLIRTensor
   MLIRTensorBufferizableOpInterfaceImpl
   MLIRTransforms
   MLIRTransformUtils
   MLIRVector
   MLIRVectorBufferizableOpInterfaceImpl
   MLIRX86VectorTransforms
   MLIRVectorToSCF
 )

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ArithInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizationInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h"
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/VectorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
registry
                 memref::MemRefDialect, tensor::TensorDialect,
                 vector::VectorDialect, scf::SCFDialect,
                 arith::ArithmeticDialect, StandardOpsDialect, AffineDialect>();
     affine_ext::registerBufferizableOpInterfaceExternalModels(registry);
     arith_ext::registerBufferizableOpInterfaceExternalModels(registry);
     bufferization_ext::registerBufferizableOpInterfaceExternalModels(registry);
     linalg_ext::registerBufferizableOpInterfaceExternalModels(registry);
     scf_ext::registerBufferizableOpInterfaceExternalModels(registry);
+    std_ext::registerModuleBufferizationExternalModels(registry);
     std_ext::registerBufferizableOpInterfaceExternalModels(registry);
     tensor_ext::registerBufferizableOpInterfaceExternalModels(registry);
     vector_ext::registerBufferizableOpInterfaceExternalModels(registry);
   }
 };
 } // namespace
 
 static void applyEnablingTransformations(ModuleOp moduleOp) {

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

func @equivalent_func_arg_2(%c0: index, %c10: index, %c1: index, %t0: tensor<?xf32>) -> tensor<?xf32> {
   // This test does not check IR. It just asserts there is no failure due to
   // non-equivalent scf.for yield values.
   %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
     %3 = call @inner_func_2(%t1) : (tensor<?xf32>) -> tensor<?xf32>
     scf.yield %3 : tensor<?xf32>
   }
   return %1: tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @write_after_select_read_one
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32> {{.*}}, %[[t2:.*]]: tensor<?xf32>
+func @write_after_select_read_one(
+    %t1 : tensor<?xf32> {linalg.inplaceable = true},
+    %t2 : tensor<?xf32> {linalg.inplaceable = true},
+    %c : i1)
+  -> (f32, tensor<?xf32>)
+{
+  %cst = arith.constant 0.0 : f32
+  %idx = arith.constant 0 : index
+
+  //      CHECK: select %{{.*}}, %[[t1]], %[[t2]]
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["none", "false", "true"]}
+  %s = std.select %c, %t1, %t2 : tensor<?xf32>
+  //      CHECK: tensor.insert
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["none", "true", "none"]}
+  %w = tensor.insert %cst into %s[%idx] : tensor<?xf32>
+  //      CHECK: tensor.extract
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["true", "none"]}
+  %f = tensor.extract %t1[%idx] : tensor<?xf32>
+
+  return %f, %w : f32, tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @write_after_select_read_both
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32> {{.*}}, %[[t2:.*]]: tensor<?xf32>
+func @write_after_select_read_both(
+    %t1 : tensor<?xf32> {linalg.inplaceable = true},
+    %t2 : tensor<?xf32> {linalg.inplaceable = true},
+    %c : i1)
+  -> (f32, f32, tensor<?xf32>)
+{
+  %cst = arith.constant 0.0 : f32
+  %idx = arith.constant 0 : index
+
+  //      CHECK: select %{{.*}}, %[[t1]], %[[t2]]
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["none", "false", "false"]}
+  %s = std.select %c, %t1, %t2 : tensor<?xf32>
+  //      CHECK: tensor.insert
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["none", "true", "none"]}
+  %w = tensor.insert %cst into %s[%idx] : tensor<?xf32>
+  //      CHECK: tensor.extract
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["true", "none"]}
+  %f = tensor.extract %t1[%idx] : tensor<?xf32>
+  //      CHECK: tensor.extract
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["true", "none"]}
+  %f2 = tensor.extract %t2[%idx] : tensor<?xf32>
+
+  return %f, %f2, %w : f32, f32, tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @write_after_select_no_conflict
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?xf32> {{.*}}, %[[t2:.*]]: tensor<?xf32>
+func @write_after_select_no_conflict(
+    %t1 : tensor<?xf32> {linalg.inplaceable = true},
+    %t2 : tensor<?xf32> {linalg.inplaceable = true},
+    %c : i1)
+  -> (f32, tensor<?xf32>)
+{
+  %cst = arith.constant 0.0 : f32
+  %idx = arith.constant 0 : index
+
+  //      CHECK: select %{{.*}}, %[[t1]], %[[t2]]
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["none", "true", "true"]}
+  %s = std.select %c, %t1, %t2 : tensor<?xf32>
+  //      CHECK: tensor.insert
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["none", "true", "none"]}
+  %w = tensor.insert %cst into %s[%idx] : tensor<?xf32>
+  //      CHECK: tensor.extract
+  // CHECK-SAME:   {__inplace_operands_attr__ = ["true", "none"]}
+  %f = tensor.extract %w[%idx] : tensor<?xf32>
+
+  return %f, %w : f32, tensor<?xf32>
+}

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

%r1 = linalg.generic #trait outs (%r0 : tensor<?xf32>) {
     } -> (tensor<?xf32>)
 
   // CHECK: return %[[ALLOC]]
   return %r1 : tensor<?xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // InitTensorOp elimination would produce SSA violations for the example below.
 //===----------------------------------------------------------------------===//
 
-func @depthwise_conv_1d_nwc_wc(%arg0: index, %arg1: index, %arg2: tensor<8x18x32xf32>) 
+func @depthwise_conv_1d_nwc_wc(%arg0: index, %arg1: index, %arg2: tensor<8x18x32xf32>)
     -> tensor<?x1x6x8xf32> {
   %c0 = arith.constant 0 : index
   %c32 = arith.constant 32 : index
   %c8 = arith.constant 8 : index
   %0 = linalg.init_tensor [4, 1, 6, 8] : tensor<4x1x6x8xf32>
   %1 = tensor.cast %0 : tensor<4x1x6x8xf32> to tensor<?x1x6x8xf32>
   %2 = linalg.init_tensor [1, 6, 8] : tensor<1x6x8xf32>
   %3 = scf.for %arg3 = %c0 to %c32 step %c8 iter_args(%arg4 = %1) -> (tensor<?x1x6x8xf32>) {
     %4 = affine.apply affine_map<(d0) -> (d0 ceildiv 8)>(%arg3)
     %5 = tensor.insert_slice %2 into %arg4[%4,0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] :
       tensor<1x6x8xf32> into tensor<?x1x6x8xf32>
     scf.yield %5 : tensor<?x1x6x8xf32>
   }
   return %3 : tensor<?x1x6x8xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @write_to_select_op_source
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32, #{{.*}}>, %[[t2:.*]]: memref<?xf32, #{{.*}}>
+func @write_to_select_op_source(
+    %t1 : tensor<?xf32> {linalg.inplaceable = true},
+    %t2 : tensor<?xf32> {linalg.inplaceable = true},
+    %c : i1)
+  -> (tensor<?xf32>, tensor<?xf32>)
+{
+  %cst = arith.constant 0.0 : f32
+  %idx = arith.constant 0 : index
+  // CHECK: %[[alloc:.*]] = memref.alloc
+  // CHECK: linalg.copy(%[[t1]], %[[alloc]])
+  // CHECK: memref.store %{{.*}}, %[[alloc]]
+  %w = tensor.insert %cst into %t1[%idx] : tensor<?xf32>
+  // CHECK: %[[select:.*]] = select %{{.*}}, %[[t1]], %[[t2]]
+  %s = std.select %c, %t1, %t2 : tensor<?xf32>
+  // CHECK: return %[[select]], %[[alloc]]
+  return %s, %w : tensor<?xf32>, tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @write_after_select_read_one
+//  CHECK-SAME:     %[[t1:.*]]: memref<?xf32, #{{.*}}>, %[[t2:.*]]: memref<?xf32, #{{.*}}>
+func @write_after_select_read_one(
+    %t1 : tensor<?xf32> {linalg.inplaceable = true},
+    %t2 : tensor<?xf32> {linalg.inplaceable = true},
+    %c : i1)
+  -> (f32, tensor<?xf32>)
+{
+  %cst = arith.constant 0.0 : f32
+  %idx = arith.constant 0 : index
+
+  // CHECK: %[[alloc:.*]] = memref.alloc
+  // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
+  // CHECK: linalg.copy(%[[t1]], %[[alloc]])
+  // CHECK: %[[select:.*]] = select %{{.*}}, %[[casted]], %[[t2]]
+  %s = std.select %c, %t1, %t2 : tensor<?xf32>
+
+  // CHECK: memref.store %{{.*}}, %[[select]]
+  %w = tensor.insert %cst into %s[%idx] : tensor<?xf32>
+
+  // CHECK: %[[f:.*]] = memref.load %[[t1]]
+  %f = tensor.extract %t1[%idx] : tensor<?xf32>
+
+  // CHECK: return %[[f]], %[[select]]
+  return %f, %w : f32, tensor<?xf32>
+}

mlir/test/lib/Dialect/Linalg/CMakeLists.txt

add_mlir_library(MLIRLinalgTestPasses
   MLIRLinalg
   MLIRLinalgBufferizableOpInterfaceImpl
   MLIRLinalgTransforms
   MLIRLLVMToLLVMIRTranslation
   MLIRMemRef
   MLIRPass
   MLIRSCF
   MLIRSCFBufferizableOpInterfaceImpl
+  MLIRStdBufferizableOpInterfaceImpl
   MLIRStandard
   MLIRTensor
   MLIRTensorBufferizableOpInterfaceImpl
   MLIRTransformUtils
   MLIRVector
   MLIRVectorBufferizableOpInterfaceImpl
   MLIRVectorToSCF
   )

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/AffineInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ArithInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizationInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/SCFInterfaceImpl.h"
+#include "mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/VectorInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
struct TestComprehensiveFunctionBufferize
   TestComprehensiveFunctionBufferize() = default;
   TestComprehensiveFunctionBufferize(
       const TestComprehensiveFunctionBufferize &pass)
       : PassWrapper(pass) {}
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect, linalg::LinalgDialect,
                     memref::MemRefDialect, tensor::TensorDialect,
-                    vector::VectorDialect, scf::SCFDialect,
+                    vector::VectorDialect, scf::SCFDialect, StandardOpsDialect,
                     arith::ArithmeticDialect, AffineDialect>();
     affine_ext::registerBufferizableOpInterfaceExternalModels(registry);
     arith_ext::registerBufferizableOpInterfaceExternalModels(registry);
     bufferization_ext::registerBufferizableOpInterfaceExternalModels(registry);
     linalg_ext::registerBufferizableOpInterfaceExternalModels(registry);
     scf_ext::registerBufferizableOpInterfaceExternalModels(registry);
+    std_ext::registerBufferizableOpInterfaceExternalModels(registry);
     tensor_ext::registerBufferizableOpInterfaceExternalModels(registry);
     vector_ext::registerBufferizableOpInterfaceExternalModels(registry);
   }
 
   void runOnFunction() override;
 
   Option<bool> allowReturnMemref{
       *this, "allow-return-memref",

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":IR",
         ":SCFDialect",
         ":Support",
         "//llvm:Support",
     ],
 )
 
 cc_library(
+    name = "StdBufferizableOpInterfaceImpl",
+    srcs = [
+        "lib/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.cpp",
+    ],
+    hdrs = [
+        "include/mlir/Dialect/Linalg/ComprehensiveBufferize/StdInterfaceImpl.h",
+    ],
+    includes = ["include"],
+    deps = [
+        ":BufferizableOpInterface",
+        ":IR",
+        ":StandardOps",
+        ":Support",
+        "//llvm:Support",
+    ],
+)
+
+cc_library(
     name = "TensorBufferizableOpInterfaceImpl",
     srcs = [
         "lib/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.cpp",
     ],
     hdrs = [
         "include/mlir/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.h",
     ],
     includes = ["include"],
deps = [
         ":MathDialect",
         ":MemRefDialect",
         ":Pass",
         ":SCFBufferizableOpInterfaceImpl",
         ":SCFDialect",
         ":SCFTransforms",
         ":StandardOps",
         ":StandardOpsTransforms",
+        ":StdBufferizableOpInterfaceImpl",
         ":Support",
         ":TensorBufferizableOpInterfaceImpl",
         ":TensorDialect",
         ":TransformUtils",
         ":VectorBufferizableOpInterfaceImpl",
         ":VectorOps",
         ":VectorToSCF",
         ":X86VectorTransforms",

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

deps = [
         "//mlir:LinalgOps",
         "//mlir:LinalgTransforms",
         "//mlir:MemRefDialect",
         "//mlir:Pass",
         "//mlir:SCFBufferizableOpInterfaceImpl",
         "//mlir:SCFDialect",
         "//mlir:SCFTransforms",
         "//mlir:StandardOps",
+        "//mlir:StdBufferizableOpInterfaceImpl",
         "//mlir:TensorBufferizableOpInterfaceImpl",
         "//mlir:TensorDialect",
         "//mlir:TransformUtils",
         "//mlir:VectorBufferizableOpInterfaceImpl",
         "//mlir:VectorOps",
         "//mlir:VectorToSCF",
     ],
 )