[mlir][linalg][bufferize] Add dialect filter to BufferizationOptions

[springerm]

Summary

This change adds a new option dialectFilter to BufferizationOptions. Only ops from dialects that are allow-listed in the filter are bufferized. Other ops are left unbufferized. Note: This option requires allowUnknownOps = true.

To make use of dialectFilter, BufferizationOptions or BufferizationState must be passed to various helper functions.

The purpose of this change is to provide a better infrastructure for partial bufferization, which will be fully activated in a subsequent change.

Depends On D114671

[springerm]

update

[springerm]

update

[springerm]

update

[nicolasvasilache]

Please move this to test passes, similar to the parent commit.
Once this is more fleshed out they can replace the existing "partial" bufferization passes (and shoule be renamed to partial at that point).

[springerm]

update

[springerm]

update

[springerm]

rebase

[springerm]

address comments

[springerm]

add test case

[rriddle]

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

97

I've always found AllowListed to be an awkward mouthful, why not just isOpAllowed/isOpSupported/etc.

98–100

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

193

Leftover debugging here.

[springerm]

update

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

 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 // TODO: from some HW description.
 static constexpr int64_t kBufferAlignments = 128;
 
 class BufferizationAliasInfo;
+class BufferizableOpInterface;
 struct BufferizationOptions;
 class BufferizationState;
 struct PostAnalysisStep;
 
 /// Callback functions that are used to allocate/deallocate/copy memory buffers.
 /// Comprehensive Bufferize provides default implementations of these functions.
 // TODO: Could be replaced with a "bufferization strategy" object with virtual
 // functions in the future.
struct BufferizationOptions {
   /// Register a "post analysis" step. Such steps are executed after the
   /// analysis, but before bufferization.
   template <typename Step, typename... Args>
   void addPostAnalysisStep(Args... args) {
     postAnalysisSteps.emplace_back(
         std::make_unique<Step>(std::forward<Args>(args)...));
   }
 
+  /// Return `true` if the op is allowed to be bufferized.
+  bool isOpAllowed(Operation *op) const {
+    if (!dialectFilter.hasValue())
+      return true;
+    return dialectFilter->contains(op->getDialect()->getNamespace());
+  }
+
+  /// Allow-list the given dialects in the dialect filter. Only ops from
+  /// allow-listed dialects will be bufferized.
+  template <typename... DialectTs>
+  void addToDialectFilter() {
+    // The following expands a call to addToDialectFilterImpl for each dialect
+    // in 'DialectTs'. This magic is necessary due to a limitation in the places
+    // that a parameter pack can be expanded in c++11.
+    // FIXME: In c++17 this can be simplified by using 'fold expressions'.
+    (void)std::initializer_list<int>{
+        0, (addToDialectFilterImpl<DialectTs>(), 0)...};
+  }
+
+  /// Try to cast the given op to BufferizableOpInterface if the op is allow
+  /// listed.
+  BufferizableOpInterface dynCastBufferizableOp(Operation *op) const;
+
+  /// Try to cast the given value to BufferizableOpInterface if the op is allow
+  /// listed.
+  BufferizableOpInterface dynCastBufferizableOp(Value value) const;
+
   /// Helper functions for allocation, deallocation, memory copying.
   std::unique_ptr<AllocationCallbacks> allocationFns;
 
   /// Specifies whether returning newly allocated memrefs should be allowed.
   /// Otherwise, a pass failure is triggered.
   bool allowReturnMemref = false;
 
   /// Specifies whether not bufferizable ops are allowed in the input. If so,
   /// bufferization.to_memref and bufferization.to_tensor ops are inserted at
   /// the boundaries.
   bool allowUnknownOps = false;
 
   /// Seed for the analysis fuzzer. If set to `0`, the fuzzer is deactivated.
   /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
 
