[mlir][linalg][bufferize] Allow non-tensor mappings in BufferizationState

[springerm]

Summary

This change makes it possible to set up custom mappings in a PostAnalysisStep. Some users of Comprehensive Bufferize have custom tensor types and it is most convenient to just reuse the same bvm.

Also add some more assertions and add a mechanism to mark ops as obsolete.

Depends On D113723

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

 //===- BufferizableOpInterface.h - Comprehensive Bufferize ------*- C++ -*-===//
 //
 // 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_BUFFERIZABLEOPINTERFACE_H_
 #define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_
 
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {

 };
 
 /// The BufferizationAliasInfo class maintains a list of buffer aliases and
 /// equivalence classes to support bufferization.
 class BufferizationAliasInfo {
 public:
   explicit BufferizationAliasInfo(Operation *rootOp);
 
+  // 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);
 
   /// Insert an info entry for `newValue` and merge its alias set with that of
   /// `alias`.
   void insertNewBufferAlias(Value newValue, Value alias);
 
struct AllocationCallbacks {
   /// calls `allocationFn` and `deallocationFn` to create (de)allocations.
   CreateAllocDeallocFn createAllocDeallocFn;
 };
 
 /// BufferizationState keeps track of bufferization state and provides access to
 /// the results of the analysis.
 struct BufferizationState {
   BufferizationState(BufferizationAliasInfo &aliasInfo,
-                     AllocationCallbacks &allocationFns,
-                     BlockAndValueMapping &tensorToBufferMap)
-      : aliasInfo(aliasInfo), allocationFns(allocationFns),
-        tensorToBufferMap(tensorToBufferMap) {}
+                     AllocationCallbacks &allocationFns)
+      : aliasInfo(aliasInfo), allocationFns(allocationFns) {}
+
+  // BufferizationState should be passed as a reference.
+  BufferizationState(const BufferizationState &) = delete;
 
   /// 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) const;
 
+  /// 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;
+
+  /// Mark `op` as obsolete, so that it is deleted after bufferization.
+  void markOpObsolete(Operation *op);
+
   /// `aliasInfo` keeps track of aliasing and equivalent values.
   BufferizationAliasInfo &aliasInfo;
 
   /// `allocationFns` contains helper functions for creating alloc ops, dealloc
   /// ops and memcpy ops.
   AllocationCallbacks &allocationFns;
 
-  /// The mapping of tensors to buffers.
-  BlockAndValueMapping &tensorToBufferMap;
+  /// The mapping of tensors to buffers. May also contain mappings of non-tensor
+  /// values.
+  BlockAndValueMapping mapping;
+
+  /// Obsolete ops that should be deleted after bufferization.
+  SmallVector<Operation *> obsoleteOps;
 };
 
 /// 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(OpBuilder &b, OpResult result, BufferizationState &state);
 
 /// PostAnalysisSteps can be registered with `BufferizationOptions` and are

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

 //===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
 /// Wrapper for better debugging.
 void mlir::linalg::comprehensive_bufferize::BufferizationState::mapBuffer(
     ValueRange tensors, ValueRange buffers) {
   assert(!tensors.empty() && "unexpected empty tensors");
-  return tensorToBufferMap.map(tensors, buffers);
+#ifndef NDEBUG
+  for (Value tensor : tensors) {
+    assert(tensor && "unexpected empty tensor");
+    assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
+  }
+  for (Value buffer : buffers) {
+    assert(buffer && "unexpected empty buffer");
+    assert((buffer.getType().isa<MemRefType>() ||
+            buffer.getType().isa<UnrankedMemRefType>()) &&
+           "expected that tensor is mapped to memref");
+  }
+#endif // NDEBUG
+  return mapping.map(tensors, buffers);
 }
 
 /// Wrapper for better debugging.
 void mlir::linalg::comprehensive_bufferize::BufferizationState::mapBuffer(
     Value tensor, Value buffer) {
   assert(tensor && "unexpected empty tensor");
   assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
-  return tensorToBufferMap.map(tensor, buffer);
+  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) const {
   // TODO: if key comes from bbArg, forward.
   assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
-  Value v = tensorToBufferMap.lookupOrNull(tensor);
+  Value v = mapping.lookupOrNull(tensor);
 
   if (!v) {
     // Dump tensor for easier debugging.
     tensor.dump();
     llvm_unreachable("tensor is not mapped");
     return Value();
   }
 
+  assert((v.getType().isa<MemRefType>() ||
+          v.getType().isa<UnrankedMemRefType>()) &&
+         "expected that tensor is mapped to memref");
   return v;
 }
+
+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);
+}

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

for (const std::unique_ptr<PostAnalysisStep> &step :
         return failure();
       // Analyze ops that were created by the PostAnalysisStep.
       if (failed(inPlaceAnalysis(newOps, aliasInfo, domInfo)))
         return failure();
     }
 
     // Bufferization phase.
     if (!options.testAnalysisOnly) {
-      BlockAndValueMapping tensorToBufferMap;
-      BufferizationState state(aliasInfo, *options.allocationFns,
-                               tensorToBufferMap);
+      BufferizationState state(aliasInfo, *options.allocationFns);
+
+      // Bufferize all ops in funcOp.
       if (failed(
               bufferizeFuncOpInternals(funcOp, state, bufferizedFunctionTypes)))
         return failure();
+
+      // Erase all obsolete ops.
+      for (Operation *op : state.obsoleteOps)
+        op->erase();
     }
   }
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly) {
     annotateOpsWithBufferizationMarkers(moduleOp, aliasInfo);
     return success();
   }