MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / HasXLAKernel

Method HasXLAKernel

tensorflow/compiler/jit/compilability_check_util.cc:128–156  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

126}
127
128bool RecursiveCompilabilityChecker::HasXLAKernel(const Node& node) const {
129 // There is a SymbolicGradient kernel on the XLA_JIT device, but the gradient
130 // is really a kind of function call and will be handled by
131 // IsCompilableCall().
132 if (node.type_string() == "SymbolicGradient") return false;
133 if (node.type_string() == "Const") {
134 // Skip Const op with type DT_STRING, since XLA doesn't support it, but the
135 // registered Const KernelDef says that it does, to support no-op Assert for
136 // tfcompile.
137 const AttrValue* attr = node.attrs().Find("dtype");
138 if (attr != nullptr && attr->type() == DT_STRING) {
139 VLOG(2) << "Rejecting " << node.name() << ": Const with STRING attribute.";
140 return false;
141 }
142 }
143
144 // XLA does not offer guaranteed aliasing between the input and output of the
145 // XLA cluster so it can't implement the forward-tensor-ref semantic. Leave
146 // such nodes out of XLA clusters.
147 if (HasForwardedRefInput(node)) {
148 VLOG(2) << "Rejecting " << node.name() << ": Identity with unsafe cast.";
149 return false;
150 }
151 Status stat = FindKernelDef(jit_device_type_, node.def(), nullptr, nullptr);
152 if (!stat.ok()) {
153 VLOG(3) << "Rejecting " << node.name() << ": Can't find a matching kernel. " << stat.error_message();
154 }
155 return stat.ok();
156}
157
158// Tests whether 'if_node' is compilable. Every operator in the then_branch and
159// else_branch functions must be compilable for 'if_node' to be compilable.

Callers

nothing calls this directly

Calls 7

HasForwardedRefInputFunction · 0.85
FindKernelDefFunction · 0.85
typeMethod · 0.65
nameMethod · 0.65
FindMethod · 0.45
attrsMethod · 0.45
okMethod · 0.45

Tested by

no test coverage detected