   /// If set to `true`, does not modify the IR apart from adding attributes (for
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
   /// Registered post analysis steps.
   PostAnalysisStepList postAnalysisSteps;
+
+  /// Only bufferize ops from dialects that are allowed-listed by the filter.
+  /// All other ops are ignored. This option controls the scope of partial
+  /// bufferization.
+  ///
+  /// Note: If no filter is specified, all ops are bufferized (as long as they
+  /// implement BufferizableOpInterface). If a filter is specified,
+  /// `allowUnknownOps` should be enabled. Otherwise, bufferization would fail
+  /// when encountering an op that is forbidden by the filter.
+  Optional<DenseSet<StringRef>> dialectFilter;
+
+private:
+  /// Allow-list a dialect in the dialect filter.
+  template <typename DialectT>
+  void addToDialectFilterImpl() {
+    if (!dialectFilter.hasValue())
+      dialectFilter.emplace();
+    dialectFilter->insert(DialectT::getDialectNamespace());
+  }
 };
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
   None,
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,
   Equivalent
 };
 
 /// The BufferizationAliasInfo class maintains a list of buffer aliases and
 /// equivalence classes to support bufferization.
 class BufferizationAliasInfo {
 public:
-  explicit BufferizationAliasInfo(Operation *rootOp);
+  explicit BufferizationAliasInfo(Operation *rootOp,
+                                  const BufferizationOptions &options);
 
   // BufferizationAliasInfo should be passed as a reference.
   BufferizationAliasInfo(const BufferizationAliasInfo &) = delete;
 
   /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
   /// beginning the alias and equivalence sets only contain `v` itself.
   void createAliasInfoEntry(Value v);
 

 ///   +----------+----------+----------+
 ///              |
 ///              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(Value value,
+findValueInReverseUseDefChain(Value value, const BufferizationOptions &options,
                               std::function<bool(Value)> condition);
 
 /// Find the Value 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);
+Value findLastPrecedingWrite(Value value, const BufferizationOptions &options);
 
 /// Dialect-specific bufferization state. Analysis/bufferization information
 /// that is specific to ops from a certain dialect can be stored in derived
 /// variants of this struct.
 struct DialectBufferizationState {
   virtual ~DialectBufferizationState() = default;
 };
 
 /// BufferizationState keeps track of memory buffers and provides a variety of
 /// helper functions for dealing with them. In particular,
 /// `BufferizableOpInterface::bufferize` implementation should utilize the
 /// following helper functions.
 ///
 /// * `createAlloc` / `createDealloc` / `createAllocDeallocPair` creates ops
 ///   that allocate and/or deallocate memref buffers.
 /// * `mapBuffer` maps a tensor value to a memref buffer during bufferization.
 /// * `lookupBuffer` returns the mapped memref buffer of a given tensor value.
 /// * `getResultBuffer` returns the memref buffer for a given tensor OpResult.
 ///   Based on inplace bufferization decisions of the analysis, it may either
 ///   directly return a mapped buffer or allocate a new brand new buffer.
 class BufferizationState {
 public:
   BufferizationState(Operation *op, const BufferizationOptions &options)
-      : aliasInfo(op), options(options), builder(op->getContext()) {}
+      : aliasInfo(op, options), options(options), builder(op->getContext()) {}
 
   // BufferizationState should be passed as a reference.
   BufferizationState(const BufferizationState &) = delete;
 
   /// Creates a memref allocation.
   Optional<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
                               ArrayRef<Value> dynShape);
 

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

let methods = [
           If this method returns `true` the given OpOperands are not considered
           to be conflicting and do not force out-of-place bufferization. (There
           may still be other conflicts that do.)
         }],
         /*retType=*/"bool",
         /*methodName=*/"isNotConflicting",
         /*args=*/(ins "OpOperand *":$uRead,
                       "OpOperand *":$uWrite,
+                      "BufferizationState &":$state,
                       "const BufferizationAliasInfo &":$aliasInfo),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return false;
         }]
       >
   ];
 

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

 // callbacks to their default functions.
 BufferizationOptions::BufferizationOptions()
     : allocationFns(defaultAllocationCallbacks()) {}
 
 //===----------------------------------------------------------------------===//
 // BufferizationAliasInfo
 //===----------------------------------------------------------------------===//
 
