[Mlir-commits] [mlir] 4a0b8c3 - [MLIR][Affine] Drop assumptions of surrounding builtin.func op in Utils/LoopUtils (#116324)

[llvmlistbot at llvm.org]

Author: Uday Bondhugula
Date: 2024-11-23T07:05:03+05:30
New Revision: 4a0b8c3e7d4e9031181400d48ffcdf703bccce35

URL: https://github.com/llvm/llvm-project/commit/4a0b8c3e7d4e9031181400d48ffcdf703bccce35
DIFF: https://github.com/llvm/llvm-project/commit/4a0b8c3e7d4e9031181400d48ffcdf703bccce35.diff

LOG: [MLIR][Affine] Drop assumptions of surrounding builtin.func op in Utils/LoopUtils (#116324)

Drop assumptions of surrounding builtin.func op in affine LoopUtils and
Utils. There are use cases of affine fusion or affine transformation in
other func-like ops.
In the context of https://github.com/llvm/llvm-project/issues/116042

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
    mlir/lib/Dialect/Affine/Utils/Utils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 7578ce61553be7..d0e0f72ab2c46f 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -1941,8 +1941,8 @@ static LogicalResult generateCopy(
   *nBegin = begin;
   *nEnd = end;
 
-  func::FuncOp f = begin->getParentOfType<func::FuncOp>();
-  OpBuilder topBuilder(f.getBody());
+  auto f = begin->getParentOfType<FunctionOpInterface>();
+  OpBuilder topBuilder(f.getFunctionBody());
   Value zeroIndex = topBuilder.create<arith::ConstantIndexOp>(f.getLoc(), 0);
 
   *sizeInBytes = 0;
@@ -1961,8 +1961,9 @@ static LogicalResult generateCopy(
   OpBuilder &b = region.isWrite() ? epilogue : prologue;
 
   // Builder to create constants at the top level.
-  auto func = copyPlacementBlock->getParent()->getParentOfType<func::FuncOp>();
-  OpBuilder top(func.getBody());
+  auto func =
+      copyPlacementBlock->getParent()->getParentOfType<FunctionOpInterface>();
+  OpBuilder top(func.getFunctionBody());
 
   auto loc = region.loc;
   auto memref = region.memref;

diff  --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 3420db771ef426..07d399adae0cd4 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -1391,11 +1391,11 @@ LogicalResult mlir::affine::replaceAllMemRefUsesWith(
   std::unique_ptr<PostDominanceInfo> postDomInfo;
   if (domOpFilter)
     domInfo = std::make_unique<DominanceInfo>(
-        domOpFilter->getParentOfType<func::FuncOp>());
+        domOpFilter->getParentOfType<FunctionOpInterface>());
 
   if (postDomOpFilter)
     postDomInfo = std::make_unique<PostDominanceInfo>(
-        postDomOpFilter->getParentOfType<func::FuncOp>());
+        postDomOpFilter->getParentOfType<FunctionOpInterface>());
 
   // Walk all uses of old memref; collect ops to perform replacement. We use a
   // DenseSet since an operation could potentially have multiple uses of a