diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEBufferizeToAllocation.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEBufferizeToAllocation.cpp
index 2f2ee3297..e42b6d597 100644
--- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEBufferizeToAllocation.cpp
+++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEBufferizeToAllocation.cpp
@@ -8,7 +8,6 @@
 #include "iree-amd-aie/Transforms/AMDAIEUtils.h"
 #include "iree-amd-aie/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/IR/Iterators.h"
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDmaUtils.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDmaUtils.cpp
index bc3cf8044..f3c734f8f 100644
--- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDmaUtils.cpp
+++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDmaUtils.cpp
@@ -6,23 +6,13 @@
 
 #include "AMDAIEDmaUtils.h"
 
-#include <cstdlib>
-
+#include "AMDAIEUtils.h"
 #include "iree-amd-aie/Transforms/AMDAIEUtils.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 
 namespace mlir::iree_compiler::AMDAIE {
 
-/// Return an ancestor of 'op' in 'block', or nullptr if no such ancestor.
-Operation *getAncestorInBlock(Operation *op, Block *block) {
-  if (!op || !block) return nullptr;
-  auto parent = op;
-  while (parent && (parent->getBlock() != block))
-    parent = parent->getParentOp();
-  return parent;
-}
-
 bool areAccessPatternsCombinable(const SmallVector<OpFoldResult> &offsetsA,
                                  const SmallVector<OpFoldResult> &sizesA,
                                  const SmallVector<OpFoldResult> &stridesA,
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEPeelForLoop.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEPeelForLoop.cpp
index 0a6118ab3..6cd8d86b4 100644
--- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEPeelForLoop.cpp
+++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEPeelForLoop.cpp
@@ -4,10 +4,9 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree-amd-aie/Transforms/AMDAIEUtils.h"
 #include "iree-amd-aie/Transforms/Passes.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
-#include "mlir/IR/Iterators.h"
 #include "mlir/Pass/Pass.h"
 
 #define DEBUG_TYPE "iree-amdaie-peel-for-loop"
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEUtils.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEUtils.cpp
index ae6081349..2cd57beb2 100644
--- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEUtils.cpp
+++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEUtils.cpp
@@ -6,10 +6,10 @@
 
 #include "AMDAIEUtils.h"
 
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringExtras.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/Iterators.h"
 
 namespace mlir::iree_compiler::AMDAIE {
 
@@ -251,7 +251,7 @@ bool isMatmulInDefChain(Value operand) {
 /// Utility to identify if `linalgOp` is an elementwise operation with a
 /// matmul-like op upstream in its computation tree.
 bool isMatmulProducerOfElementwise(linalg::LinalgOp linalgOp) {
-  if (!isElementwise(linalgOp) || isa<linalg::FillOp>(linalgOp)) {
+  if (!linalg::isElementwise(linalgOp) || isa<linalg::FillOp>(linalgOp)) {
     return false;
   }
   // Check if any of the defining op is a matmul-like op.
@@ -272,6 +272,13 @@ std::string utohexstr(uint32_t value, size_t width, bool header,
   return res + prefix + hexStr;
 }
 
+/// Return an ancestor of 'op' in 'block', or nullptr if no such ancestor.
+Operation *getAncestorInBlock(Operation *op, Block *block) {
+  if (!op || !block) return nullptr;
+  while (op && (op->getBlock() != block)) op = op->getParentOp();
+  return op;
+}
+
 /// Find the largest factor of 'num' which is not larger than 'max'.
 int detail::findLargestFactor(int num, int max) {
   assert(max > 0 && "No factors less than or equal to 0 exist");
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEUtils.h b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEUtils.h
index e78ebc281..3daac3cd7 100644
--- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEUtils.h
+++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEUtils.h
@@ -7,11 +7,9 @@
 #ifndef IREE_AMD_AIE_TRANSFORMS_AMDAIEUTILS_H_
 #define IREE_AMD_AIE_TRANSFORMS_AMDAIEUTILS_H_
 
-#include <array>
-
-#include "iree-amd-aie/IR/AMDAIEAttrs.h"
-#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
-#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "iree-amd-aie/aie_runtime/AMDAIEEnums.h"
+#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
+#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
 #include "mlir/IR/Types.h"
 
 namespace mlir::iree_compiler::AMDAIE {
@@ -62,15 +60,22 @@ bool isMatmulProducerOfElementwise(linalg::LinalgOp linalgOp);
 std::string utohexstr(uint32_t value, size_t width, bool header = true,
                       bool lowercase = false);
 
+/// If `op` is in `block`, then return `op`. Otherwise traverse through parents
+/// to the first ancestor of `op` that is in `block`, and return that
+/// ancestor. If `op` has no ancestor in `block`, or if `op` is nullptr or
+/// `block` is nullptr, return nullptr.
+Operation *getAncestorInBlock(Operation *op, Block *block);
+
 namespace detail {
 
-// Returns the largest number that perfectly divides `num` that
-// is less than or equal to max
+/// Returns the largest number that perfectly divides `num` that
+/// is less than or equal to max
 int findLargestFactor(int num, int max);
 
-// A variant where we prefer factors to also be a multiple of `multiple`
+/// A variant where we prefer factors to also be a multiple of `multiple`
 int findLargestFactor(int num, int max, int multiple);
 
+
 }  // namespace detail
 
 }  // namespace mlir::iree_compiler::AMDAIE