[Mlir-commits] [mlir] bfefa15 - [mlir][bufferization] Use original type when convert arg for users (#124826)

[llvmlistbot at llvm.org]

Author: Yi Zhang
Date: 2025-01-28T15:30:16-05:00
New Revision: bfefa15cc18f6f4b0b07849c619989c1a8c5aef9

URL: https://github.com/llvm/llvm-project/commit/bfefa15cc18f6f4b0b07849c619989c1a8c5aef9
DIFF: https://github.com/llvm/llvm-project/commit/bfefa15cc18f6f4b0b07849c619989c1a8c5aef9.diff

LOG: [mlir][bufferization] Use original type when convert arg for users (#124826)

This change will keep the memory space information for the tensor if
there is any.

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 64d79f5b5d60cf..313dbbd16984b0 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);
     }