[NFC] Remove type-inferring builder of ToTensorOp

The builder is ambiguous given customizable tensor-like -> buffer-like
conversion and is thus removed. The places where reverse bufferization
has to happen rely on the pre-existing functionality.

Author: andrey-golubev

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

@@ -748,10 +748,6 @@ FailureOr<BaseMemRefType> castToMemRef(FailureOr<BufferLikeType> bufferType);
 /// bufferization::ConversionInterface to verify the types in tensor and buffer
 /// worlds match.
 bool typesMatchAfterBufferization(Operation &op, Value tensor, Value buffer);
-
-/// This function is a free-standing helper that relies on
-/// bufferization::ConversionInterface to perform the conversion.
-Type getTensorFromBuffer(Type buffer);
 } // namespace detail
 
 } // namespace bufferization

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

@@ -35,10 +35,6 @@ struct ConversionDialectInterface
   virtual LogicalResult typesMatch(
       TensorLikeType tensor, BufferLikeType buffer,
       function_ref<InFlightDiagnostic(const Twine &)> emitError) const = 0;
-
-  /// Hook to customize buffer-like -> tensor-like conversion, which is the
-  /// opposite of bufferization.
-  virtual TensorLikeType getTensorFromBuffer(BufferLikeType buffer) const = 0;
 };
 
 /// Interface collection for conversion between tensor-like and buffer-like
@@ -60,10 +56,6 @@ struct ConversionInterface
   LogicalResult
   typesMatch(TensorLikeType tensor, BufferLikeType buffer,
              function_ref<InFlightDiagnostic(const Twine &)> emitError) const;
-
-  /// Dispatches to ConversionDialectInterface::getTensorFromBuffer() of the
-  /// dialect associated with the value type.
-  TensorLikeType getTensorFromBuffer(BufferLikeType buffer) const;
 };
 
 } // namespace bufferization

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

@@ -490,13 +490,6 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
       `:` type($buffer) `to` type($result)
   }];
 
-  let builders = [
-    OpBuilder<(ins "Value":$buffer, CArg<"bool", "false">:$restrict, CArg<"bool", "false">:$writeable), [{
-      auto rtt = bufferization::detail::getTensorFromBuffer(buffer.getType());
-      build($_builder, $_state, rtt, buffer, restrict, writeable);
-    }]>
-  ];
-
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }

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

@@ -172,7 +172,9 @@ FailureOr<Value> bufferization::allocateTensorForShapedValue(
   if (llvm::isa<RankedTensorType>(shapedValue.getType())) {
     tensor = shapedValue;
   } else if (llvm::isa<MemRefType>(shapedValue.getType())) {
-    tensor = b.create<ToTensorOp>(loc, shapedValue);
+    tensor = b.create<ToTensorOp>(
+        loc, memref::getTensorTypeFromMemRefType(shapedValue.getType()),
+        shapedValue);
   } else if (llvm::isa<UnrankedTensorType>(shapedValue.getType()) ||
              llvm::isa<UnrankedMemRefType>(shapedValue.getType())) {
     return getOwnerOfValue(shapedValue)
@@ -1064,9 +1066,3 @@ bool bufferization::detail::typesMatchAfterBufferization(Operation &op,
       cast<BufferLikeType>(buffer.getType()),
       [&](const Twine &message) { return op.emitError(message); }));
 }
-
-Type bufferization::detail::getTensorFromBuffer(Type buffer) {
-  assert(isa<BufferLikeType>(buffer) && "expected BufferLikeType");
-  bufferization::ConversionInterface iface(buffer.getContext());
-  return iface.getTensorFromBuffer(cast<BufferLikeType>(buffer));
-}

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

@@ -54,14 +54,5 @@ LogicalResult ConversionInterface::typesMatch(
   return success();
 }
 
-TensorLikeType
-ConversionInterface::getTensorFromBuffer(BufferLikeType buffer) const {
-  Dialect *dialect = &buffer.getDialect();
-  if (const ConversionDialectInterface *iface = getInterfaceFor(dialect))
-    return iface->getTensorFromBuffer(buffer);
-
-  return cast<TensorLikeType>(memref::getTensorTypeFromMemRefType(buffer));
-}
-
 } // namespace bufferization
 } // namespace mlir

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