-BufferizationAliasInfo::BufferizationAliasInfo(Operation *rootOp) {
+BufferizationAliasInfo::BufferizationAliasInfo(
+    Operation *rootOp, const BufferizationOptions &options) {
   rootOp->walk([&](Operation *op) {
     for (Value v : op->getResults())
       if (v.getType().isa<TensorType>())
         createAliasInfoEntry(v);
     for (Region &r : op->getRegions())
       for (Block &b : r.getBlocks())
         for (auto bbArg : b.getArguments())
           if (bbArg.getType().isa<TensorType>())
             createAliasInfoEntry(bbArg);
   });
 
   // Set up alias sets for OpResults that must bufferize in-place. This should
   // be done before making any other bufferization decisions.
   rootOp->walk([&](BufferizableOpInterface bufferizableOp) {
+    if (!options.isOpAllowed(bufferizableOp))
+      return WalkResult::skip();
     for (OpResult opResult : bufferizableOp->getOpResults()) {
       if (opResult.getType().isa<TensorType>())
         if (bufferizableOp.mustBufferizeInPlace(opResult)) {
           SmallVector<OpOperand *> operands =
               bufferizableOp.getAliasingOpOperand(opResult);
           assert(!operands.empty() &&
                  "expected that OpResult has aliasing OpOperand");
           for (OpOperand *operand : operands)
             aliasInfo.unionSets(operand->get(), opResult);
           markInPlace(opResult);
         }
     }
+    return WalkResult::advance();
   });
 }
 
 /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
 /// beginning the alias and equivalence sets only contain `v` itself.
 void BufferizationAliasInfo::createAliasInfoEntry(Value v) {
   aliasInfo.insert(v);
   equivalentInfo.insert(v);

 static void setInsertionPointAfter(OpBuilder &b, Value value) {
   if (auto bbArg = value.dyn_cast<BlockArgument>()) {
     b.setInsertionPointToStart(bbArg.getOwner());
   } else {
     b.setInsertionPointAfter(value.getDefiningOp());
   }
 }
 
+BufferizableOpInterface mlir::linalg::comprehensive_bufferize::
+    BufferizationOptions::dynCastBufferizableOp(Operation *op) const {
+  if (isOpAllowed(op))
+    return dyn_cast<BufferizableOpInterface>(op);
+  return nullptr;
+}
+
+BufferizableOpInterface mlir::linalg::comprehensive_bufferize::
+    BufferizationOptions::dynCastBufferizableOp(Value value) const {
+  if (auto bufferizableOp = value.getDefiningOp<BufferizableOpInterface>())
+    if (isOpAllowed(bufferizableOp.getOperation()))
+      return bufferizableOp;
+  return nullptr;
+}
+
 /// Determine which OpOperand* will alias with `result` if the op is bufferized
 /// in place. Return an empty vector if the op is not bufferizable.
 SmallVector<OpOperand *>
 mlir::linalg::comprehensive_bufferize::getAliasingOpOperand(OpResult result) {
   if (Operation *op = result.getDefiningOp())
     if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
       return bufferizableOp.getAliasingOpOperand(result);
   return {};

 }
 
 // Starting from `value`, follow the use-def chain in reverse, always selecting
 // the aliasing OpOperands. Find and return Values for which `condition`
 // evaluates to true. OpOperands of such matching Values are not traversed any
 // further.
 llvm::SetVector<Value>
 mlir::linalg::comprehensive_bufferize::findValueInReverseUseDefChain(
-    Value value, std::function<bool(Value)> condition) {
+    Value value, const BufferizationOptions &options,
+    std::function<bool(Value)> condition) {
   llvm::SetVector<Value> result, workingSet;
   workingSet.insert(value);
 
   while (!workingSet.empty()) {
     Value value = workingSet.pop_back_val();
     if (condition(value) || value.isa<BlockArgument>()) {
       result.insert(value);
       continue;
     }
 
     OpResult opResult = value.cast<OpResult>();
     SmallVector<OpOperand *> opOperands = getAliasingOpOperand(opResult);
-    if (opOperands.empty()) {
+    if (opOperands.empty() || !options.isOpAllowed(value.getDefiningOp())) {
       result.insert(value);
       continue;
     }
 
     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::findLastPrecedingWrite(
-    Value value) {
+    Value value, const BufferizationOptions &options) {
   SetVector<Value> result =
-      findValueInReverseUseDefChain(value, [](Value value) {
+      findValueInReverseUseDefChain(value, options, [&](Value value) {
         Operation *op = value.getDefiningOp();
         if (!op)
           return true;
-        auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
+        auto bufferizableOp = options.dynCastBufferizableOp(op);
         if (!bufferizableOp)
           return true;
         return bufferizableOp.isMemoryWrite(value.cast<OpResult>());
       });
 
   // 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
if (!aliasInfo.isInPlace(result)) {
     // Allocate the result buffer.
     Value resultBuffer = createAllocDeallocPair(builder, loc, operandBuffer);
     bool skipCopy = false;
     // Do not copy if the last preceding write of `operand` is an op that does
     // 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);
-    if (auto bufferizableOp =
-            lastWrite.getDefiningOp<BufferizableOpInterface>())
+    Value lastWrite = findLastPrecedingWrite(operand, options);
+    if (auto bufferizableOp = options.dynCastBufferizableOp(lastWrite))
       if (!bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>()))
         skipCopy = true;
     // Do not copy if the copied data is never read.
     if (!isValueRead(result))
       skipCopy = true;
     // Do not copy if this op does not read the data, but writes it.
     if (bufferizesToMemoryWrite(*opOperand) &&
         !bufferizesToMemoryRead(*opOperand))
mlir::linalg::comprehensive_bufferize::bufferize(Operation *op,
   bool hasRegions = !op->getRegions().empty();
 
   // No tensor results/operands or regions. We are done.
   if (!hasTensorResult && !hasTensorOperand && !hasRegions)
     return success();
 
   // Bufferize using `BufferizableOpInterface`. Interface implementations are
   // responsible for bufferizing nested ops.
-  if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op)) {
+  if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op)) {
     b.setInsertionPoint(op);
     return bufferizableOp.bufferize(b, state);
   }
 
   // `op` is an unbufferizable tensor op.
   if (!state.getOptions().allowUnknownOps)
     return op->emitError() << "unsupported op with tensors";
 
Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupBuffer(
   // TODO: if key comes from bbArg, forward.
   assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
   Value buffer = mapping.lookupOrNull(tensor);
 
   if (!buffer) {
     if (options.allowUnknownOps) {
       // `tensor` was not bufferized yet. This should never happen with
       // bufferizable ops.
-      assert(!tensor.getDefiningOp<BufferizableOpInterface>() &&
-             "tensor is not mapped");
+      assert(!options.dynCastBufferizableOp(tensor) && "tensor is not mapped");
       // Insert to_memref op.
       OpBuilder b(tensor.getContext());
       setInsertionPointAfter(b, tensor);
       return b.create<bufferization::ToMemrefOp>(
           tensor.getLoc(),
           getDynamicMemRefType(tensor.getType().cast<RankedTensorType>()),
           tensor);
     }

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

 static bool aliasesNonWritableBuffer(Value value,
                                      const BufferizationAliasInfo &aliasInfo,
                                      BufferizationState &state) {
   LDBG("WRITABILITY ANALYSIS FOR " << printValueInfo(value) << "\n");
   bool foundNonWritableBuffer = false;
   aliasInfo.applyOnAliases(value, [&](Value v) {
     // Query BufferizableOpInterface to see if the OpResult is writable.
     // TODO: Out-of-place bufferized OpResult could be considered writable.
-    if (auto bufferizableOp = v.getDefiningOp<BufferizableOpInterface>())
+    if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(v))
       if (bufferizableOp && bufferizableOp.isWritable(v, state))
         return;
 
     // Query BufferizableOpInterface to see if the BlockArgument is writable.
     if (auto bbArg = v.dyn_cast<BlockArgument>())
-      if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(
+      if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(
               bbArg.getOwner()->getParentOp()))
         if (bufferizableOp.isWritable(bbArg, state))
           return;
 
     foundNonWritableBuffer = true;
   });
 
   if (foundNonWritableBuffer)

 
 /// Given sets of uses and writes, return true if there is a RaW conflict under
 /// the assumption that all given reads/writes alias the same buffer and that
 /// all given writes bufferize inplace.
 ///
 /// A conflict is: According to SSA use-def chains, a read R is supposed to read
 /// the result of a write W1. But because of bufferization decisions, R actually
 /// reads another write W2.
-static bool
-hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,
-                              const DenseSet<OpOperand *> &usesWrite,
-                              const DominanceInfo &domInfo,
-                              const BufferizationAliasInfo &aliasInfo) {
+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.:
     //
     // %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 = findLastPrecedingWrite(uRead->get());
+    Value lastWrite = findLastPrecedingWrite(uRead->get(), options);
 
     // 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) {
 
       // No conflict if the reading use equals the use of the conflicting write.
       // A use cannot conflict with itself. Note: Just being the same op is not
       // enough. It has to be the same use.
       if (uConflictingWrite == uRead)
         continue;
 
       // No conflict if the op interface says so.
-      if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(readingOp))
-        if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite,
+      if (auto bufferizableOp = options.dynCastBufferizableOp(readingOp))
+        if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state,
                                             aliasInfo))
           continue;
 
       if (conflictingWritingOp != readingOp)
         if (auto bufferizableOp =
-                dyn_cast<BufferizableOpInterface>(conflictingWritingOp))
-          if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite,
+                options.dynCastBufferizableOp(conflictingWritingOp))
+          if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state,
                                               aliasInfo))
             continue;
 
       // Ops are not conflicting if they are in mutually exclusive regions.
       if (insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp))
         continue;
 
       LDBG("WRITE = #" << printValueInfo(lastWrite) << "\n");

 /// read-after-write conflict without bufferizing `operand` inplace. This would
 /// indicate a problem with the current inplace bufferization decisions.
 ///
 /// Note: If `checkConsistencyOnly`, this function may be called with a null
 /// OpResult. In that case, only the consistency of bufferization decisions
 /// involving aliases of the given OpOperand are checked.
 bool wouldCreateReadAfterWriteInterference(
     OpOperand &operand, OpResult result, const DominanceInfo &domInfo,
-    const BufferizationAliasInfo &aliasInfo,
+    BufferizationState &state, const BufferizationAliasInfo &aliasInfo,
     bool checkConsistencyOnly = false) {
 #ifndef NDEBUG
   if (result) {
     SmallVector<OpOperand *> opOperands = getAliasingOpOperand(result);
     assert(llvm::find(opOperands, &operand) != opOperands.end() &&
            "operand and result do not match");
   } else {
     assert(checkConsistencyOnly &&
#endif // NDEBUG
   if (result)
     getAliasingReads(usesRead, result);
   getAliasingInplaceWrites(usesWrite, operand.get());
   if (result)
     getAliasingInplaceWrites(usesWrite, result);
   if (!checkConsistencyOnly && bufferizesToMemoryWrite(operand))
     usesWrite.insert(&operand);
 
-  return hasReadAfterWriteInterference(usesRead, usesWrite, domInfo, aliasInfo);
+  return hasReadAfterWriteInterference(usesRead, usesWrite, domInfo, state,
+                                       aliasInfo);
 }
 
 /// Return true if bufferizing `opOperand` inplace with `opResult` would create
 /// a write to a non-writable buffer.
 static bool
 wouldCreateWriteToNonWritableBuffer(OpOperand &opOperand, OpResult opResult,
                                     const BufferizationAliasInfo &aliasInfo,
                                     BufferizationState &state) {
#endif // NDEBUG
   (void)resultNumber;
   LDBG('\n');
   LDBG("Inplace analysis for <- #" << resultNumber << " -> #"
                                    << operand.getOperandNumber() << " in "
                                    << printValueInfo(result) << '\n');
 
   bool foundInterference =
       wouldCreateWriteToNonWritableBuffer(operand, result, aliasInfo, state) ||
-      wouldCreateReadAfterWriteInterference(operand, result, domInfo,
+      wouldCreateReadAfterWriteInterference(operand, result, domInfo, state,
                                             aliasInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(result);
   else
     aliasInfo.bufferizeInPlace(result, operand);
 
   LDBG("Done inplace analysis for result #" << resultNumber << '\n');
if (analysisFuzzerSeed) {
     std::mt19937 g(analysisFuzzerSeed);
     llvm::shuffle(ops.begin(), ops.end(), g);
   }
 
   // Walk ops in reverse for better interference analysis.
   for (Operation *op : reverse(ops))
     for (OpOperand &opOperand : op->getOpOperands())
       if (opOperand.get().getType().isa<TensorType>())
-        if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+        if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
           if (OpResult opResult = bufferizableOp.getAliasingOpResult(opOperand))
             if (failed(bufferizableInPlaceAnalysisImpl(
                     opOperand, opResult, aliasInfo, state, domInfo)))
               return failure();
 
   return success();
 }
 
op->walk([&](Operation *op) {
     ops.push_back(op);
   });
 
   return inPlaceAnalysis(ops, aliasInfo, state, domInfo, analysisFuzzerSeed);
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.
 static void equivalenceAnalysis(SmallVector<Operation *> &ops,
-                                BufferizationAliasInfo &aliasInfo) {
+                                BufferizationAliasInfo &aliasInfo,
+                                const BufferizationOptions &options) {
   for (Operation *op : ops)
-    if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+    if (auto bufferizableOp = options.dynCastBufferizableOp(op))
       for (OpResult opResult : op->getOpResults())
         if (opResult.getType().isa<TensorType>())
           if (aliasInfo.isInPlace(opResult)) {
             SmallVector<OpOperand *> opOperands =
                 bufferizableOp.getAliasingOpOperand(opResult);
             if (!opOperands.empty())
               if (bufferizableOp.bufferRelation(opResult, aliasInfo) ==
                   BufferRelation::Equivalent)
                 for (OpOperand *opOperand : opOperands)
                   aliasInfo.unionEquivalenceClasses(opResult, opOperand->get());
           }
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of all ops contained
 /// in `op`.
 static void equivalenceAnalysis(Operation *op,
-                                BufferizationAliasInfo &aliasInfo) {
+                                BufferizationAliasInfo &aliasInfo,
+                                const BufferizationOptions &options) {
   // Traverse ops in PostOrder: Nested ops first, then enclosing ops.
   SmallVector<Operation *> ops;
   op->walk<WalkOrder::PostOrder>([&](Operation *op) {
     // No tensors => no buffers.
     if (none_of(op->getResultTypes(), isaTensor))
       return;
     ops.push_back(op);
   });
 
-  equivalenceAnalysis(ops, aliasInfo);
+  equivalenceAnalysis(ops, aliasInfo, options);
 }
 
 /// Assert that the current bufferization decisions are consistent.
 static LogicalResult
 checkAliasInfoConsistency(Operation *op, const DominanceInfo &domInfo,
+                          BufferizationState &state,
                           const BufferizationAliasInfo &aliasInfo) {
+  const BufferizationOptions &options = state.getOptions();
   Operation *inconsistentOp = nullptr;
   WalkResult walkResult = op->walk([&](Operation *op) {
-    if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+    if (auto bufferizableOp = options.dynCastBufferizableOp(op))
       for (OpOperand &opOperand : op->getOpOperands())
         if (opOperand.get().getType().isa<TensorType>()) {
           OpResult opResult = bufferizableOp.getAliasingOpResult(opOperand);
           if (wouldCreateReadAfterWriteInterference(
-                  opOperand, opResult, domInfo, aliasInfo,
+                  opOperand, opResult, domInfo, state, aliasInfo,
                   /*checkConsistencyOnly=*/true)) {
             // This error can happen for two reasons. Either the input IR
             // already has a read-after-write conflict. Or certain
             // "mustBufferizeInPlace" interface methods are implemented
             // incorrectly.
             inconsistentOp = op;
             return WalkResult::interrupt();
           }

 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     Operation *op, const BufferizationOptions &options,
     BufferizationState &state, const PostAnalysisStepList &extraSteps) {
 
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
-  if (failed(checkAliasInfoConsistency(op, domInfo, aliasInfo)))
+  if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
   if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
                              options.analysisFuzzerSeed)))
     return failure();
-  equivalenceAnalysis(op, aliasInfo);
+  equivalenceAnalysis(op, aliasInfo, options);
 
   auto runPostAnalysisSteps = [&](const PostAnalysisStepList &steps) {
     for (const std::unique_ptr<PostAnalysisStep> &step : steps) {
       SmallVector<Operation *> newOps;
       if (failed(step->run(op, state, aliasInfo, newOps)))
         return failure();
       // Analyze ops that were created by the PostAnalysisStep.
       if (failed(inPlaceAnalysis(newOps, aliasInfo, state, domInfo)))
         return failure();
-      equivalenceAnalysis(newOps, aliasInfo);
+      equivalenceAnalysis(newOps, aliasInfo, options);
     }
     return success();
   };
 
   if (failed(runPostAnalysisSteps(extraSteps)))
     return failure();
   if (failed(runPostAnalysisSteps(options.postAnalysisSteps)))
     return failure();

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

 LogicalResult mlir::linalg::comprehensive_bufferize::linalg_ext::
     InitTensorEliminationStep::eliminateInitTensors(
         Operation *op, BufferizationState &state,
         BufferizationAliasInfo &aliasInfo,
         std::function<bool(OpOperand &)> anchorMatchFunc,
         std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
         SmallVector<Operation *> &newOps) {
   OpBuilder b(op->getContext());
+  const BufferizationOptions &options = state.getOptions();
 
   WalkResult status = op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Is this a matching OpOperand?
       if (!anchorMatchFunc(operand))
         continue;
 
       SetVector<Value> maybeInitTensor =
-          findValueInReverseUseDefChain(operand.get(), [&](Value val) {
+          findValueInReverseUseDefChain(operand.get(), options, [&](Value val) {
             // Continue traversal until this function returns true.
             OpResult opResult = val.dyn_cast<OpResult>();
             if (!opResult)
               return true;
             if (!aliasInfo.isInPlace(opResult))
               return true;
             // Only equivalent tensors are supported at the moment.
             // TODO: Support cases such as extract_slice(init_tensor).

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

struct ExtractOpInterface
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto extractOp = cast<tensor::ExtractOp>(op);
     Location loc = extractOp.getLoc();
     Value srcMemref = state.lookupBuffer(extractOp.tensor());
     Value l = b.create<memref::LoadOp>(loc, srcMemref, extractOp.indices());
     extractOp.replaceAllUsesWith(l);
     return success();
   }
 };
 
aliasInfo.applyOnEquivalenceClass(insertSliceOp.source(), [&](Value value) {
     }
   });
   return foundOp;
 }
 
 /// Return true if `value` is originating from an ExtractSliceOp that matches
 /// the given InsertSliceOp.
 static bool hasMatchingExtractSliceOp(const BufferizationAliasInfo &aliasInfo,
+                                      const BufferizationOptions &options,
                                       Value value, InsertSliceOp insertOp) {
   auto condition = [&](Value val) {
     if (auto extractOp = val.getDefiningOp<ExtractSliceOp>())
       if (areEquivalentExtractSliceOps(aliasInfo, extractOp, insertOp))
         return true;
     return false;
   };
 
-  return llvm::all_of(findValueInReverseUseDefChain(value, condition),
+  return llvm::all_of(findValueInReverseUseDefChain(value, options, condition),
                       condition);
 }
 
 struct InsertSliceOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertSliceOpInterface,
                                                     tensor::InsertSliceOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;
struct InsertSliceOpInterface
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const BufferizationAliasInfo &aliasInfo) const {
     return BufferRelation::Equivalent;
   }
 
   bool isNotConflicting(Operation *op, OpOperand *uRead,
-                        OpOperand *uConflictingWrite,
+                        OpOperand *uConflictingWrite, BufferizationState &state,
                         const BufferizationAliasInfo &aliasInfo) const {
     Operation *readingOp = uRead->getOwner();
     Operation *conflictingWritingOp = uConflictingWrite->getOwner();
 
     // Special rules for matching ExtractSliceOp/InsertSliceOp pairs. If
     // uRead is an InsertSliceOp...
     if (auto insertSliceOp = dyn_cast<InsertSliceOp>(readingOp)) {
       // As an example, consider the following IR.
       //
       // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
       // %1 = linalg.fill %cst, %0 {inplace= [true] }
       // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
       //     {inplace= [true] }
 
       // TODO: Use insertSliceOp.getDestOpOperand etc. when available.
       if (uRead == &insertSliceOp->getOpOperand(1) /*dest*/ &&
-          hasMatchingExtractSliceOp(aliasInfo, uConflictingWrite->get(),
-                                    insertSliceOp))
+          hasMatchingExtractSliceOp(aliasInfo, state.getOptions(),
+                                    uConflictingWrite->get(), insertSliceOp))
         // Case 1: The main insight is that InsertSliceOp reads only part of
         // the destination tensor. The overwritten area is not read. If
         // uConflictingWrite writes into exactly the memory location that is
         // being read by uRead, this is not a conflict.
         //
         // In the above example:
         // uRead             = OpOperand 1 (%t) of tensor.insert_slice
         // uConflictingWrite = OpOperand 1 (%0) of linalg.fill
         //
         // The read of %t does not conflict with the write of the FillOp
         // (same aliases!) because the area that the FillOp operates on is
         // exactly the one that is *not* read via %t.
         return true;
 
       if (uRead == &insertSliceOp->getOpOperand(0) /*source*/ &&
           uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
-          hasMatchingExtractSliceOp(aliasInfo, uRead->get(), insertSliceOp))
+          hasMatchingExtractSliceOp(aliasInfo, state.getOptions(), uRead->get(),
+                                    insertSliceOp))
         // Case 2: The read of the source tensor and the write to the dest
         // tensor via an InsertSliceOp is not a conflict if the read is
         // reading exactly that part of an equivalent tensor that the
         // InsertSliceOp is writing.
         //
         // In the above example:
         // uRead             = OpOperand 0 (%1) of tensor.insert_slice
         // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
if (auto insertSliceOp = dyn_cast<InsertSliceOp>(conflictingWritingOp))
       // lastWrite         = %1
       //
       // This is not a conflict because the InsertSliceOp overwrites the
       // memory segment of %1 with the exact same data. (Effectively, there
       // is no memory write here.)
       if (uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
           aliasInfo.areEquivalentBufferizedValues(uRead->get(),
                                                   insertSliceOp.source()) &&
-          hasMatchingExtractSliceOp(aliasInfo, insertSliceOp.source(),
-                                    insertSliceOp))
+          hasMatchingExtractSliceOp(aliasInfo, state.getOptions(),
+                                    insertSliceOp.source(), insertSliceOp))
         return true;
 
     return false;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     // insert_slice ops arise from tiling and bufferizing them out-of-place is

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

 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
 
 // TODO: Bufferize result IR of bufferization.
 // TODO: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops" -linalg-comprehensive-module-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
 
 // Run fuzzer with different seeds.
 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
 // RUN: mlir-opt %s -allow-unregistered-dialect -linalg-comprehensive-module-bufferize="test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
 
+// RUN: mlir-opt %s -allow-unregistered-dialect -test-comprehensive-function-bufferize="dialect-filter=tensor allow-unknown-ops allow-return-memref" -canonicalize -split-input-file | FileCheck %s --check-prefix=CHECK-TENSOR
+
 // CHECK-LABEL: func @use_of_unknown_op_1(
 //  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
 func @use_of_unknown_op_1(%t1: tensor<?xf32> {linalg.inplaceable = true})
     -> vector<5xf32> {
   // ToTensorOp is generated because the function is bufferized and has a
   // memref block argument.
   // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
   // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
func @unknown_op_not_writable(
   // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
   // CHECK: linalg.copy(%[[dummy_memref]], %[[alloc]])
   // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
   %1 = vector.transfer_write %v, %0[%idx] : vector<5xf32>, tensor<?xf32>
 
   // CHECK: return %[[alloc]]
   return %1 : tensor<?xf32>
 }
+
+// -----
+
+// CHECK-TENSOR-LABEL: func @simple_tensor_test(
+//  CHECK-TENSOR-SAME:     %[[t1:.*]]: tensor<?xf32>
+func @simple_tensor_test(%t1 : tensor<?xf32>, %f : f32) -> tensor<?xf32> {
+  // CHECK-TENSOR: %[[t1_memref:.*]] = bufferization.to_memref %[[t1]]
+  %c0 = arith.constant 0 : index
+  // CHECK-TENSOR: %[[alloc:.*]] = memref.alloc
+  // CHECK-TENSOR: %[[casted:.*]] = memref.cast %[[alloc]]
+  // CHECK-TENSOR: memref.copy %[[t1_memref]], %[[casted]]
+  // CHECK-TENSOR: memref.store %{{.*}}, %[[alloc]]
+  %0 = tensor.insert %f into %t1[%c0] : tensor<?xf32>
+  // CHECK-TENSOR: %[[casted_tensor:.*]] = bufferization.to_tensor %[[casted]]
+  // CHECK-TENSOR: return %[[casted_tensor]]
+  return %0 : tensor<?xf32>
+}

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

Option<bool> testAnalysisOnly{
       *this, "test-analysis-only",
       llvm::cl::desc(
           "Only runs inplaceability analysis (for testing purposes only)"),
       llvm::cl::init(false)};
   Option<unsigned> analysisFuzzerSeed{
       *this, "analysis-fuzzer-seed",
       llvm::cl::desc("Analyze ops in random order with a given seed (fuzzer)"),
       llvm::cl::init(0)};
+  ListOption<std::string> dialectFilter{
+      *this, "dialect-filter",
+      llvm::cl::desc("Bufferize only ops from the specified dialects"),
+      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
 };
 } // namespace
 
 void TestComprehensiveFunctionBufferize::runOnFunction() {
   BufferizationOptions options;
 
   // Enable InitTensorOp elimination.
   options.addPostAnalysisStep<
       linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   // TODO: Find a way to enable this step automatically when bufferizing
   // tensor dialect ops.
   options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
   options.addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
   options.allowReturnMemref = allowReturnMemref;
   options.allowUnknownOps = allowUnknownOps;
   options.testAnalysisOnly = testAnalysisOnly;
   options.analysisFuzzerSeed = analysisFuzzerSeed;
 
+  if (dialectFilter.hasValue()) {
+    options.dialectFilter.emplace();
+    for (const std::string &dialectNamespace : dialectFilter)
+      options.dialectFilter->insert(dialectNamespace);
+  }
+
   Operation *op = getFunction().getOperation();
   if (failed(runComprehensiveBufferize(op, options)))
     return;
 
   OpPassManager cleanupPipeline("builtin.func");
   cleanupPipeline.addPass(createCanonicalizerPass());
   cleanupPipeline.addPass(createCSEPass());
   cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());