diff --git a/backends/jitrt/BUILD b/backends/jitrt/BUILD
index 514949e41a6..68481a6f03b 100644
--- a/backends/jitrt/BUILD
+++ b/backends/jitrt/BUILD
@@ -214,6 +214,7 @@ licenses(["notice"])
 #         "@llvm-project//mlir:ArithmeticTransforms",
 #         "@llvm-project//mlir:ArmNeonToLLVMIRTranslation",
 #         "@llvm-project//mlir:ArmSVEToLLVMIRTranslation",
+#         "@llvm-project//mlir:AsyncDialect",
 #         "@llvm-project//mlir:AsyncToLLVM",
 #         "@llvm-project//mlir:AsyncTransforms",
 #         "@llvm-project//mlir:ComplexToLLVM",
@@ -679,16 +680,13 @@ licenses(["notice"])
 #         "@llvm-project//mlir:ArithmeticDialect",
 #         "@llvm-project//mlir:AsyncDialect",
 #         "@llvm-project//mlir:AsyncTransforms",
-#         "@llvm-project//mlir:FuncDialect",
 #         "@llvm-project//mlir:IR",
 #         "@llvm-project//mlir:MathDialect",
-#         "@llvm-project//mlir:MathTransforms",
 #         "@llvm-project//mlir:MemRefDialect",
 #         "@llvm-project//mlir:Pass",
 #         "@llvm-project//mlir:SCFDialect",
 #         "@llvm-project//mlir:Transforms",
 #         "@llvm-project//mlir:VectorDialect",
-#         "@llvm-project//mlir:X86VectorDialect",
 #     ],
 # )
 #
diff --git a/backends/jitrt/include/tfrt/jitrt/transforms/codegen_passes.h b/backends/jitrt/include/tfrt/jitrt/transforms/codegen_passes.h
index 4ca743aa50f..260c89f0db5 100644
--- a/backends/jitrt/include/tfrt/jitrt/transforms/codegen_passes.h
+++ b/backends/jitrt/include/tfrt/jitrt/transforms/codegen_passes.h
@@ -19,12 +19,6 @@
 
 #include <memory>
 
-#include "mlir/Dialect/Async/IR/Async.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/Pass/Pass.h"
 
 namespace tfrt {
diff --git a/backends/jitrt/lib/jitrt_compiler.cc b/backends/jitrt/lib/jitrt_compiler.cc
index 334cd42fabf..334bc74c672 100644
--- a/backends/jitrt/lib/jitrt_compiler.cc
+++ b/backends/jitrt/lib/jitrt_compiler.cc
@@ -37,6 +37,7 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
+#include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
diff --git a/backends/jitrt/lib/transforms/cost_driven_async_parallel_for.cc b/backends/jitrt/lib/transforms/cost_driven_async_parallel_for.cc
index 4bf7db925a9..56bbd14ce7a 100644
--- a/backends/jitrt/lib/transforms/cost_driven_async_parallel_for.cc
+++ b/backends/jitrt/lib/transforms/cost_driven_async_parallel_for.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Async/Transforms.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"