From 5d0d0ff9c0f39611138e5227a0449612116f52a9 Mon Sep 17 00:00:00 2001
From: linyayAmd <linyay@amd.com>
Date: Tue, 10 Dec 2024 23:34:29 -0800
Subject: [PATCH] Extend affine::normalizeMemRef to return
 FailureOr<AllocLikeOp>, where AllocLikeOp is the normalized memref.alloca or
 memref.alloc (#412)

Co-authored-by: Lin-Ya Yu <lin-ya.yu@amd.com>
---
 mlir/include/mlir/Dialect/Affine/Utils.h | 6 +++---
 mlir/lib/Dialect/Affine/Utils/Utils.cpp  | 8 ++++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h
index 3ea45f235ca967..db5637c024b7f0 100644
--- 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
index cf2f084bba3153..222178904f9546 100644
--- 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) {