Address last remaining comments

Author: christopherbate

mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h

@@ -19,10 +19,6 @@
 #include "mlir/Interfaces/SubsetOpInterface.h"
 #include "llvm/Support/Debug.h"
 
-namespace mlir::bufferization::detail {
-bool tensorTypesMatchUpToEncoding(Type lhs, Type rhs);
-} // namespace mlir::bufferization::detail
-
 //===----------------------------------------------------------------------===//
 // Bufferization Dialect
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

@@ -23,16 +23,6 @@ using namespace mlir::bufferization;
 // Helper functions
 //===----------------------------------------------------------------------===//
 
-bool bufferization::detail::tensorTypesMatchUpToEncoding(Type lhs, Type rhs) {
-  auto lhsType = cast<ShapedType>(lhs);
-  auto rhsType = cast<ShapedType>(rhs);
-  if (lhsType.getElementType() != rhsType.getElementType())
-    return false;
-  if (lhsType.hasRank() && rhsType.hasRank())
-    return lhsType.getShape() == rhsType.getShape();
-  return true;
-}
-
 FailureOr<Value> mlir::bufferization::castOrReallocMemRefValue(
     OpBuilder &b, Value value, MemRefType destType,
     const BufferizationOptions &options) {

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

@@ -147,19 +147,31 @@ struct OneShotBufferizePass
       opt.dumpAliasSets = dumpAliasSets;
       opt.setFunctionBoundaryTypeConversion(
           parseLayoutMapOption(functionBoundaryTypeConversion));
+
+      if (mustInferMemorySpace && useEncodingForMemorySpace) {
+        emitError(getOperation()->getLoc())
+            << "only one of 'must-infer-memory-space' and "
+               "'use-encoding-for-memory-space' are allowed in "
+            << getArgument();
+        return signalPassFailure();
+      }
+
       if (mustInferMemorySpace) {
         opt.defaultMemorySpaceFn =
             [](TensorType t) -> std::optional<Attribute> {
           return std::nullopt;
         };
-      } else if (useEncodingForMemorySpace) {
+      }
+
+      if (useEncodingForMemorySpace) {
         opt.defaultMemorySpaceFn =
             [](TensorType t) -> std::optional<Attribute> {
           if (auto rtt = dyn_cast<RankedTensorType>(t))
             return rtt.getEncoding();
           return std::nullopt;
         };
       }
+
       opt.printConflicts = printConflicts;
       opt.bufferAlignment = bufferAlignment;
       opt.testAnalysisOnly = testAnalysisOnly;