diff --git a/tensorflow/compiler/mlir/tfrt/BUILD b/tensorflow/compiler/mlir/tfrt/BUILD
index 9a71c36..572f8ba 100644
--- a/tensorflow/compiler/mlir/tfrt/BUILD
+++ b/tensorflow/compiler/mlir/tfrt/BUILD
@@ -481,7 +481,7 @@
         "//platforms/xla/tests/gpu:__pkg__",
     ],
     deps = [
-        "//tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu:pass",
+        "//tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu:passes",
         "//tensorflow/compiler/mlir/xla:mhlo_to_lhlo_with_xla",
         "//tensorflow/compiler/xla/service:hlo",
         "//tensorflow/core/platform:errors",
diff --git a/tensorflow/compiler/mlir/tfrt/lhlo-tfrt-opt.cc b/tensorflow/compiler/mlir/tfrt/lhlo-tfrt-opt.cc
index 8f7eab0..b04ab36 100644
--- a/tensorflow/compiler/mlir/tfrt/lhlo-tfrt-opt.cc
+++ b/tensorflow/compiler/mlir/tfrt/lhlo-tfrt-opt.cc
@@ -36,7 +36,7 @@
                   mlir::mhlo::MhloDialect, tfrt::gpu::GpuDialect,
                   tfrt::gpu::conversion::GpuConversionDialect>();
   tfrt::RegisterTFRTDialects(registry);
-  tensorflow::registerTFRTGPUPasses();
+  tensorflow::registerXlirPasses();
   tfrt::gpu::registerPasses();
   return failed(mlir::MlirOptMain(argc, argv, "MHLO TFRT pass driver\n",
                                   registry,
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/BUILD b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/BUILD
index 5a4c25a..aa47f40 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/BUILD
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/BUILD
@@ -3,20 +3,8 @@
 load("//tensorflow:tensorflow.bzl", "if_google")
 
 # TF to TFRT kernels conversion.
-package(licenses = ["notice"])
-
-package_group(
-    name = "friends",
-    packages = if_google([
-        "//learning/brain/experimental/mlir/tfrt_compiler/...",
-        "//learning/brain/experimental/tfrt/...",
-        "//platforms/xla/tests/gpu/...",
-        "//third_party/tf_runtime_google/...",
-    ]) + [
-        "//tensorflow/compiler/...",
-        "//tensorflow/core/runtime_fallback/...",
-        "//tensorflow/core/tfrt/saved_model/...",
-    ],
+package(
+    licenses = ["notice"],
 )
 
 gentbl_cc_library(
@@ -26,7 +14,7 @@
         (
             [
                 "-gen-pass-decls",
-                "-name=TFRTGPU",
+                "-name=Xlir",
             ],
             "gpu_passes.h.inc",
         ),
@@ -39,43 +27,31 @@
 )
 
 cc_library(
-    name = "pass_details",
-    hdrs = [
-        "PassDetail.h",
-    ],
-    deps = [
-        ":TfrtGpuPassIncGen",
-        "//tensorflow/compiler/xla/service/gpu:xlir_opdefs",
-        "@llvm-project//mlir:GPUDialect",
-        "@llvm-project//mlir:Pass",
-        "@tf_runtime//:basic_kernels_opdefs",
-        "@tf_runtime//backends/gpu:gpu_opdefs",
-    ],
-)
-
-cc_library(
-    name = "pass",
+    name = "passes",
     srcs = [
         "ccl_pattern.cc",
-        "ccl_pattern.h",
         "custom_call_pattern.cc",
-        "custom_call_pattern.h",
         "gemm_pattern.cc",
-        "gemm_pattern.h",
         "gpu_passes.cc",
         "memcpy_pattern.cc",
-        "memcpy_pattern.h",
         "memset_pattern.cc",
-        "memset_pattern.h",
     ],
-    hdrs = ["gpu_passes.h"],
+    hdrs = [
+        "gpu_passes.h",
+        "register_passes.h",
+    ],
     tags = [
         "gpu",
         "no_oss",
     ],
-    visibility = [":friends"],
+    visibility = if_google([
+        "//platforms/xla/tests/gpu:__pkg__",
+    ]) + [
+        "//tensorflow/compiler/mlir/tfrt:__pkg__",
+        "//tensorflow/compiler/xla/service/gpu:__pkg__",
+    ],
     deps = [
-        ":pass_details",
+        ":TfrtGpuPassIncGen",
         "//tensorflow/compiler/mlir/hlo:lhlo",
         "//tensorflow/compiler/mlir/hlo:lhlo_gpu",
         "//tensorflow/compiler/mlir/tensorflow",
@@ -92,6 +68,7 @@
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Rewrite",
         "@llvm-project//mlir:StandardOps",
+        "@llvm-project//mlir:StandardOpsTransforms",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:Transforms",
         "@tf_runtime//:basic_kernels_opdefs",
@@ -101,18 +78,3 @@
     ],
     alwayslink = 1,
 )
-
-cc_library(
-    name = "passes",
-    hdrs = ["register_passes.h"],
-    tags = [
-        "gpu",
-        "no_oss",
-    ],
-    visibility = [":friends"],
-    deps = [
-        ":TfrtGpuPassIncGen",
-        ":pass",
-        "@llvm-project//mlir:Pass",
-    ],
-)
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/PassDetail.h b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/PassDetail.h
deleted file mode 100644
index 4e08993..0000000
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/PassDetail.h
+++ /dev/null
@@ -1,31 +0,0 @@
-// Copyright 2020 The TensorFlow Runtime Authors
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_PASSDETAIL_H_
-#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_PASSDETAIL_H_
-
-#include "mlir/Dialect/GPU/GPUDialect.h"
-#include "mlir/Pass/Pass.h"
-#include "tensorflow/compiler/xla/service/gpu/xlir_ops.h"
-#include "tfrt/gpu/kernels/gpu_ops.h"  // from @tf_runtime
-#include "tfrt/basic_kernels/opdefs/tfrt_base.h"  // from @tf_runtime
-
-namespace tensorflow {
-
-#define GEN_PASS_CLASSES
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.h.inc"
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_PASSDETAIL_H_
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/ccl_pattern.cc b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/ccl_pattern.cc
index 5a2433b..3717d5e 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/ccl_pattern.cc
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/ccl_pattern.cc
@@ -18,8 +18,6 @@
 // Pattern to lower lmhlo collective ops to tfrt_gpu/xlir dialect.
 //
 //===----------------------------------------------------------------------===//
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/ccl_pattern.h"
-
 #include <functional>
 #include <string>
 
@@ -311,13 +309,14 @@
 
 }  // namespace
 
-void populateCclConversionPattern(RewritePatternSet& patterns) {
+void populateCclConversionPattern(RewritePatternSet& patterns,
+                                  TypeConverter& converter) {
   patterns.add<CclRewritePattern<lmhlo::AllGatherOp>,
                CclRewritePattern<lmhlo::AllReduceOp>,
                CclRewritePattern<lmhlo::ReduceScatterOp>,
                CclRewritePattern<lmhlo::AllToAllOp>,
                CclRewritePattern<lmhlo::CollectivePermuteOp>>(
-      patterns.getContext());
+      converter, patterns.getContext());
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/ccl_pattern.h b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/ccl_pattern.h
deleted file mode 100644
index 1185511..0000000
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/ccl_pattern.h
+++ /dev/null
@@ -1,26 +0,0 @@
-// Copyright 2021 The TensorFlow Runtime Authors
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_CUDA_CCL_PATTERN_H_
-#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_CUDA_CCL_PATTERN_H_
-
-#include "mlir/IR/PatternMatch.h"
-
-namespace tensorflow {
-
-void populateCclConversionPattern(mlir::RewritePatternSet& patterns);
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_CUDA_CCL_PATTERN_H_
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/custom_call_pattern.cc b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/custom_call_pattern.cc
index e1c2a3d..e9e01c9 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/custom_call_pattern.cc
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/custom_call_pattern.cc
@@ -18,8 +18,6 @@
 // Pattern to lower lmhlo.custom_call op to tfrt_gpu/xlir dialect.
 //
 //===----------------------------------------------------------------------===//
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/custom_call_pattern.h"
-
 #include <functional>
 #include <string>
 
@@ -85,8 +83,9 @@
 
 }  // namespace
 
-void populateCustomCallConversionPattern(RewritePatternSet& patterns) {
-  patterns.add<CustomCallRewritePattern>(patterns.getContext());
+void populateCustomCallConversionPattern(RewritePatternSet& patterns,
+                                         TypeConverter& converter) {
+  patterns.add<CustomCallRewritePattern>(converter, patterns.getContext());
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/custom_call_pattern.h b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/custom_call_pattern.h
deleted file mode 100644
index 735fd24..0000000
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/custom_call_pattern.h
+++ /dev/null
@@ -1,26 +0,0 @@
-// Copyright 2021 The TensorFlow Runtime Authors
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_CUSTOM_CALL_PATTERN_H_
-#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_CUSTOM_CALL_PATTERN_H_
-
-#include "mlir/IR/PatternMatch.h"
-
-namespace tensorflow {
-
-void populateCustomCallConversionPattern(mlir::RewritePatternSet& patterns);
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_CUSTOM_CALL_PATTERN_H_
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gemm_pattern.cc b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gemm_pattern.cc
index 19c64f8..ef34270 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gemm_pattern.cc
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gemm_pattern.cc
@@ -18,8 +18,6 @@
 // Pattern to lower lhlogpu_gemm Ops to tfrt cuda dialect.
 //
 //===----------------------------------------------------------------------===//
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gemm_pattern.h"
-
 #include <assert.h>
 #include <stdint.h>
 
@@ -345,10 +343,11 @@
 
 }  // namespace
 
-void populateGemmConversionPattern(RewritePatternSet& patterns) {
+void populateGemmConversionPattern(RewritePatternSet& patterns,
+                                   TypeConverter& converter) {
   patterns.add<GemmRewritePattern<lmhlo_gpu::GEMMOp>,
                GemmRewritePattern<lmhlo_gpu::GEMM_BiasOp>>(
-      patterns.getContext());
+      converter, patterns.getContext());
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gemm_pattern.h b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gemm_pattern.h
deleted file mode 100644
index 07451af..0000000
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gemm_pattern.h
+++ /dev/null
@@ -1,30 +0,0 @@
-// Copyright 2020 The TensorFlow Runtime Authors
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_GEMM_PATTERN_H_
-#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_GEMM_PATTERN_H_
-
-#include "mlir-hlo/Dialect/mhlo/IR/lhlo_gpu_ops.h"
-#include "mlir/IR/BlockAndValueMapping.h"
-#include "mlir/Transforms/DialectConversion.h"
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/StringRef.h"
-
-namespace tensorflow {
-
-void populateGemmConversionPattern(mlir::RewritePatternSet& patterns);
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_GEMM_PATTERN_H_
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.cc b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.cc
index 22498f3..02e551f 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.cc
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.cc
@@ -21,85 +21,105 @@
 #include <memory>
 #include <utility>
 
+#include "mlir-hlo/Dialect/mhlo/IR/lhlo_gpu_ops.h"
 #include "mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Transforms/FuncConversions.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Pass/Pass.h"
-#include "mlir/Rewrite/PatternApplicator.h"
 #include "mlir/Transforms/DialectConversion.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "llvm/ADT/ArrayRef.h"
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/PassDetail.h"
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/ccl_pattern.h"
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/custom_call_pattern.h"
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gemm_pattern.h"
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memcpy_pattern.h"
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memset_pattern.h"
+#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h"
 #include "tensorflow/compiler/xla/service/gpu/xlir_ops.h"
 #include "tfrt/gpu/kernels/gpu_ops.h"  // from @tf_runtime
 #include "tfrt/gpu/passes/passes.h"  // from @tf_runtime
 #include "tfrt/basic_kernels/opdefs/basic_kernels.h"  // from @tf_runtime
 
 namespace tensorflow {
+
+void populateCclConversionPattern(RewritePatternSet&, TypeConverter&);
+void populateCustomCallConversionPattern(RewritePatternSet&, TypeConverter&);
+void populateGemmConversionPattern(RewritePatternSet&, TypeConverter&);
+void populateMemcpyConversionPattern(RewritePatternSet&, TypeConverter&);
+void populateMemsetConversionPattern(RewritePatternSet&, TypeConverter&);
+
+#define GEN_PASS_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.h.inc"
+
 namespace {
 
-struct LmhloGpuAsyncConversionPass
-    : public LmhloGpuAsyncConversionPassBase<LmhloGpuAsyncConversionPass> {
+struct ConvertLmhloToGpuPass
+    : public ConvertLmhloToGpuPassBase<ConvertLmhloToGpuPass> {
  private:
-  void runOnFunction() override {
-    auto* context = &getContext();
-
-    TypeConverter converter;
-    converter.addConversion([](Type type) { return type; });
-    auto buffer_type = tfrt::gpu::BufferType::get(context);
-    converter.addConversion([&](BaseMemRefType) { return buffer_type; });
-
-    ConversionTarget target(*context);
-    target
-        .addIllegalDialect<lmhlo_gpu::LmhloGpuDialect, mlir::gpu::GPUDialect>();
-    target.addLegalDialect<tfrt::compiler::TFRTDialect, tfrt::gpu::GpuDialect,
-                           xla::gpu::XlirDialect>();
-    target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
-      return converter.isSignatureLegal(op.getType()) &&
-             converter.isLegal(&op.body());
-    });
-    target.addDynamicallyLegalOp<tfrt::gpu::conversion::AsyncExecuteOp>(
-        [&](tfrt::gpu::conversion::AsyncExecuteOp op) {
-          return converter.isLegal(&op.body());
-        });
-
-    RewritePatternSet patterns(context);
-    populateCclConversionPattern(patterns);
-    populateCustomCallConversionPattern(patterns);
-    populateGemmConversionPattern(patterns);
-    populateMemcpyConversionPattern(patterns);
-    populateMemsetConversionPattern(patterns);
-    populateFuncOpTypeConversionPattern(patterns, converter);
-
-    ConversionTarget wrap_target(*context);
-    wrap_target
-        .addLegalDialect<lmhlo_gpu::LmhloGpuDialect, mlir::gpu::GPUDialect>();
-    wrap_target.addLegalOp<lmhlo::AllGatherOp, lmhlo::AllReduceOp,
-                           lmhlo::ReduceScatterOp, lmhlo::AllToAllOp,
-                           lmhlo::CollectivePermuteOp, lmhlo::CustomCallOp>();
-    tfrt::gpu::populateGpuAsyncConversionPatterns(patterns, converter,
-                                                  wrap_target);
-
-    if (failed(applyPartialConversion(getOperation(), target,
-                                      std::move(patterns))))
-      return signalPassFailure();
-  }
+  void runOnFunction() override;
 };
 
 }  // namespace
 
-std::unique_ptr<FunctionPass> createLmhloGpuAsyncConversionPass() {
-  return std::make_unique<LmhloGpuAsyncConversionPass>();
+static Value MaterializeCast(OpBuilder& builder, Type type, ValueRange values,
+                             Location loc) {
+  auto cast_op = builder.create<UnrealizedConversionCastOp>(loc, type, values);
+  return cast_op.getResult(0);
+}
+
+void ConvertLmhloToGpuPass::runOnFunction() {
+  auto* context = &getContext();
+
+  TypeConverter converter;
+  converter.addConversion([](Type type) { return type; });
+  auto buffer_type = tfrt::gpu::BufferType::get(context);
+  converter.addConversion([&](BaseMemRefType) { return buffer_type; });
+  converter.addArgumentMaterialization(MaterializeCast);
+  converter.addSourceMaterialization(MaterializeCast);
+  converter.addTargetMaterialization(MaterializeCast);
+
+  RewritePatternSet patterns(context);
+  populateCclConversionPattern(patterns, converter);
+  populateCustomCallConversionPattern(patterns, converter);
+  populateGemmConversionPattern(patterns, converter);
+  populateMemcpyConversionPattern(patterns, converter);
+  populateMemsetConversionPattern(patterns, converter);
+  populateFuncOpTypeConversionPattern(patterns, converter);
+  populateReturnOpTypeConversionPattern(patterns, converter);
+
+  // Set of ops that need to be wrapped in tfrt_gpu_conversion.async.execute
+  // before lowering directly to tfrt_gpu ops (and therefore require some chain
+  // and stream, which the wrapper op provides as block arguments). On the other
+  // hand, ops which lower to the gpu dialect do not need to be wrapped.
+  ConversionTarget wrap_target(*context);
+  wrap_target
+      .addLegalDialect<lmhlo_gpu::LmhloGpuDialect, mlir::gpu::GPUDialect>();
+  wrap_target.addLegalOp<lmhlo::AllGatherOp, lmhlo::AllReduceOp,
+                         lmhlo::ReduceScatterOp, lmhlo::AllToAllOp,
+                         lmhlo::CollectivePermuteOp, lmhlo::CustomCallOp>();
+  tfrt::gpu::populateGpuAsyncConversionPatterns(patterns, converter,
+                                                wrap_target);
+
+  ConversionTarget target(*context);
+  target.addIllegalOp<memref::ReinterpretCastOp, memref::ViewOp>();
+  target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
+    return converter.isSignatureLegal(op.getType()) &&
+           converter.isLegal(&op.body());
+  });
+  target.addDynamicallyLegalOp<tfrt::gpu::conversion::AsyncExecuteOp>(
+      [&](tfrt::gpu::conversion::AsyncExecuteOp op) {
+        return converter.isLegal(&op.body());
+      });
+  target.markUnknownOpDynamicallyLegal([&](Operation* op) {
+    if (op->hasTrait<OpTrait::ReturnLike>()) return converter.isLegal(op);
+    return !wrap_target.isLegal(op);  // Wrapped ops are immediately lowered.
+  });
+
+  if (failed(
+          applyPartialConversion(getOperation(), target, std::move(patterns))))
+    return signalPassFailure();
+}
+
+std::unique_ptr<FunctionPass> createConvertLmhloToGpuPass() {
+  return std::make_unique<ConvertLmhloToGpuPass>();
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.h b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.h
index 1994a34..29f6e67 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.h
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.h
@@ -24,7 +24,7 @@
 
 // Creates a pass that lowers lmhlo_gpu ops to tfrt_gpu. Prepares the function
 // to be consumed by MLIR's gpu-async-region pass.
-std::unique_ptr<mlir::FunctionPass> createLmhloGpuAsyncConversionPass();
+std::unique_ptr<mlir::FunctionPass> createConvertLmhloToGpuPass();
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.td b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.td
index 7cf1b7b..a84c662 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.td
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.td
@@ -15,13 +15,13 @@
 
 include "mlir/Pass/PassBase.td"
 
-def LmhloGpuAsyncConversionPass : FunctionPass<"lmhlo-to-gpu"> {
+def ConvertLmhloToGpuPass : FunctionPass<"lmhlo-to-gpu"> {
   let summary = "Convert lmhlo_gpu ops to tfrt_gpu dialect.";
   let description = [{
     Move lmhlo_gpu ops inside tfrt_gpu_conversion.execute ops and convert them
     to tfrt_gpu dialect.
   }];
-  let constructor = "createLmhloGpuAsyncConversionPass()";
+  let constructor = "createConvertLmhloToGpuPass()";
   let dependentDialects = [
     "::mlir::gpu::GPUDialect",
     "::tfrt::compiler::TFRTDialect",
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memcpy_pattern.cc b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memcpy_pattern.cc
index a1b55d7..10d53a3 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memcpy_pattern.cc
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memcpy_pattern.cc
@@ -13,8 +13,6 @@
 // limitations under the License.
 
 // Pattern to lower gpu.memcpy ops to tfrt_gpu.mem.copy.
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memcpy_pattern.h"
-
 #include "mlir/Dialect/GPU/GPUDialect.h"  // from @llvm-project
 #include "tfrt/gpu/kernels/gpu_ops.h"  // from @tf_runtime
 #include "tfrt/gpu/passes/passes.h"  // from @tf_runtime
@@ -44,8 +42,9 @@
 
 }  // namespace
 
-void populateMemcpyConversionPattern(RewritePatternSet& patterns) {
-  patterns.add<MemcpyRewritePattern>(patterns.getContext());
+void populateMemcpyConversionPattern(RewritePatternSet& patterns,
+                                     TypeConverter& converter) {
+  patterns.add<MemcpyRewritePattern>(converter, patterns.getContext());
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memcpy_pattern.h b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memcpy_pattern.h
deleted file mode 100644
index 27a9364..0000000
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memcpy_pattern.h
+++ /dev/null
@@ -1,28 +0,0 @@
-// Copyright 2021 The TensorFlow Runtime Authors
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_MEMCPY_PATTERN_H_
-#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_MEMCPY_PATTERN_H_
-
-#include "mlir/IR/PatternMatch.h"  // from @llvm-project
-
-namespace tensorflow {
-
-// Add a pattern to the given pattern list to convert from mlir::gpu::MemcpyOp
-// to tfrt::gpu::MemCopyOp.
-void populateMemcpyConversionPattern(mlir::RewritePatternSet& patterns);
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_MEMCPY_PATTERN_H_
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memset_pattern.cc b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memset_pattern.cc
index 117c086..e3b5965 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memset_pattern.cc
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memset_pattern.cc
@@ -13,8 +13,6 @@
 // limitations under the License.
 
 // Pattern to lower mlir::gpu::memset Ops to tfrt cuda dialect.
-#include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memset_pattern.h"
-
 #include "mlir/Dialect/GPU/GPUDialect.h"  // from @llvm-project
 #include "tfrt/gpu/kernels/gpu_ops.h"  // from @tf_runtime
 #include "tfrt/gpu/passes/passes.h"  // from @tf_runtime
@@ -45,8 +43,9 @@
 
 }  // namespace
 
-void populateMemsetConversionPattern(RewritePatternSet& patterns) {
-  patterns.add<MemsetRewritePattern>(patterns.getContext());
+void populateMemsetConversionPattern(RewritePatternSet& patterns,
+                                     TypeConverter& converter) {
+  patterns.add<MemsetRewritePattern>(converter, patterns.getContext());
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memset_pattern.h b/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memset_pattern.h
deleted file mode 100644
index f4e354b..0000000
--- a/tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/memset_pattern.h
+++ /dev/null
@@ -1,28 +0,0 @@
-// Copyright 2020 The TensorFlow Runtime Authors
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_MEMSET_PATTERN_H_
-#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_MEMSET_PATTERN_H_
-
-#include "mlir/IR/PatternMatch.h"  // from @llvm-project
-
-namespace tensorflow {
-
-// Add a pattern to the given pattern list to convert from mlir::gpu::MemsetOp
-// to tfrt::gpu::MemSetOp.
-void populateMemsetConversionPattern(mlir::RewritePatternSet& patterns);
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_LHLO_GPU_TO_TFRT_GPU_MEMSET_PATTERN_H_
diff --git a/tensorflow/compiler/mlir/tfrt/translate/convert_xla_gpu.cc b/tensorflow/compiler/mlir/tfrt/translate/convert_xla_gpu.cc
index 376700b..4fee874 100644
--- a/tensorflow/compiler/mlir/tfrt/translate/convert_xla_gpu.cc
+++ b/tensorflow/compiler/mlir/tfrt/translate/convert_xla_gpu.cc
@@ -14,12 +14,13 @@
 ==============================================================================*/
 #include "tensorflow/compiler/mlir/tfrt/translate/convert_xla_gpu.h"
 
+#include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
+#include "mlir/InitAllDialects.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "mlir/Dialect/GPU/Passes.h"  // from @llvm-project
-#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"
 #include "tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu/gpu_passes.h"
 #include "tensorflow/compiler/mlir/xla/transforms/mhlo_to_lhlo_with_xla.h"
 #include "tensorflow/core/platform/errors.h"
@@ -55,20 +56,15 @@
 
   // LHLO -> TFRT Dialect (gpu kernels)
   mlir::PassManager pm(&context, mlir::PassManager::Nesting::Implicit);
-  pm.addPass(tensorflow::createLmhloGpuAsyncConversionPass());
+  pm.addPass(tensorflow::createConvertLmhloToGpuPass());
   pm.addPass(mlir::createGpuAsyncRegionPass());
   tfrt::gpu::populateGpuToTfrtGpuPasses(pm);
+  pm.addPass(mlir::createCanonicalizerPass());
   if (pm.run(*module).failed()) {
     return errors::Internal(
         "Failed to lower LHLO to TFRT Dialect with gpu kernels.");
   }
 
-  // Perform DCE with empty pattern set.
-  if (failed(mlir::applyPatternsAndFoldGreedily(*module,
-                                                RewritePatternSet(&context)))) {
-    return errors::Internal("Failed to remove dead ops.");
-  }
-
   // TFRT Dialect -> BEF
   std::string bef;
   llvm::raw_string_ostream bef_ostream(bef);
diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD
index 4afdae4..7fcb6ad 100644
--- a/tensorflow/compiler/xla/service/gpu/BUILD
+++ b/tensorflow/compiler/xla/service/gpu/BUILD
@@ -846,7 +846,7 @@
             "@llvm-project//mlir:GPUTransforms",
             "@llvm-project//mlir:Pass",
             "//tensorflow/compiler/mlir/hlo:lhlo",
-            "//tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu:pass",
+            "//tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu:passes",
             "//tensorflow/compiler/mlir/xla:attribute_exporter",
             "//tensorflow/compiler/xla/service:collective_ops_utils",
             "//tensorflow/core/tfrt/runtime:work_queue_interface",
@@ -1719,7 +1719,7 @@
         "//tensorflow/stream_executor:stream_executor_headers",
     ] + select({
         ":is_bef_executable_enabled": [
-            "//tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu:pass",
+            "//tensorflow/compiler/mlir/tfrt/transforms/lhlo_gpu_to_tfrt_gpu:passes",
             "@tf_runtime//:mlirtobef_translate",
             "@tf_runtime//:support",
             "@tf_runtime//:bef",
diff --git a/tensorflow/compiler/xla/service/gpu/bef_thunk.cc b/tensorflow/compiler/xla/service/gpu/bef_thunk.cc
index 42b9cb8..541625a 100644
--- a/tensorflow/compiler/xla/service/gpu/bef_thunk.cc
+++ b/tensorflow/compiler/xla/service/gpu/bef_thunk.cc
@@ -189,7 +189,7 @@
 static Status RunLmhloGpuToTfrtConversionPipeline(mlir::ModuleOp module) {
   mlir::PassManager pass_manager(module->getContext(),
                                  mlir::PassManager::Nesting::Implicit);
-  pass_manager.addPass(tensorflow::createLmhloGpuAsyncConversionPass());
+  pass_manager.addPass(tensorflow::createConvertLmhloToGpuPass());
   pass_manager.addPass(mlir::createGpuAsyncRegionPass());
   tfrt::gpu::populateGpuToTfrtGpuPasses(pass_manager);
 
diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
index 3c8cd8d..88b37f6 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
+++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc
@@ -40,6 +40,7 @@
 #include "mlir/InitAllDialects.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/utils/name_utils.h"
 #include "tensorflow/compiler/mlir/xla/hlo_utils.h"
 #include "tensorflow/compiler/mlir/xla/type_to_shape.h"
@@ -749,20 +750,15 @@
   // LHLO -> TFRT Dialect (gpu kernels)
   mlir::PassManager pm(mlir_module.getContext(),
                        mlir::PassManager::Nesting::Implicit);
-  pm.addPass(tensorflow::createLmhloGpuAsyncConversionPass());
+  pm.addPass(tensorflow::createConvertLmhloToGpuPass());
   pm.addPass(mlir::createGpuAsyncRegionPass());
   tfrt::gpu::populateGpuToTfrtGpuPasses(pm);
+  pm.addPass(mlir::createCanonicalizerPass());
   if (pm.run(mlir_module).failed()) {
     return InternalError(
         "Failed to lower LHLO to TFRT Dialect with gpu kernels.");
   }
 
-  // Perform DCE with empty pattern set.
-  if (failed(mlir::applyPatternsAndFoldGreedily(
-          mlir_module, mlir::RewritePatternSet(mlir_module.getContext())))) {
-    return InternalError("Failed to remove dead ops.");
-  }
-
   // TFRT Dialect -> BEF
   std::string bef;
   llvm::raw_string_ostream bef_ostream(bef);