Extend affine::normalizeMemRef to return FailureOr<AllocLikeOp>, where AllocLikeOp is the normalized memref.alloca or memref.alloc (#412)

Co-authored-by: Lin-Ya Yu <[email protected]>

Author: linyayAmd

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

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) {