message: Check sub-computations for call sites too
diff: @@ -184,24 +184,46 @@ public:
  class CallTargetFinder : public DfsHloVisitorWithDefault {
  public:
  - CallTargetFinder() {}
  + CallTargetFinder(HloComputation* entry) {
  + todo.insert(entry);
  + }
  Status DefaultAction(HloInstruction*) override { return Status::OK(); }
  Status HandleCall(HloInstruction* inst) override {
  - targets[inst->to_apply()]++;
  + CallSiteFound(inst->to_apply(), 1);
  return Status::OK();
  }
  Status HandleWhile(HloInstruction* inst) override {
  - // While targets should always be compiled outlined - so we
  - // add 2 to ensure that they will be.
  - targets[inst->while_condition()] += 2;
  - targets[inst->while_body()] += 2;
  + CallSiteFound(inst->while_condition(), 2);
  + CallSiteFound(inst->while_body(), 2);
  + return Status::OK();
  + }
  +
  + Status Run() {
  + while (!todo.empty()) {
  + auto it = todo.begin();
  + HloComputation *comp = *it;
  + todo.erase(it);
  + done.insert(comp);
  + TF_RETURN_IF_ERROR(comp->Accept(this));
  + }
  return Status::OK();
  }
  std::map<HloComputation*,int> targets;
  +
  +private:
  + void CallSiteFound(HloComputation* comp, int count) {
  + if (done.find(comp) != done.end()) {
  + todo.insert(comp);
  + }
  + targets[comp] += count;
  + }
  +
  + std::set<HloComputation*> todo;
  + std::set<HloComputation*> done;
  };
  Status PoplarCompiler::RunHloOptimization(HloModule* hlo_module) {
  @@ -258,8 +280,8 @@ StatusOr<std::unique_ptr<Executable>> PoplarCompiler::Compile(
  VLOG(2) << "Running poplar call site finder";
  // Find all Call instructions
  - CallTargetFinder call_finder;
  - TF_RETURN_IF_ERROR(entry->Accept(&call_finder));
  + CallTargetFinder call_finder(entry);
  + TF_RETURN_IF_ERROR(call_finder.Run());
  VLOG(2) << "Running tensor allocation tracker";
  