Fix for OOM in EfficientDet

Summary:
The changes for a unified instruction cache was causing EfficientDet to go OOM (T66916). This was caused by the input layout flag not being provided correctly to the SubcomputationGraphCache.

FIX T68951

Reviewers: alfiee, jackh, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, zigmasb

Reviewed By: alfiee, jackh, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, zigmasb

Maniphest Tasks: T68951

Differential Revision: https://phabricator.sourcevertex.net/D79039

Author: Gautham Ganapathy

tensorflow/compiler/plugin/poplar/driver/ops/map_ops.cc

@@ -666,7 +666,7 @@ StatusOr<DriverProgramSequence> CreateFunctionOp(
   HloComputation* comp = inst->to_apply();
 
   bool keep_input_layouts = false;
-  if (IsFunction(inst)) {
+  if (IsFunction(inst) || IsCall(inst)) {
     keep_input_layouts = GetFunctionKeepInputLayouts(inst);
   }
 

tensorflow/compiler/plugin/poplar/driver/tools/util.cc

@@ -457,6 +457,10 @@ bool IsFunction(const HloInstruction* inst) {
   return CallConfigHasType(inst, PoplarBackendConfig::CallConfig::Function);
 }
 
+bool IsCall(const HloInstruction* inst) {
+  return CallConfigHasType(inst, PoplarBackendConfig::CallConfig::Call);
+}
+
 bool IsMultiConv(const HloInstruction* inst) {
   return CallConfigHasType(inst, PoplarBackendConfig::CallConfig::MultiConv);
 }

tensorflow/compiler/plugin/poplar/driver/tools/util.h

@@ -173,6 +173,7 @@ bool IsPipelineStageBackward(const HloInstruction*);
 bool IsPipelineStageRecomputation(const HloInstruction*);
 bool IsResourceUpdate(const HloInstruction*);
 bool IsFunction(const HloInstruction*);
+bool IsCall(const HloInstruction*);
 bool IsMultiConv(const HloInstruction*);
 bool IsPipelineOp(const HloInstruction*);
 bool IsBatchSerializedPipelineOp(const HloInstruction*);