From 14b3a42a42a488a32eae0951b9d4a92da2cd56a2 Mon Sep 17 00:00:00 2001 From: Sanjoy Das Date: Tue, 13 Feb 2018 12:45:27 -0800 Subject: [PATCH] Make LLVM IR dumps more readable Before this the IR we dumped out would contain constant tensors that were sometimes huge resulting in unweildy IR files. With this change we dump out a "noconst" IR file that has the constant initializers stripped out. PiperOrigin-RevId: 185572700 --- tensorflow/compiler/xla/service/llvm_ir/BUILD | 1 + .../compiler/xla/service/llvm_ir/llvm_util.cc | 41 ++++++++++++++++++---- 2 files changed, 35 insertions(+), 7 deletions(-) diff --git a/tensorflow/compiler/xla/service/llvm_ir/BUILD b/tensorflow/compiler/xla/service/llvm_ir/BUILD index ffc78bd..37261ed 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/BUILD +++ b/tensorflow/compiler/xla/service/llvm_ir/BUILD @@ -54,6 +54,7 @@ cc_library( "@llvm//:core", "@llvm//:support", "@llvm//:target", + "@llvm//:transform_utils", ], ) diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc index 8d1e633..15a4cf8 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc @@ -23,6 +23,7 @@ limitations under the License. #include "llvm/IR/MDBuilder.h" #include "llvm/IR/Operator.h" #include "llvm/Target/TargetOptions.h" +#include "llvm/Transforms/Utils/Cloning.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/name_uniquer.h" @@ -61,6 +62,15 @@ llvm::StringRef AsStringRef(tensorflow::StringPiece str) { return llvm::StringRef(str.data(), str.size()); } +std::unique_ptr DropConstantInitializers( + const llvm::Module& module) { + std::unique_ptr cloned_module = CloneModule(&module); + for (llvm::GlobalVariable& global_var : cloned_module->globals()) { + global_var.setInitializer(nullptr); + } + return cloned_module; +} + string DumpModuleToString(const llvm::Module& module) { std::string buffer_string; llvm::raw_string_ostream ostream(buffer_string); @@ -672,6 +682,19 @@ static string GetProcessUniqueIrFileName(tensorflow::StringPiece prefix) { return uniquer->GetUniqueName(prefix); } +static Status CreateAndWriteStringToFile(const string& directory_name, + const string& file_name, + const string& text) { + std::unique_ptr f; + TF_RETURN_IF_ERROR( + tensorflow::Env::Default()->RecursivelyCreateDir(directory_name)); + TF_RETURN_IF_ERROR( + tensorflow::Env::Default()->NewWritableFile(file_name, &f)); + TF_RETURN_IF_ERROR(f->Append(text)); + TF_RETURN_IF_ERROR(f->Close()); + return Status::OK(); +} + Status DumpIRToDirectory(const string& directory_name, const string& hlo_module_name, const llvm::Module& llvm_module, bool optimized) { @@ -686,13 +709,17 @@ Status DumpIRToDirectory(const string& directory_name, directory_name, tensorflow::strings::StrCat(unique_and_safe_file_name, ".ll")); - std::unique_ptr f; - TF_RETURN_IF_ERROR( - tensorflow::Env::Default()->RecursivelyCreateDir(directory_name)); - TF_RETURN_IF_ERROR( - tensorflow::Env::Default()->NewWritableFile(ir_file_name, &f)); - TF_RETURN_IF_ERROR(f->Append(DumpModuleToString(llvm_module))); - return f->Close(); + // For some models the embedded constants can be huge, so also dump the module + // with the constants stripped to get IR that is easier to manipulate. + string ir_no_constant_initializers_file_name = tensorflow::io::JoinPath( + directory_name, + tensorflow::strings::StrCat(unique_and_safe_file_name, "-noconst.ll")); + + TF_RETURN_IF_ERROR(CreateAndWriteStringToFile( + directory_name, ir_file_name, DumpModuleToString(llvm_module))); + return CreateAndWriteStringToFile( + directory_name, ir_no_constant_initializers_file_name, + DumpModuleToString(*DropConstantInitializers(llvm_module))); } llvm::Function* CreateFunction(llvm::FunctionType* function_type, -- 2.7.4