[mlir][linalg][bufferize][NFC] Map only tensors in BufferizationState

[springerm]

Summary

BufferizationState had map/lookup overloads for non-tensor values. This was necessary for IREE. There is now a better way to do this, so these overloads can be removed.

Depends On D114928

[springerm]

update

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

struct BufferizationState {
   /// additional optimizations such as buffer allocation hoisting. This function
   /// calls `allocationFn` and `deallocationFn` to create (de)allocations.
   Value createAllocDeallocFn(OpBuilder &builder, Location loc,
                              Value shapedValue);
 
   /// Map tensor values to memref buffers.
   void mapBuffer(ValueRange tensors, ValueRange buffers);
 
-  /// Map a value to another value.
-  void mapValue(Value from, Value to);
-
   /// Map a tensor value to a memref buffer.
   void mapBuffer(Value tensor, Value buffer);
 
   /// Lookup the memref buffer that is associated to the given tensor value.
   /// Asserts if no buffer is associated.
   Value lookupBuffer(Value tensor);
 
-  /// Lookup the value that is associated to the given value. Asserts if no
-  /// value is associated.
-  Value lookupValue(Value value) const;
-
   /// Return `true` if the given value is mapped.
   bool isMapped(Value value) const;
 
   /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
   /// a new buffer and copy over data from the existing buffer if out-of-place
   /// bufferization is necessary.
   Value getResultBuffer(OpResult result);
 
template <typename StateT> StateT &getDialectState(StringRef name) {
     if (!dialectState.count(name))
       dialectState[name] = std::make_unique<StateT>();
     return static_cast<StateT &>(*dialectState[name]);
   }
 
   /// `aliasInfo` keeps track of aliasing and equivalent values.
   BufferizationAliasInfo aliasInfo;
 
-  /// The mapping of tensors to buffers. May also contain mappings of non-tensor
-  /// values.
+  /// The mapping of tensors to buffers.
   BlockAndValueMapping mapping;
 
   /// Obsolete ops that should be deleted after bufferization.
   SmallVector<Operation *> obsoleteOps;
 
   /// Dialect-specific bufferization state.
   DenseMap<StringRef, std::unique_ptr<DialectBufferizationState>> dialectState;
 

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

void mlir::linalg::comprehensive_bufferize::BufferizationState::mapBuffer(
   assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
   assert(buffer && "unexpected empty buffer");
   assert((buffer.getType().isa<MemRefType>() ||
           buffer.getType().isa<UnrankedMemRefType>()) &&
          "expected that tensor is mapped to memref");
   return mapping.map(tensor, buffer);
 }
 
-void mlir::linalg::comprehensive_bufferize::BufferizationState::mapValue(
-    Value from, Value to) {
-  assert(from && "unexpected empty value");
-  return mapping.map(from, to);
-}
-
 /// Wrapper for better debugging.
 Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupBuffer(
     Value tensor) {
   // TODO: if key comes from bbArg, forward.
   assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
   Value buffer = mapping.lookupOrNull(tensor);
 
   if (!buffer) {
Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupBuffer(
   }
 
   assert((buffer.getType().isa<MemRefType>() ||
           buffer.getType().isa<UnrankedMemRefType>()) &&
          "expected that tensor is mapped to memref");
   return buffer;
 }
 
-Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupValue(
-    Value value) const {
-  Value v = mapping.lookupOrNull(value);
-  if (!v) {
-    llvm_unreachable("tensor is not mapped");
-    return Value();
-  }
-  return v;
-}
-
 bool mlir::linalg::comprehensive_bufferize::BufferizationState::isMapped(
     Value value) const {
   return mapping.contains(value);
 }
 
 void mlir::linalg::comprehensive_bufferize::BufferizationState::markOpObsolete(
     Operation *op) {
   obsoleteOps.push_back(op);