Extend affine::normalizeMemRef to return FailureOr<AllocLikeOp>, wher…

…e AllocLikeOp is the normalized memref.alloca or memref.alloc (#412) Co-authored-by: Lin-Ya Yu <lin-ya.yu@amd.com>
Xilinx · Dec 11, 2024 · 5d0d0ff · 5d0d0ff
1 parent 2babd7e
commit 5d0d0ff
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h
@@ -247,10 +247,10 @@ LogicalResult replaceAllMemRefUsesWith(Value oldMemRef, Value newMemRef,
 /// and updates all its indexing uses. Returns failure if any of its uses
 /// escape (while leaving the IR in a valid state).
 template <typename AllocLikeOp>
-LogicalResult normalizeMemRef(AllocLikeOp *op);
-extern template LogicalResult
+FailureOr<AllocLikeOp> normalizeMemRef(AllocLikeOp *op);
+extern template FailureOr<memref::AllocaOp>
 normalizeMemRef<memref::AllocaOp>(memref::AllocaOp *op);
-extern template LogicalResult
+extern template FailureOr<memref::AllocOp>
 normalizeMemRef<memref::AllocOp>(memref::AllocOp *op);
 
 /// Normalizes `memrefType` so that the affine layout map of the memref is

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1689,7 +1689,7 @@ static void createNewDynamicSizes(MemRefType oldMemRefType,
 
 // TODO: Currently works for static memrefs with a single layout map.
 template <typename AllocLikeOp>
-LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
+FailureOr<AllocLikeOp> mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
   MemRefType memrefType = allocOp->getType();
   OpBuilder b(*allocOp);
 
@@ -1743,12 +1743,12 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
   }));
   oldMemRef.replaceAllUsesWith(newAlloc);
   allocOp->erase();
-  return success();
+  return newAlloc;
 }
 
-template LogicalResult
+template FailureOr<memref::AllocaOp>
 mlir::affine::normalizeMemRef<memref::AllocaOp>(memref::AllocaOp *op);
-template LogicalResult
+template FailureOr<memref::AllocOp>
 mlir::affine::normalizeMemRef<memref::AllocOp>(memref::AllocOp *op);
 
 MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {