8000 Add uncompilable node info to error message if function is uncompilable. · compnerd/tensorflow@d7a0333 · GitHub
[go: up one dir, main page]
Skip to content

Commit d7a0333

Browse files
er91tensorflower-gardener
er91
authored and
tensorflower-gardener
committed
Add uncompilable node info to error message if function is uncompilable.
PiperOrigin-RevId: 252750929
1 parent fdbe59e commit d7a0333

File tree

4 files changed

+46
-11
lines changed

4 files changed

+46
-11
lines changed

‎tensorflow/compiler/jit/BUILD

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -336,6 +336,7 @@ cc_library(
336336
],
337337
deps = [
338338
":common",
339+
":compilability_check_util",
339340
":compilation_passes",
340341
"//tensorflow/compiler/jit/kernels:xla_ops",
341342
"//tensorflow/compiler/tf2xla:xla_compiler",
@@ -344,6 +345,8 @@ cc_library(
344345
"//tensorflow/core:lib",
345346
"//tensorflow/core:protos_all_cc",
346347
"@com_google_absl//absl/memory",
348+
"@com_google_absl//absl/strings",
349+
"@com_google_absl//absl/strings:str_format",
347350
],
348351
alwayslink = 1,
349352
)

‎tensorflow/compiler/jit/mark_for_compilation_pass.cc

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1559,7 +1559,10 @@ std::atomic<int64>* GetPointerToFuel(int64 initial_value) {
15591559
}
15601560
} // anonymous namespace
15611561

1562-
bool IsCompilable(FunctionLibraryRuntime* flr, const NodeDef& ndef) {
1562+
bool IsCompilable(
1563+
FunctionLibraryRuntime* flr, const NodeDef& ndef,
1564+
std::vector<RecursiveCompilabilityChecker::UncompilableNodeInfo>*
1565+
uncompilable_node_info) {
15631566
Device* device = flr->device();
15641567
const XlaOpRegistry::DeviceRegistration* registration;
15651568
CHECK(XlaOpRegistry::GetCompilationDevice(device->device_type(),
@@ -1579,8 +1582,16 @@ bool IsCompilable(FunctionLibraryRuntime* flr, const NodeDef& ndef) {
15791582
op_filter.allow_slow_ops = true;
15801583
op_filter.allow_inaccurate_ops = true;
15811584

1582-
return RecursiveCompilabilityChecker{&op_filter, &jit_device_type}
1583-
.IsCompilableCall(ndef, flr);
1585+
RecursiveCompilabilityChecker checker{&op_filter, &jit_device_type};
1586+
if (!uncompilable_node_info) {
1587+
// We do not need uncompilable node info. Just return the result.
1588+
return checker.IsCompilableCall(ndef, flr);
1589+
}
1590+
1591+
std::vector<RecursiveCompilabilityChecker::UncompilableNodeInfo>
1592+
uncompilable_node_result = checker.FindUncompilableNodes(ndef, flr);
1593+
uncompilable_node_info->swap(uncompilable_node_result);
1594+
return uncompilable_node_info->empty();
15841595
}
15851596

15861597
Status MarkForCompilationPass::Run(

‎tensorflow/compiler/jit/mark_for_compilation_pass.h

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ limitations under the License.
2020
#ifndef TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_H_
2121
#define TENSORFLOW_COMPILER_JIT_MARK_FOR_COMPILATION_PASS_H_
2222

23+
#include "tensorflow/compiler/jit/compilability_check_util.h"
2324
#include "tensorflow/core/common_runtime/optimization_registry.h"
2425

2526
namespace tensorflow {
@@ -49,8 +50,12 @@ class MarkForCompilationPass : public GraphOptimizationPass {
4950

5051
// Returns true iff 'ndef' is a call to a function that is compilable. A
5152
// function is compilable iff every operator in the function body is
52-
// compilable.
53-
bool IsCompilable(FunctionLibraryRuntime* flr, const NodeDef& ndef);
53+
// compilable. If 'ndef' is not compilable and 'uncompilable_node_info' is not
54+
// null, we will populate 'uncompilable_node_info' with uncompilable node info.
55+
bool IsCompilable(
56+
FunctionLibraryRuntime* flr, const NodeDef& ndef,
57+
std::vector<RecursiveCompilabilityChecker::UncompilableNodeInfo>*
58+
uncompilable_node_info = nullptr);
5459

5560
namespace testing {
5661
// DO NOT USE IN PRODUCTION.

‎tensorflow/compiler/jit/xla_kernel_creator.cc

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@ limitations under the License.
1515
#include "tensorflow/compiler/jit/xla_kernel_creator.h"
1616

1717
#include "absl/memory/memory.h"
18+
#include "absl/strings/str_cat.h"
19+
#include "absl/strings/str_format.h"
20+
#include "tensorflow/compiler/jit/compilability_check_util.h"
1821
#include "tensorflow/compiler/jit/defs.h"
1922
#include "tensorflow/compiler/jit/kernels/xla_ops.h"
2023
#include "tensorflow/compiler/jit/mark_for_compilation_pass.h"
@@ -156,13 +159,26 @@ Status XlaKernelCreator::CreateKernel(FunctionLibraryRuntime* flr,
156159

157160
// Make sure that kernels have been registered on the JIT device.
158161
XlaOpRegistry::RegisterCompilationKernels();
159-
if (!IsCompilable(flr, node_def)) {
160-
VLOG(1) << "Not creating XlaLaunchOp because function invoked by the "
161-
"following node is not compilable: "
162-
<< node_def.DebugString();
162+
std::vector<RecursiveCompilabilityChecker::UncompilableNodeInfo>
163+
uncompilable_node_info;
164+
if (!IsCompilable(flr, node_def, &uncompilable_node_info)) {
165+
string message = absl::StrCat(
166+
"Function invoked by the following node is not compilable: ",
167+
node_def.ShortDebugString(), ".\n");
168+
absl::StrAppend(&message, "Uncompilable nodes:\n");
169+
for (const auto& node_info : uncompilable_node_info) {
170+
string node_message =
171+
absl::StrCat("\t", node_info.name, ": ",
172+
node_info.uncompilable_reason, "\n", "\tStacktrace:\n");
173+
for (const auto& stack_frame : node_info.stack_trace) {
174+
absl::StrAppendFormat(&node_message, "\t\tNode: %s, function: %s\n",
175+
stack_frame.name, stack_frame.function_name);
176+
}
177+
absl::StrAppend(&message, node_message);
178+
}
179+
VLOG(1) << message;
163180
// node_def is calling a function that XLA can't compile.
164-
return errors::InvalidArgument("Not compilable: ",
165-
node_def.ShortDebugString());
181+
return errors::InvalidArgument(message);
166182
}
167183

168184
// Get function body, constant args, and resource args.

0 commit comments

Comments
 (0)
0