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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-bufferization

Author: Yi Zhang (cathyzhyi)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/124826.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp (+3-2) 


``````````diff
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);
     }

``````````

</details>


https://github.com/llvm/llvm-project/pull/124826