diff --git a/xla/codegen/llvm_ir_kernel_source.cc b/xla/codegen/llvm_ir_kernel_source.cc
index bc9af24f45cfce..c7a0ad85bc6463 100644
--- a/xla/codegen/llvm_ir_kernel_source.cc
+++ b/xla/codegen/llvm_ir_kernel_source.cc
@@ -17,12 +17,14 @@ limitations under the License.
 
 #include <string>
 
+#include "llvm/IR/Module.h"
 #include "xla/service/llvm_ir/llvm_util.h"
 
 namespace xla {
 
 std::string LlvmIrKernelSource::ToString() const {
-  return llvm_ir::DumpToString(module_.get());
+  return module_.withModuleDo(
+      [&](llvm::Module& m) { return llvm_ir::DumpToString(&m); });
 }
 
 }  // namespace xla
diff --git a/xla/codegen/llvm_ir_kernel_source.h b/xla/codegen/llvm_ir_kernel_source.h
index be3a18927e69ac..872df7567b2548 100644
--- a/xla/codegen/llvm_ir_kernel_source.h
+++ b/xla/codegen/llvm_ir_kernel_source.h
@@ -35,20 +35,19 @@ class LlvmIrKernelSource final : public KernelSource {
  public:
   LlvmIrKernelSource(llvm::orc::ThreadSafeContext context,
                      std::unique_ptr<llvm::Module> module)
-      : context_(std::move(context)), module_(std::move(module)) {}
+      : module_(std::move(module), std::move(context)) {}
 
   LlvmIrKernelSource(LlvmIrKernelSource&& other) = default;
   LlvmIrKernelSource& operator=(LlvmIrKernelSource&& other) = default;
 
   llvm::orc::ThreadSafeModule thread_safe_module() && {
-    return llvm::orc::ThreadSafeModule(std::move(module_), context_);
+    return std::move(module_);
   }
 
   std::string ToString() const final;
 
  private:
-  llvm::orc::ThreadSafeContext context_;
-  std::unique_ptr<llvm::Module> module_;
+  llvm::orc::ThreadSafeModule module_;
 };
 
 }  // namespace xla