@@ -643,8 +643,9 @@ Value MaterializeInDestinationOp::buildSubsetExtraction(OpBuilder &builder,
   assert(getRestrict() &&
          "expected that ops with memrefs dest have 'restrict'");
   setRestrict(false);
-  return builder.create<ToTensorOp>(loc, getDest(), /*restrict=*/true,
-                                    getWritable());
+  return builder.create<ToTensorOp>(
+      loc, memref::getTensorTypeFromMemRefType(getDest().getType()), getDest(),
+      /*restrict=*/true, getWritable());
 }
 
 bool MaterializeInDestinationOp::isEquivalentSubset(

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

@@ -252,7 +252,8 @@ Value linalg::bufferizeToAllocation(
   // Create bufferization.to_tensor with "restrict" and "writable". The returned
   // tensor is a new buffer allocation, so it does not alias with any buffer.
   Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
-      loc, alloc, /*restrict=*/true, /*writable=*/true);
+      loc, padOp.getResult().getType(), alloc, /*restrict=*/true,
+      /*writable=*/true);
   rewriter.replaceOp(padOp, toTensorOp);
   return alloc;
 }
@@ -340,7 +341,8 @@ Value linalg::bufferizeToAllocation(
   // Create bufferization.to_tensor with "restrict" and "writable". The returned
   // tensor is a new buffer allocation, so it does not alias with any buffer.
   Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
-      loc, alloc, /*restrict=*/true, /*writable=*/true);
+      loc, allocTensorOp.getResult().getType(), alloc, /*restrict=*/true,
+      /*writable=*/true);
   rewriter.replaceOp(allocTensorOp, toTensorOp);
   return alloc;
 }
@@ -567,7 +569,8 @@ Value linalg::bufferizeToAllocation(
       createMemcpy(rewriter, op->getLoc(), operand->get(), alloc, options);
     }
     rewriter.modifyOpInPlace(op, [&]() {
-      auto toTensorOp = rewriter.create<ToTensorOp>(op->getLoc(), alloc);
+      auto toTensorOp = rewriter.create<ToTensorOp>(
+          op->getLoc(), operand->get().getType(), alloc);
       operand->set(toTensorOp);
       if (options.bufferizeDestinationOnly) {
         rewriter.modifyOpInPlace(toTensorOp, [&]() {

mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp

@@ -67,7 +67,7 @@ struct AssumingOpInterface
     for (const auto &it : llvm::enumerate(assumingOp->getResultTypes())) {
       if (isa<TensorType>(it.value())) {
         newResults.push_back(rewriter.create<bufferization::ToTensorOp>(
-            assumingOp.getLoc(), newOp->getResult(it.index())));
+            assumingOp.getLoc(), it.value(), newOp->getResult(it.index())));
       } else {
         newResults.push_back(newOp->getResult(it.index()));
       }

mlir/lib/Dialect/SparseTensor/Transforms/SparseGPUCodegen.cpp

@@ -651,7 +651,7 @@ static LogicalResult rewriteSpMV(PatternRewriter &rewriter,
   tokens.clear();
 
   // Done.
-  rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, memY);
+  rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, y.getType(), memY);
   return success();
 }
 
@@ -752,7 +752,7 @@ static LogicalResult rewriteSpMM(PatternRewriter &rewriter,
   tokens.clear();
 
   // Done.
-  rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, bufC);
+  rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, c.getType(), bufC);
   return success();
 }
 
@@ -925,9 +925,12 @@ static LogicalResult rewriteSpGEMM(PatternRewriter &rewriter,
   tokens.clear();
 
   // Done.
-  Value vt = rewriter.create<bufferization::ToTensorOp>(loc, valH);
-  Value rt = rewriter.create<bufferization::ToTensorOp>(loc, rowH);
-  Value ct = rewriter.create<bufferization::ToTensorOp>(loc, colH);
+  Value vt = rewriter.create<bufferization::ToTensorOp>(
+      loc, memref::getTensorTypeFromMemRefType(valH.getType()), valH);
+  Value rt = rewriter.create<bufferization::ToTensorOp>(
+      loc, memref::getTensorTypeFromMemRefType(rowH.getType()), rowH);
+  Value ct = rewriter.create<bufferization::ToTensorOp>(
+      loc, memref::getTensorTypeFromMemRefType(colH.getType()), colH);
   rewriter.replaceOpWithNewOp<AssembleOp>(op, c.getType(), ValueRange{rt, ct},
                                           vt);
   return success();
@@ -1043,7 +1046,7 @@ static LogicalResult rewrite2To4SpMM(PatternRewriter &rewriter,
   tokens.clear();
 
   // Done.
-  rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, bufC);
+  rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, C.getType(), bufC);
   return success();
 }
 

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp

@@ -1471,7 +1471,8 @@ struct SparseDisassembleOpConverter
     // Converts MemRefs back to Tensors.
     SmallVector<Value> retValues = llvm::to_vector(
         llvm::map_range(retMem, [&rewriter, loc](Value v) -> Value {
-          return rewriter.create<bufferization::ToTensorOp>(loc, v);
+          return rewriter.create<bufferization::ToTensorOp>(
+              loc, memref::getTensorTypeFromMemRefType(v.getType()), v);
         }));
     // Appends the actual memory length used in each buffer returned.
     retValues.append(retLen.begin(), retLen.end());