[MLIR][Affine] Fix signatures of normalize memref utilities

[bondhugula]

These methods were passing derived op types by pointers, which deviates
from the style. While on this, fix obsolete comments on those methods.

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir-memref

Author: Uday Bondhugula (bondhugula)

Changes

These methods were passing derived op types by pointers, which deviates
from the style. While on this, fix obsolete comments on those methods.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/Affine/Utils.h (+3-3)
• (modified) mlir/lib/Dialect/Affine/Utils/Utils.cpp (+17-18)
• (modified) mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp (+2-2)

diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h
index ff1900bc8f2eb..3b4bb34105581 100644
--- a/mlir/include/mlir/Dialect/Affine/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Utils.h
@@ -247,11 +247,11 @@ 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);
+LogicalResult normalizeMemRef(AllocLikeOp op);
 extern template LogicalResult
-normalizeMemRef<memref::AllocaOp>(memref::AllocaOp *op);
+normalizeMemRef<memref::AllocaOp>(memref::AllocaOp op);
 extern template LogicalResult
-normalizeMemRef<memref::AllocOp>(memref::AllocOp *op);
+normalizeMemRef<memref::AllocOp>(memref::AllocOp op);
 
 /// Normalizes `memrefType` so that the affine layout map of the memref is
 /// transformed to an identity map with a new shape being computed for the
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 2723cff6900d0..2925aa918cb1c 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1741,7 +1741,7 @@ static AffineExpr createDimSizeExprForTiledLayout(AffineExpr oldMapOutput,
 template <typename AllocLikeOp>
 static void createNewDynamicSizes(MemRefType oldMemRefType,
                                   MemRefType newMemRefType, AffineMap map,
-                                  AllocLikeOp *allocOp, OpBuilder b,
+                                  AllocLikeOp allocOp, OpBuilder b,
                                   SmallVectorImpl<Value> &newDynamicSizes) {
   // Create new input for AffineApplyOp.
   SmallVector<Value, 4> inAffineApply;
@@ -1750,13 +1750,13 @@ static void createNewDynamicSizes(MemRefType oldMemRefType,
   for (unsigned d = 0; d < oldMemRefType.getRank(); ++d) {
     if (oldMemRefShape[d] < 0) {
       // Use dynamicSizes of allocOp for dynamic dimension.
-      inAffineApply.emplace_back(allocOp->getDynamicSizes()[dynIdx]);
+      inAffineApply.emplace_back(allocOp.getDynamicSizes()[dynIdx]);
       dynIdx++;
     } else {
       // Create ConstantOp for static dimension.
       auto constantAttr = b.getIntegerAttr(b.getIndexType(), oldMemRefShape[d]);
       inAffineApply.emplace_back(
-          b.create<arith::ConstantOp>(allocOp->getLoc(), constantAttr));
+          b.create<arith::ConstantOp>(allocOp.getLoc(), constantAttr));
     }
   }
 
@@ -1780,18 +1780,17 @@ static void createNewDynamicSizes(MemRefType oldMemRefType,
       AffineMap newMap =
           AffineMap::get(map.getNumInputs(), map.getNumSymbols(), newMapOutput);
       Value affineApp =
-          b.create<AffineApplyOp>(allocOp->getLoc(), newMap, inAffineApply);
+          b.create<AffineApplyOp>(allocOp.getLoc(), newMap, inAffineApply);
       newDynamicSizes.emplace_back(affineApp);
     }
     newDimIdx++;
   }
 }
 
-// TODO: Currently works for static memrefs with a single layout map.
 template <typename AllocLikeOp>
-LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
-  MemRefType memrefType = allocOp->getType();
-  OpBuilder b(*allocOp);
+LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp allocOp) {
+  MemRefType memrefType = allocOp.getType();
+  OpBuilder b(allocOp);
 
   // Fetch a new memref type after normalizing the old memref to have an
   // identity map layout.
@@ -1801,9 +1800,9 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
     // transformed to an identity map.
     return failure();
 
-  Value oldMemRef = allocOp->getResult();
+  Value oldMemRef = allocOp.getResult();
 
-  SmallVector<Value, 4> symbolOperands(allocOp->getSymbolOperands());
+  SmallVector<Value, 4> symbolOperands(allocOp.getSymbolOperands());
   AffineMap layoutMap = memrefType.getLayout().getAffineMap();
   AllocLikeOp newAlloc;
   // Check if `layoutMap` is a tiled layout. Only single layout map is
@@ -1811,17 +1810,17 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
   SmallVector<std::tuple<AffineExpr, unsigned, unsigned>> tileSizePos;
   (void)getTileSizePos(layoutMap, tileSizePos);
   if (newMemRefType.getNumDynamicDims() > 0 && !tileSizePos.empty()) {
-    MemRefType oldMemRefType = cast<MemRefType>(oldMemRef.getType());
+    auto oldMemRefType = cast<MemRefType>(oldMemRef.getType());
     SmallVector<Value, 4> newDynamicSizes;
     createNewDynamicSizes(oldMemRefType, newMemRefType, layoutMap, allocOp, b,
                           newDynamicSizes);
     // Add the new dynamic sizes in new AllocOp.
     newAlloc =
-        b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType, newDynamicSizes,
-                              allocOp->getAlignmentAttr());
+        b.create<AllocLikeOp>(allocOp.getLoc(), newMemRefType, newDynamicSizes,
+                              allocOp.getAlignmentAttr());
   } else {
-    newAlloc = b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType,
-                                     allocOp->getAlignmentAttr());
+    newAlloc = b.create<AllocLikeOp>(allocOp.getLoc(), newMemRefType,
+                                     allocOp.getAlignmentAttr());
   }
   // Replace all uses of the old memref.
   if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
@@ -1842,14 +1841,14 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
     return hasSingleEffect<MemoryEffects::Free>(op, oldMemRef);
   }));
   oldMemRef.replaceAllUsesWith(newAlloc);
-  allocOp->erase();
+  allocOp.erase();
   return success();
 }
 
 template LogicalResult
-mlir::affine::normalizeMemRef<memref::AllocaOp>(memref::AllocaOp *op);
+mlir::affine::normalizeMemRef<memref::AllocaOp>(memref::AllocaOp op);
 template LogicalResult
-mlir::affine::normalizeMemRef<memref::AllocOp>(memref::AllocOp *op);
+mlir::affine::normalizeMemRef<memref::AllocOp>(memref::AllocOp op);
 
 MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
   unsigned rank = memrefType.getRank();
diff --git a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
index 08b853fe65b85..95fed04a7864e 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp
@@ -356,12 +356,12 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(func::FuncOp funcOp,
   SmallVector<memref::AllocOp, 4> allocOps;
   funcOp.walk([&](memref::AllocOp op) { allocOps.push_back(op); });
   for (memref::AllocOp allocOp : allocOps)
-    (void)normalizeMemRef(&allocOp);
+    (void)normalizeMemRef(allocOp);
 
   SmallVector<memref::AllocaOp> allocaOps;
   funcOp.walk([&](memref::AllocaOp op) { allocaOps.push_back(op); });
   for (memref::AllocaOp allocaOp : allocaOps)
-    (void)normalizeMemRef(&allocaOp);
+    (void)normalizeMemRef(allocaOp);
 
   // We use this OpBuilder to create new memref layout later.
   OpBuilder b(funcOp);