From 80ce71187bf7d8e07383677eb2607788f2902f26 Mon Sep 17 00:00:00 2001
From: Yi Zhang <cathyzhyi@google.com>
Date: Tue, 28 Jan 2025 12:13:45 -0500
Subject: [PATCH] Use original type when convert arg for users

This will keep the memory space information for the tensor.
---
 mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 64d79f5b5d60c..313dbbd16984b 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -453,14 +453,15 @@ bufferization::bufferizeBlockSignature(Block *block, RewriterBase &rewriter,
     for (OpOperand &use : bbArg.getUses())
       bbArgUses.push_back(&use);
 
+    Type tensorType = bbArg.getType();
     // Change the bbArg type to memref.
     bbArg.setType(type);
 
     // Replace all uses of the original tensor bbArg.
     rewriter.setInsertionPointToStart(block);
     if (!bbArgUses.empty()) {
-      Value toTensorOp =
-          rewriter.create<bufferization::ToTensorOp>(bbArg.getLoc(), bbArg);
+      Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
+          bbArg.getLoc(), tensorType, bbArg);
       for (OpOperand *use : bbArgUses)
         use->set(toTensorOp);
     }