diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/BUILD b/tensorflow/compiler/mlir/tools/kernel_gen/BUILD index b4cbf765c79..5949ee6498b 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/BUILD +++ b/tensorflow/compiler/mlir/tools/kernel_gen/BUILD @@ -85,7 +85,10 @@ cc_library( tf_cc_binary( name = "tf_to_gpu_binary", - srcs = ["tf_to_gpu_binary.cc"], + srcs = [ + "crash_handler.h", + "tf_to_gpu_binary.cc", + ], visibility = [ "//tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_to_gpu_binary:__pkg__", "//tensorflow/core/kernels/mlir_generated:__pkg__", @@ -95,6 +98,7 @@ tf_cc_binary( "//tensorflow/compiler/mlir:init_mlir", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/core:lib", + "//tensorflow/core/platform", "//tensorflow/stream_executor/lib", "@com_google_absl//absl/strings", "@llvm-project//llvm:Support", diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/crash_handler.h b/tensorflow/compiler/mlir/tools/kernel_gen/crash_handler.h new file mode 100644 index 00000000000..9ecaa40e567 --- /dev/null +++ b/tensorflow/compiler/mlir/tools/kernel_gen/crash_handler.h @@ -0,0 +1,38 @@ +// 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_TOOLS_KERNEL_GEN_CRASH_HANDLER_H_ +#define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_CRASH_HANDLER_H_ + +#include "llvm/Support/PrettyStackTrace.h" +#include "tensorflow/core/platform/platform.h" + +namespace tensorflow { +namespace kernel_gen { + +inline void SetCrashReportMessage() { +#if defined(PLATFORM_GOOGLE) + llvm::setBugReportMsg( + "The TensorFlow Kernel Generator crashed, see the docs at " + "go/tf-kernel-gen for debug hints and contact information.\n"); +#else + llvm::setBugReportMsg( + "The TensorFlow Kernel Generator crashed, please report a bug with the " + "trace below on https://github.com/tensorflow/tensorflow/issues.\n"); +#endif +} +} // namespace kernel_gen +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_CRASH_HANDLER_H_ diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/tf_to_gpu_binary.cc b/tensorflow/compiler/mlir/tools/kernel_gen/tf_to_gpu_binary.cc index 84c2bf46b55..7ecae51c194 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/tf_to_gpu_binary.cc +++ b/tensorflow/compiler/mlir/tools/kernel_gen/tf_to_gpu_binary.cc @@ -25,6 +25,7 @@ #include "llvm/Support/CommandLine.h" #include "mlir/Pass/PassManager.h" // from @llvm-project #include "tensorflow/compiler/mlir/init_mlir.h" +#include "tensorflow/compiler/mlir/tools/kernel_gen/crash_handler.h" #include "tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" @@ -63,6 +64,7 @@ xla::Status Run(llvm::StringRef input_file, llvm::StringRef output_file, } // namespace tensorflow int main(int argc, char** argv) { + tensorflow::kernel_gen::SetCrashReportMessage(); llvm::cl::opt input_file("input", llvm::cl::desc("input file"), llvm::cl::value_desc("filename"), llvm::cl::init("foo.mlir"));