From e3614c4abf383f66093f790ce313c17edb5eda6f Mon Sep 17 00:00:00 2001
From: Lin-Ya Yu <lin-ya.yu@amd.com>
Date: Tue, 19 Nov 2024 23:13:56 -0800
Subject: [PATCH] Extend affine::normalizeMemRef to return
 FailureOr<AllocLikeOp>, where AllocLikeOp is the normalized memref.alloca or
 memref.alloc

---
 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 ef3081f75608a8..36e817412acabf 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1690,7 +1690,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);
 
@@ -1744,12 +1744,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) {