[XLA:CPU] Store ThreadSafeModule in LlvmIrKernelSource

PiperOrigin-RevId: 719321299
openxla · Jan 27, 2025 · 7ea5233 · 7ea5233
1 parent fd2988c
commit 7ea5233
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git 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
@@ -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