<commit_before>#include "wayfire/nonstd/observer_ptr.h"
#include "wayfire/object.hpp"
#include "wayfire/output.hpp"
#include "wayfire/util.hpp"
#include <sys/types.h>
#include <chrono>
#include <wayfire/plugin.hpp>
#include <wayfire/view.hpp>
#include <wayfire/core.hpp>
#include <wayfire/workspace-manager.hpp>
#include <wayfire/signal-definitions.hpp>
#include <wayfire/output-layout.hpp>
#include <wayfire/nonstd/wlroots-full.hpp>
#include <wayfire/util/log.hpp>
#include <wlr/util/edges.h>

/**
 * View last output info
 */

class last_output_info_t : public wf::custom_data_t
{
  public:
    std::string output_identifier;
    wf::geometry_t geometry;
    bool fullscreen = false;
    bool minimized  = false;
    uint32_t tiled_edges = 0;
    uint z_order;
    bool focused = false;
};

std::string make_output_identifier(wf::output_t *output)
{
    std::string identifier = "";
    identifier += output->handle->make;
    identifier += "|";
    identifier += output->handle->model;
    identifier += "|";
    identifier += output->handle->serial;
    return identifier;
}

void view_store_data(wayfire_view view, wf::output_t *output, int z_order)
{
    auto view_data = view->get_data_safe<last_output_info_t>();
    view_data->output_identifier = make_output_identifier(output);
    view_data->geometry    = view->get_wm_geometry();
    view_data->fullscreen  = view->fullscreen;
    view_data->minimized   = view->minimized;
    view_data->tiled_edges = view->tiled_edges;
    view_data->z_order     = z_order;
    if (view == output->get_active_view())
    {
        view_data->focused = true;
    }
}

nonstd::observer_ptr<last_output_info_t> view_get_data(wayfire_view view)
{
    return view->get_data<last_output_info_t>();
}

bool view_has_data(wayfire_view view)
{
    return view->has_data<last_output_info_t>();
}

void view_erase_data(wayfire_view view)
{
    view->erase_data<last_output_info_t>();
}

/**
 * Core preserve-output info
 */

wf::option_wrapper_t<int> last_output_focus_timeout{
    "preserve-output/last_output_focus_timeout"};

class preserve_output_t : public wf::custom_data_t
{
  public:
    int instances = 0;
    std::string last_focused_output_identifier = "";
    std::chrono::time_point<std::chrono::steady_clock> last_focused_output_timestamp;

    std::map<std::string, wf::point_t> output_saved_workspace;
};

nonstd::observer_ptr<preserve_output_t> get_preserve_output_data()
{
    return wf::get_core().get_data_safe<preserve_output_t>();
}

bool core_focused_output_expired()
{
    using namespace std::chrono;
    const auto now = steady_clock::now();
    const auto last_focus_ts =
        get_preserve_output_data()->last_focused_output_timestamp;
    const auto elapsed_since_focus =
        duration_cast<milliseconds>(now - last_focus_ts).count();

    return elapsed_since_focus > last_output_focus_timeout;
}

void core_store_focused_output(wf::output_t *output)
{
    auto& last_focused_output =
        get_preserve_output_data()->last_focused_output_identifier;
    // Store the output as last focused if no other output has been stored as last
    // focused in the last 10 seconds
    if ((last_focused_output == "") || core_focused_output_expired())
    {
        LOGD("Setting last focused output to: ", output->to_string());
        last_focused_output = make_output_identifier(output);
        get_preserve_output_data()->last_focused_output_timestamp =
            std::chrono::steady_clock::now();
    }
}

std::string core_get_focused_output()
{
    return wf::get_core().get_data_safe<preserve_output_t>()->
           last_focused_output_identifier;
}

void core_erase_focused_output()
{
    wf::get_core().erase_data<preserve_output_t>();
}

/**
 * preserve-output plugin
 */

class wayfire_preserve_output : public wf::plugin_interface_t
{
    bool outputs_being_removed = false;

    wf::signal_connection_t output_pre_remove = [=] (wf::signal_data_t *data)
    {
        auto signal_data = (wf::output_pre_remove_signal*)data;
        LOGD("Received pre-remove event: ", signal_data->output->to_string());
        outputs_being_removed = true;

        if (signal_data->output != output)
        {
            // This event is not for this output
            return;
        }

        // This output is being destroyed
        std::string identifier = make_output_identifier(output);

        // Store this output as the focused one
        if (wf::get_core().get_active_output() == output)
        {
            core_store_focused_output(output);
        }

        get_preserve_output_data()->output_saved_workspace[identifier] =
            output->workspace->get_current_workspace();

        auto views = output->workspace->get_views_in_layer(wf::LAYER_WORKSPACE);
        for (size_t i = 0; i < views.size(); i++)
        {
            auto view = views[i];
            if ((view->role != wf::VIEW_ROLE_TOPLEVEL) || !view->is_mapped())
            {
                continue;
            }

            // Set current output and geometry in the view's last output data
            if (!view_has_data(view))
            {
                view_store_data(view, output, i);
            }
        }
    };

    wf::signal_connection_t output_removed = [=] (wf::signal_data_t *data)
    {
        auto signal_data = (wf::output_removed_signal*)data;
        LOGD("Received output-removed event: ", signal_data->output->to_string());
        outputs_being_removed = false;
    };

    void restore_views_to_output()
    {
        std::string identifier = make_output_identifier(output);

        // Restore active workspace on the output
        // We do this first so that when restoring view's geometries, they land
        // directly on the correct workspace.
        auto core_data = get_preserve_output_data();
        if (core_data->output_saved_workspace.count(identifier))
        {
            output->workspace->set_workspace(
                core_data->output_saved_workspace[identifier]);
        }

        // Focus this output if it was the last one focused
        if (core_get_focused_output() == identifier)
        {
            LOGD("This is last focused output, refocusing: ", output->to_string());
            wf::get_core().focus_output(output);
            core_erase_focused_output();
        }

        // Make a list of views to move to this output
        auto views = std::vector<wayfire_view>();
        for (auto& view : wf::get_core().get_all_views())
        {
            if (!view->is_mapped())
            {
                continue;
            }

            if (!view_has_data(view))
            {
                continue;
            }

            auto last_output_info = view_get_data(view);
            if (last_output_info->output_identifier == identifier)
            {
                views.push_back(view);
            }
        }

        // Sorts with the views closest to front last
        std::sort(views.begin(), views.end(),
            [=] (wayfire_view & view1, wayfire_view & view2)
        {
            return view_get_data(view1)->z_order > view_get_data(view2)->z_order;
        });

        // Move views to this output
        for (auto& view : views)
        {
            auto last_output_info = view_get_data(view);
            LOGD("Restoring view: ",
                view->get_title(), " to: ", output->to_string());

            wf::get_core().move_view_to_output(view, output, false);
            view->set_fullscreen(last_output_info->fullscreen);
            view->set_minimized(last_output_info->minimized);
            if (last_output_info->tiled_edges != 0)
            {
                view->tile_request(last_output_info->tiled_edges);
            }

            view->set_geometry(last_output_info->geometry);

            // Focus
            if (last_output_info->focused)
            {
                LOGD("Focusing view: ", view->get_title());
                output->focus_view(view, false);
            }

            // Z Order
            output->workspace->bring_to_front(view);

            // Remove all last output info from views
            view_erase_data(view);
        }

        // Start listening for view resize events AFTER this callback has finished
        output->connect_signal("view-geometry-changed", &view_moved);
    }

    wf::signal_connection_t view_moved = [=] (wf::signal_data_t *data)
    {
        // Triggered whenever a view on this output's geometry changed
        auto signal_data = (wf::view_geometry_changed_signal*)data;
        auto view = signal_data->view;

        // Ignore event if geometry didn't actually change
        if (signal_data->old_geometry == view->get_wm_geometry())
        {
            return;
        }

        if (view_has_data(view))
        {
            // Remove a view's last output info if it is deliberately moved
            // by user
            if (!outputs_being_removed)
            {
                LOGD("View moved, deleting last output info for: ",
                    view->get_title());
                view_erase_data(view);
            }
        }
    };

    wf::wl_idle_call idle_restore_views;

  public:
    void init() override
    {
        // Increment number of instances of this plugin
        wf::get_core().get_data_safe<preserve_output_t>()->instances++;

        if (wlr_output_is_noop(output->handle))
        {
            // Don't do anything for NO-OP outputs
            return;
        }

        // Call restore_views_to_output() after returning to main loop
        idle_restore_views.run_once([=] ()
        {
            restore_views_to_output();
        });

        wf::get_core().output_layout->connect_signal("output-pre-remove",
            &output_pre_remove);
        wf::get_core().output_layout->connect_signal("output-removed",
            &output_removed);
    }

    void fini() override
    {
        // Decrement number of instances of this plugin
        wf::get_core().get_data_safe<preserve_output_t>()->instances--;
        LOGD("Destroying instance, ",
            wf::get_core().get_data_safe<preserve_output_t>()->instances,
            " remaining");

        // If this is the last instance, delete all data related to this plugin
        if (wf::get_core().get_data_safe<preserve_output_t>()->instances == 0)
        {
            LOGD("This is last instance - deleting all data");
            // Delete data from all views
            for (auto& view : wf::get_core().get_all_views())
            {
                view_erase_data(view);
            }

            // Delete data from core
            core_erase_focused_output();
        }
    }
};

DECLARE_WAYFIRE_PLUGIN(wayfire_preserve_output);
<commit_msg>preserve-output: fix reference counting by using shared_data helpers<commit_after>#include "wayfire/nonstd/observer_ptr.h"
#include "wayfire/object.hpp"
#include "wayfire/output.hpp"
#include "wayfire/util.hpp"
#include <sys/types.h>
#include <chrono>
#include <wayfire/plugin.hpp>
#include <wayfire/view.hpp>
#include <wayfire/core.hpp>
#include <wayfire/workspace-manager.hpp>
#include <wayfire/signal-definitions.hpp>
#include <wayfire/output-layout.hpp>
#include <wayfire/nonstd/wlroots-full.hpp>
#include <wayfire/util/log.hpp>
#include <wlr/util/edges.h>

#include <wayfire/plugins/common/shared-core-data.hpp>

/**
 * View last output info
 */

class last_output_info_t : public wf::custom_data_t
{
  public:
    std::string output_identifier;
    wf::geometry_t geometry;
    bool fullscreen = false;
    bool minimized  = false;
    uint32_t tiled_edges = 0;
    uint z_order;
    bool focused = false;
};

std::string make_output_identifier(wf::output_t *output)
{
    std::string identifier = "";
    identifier += output->handle->make;
    identifier += "|";
    identifier += output->handle->model;
    identifier += "|";
    identifier += output->handle->serial;
    return identifier;
}

void view_store_data(wayfire_view view, wf::output_t *output, int z_order)
{
    auto view_data = view->get_data_safe<last_output_info_t>();
    view_data->output_identifier = make_output_identifier(output);
    view_data->geometry    = view->get_wm_geometry();
    view_data->fullscreen  = view->fullscreen;
    view_data->minimized   = view->minimized;
    view_data->tiled_edges = view->tiled_edges;
    view_data->z_order     = z_order;
    if (view == output->get_active_view())
    {
        view_data->focused = true;
    }
}

nonstd::observer_ptr<last_output_info_t> view_get_data(wayfire_view view)
{
    return view->get_data<last_output_info_t>();
}

bool view_has_data(wayfire_view view)
{
    return view->has_data<last_output_info_t>();
}

void view_erase_data(wayfire_view view)
{
    view->erase_data<last_output_info_t>();
}

/**
 * Core preserve-output info
 */

wf::option_wrapper_t<int> last_output_focus_timeout{
    "preserve-output/last_output_focus_timeout"};

class preserve_output_t : public noncopyable_t
{
  public:
    int instances = 0;
    std::string last_focused_output_identifier = "";
    std::chrono::time_point<std::chrono::steady_clock> last_focused_output_timestamp;

    std::map<std::string, wf::point_t> output_saved_workspace;

    ~preserve_output_t()
    {
        LOGD("This is last instance - deleting all data");
        // Delete data from all views
        for (auto& view : wf::get_core().get_all_views())
        {
            view_erase_data(view);
        }
    }
};

/**
 * preserve-output plugin
 */

class wayfire_preserve_output : public wf::plugin_interface_t
{
    bool outputs_being_removed = false;
    wf::shared_data::ref_ptr_t<preserve_output_t> core_data;

    bool focused_output_expired()
    {
        using namespace std::chrono;
        const auto now = steady_clock::now();
        const auto last_focus_ts = core_data->last_focused_output_timestamp;
        const auto elapsed_since_focus =
            duration_cast<milliseconds>(now - last_focus_ts).count();

        return elapsed_since_focus > last_output_focus_timeout;
    }

    void store_focused_output(wf::output_t *output)
    {
        auto& last_focused_output = core_data->last_focused_output_identifier;
        // Store the output as last focused if no other output has been stored as
        // last
        // focused in the last 10 seconds
        if ((last_focused_output == "") || focused_output_expired())
        {
            LOGD("Setting last focused output to: ", output->to_string());
            last_focused_output = make_output_identifier(output);
            core_data->last_focused_output_timestamp =
                std::chrono::steady_clock::now();
        }
    }

    wf::signal_connection_t output_pre_remove = [=] (wf::signal_data_t *data)
    {
        auto signal_data = (wf::output_pre_remove_signal*)data;
        LOGD("Received pre-remove event: ", signal_data->output->to_string());
        outputs_being_removed = true;

        if (signal_data->output != output)
        {
            // This event is not for this output
            return;
        }

        // This output is being destroyed
        std::string identifier = make_output_identifier(output);

        // Store this output as the focused one
        if (wf::get_core().get_active_output() == output)
        {
            store_focused_output(output);
        }

        core_data->output_saved_workspace[identifier] =
            output->workspace->get_current_workspace();

        auto views = output->workspace->get_views_in_layer(wf::LAYER_WORKSPACE);
        for (size_t i = 0; i < views.size(); i++)
        {
            auto view = views[i];
            if ((view->role != wf::VIEW_ROLE_TOPLEVEL) || !view->is_mapped())
            {
                continue;
            }

            // Set current output and geometry in the view's last output data
            if (!view_has_data(view))
            {
                view_store_data(view, output, i);
            }
        }
    };

    wf::signal_connection_t output_removed = [=] (wf::signal_data_t *data)
    {
        auto signal_data = (wf::output_removed_signal*)data;
        LOGD("Received output-removed event: ", signal_data->output->to_string());
        outputs_being_removed = false;
    };

    void restore_views_to_output()
    {
        std::string identifier = make_output_identifier(output);

        // Restore active workspace on the output
        // We do this first so that when restoring view's geometries, they land
        // directly on the correct workspace.
        if (core_data->output_saved_workspace.count(identifier))
        {
            output->workspace->set_workspace(
                core_data->output_saved_workspace[identifier]);
        }

        // Focus this output if it was the last one focused
        if (core_data->last_focused_output_identifier == identifier)
        {
            LOGD("This is last focused output, refocusing: ", output->to_string());
            wf::get_core().focus_output(output);
            core_data->last_focused_output_identifier.clear();
        }

        // Make a list of views to move to this output
        auto views = std::vector<wayfire_view>();
        for (auto& view : wf::get_core().get_all_views())
        {
            if (!view->is_mapped())
            {
                continue;
            }

            if (!view_has_data(view))
            {
                continue;
            }

            auto last_output_info = view_get_data(view);
            if (last_output_info->output_identifier == identifier)
            {
                views.push_back(view);
            }
        }

        // Sorts with the views closest to front last
        std::sort(views.begin(), views.end(),
            [=] (wayfire_view & view1, wayfire_view & view2)
        {
            return view_get_data(view1)->z_order > view_get_data(view2)->z_order;
        });

        // Move views to this output
        for (auto& view : views)
        {
            auto last_output_info = view_get_data(view);
            LOGD("Restoring view: ",
                view->get_title(), " to: ", output->to_string());

            wf::get_core().move_view_to_output(view, output, false);
            view->set_fullscreen(last_output_info->fullscreen);
            view->set_minimized(last_output_info->minimized);
            if (last_output_info->tiled_edges != 0)
            {
                view->tile_request(last_output_info->tiled_edges);
            }

            view->set_geometry(last_output_info->geometry);

            // Focus
            if (last_output_info->focused)
            {
                LOGD("Focusing view: ", view->get_title());
                output->focus_view(view, false);
            }

            // Z Order
            output->workspace->bring_to_front(view);

            // Remove all last output info from views
            view_erase_data(view);
        }

        // Start listening for view resize events AFTER this callback has finished
        output->connect_signal("view-geometry-changed", &view_moved);
    }

    wf::signal_connection_t view_moved = [=] (wf::signal_data_t *data)
    {
        // Triggered whenever a view on this output's geometry changed
        auto signal_data = (wf::view_geometry_changed_signal*)data;
        auto view = signal_data->view;

        // Ignore event if geometry didn't actually change
        if (signal_data->old_geometry == view->get_wm_geometry())
        {
            return;
        }

        if (view_has_data(view))
        {
            // Remove a view's last output info if it is deliberately moved
            // by user
            if (!outputs_being_removed)
            {
                LOGD("View moved, deleting last output info for: ",
                    view->get_title());
                view_erase_data(view);
            }
        }
    };

    wf::wl_idle_call idle_restore_views;

  public:
    void init() override
    {
        if (wlr_output_is_noop(output->handle))
        {
            // Don't do anything for NO-OP outputs
            return;
        }

        // Call restore_views_to_output() after returning to main loop
        idle_restore_views.run_once([=] ()
        {
            restore_views_to_output();
        });

        wf::get_core().output_layout->connect_signal("output-pre-remove",
            &output_pre_remove);
        wf::get_core().output_layout->connect_signal("output-removed",
            &output_removed);
    }

    void fini() override
    {
        // Nothing to do
    }
};

DECLARE_WAYFIRE_PLUGIN(wayfire_preserve_output);
<|endoftext|>


<commit_before>/* Copyright 2015 Google Inc. All Rights Reserved.

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.
==============================================================================*/

#include <algorithm>
#include <atomic>
#include <set>
#include <unordered_map>
#include <vector>

#include "tensorflow/core/common_runtime/constant_folding.h"

#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/executor.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/rendezvous_mgr.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/public/session_options.h"

namespace tensorflow {

namespace {

bool IsConstantFoldable(const Node* n,
                        std::function<bool(const Node*)> consider) {
  if (n->op_def().is_stateful()) {
    return false;
  }
  if (consider && !consider(n)) {
    return false;
  }
  if (n->IsControlFlow() || n->IsSend() || n->IsRecv()) {
    return false;
  }
  return true;
}

// Returns the constant foldable nodes in `nodes_result` in data flow order.
void FindConstantFoldableNodes(const Graph* graph, ConstantFoldingOptions opts,
                               std::vector<Node*>* nodes_result) {
  std::set<const Node*> node_set;
  std::vector<Node*>& nodes = *nodes_result;
  bool internal_node_inserted = false;
  // Walk the nodes in data flow order
  ReverseDFS(*graph, nullptr,
             [&nodes, &node_set, &internal_node_inserted, opts](Node* n) {
               if (n->IsConstant()) {
                 // Constants with no control inputs (except from _SOURCE node)
                 // are definitely constant foldable.
                 if (n->in_edges().size() == 0 ||
                     (n->in_edges().size() == 1 &&
                      (*n->in_edges().begin())->src()->IsSource())) {
                   node_set.insert(n);
                   nodes.push_back(n);
                 }
               } else if (IsConstantFoldable(n, opts.consider)) {
                 // Check whether the set of this node's in_nodes is completely
                 // included in the set of constant foldable nodes. If true,
                 // then this node is also constant foldable.
                 bool all_parents_constant = n->num_inputs() > 0;
                 for (const Node* parent : n->in_nodes()) {
                   if (node_set.count(parent) == 0) {
                     all_parents_constant = false;
                     break;
                   }
                 }
                 if (all_parents_constant) {
                   node_set.insert(n);
                   nodes.push_back(n);
                   internal_node_inserted = true;
                 }
               }
             });
  // If we have inserted just leaf level nodes, then there is nothing to fold.
  if (!internal_node_inserted) {
    nodes.clear();
  }
}

typedef std::pair<Node*, int> NodeAndOutput;

// Given the constant foldable nodes in 'nodes', returns a new graph 'g'. 'g'
// will contain copies of the nodes in 'nodes'. In addition, if there is an edge
// going from a node 'n' in 'nodes' to another node in 'orig_graph' but not in
// 'nodes', then 'tensors_to_fetch' will contain the mapping from the
// corresponding copy of 'n' and the edge number in 'g' to 'n'.
Graph* GetConstantGraph(const Graph* orig_graph,
                        const std::vector<Node*>& nodes,
                        std::map<NodeAndOutput, Node*>* tensors_to_fetch) {
  Graph* constant_graph = new Graph(orig_graph->op_registry());
  std::unordered_map<Node*, Node*> node_map;
  std::set<Node*> already_added;
  already_added.insert(constant_graph->source_node());
  already_added.insert(constant_graph->sink_node());
  node_map[orig_graph->source_node()] = constant_graph->source_node();
  node_map[orig_graph->sink_node()] = constant_graph->sink_node();
  for (Node* n : nodes) {
    Node* added = constant_graph->CopyNode(n);
    node_map[n] = added;
    already_added.insert(added);
    for (const Edge* in_edge : n->in_edges()) {
      Node* in = in_edge->src();
      CHECK_GT(node_map.count(in), size_t{0}) << n->DebugString() << " <-"
                                              << in->DebugString();
      CHECK_GT(already_added.count(node_map[in]), size_t{0})
          << in->DebugString();
      constant_graph->AddEdge(node_map[in], in_edge->src_output(), added,
                              in_edge->dst_input());
    }
  }

  for (auto const& added_nodes : node_map) {
    for (const Edge* out_edge : added_nodes.first->out_edges()) {
      if (node_map.count(out_edge->dst()) == 0) {
        if (out_edge->IsControlEdge()) continue;
        tensors_to_fetch->insert(
            {{added_nodes.second, out_edge->src_output()}, added_nodes.first});
      }
    }
  }

  return constant_graph;
}

int64 UniqueConstantId() {
  static std::atomic_int_fast64_t id;
  return id.fetch_add(1);
}

Device* GetCPUDevice() {
  static Device* device = nullptr;
  if (!device) {
    std::vector<Device*> devices;
    DeviceFactory::GetFactory(DEVICE_CPU)
        ->CreateDevices(SessionOptions{}, "", &devices);
    if (devices.size() > 0) {
      device = devices[0];
    }
  }
  return device;
}

thread::ThreadPool* GetThreadPool() {
  static thread::ThreadPool* thread_pool =
      new thread::ThreadPool(Env::Default(), "Compute", 1);
  return thread_pool;
}

// A simple rendezvous class.
// Assumes a single sender and a single receiver, no duplicate sends, and no
// sends of dead tensors.
class SimpleRendezvous : public Rendezvous {
 public:
  explicit SimpleRendezvous() {}

  Status Send(const string& key, const Args& send_args, const Tensor& val,
              const bool is_dead) override {
    if (is_dead) {
      return errors::Internal("Send of a dead tensor");
    }
    ParsedKey parsed;
    TF_RETURN_IF_ERROR(ParseKey(key, &parsed));

    mutex_lock l(mu_);
    if (table_.count(parsed.edge_name) > 0) {
      return errors::Internal("Send of an already sent tensor");
    }
    table_[parsed.edge_name] = val;
    return Status::OK();
  }

  void RecvAsync(const string& key, const Args& recv_args,
                 DoneCallback done) override {
    Tensor tensor;
    Status status = Status::OK();
    {
      mutex_lock l(mu_);
      if (table_.count(key) <= 0) {
        status = errors::Internal("Did not find key ", key);
      } else {
        tensor = table_[key];
      }
    }
    done(status, Args{}, recv_args, tensor, false);
  }

  void StartAbort(const Status& status) override {}

 private:
  typedef std::unordered_map<string, Tensor> Table;

  mutex mu_;
  Table table_ GUARDED_BY(mu_);
};

}  // namespace

void ReplaceTensorWithConstant(Graph* graph, NodeAndOutput tensor,
                               const Tensor& constant) {
  Node* n = tensor.first;
  std::vector<const Edge*> edges_to_remove;
  for (const Edge* out_edge : n->out_edges()) {
    if (out_edge->src_output() == tensor.second) {
      edges_to_remove.push_back(out_edge);
    }
  }
  string node_name = n->name();
  Node* constant_node;
  TF_CHECK_OK(NodeBuilder(strings::StrCat(graph->NewName(node_name), "__cf__",
                                          UniqueConstantId()),
                          "Const")
                  .Attr("dtype", constant.dtype())
                  .Attr("value", constant)
                  .Finalize(graph, &constant_node));
  for (auto edge : edges_to_remove) {
    graph->AddEdge(constant_node, 0, edge->dst(), edge->dst_input());
    graph->RemoveEdge(edge);
  }
  graph->AddEdge(graph->source_node(), -1, constant_node, -1);
}

bool DoConstantFolding(const ConstantFoldingOptions& opts, Graph* graph) {
  DumpGraph("Before", graph);
  Device* device = GetCPUDevice();
  thread::ThreadPool* thread_pool = GetThreadPool();
  if (!device || !thread_pool) {
    VLOG(1) << "Cannot find a device and/or a thread pool to do constant "
               "folding on";
    return false;
  }

  std::vector<Node*> constant_foldable_nodes;
  FindConstantFoldableNodes(graph, opts, &constant_foldable_nodes);
  if (constant_foldable_nodes.empty()) {
    VLOG(1) << "No constant foldable nodes found";
    return false;
  }

  std::map<NodeAndOutput, Node*> tensors_to_fetch;
  Graph* constant_graph =
      GetConstantGraph(graph, constant_foldable_nodes, &tensors_to_fetch);
  DumpGraph("Constant graph", constant_graph);

  if (tensors_to_fetch.empty()) {
    VLOG(1) << "No constant nodes found that feed into the original graph.";
    delete constant_graph;
    return false;
  }
  VLOG(1) << "Constant foldable " << constant_graph->num_node_ids() << " : "
          << graph->num_node_ids();

  // Create a local executor and evaluate the constant foldable nodes.
  subgraph::NameIndex name_index;
  for (Node* n : constant_graph->nodes()) {
    name_index[n->name()] = n;
  }

  std::vector<Node*> fetch_nodes;
  std::vector<string> tensors_to_fetch_names;
  std::vector<NodeAndOutput> tensors_to_replace;
  for (auto n : tensors_to_fetch) {
    tensors_to_fetch_names.push_back(
        strings::StrCat(n.first.first->name(), ":", n.first.second));
    tensors_to_replace.push_back({n.second, n.first.second});
  }
  // For nodes that need to be fetched back from the constant_graph, attach Send
  // nodes.
  Status s =
      subgraph::FetchOutputs(constant_graph, device->attributes(),
                             tensors_to_fetch_names, &name_index, &fetch_nodes);
  if (!s.ok()) {
    delete constant_graph;
    VLOG(1) << "Could not fetch constants: " << s;
    return false;
  }

  CHECK_EQ(fetch_nodes.size(), tensors_to_fetch.size());

  // Create the local executor and the Rendezvous for fetching back the
  // constants.
  auto runner = [thread_pool](Executor::Args::Closure c) {
    thread_pool->Schedule(c);
  };
  LocalExecutorParams params;
  params.device = device;
  params.create_kernel = [device, constant_graph](const NodeDef& ndef,
                                                  OpKernel** kernel) {
    return CreateNonCachedKernel(device, nullptr, ndef,
                                 constant_graph->versions().producer(), kernel);
  };
  params.delete_kernel = [](OpKernel* kernel) { delete kernel; };
  Executor* executor;
  if (!NewLocalExecutor(params, constant_graph, &executor).ok()) {
    return false;
  }

  std::unique_ptr<Executor> executor_unref(executor);

  SimpleRendezvous* rendez = new SimpleRendezvous;
  core::ScopedUnref rendez_unref(rendez);

  Executor::Args args;
  args.runner = runner;
  args.rendezvous = rendez;

  // Run the constant_graph.
  Notification executor_done;
  Status executor_done_status;
  ExecutorBarrier* barrier = new ExecutorBarrier(
      1, rendez, [&executor_done, &executor_done_status](const Status& ret) {
        executor_done_status = ret;
        executor_done.Notify();
      });

  executor->RunAsync(args, barrier->Get());

  if (!executor_done_status.ok()) {
    return false;
  }
  executor_done.WaitForNotification();

  // Fetch the constant tensors and replace the corresponding tensors in the
  // original graph with those constants.
  for (size_t c = 0; c < fetch_nodes.size(); ++c) {
    Tensor output;
    bool is_dead;
    string tensor_name;
    if (!GetNodeAttr(fetch_nodes[c]->def(), "tensor_name", &tensor_name).ok()) {
      // We successfully replaced some nodes previously, but had a problem with
      // this node. Don't bother processing the rest of the nodes.
      return c > 0;
    }
    Status s = rendez->Recv(tensor_name, Rendezvous::Args(), &output, &is_dead);
    if (!s.ok() || is_dead) {
      return c > 0;
    }
    VLOG(1) << "Replacing " << tensors_to_replace[c].first->DebugString()
            << " :: " << tensors_to_replace[c].second << " with constant "
            << output.DebugString();
    ReplaceTensorWithConstant(graph, tensors_to_replace[c], output);
  }

  DumpGraph("After", graph);

  return true;
}

}  // namespace tensorflow
<commit_msg>remove duplicate typedef<commit_after>/* Copyright 2015 Google Inc. All Rights Reserved.

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.
==============================================================================*/

#include <algorithm>
#include <atomic>
#include <set>
#include <unordered_map>
#include <vector>

#include "tensorflow/core/common_runtime/constant_folding.h"

#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/executor.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/rendezvous_mgr.h"
#include "tensorflow/core/graph/algorithm.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/public/session_options.h"

namespace tensorflow {

namespace {

bool IsConstantFoldable(const Node* n,
                        std::function<bool(const Node*)> consider) {
  if (n->op_def().is_stateful()) {
    return false;
  }
  if (consider && !consider(n)) {
    return false;
  }
  if (n->IsControlFlow() || n->IsSend() || n->IsRecv()) {
    return false;
  }
  return true;
}

// Returns the constant foldable nodes in `nodes_result` in data flow order.
void FindConstantFoldableNodes(const Graph* graph, ConstantFoldingOptions opts,
                               std::vector<Node*>* nodes_result) {
  std::set<const Node*> node_set;
  std::vector<Node*>& nodes = *nodes_result;
  bool internal_node_inserted = false;
  // Walk the nodes in data flow order
  ReverseDFS(*graph, nullptr,
             [&nodes, &node_set, &internal_node_inserted, opts](Node* n) {
               if (n->IsConstant()) {
                 // Constants with no control inputs (except from _SOURCE node)
                 // are definitely constant foldable.
                 if (n->in_edges().size() == 0 ||
                     (n->in_edges().size() == 1 &&
                      (*n->in_edges().begin())->src()->IsSource())) {
                   node_set.insert(n);
                   nodes.push_back(n);
                 }
               } else if (IsConstantFoldable(n, opts.consider)) {
                 // Check whether the set of this node's in_nodes is completely
                 // included in the set of constant foldable nodes. If true,
                 // then this node is also constant foldable.
                 bool all_parents_constant = n->num_inputs() > 0;
                 for (const Node* parent : n->in_nodes()) {
                   if (node_set.count(parent) == 0) {
                     all_parents_constant = false;
                     break;
                   }
                 }
                 if (all_parents_constant) {
                   node_set.insert(n);
                   nodes.push_back(n);
                   internal_node_inserted = true;
                 }
               }
             });
  // If we have inserted just leaf level nodes, then there is nothing to fold.
  if (!internal_node_inserted) {
    nodes.clear();
  }
}

// Given the constant foldable nodes in 'nodes', returns a new graph 'g'. 'g'
// will contain copies of the nodes in 'nodes'. In addition, if there is an edge
// going from a node 'n' in 'nodes' to another node in 'orig_graph' but not in
// 'nodes', then 'tensors_to_fetch' will contain the mapping from the
// corresponding copy of 'n' and the edge number in 'g' to 'n'.
Graph* GetConstantGraph(const Graph* orig_graph,
                        const std::vector<Node*>& nodes,
                        std::map<NodeAndOutput, Node*>* tensors_to_fetch) {
  Graph* constant_graph = new Graph(orig_graph->op_registry());
  std::unordered_map<Node*, Node*> node_map;
  std::set<Node*> already_added;
  already_added.insert(constant_graph->source_node());
  already_added.insert(constant_graph->sink_node());
  node_map[orig_graph->source_node()] = constant_graph->source_node();
  node_map[orig_graph->sink_node()] = constant_graph->sink_node();
  for (Node* n : nodes) {
    Node* added = constant_graph->CopyNode(n);
    node_map[n] = added;
    already_added.insert(added);
    for (const Edge* in_edge : n->in_edges()) {
      Node* in = in_edge->src();
      CHECK_GT(node_map.count(in), size_t{0}) << n->DebugString() << " <-"
                                              << in->DebugString();
      CHECK_GT(already_added.count(node_map[in]), size_t{0})
          << in->DebugString();
      constant_graph->AddEdge(node_map[in], in_edge->src_output(), added,
                              in_edge->dst_input());
    }
  }

  for (auto const& added_nodes : node_map) {
    for (const Edge* out_edge : added_nodes.first->out_edges()) {
      if (node_map.count(out_edge->dst()) == 0) {
        if (out_edge->IsControlEdge()) continue;
        tensors_to_fetch->insert(
            {{added_nodes.second, out_edge->src_output()}, added_nodes.first});
      }
    }
  }

  return constant_graph;
}

int64 UniqueConstantId() {
  static std::atomic_int_fast64_t id;
  return id.fetch_add(1);
}

Device* GetCPUDevice() {
  static Device* device = nullptr;
  if (!device) {
    std::vector<Device*> devices;
    DeviceFactory::GetFactory(DEVICE_CPU)
        ->CreateDevices(SessionOptions{}, "", &devices);
    if (devices.size() > 0) {
      device = devices[0];
    }
  }
  return device;
}

thread::ThreadPool* GetThreadPool() {
  static thread::ThreadPool* thread_pool =
      new thread::ThreadPool(Env::Default(), "Compute", 1);
  return thread_pool;
}

// A simple rendezvous class.
// Assumes a single sender and a single receiver, no duplicate sends, and no
// sends of dead tensors.
class SimpleRendezvous : public Rendezvous {
 public:
  explicit SimpleRendezvous() {}

  Status Send(const string& key, const Args& send_args, const Tensor& val,
              const bool is_dead) override {
    if (is_dead) {
      return errors::Internal("Send of a dead tensor");
    }
    ParsedKey parsed;
    TF_RETURN_IF_ERROR(ParseKey(key, &parsed));

    mutex_lock l(mu_);
    if (table_.count(parsed.edge_name) > 0) {
      return errors::Internal("Send of an already sent tensor");
    }
    table_[parsed.edge_name] = val;
    return Status::OK();
  }

  void RecvAsync(const string& key, const Args& recv_args,
                 DoneCallback done) override {
    Tensor tensor;
    Status status = Status::OK();
    {
      mutex_lock l(mu_);
      if (table_.count(key) <= 0) {
        status = errors::Internal("Did not find key ", key);
      } else {
        tensor = table_[key];
      }
    }
    done(status, Args{}, recv_args, tensor, false);
  }

  void StartAbort(const Status& status) override {}

 private:
  typedef std::unordered_map<string, Tensor> Table;

  mutex mu_;
  Table table_ GUARDED_BY(mu_);
};

}  // namespace

void ReplaceTensorWithConstant(Graph* graph, NodeAndOutput tensor,
                               const Tensor& constant) {
  Node* n = tensor.first;
  std::vector<const Edge*> edges_to_remove;
  for (const Edge* out_edge : n->out_edges()) {
    if (out_edge->src_output() == tensor.second) {
      edges_to_remove.push_back(out_edge);
    }
  }
  string node_name = n->name();
  Node* constant_node;
  TF_CHECK_OK(NodeBuilder(strings::StrCat(graph->NewName(node_name), "__cf__",
                                          UniqueConstantId()),
                          "Const")
                  .Attr("dtype", constant.dtype())
                  .Attr("value", constant)
                  .Finalize(graph, &constant_node));
  for (auto edge : edges_to_remove) {
    graph->AddEdge(constant_node, 0, edge->dst(), edge->dst_input());
    graph->RemoveEdge(edge);
  }
  graph->AddEdge(graph->source_node(), -1, constant_node, -1);
}

bool DoConstantFolding(const ConstantFoldingOptions& opts, Graph* graph) {
  DumpGraph("Before", graph);
  Device* device = GetCPUDevice();
  thread::ThreadPool* thread_pool = GetThreadPool();
  if (!device || !thread_pool) {
    VLOG(1) << "Cannot find a device and/or a thread pool to do constant "
               "folding on";
    return false;
  }

  std::vector<Node*> constant_foldable_nodes;
  FindConstantFoldableNodes(graph, opts, &constant_foldable_nodes);
  if (constant_foldable_nodes.empty()) {
    VLOG(1) << "No constant foldable nodes found";
    return false;
  }

  std::map<NodeAndOutput, Node*> tensors_to_fetch;
  Graph* constant_graph =
      GetConstantGraph(graph, constant_foldable_nodes, &tensors_to_fetch);
  DumpGraph("Constant graph", constant_graph);

  if (tensors_to_fetch.empty()) {
    VLOG(1) << "No constant nodes found that feed into the original graph.";
    delete constant_graph;
    return false;
  }
  VLOG(1) << "Constant foldable " << constant_graph->num_node_ids() << " : "
          << graph->num_node_ids();

  // Create a local executor and evaluate the constant foldable nodes.
  subgraph::NameIndex name_index;
  for (Node* n : constant_graph->nodes()) {
    name_index[n->name()] = n;
  }

  std::vector<Node*> fetch_nodes;
  std::vector<string> tensors_to_fetch_names;
  std::vector<NodeAndOutput> tensors_to_replace;
  for (auto n : tensors_to_fetch) {
    tensors_to_fetch_names.push_back(
        strings::StrCat(n.first.first->name(), ":", n.first.second));
    tensors_to_replace.push_back({n.second, n.first.second});
  }
  // For nodes that need to be fetched back from the constant_graph, attach Send
  // nodes.
  Status s =
      subgraph::FetchOutputs(constant_graph, device->attributes(),
                             tensors_to_fetch_names, &name_index, &fetch_nodes);
  if (!s.ok()) {
    delete constant_graph;
    VLOG(1) << "Could not fetch constants: " << s;
    return false;
  }

  CHECK_EQ(fetch_nodes.size(), tensors_to_fetch.size());

  // Create the local executor and the Rendezvous for fetching back the
  // constants.
  auto runner = [thread_pool](Executor::Args::Closure c) {
    thread_pool->Schedule(c);
  };
  LocalExecutorParams params;
  params.device = device;
  params.create_kernel = [device, constant_graph](const NodeDef& ndef,
                                                  OpKernel** kernel) {
    return CreateNonCachedKernel(device, nullptr, ndef,
                                 constant_graph->versions().producer(), kernel);
  };
  params.delete_kernel = [](OpKernel* kernel) { delete kernel; };
  Executor* executor;
  if (!NewLocalExecutor(params, constant_graph, &executor).ok()) {
    return false;
  }

  std::unique_ptr<Executor> executor_unref(executor);

  SimpleRendezvous* rendez = new SimpleRendezvous;
  core::ScopedUnref rendez_unref(rendez);

  Executor::Args args;
  args.runner = runner;
  args.rendezvous = rendez;

  // Run the constant_graph.
  Notification executor_done;
  Status executor_done_status;
  ExecutorBarrier* barrier = new ExecutorBarrier(
      1, rendez, [&executor_done, &executor_done_status](const Status& ret) {
        executor_done_status = ret;
        executor_done.Notify();
      });

  executor->RunAsync(args, barrier->Get());

  if (!executor_done_status.ok()) {
    return false;
  }
  executor_done.WaitForNotification();

  // Fetch the constant tensors and replace the corresponding tensors in the
  // original graph with those constants.
  for (size_t c = 0; c < fetch_nodes.size(); ++c) {
    Tensor output;
    bool is_dead;
    string tensor_name;
    if (!GetNodeAttr(fetch_nodes[c]->def(), "tensor_name", &tensor_name).ok()) {
      // We successfully replaced some nodes previously, but had a problem with
      // this node. Don't bother processing the rest of the nodes.
      return c > 0;
    }
    Status s = rendez->Recv(tensor_name, Rendezvous::Args(), &output, &is_dead);
    if (!s.ok() || is_dead) {
      return c > 0;
    }
    VLOG(1) << "Replacing " << tensors_to_replace[c].first->DebugString()
            << " :: " << tensors_to_replace[c].second << " with constant "
            << output.DebugString();
    ReplaceTensorWithConstant(graph, tensors_to_replace[c], output);
  }

  DumpGraph("After", graph);

  return true;
}

}  // namespace tensorflow
<|endoftext|>


<commit_before>/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/

#include "tensorflow/core/common_runtime/executor_factory.h"

#include <unordered_map>

#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"

namespace tensorflow {
namespace {

mutex& executor_factory_lock() {
  static mutex mu(LINKER_INITIALIZED);
  return mu;
}

typedef std::unordered_map<string, ExecutorFactory*> ExecutorFactories;
ExecutorFactories* executor_factories() {
  static ExecutorFactories* factories = new ExecutorFactories;
  return factories;
}

}  // namespace

void ExecutorFactory::Register(const string& executor_type,
                               ExecutorFactory* factory) {
  mutex_lock l(executor_factory_lock());
  if (!executor_factories()->insert({executor_type, factory}).second) {
    LOG(FATAL) << "Two executor factories are being registered "
               << "under" << executor_type;
  }
}

namespace {
const string RegisteredFactoriesErrorMessageLocked()
    TF_SHARED_LOCKS_REQUIRED(executor_factory_lock()) {
  std::vector<string> factory_types;
  for (const auto& executor_factory : *executor_factories()) {
    factory_types.push_back(executor_factory.first);
  }
  return strings::StrCat("Registered factories are {",
                         absl::StrJoin(factory_types, ", "), "}.");
}
}  // namespace

Status ExecutorFactory::GetFactory(const string& executor_type,
                                   ExecutorFactory** out_factory) {
  tf_shared_lock l(executor_factory_lock());

  auto iter = executor_factories()->find(executor_type);
  if (iter == executor_factories()->end()) {
    return errors::NotFound(
        "No executor factory registered for the given executor type: ",
        executor_type, " ", RegisteredFactoriesErrorMessageLocked());
  }

  *out_factory = iter->second;
  return OkStatus();
}

Status NewExecutor(const string& executor_type,
                   const LocalExecutorParams& params, const Graph& graph,
                   std::unique_ptr<Executor>* out_executor) {
  ExecutorFactory* factory = nullptr;
  TF_RETURN_IF_ERROR(ExecutorFactory::GetFactory(executor_type, &factory));
  return factory->NewExecutor(params, std::move(graph), out_executor);
}

}  // namespace tensorflow
<commit_msg>Fix static variable initialization order issue<commit_after>/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/

#include "tensorflow/core/common_runtime/executor_factory.h"

#include <unordered_map>

#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"

namespace tensorflow {
namespace {

static mutex executor_factory_lock(LINKER_INITIALIZED);

typedef std::unordered_map<string, ExecutorFactory*> ExecutorFactories;
ExecutorFactories* executor_factories() {
  static ExecutorFactories* factories = new ExecutorFactories;
  return factories;
}

}  // namespace

void ExecutorFactory::Register(const string& executor_type,
                               ExecutorFactory* factory) {
  mutex_lock l(executor_factory_lock);
  if (!executor_factories()->insert({executor_type, factory}).second) {
    LOG(FATAL) << "Two executor factories are being registered "
               << "under" << executor_type;
  }
}

namespace {
const string RegisteredFactoriesErrorMessageLocked()
    TF_SHARED_LOCKS_REQUIRED(executor_factory_lock) {
  std::vector<string> factory_types;
  for (const auto& executor_factory : *executor_factories()) {
    factory_types.push_back(executor_factory.first);
  }
  return strings::StrCat("Registered factories are {",
                         absl::StrJoin(factory_types, ", "), "}.");
}
}  // namespace

Status ExecutorFactory::GetFactory(const string& executor_type,
                                   ExecutorFactory** out_factory) {
  tf_shared_lock l(executor_factory_lock);

  auto iter = executor_factories()->find(executor_type);
  if (iter == executor_factories()->end()) {
    return errors::NotFound(
        "No executor factory registered for the given executor type: ",
        executor_type, " ", RegisteredFactoriesErrorMessageLocked());
  }

  *out_factory = iter->second;
  return OkStatus();
}

Status NewExecutor(const string& executor_type,
                   const LocalExecutorParams& params, const Graph& graph,
                   std::unique_ptr<Executor>* out_executor) {
  ExecutorFactory* factory = nullptr;
  TF_RETURN_IF_ERROR(ExecutorFactory::GetFactory(executor_type, &factory));
  return factory->NewExecutor(params, std::move(graph), out_executor);
}

}  // namespace tensorflow
<|endoftext|>


<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/
#include "tensorflow/core/tpu/kernels/tpu_compile_op_impl.h"

#include "tensorflow/compiler/xla/status.h"
#include "tensorflow/core/tpu/kernels/tpu_compile.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
#include "tensorflow/core/tpu/kernels/tpu_program_group.h"
#include "tensorflow/core/tpu/kernels/tpu_program_group_interface.h"

namespace tensorflow {
namespace tpu {
Status TpuCompileOpKernelImpl::Compile(
    const std::variant<MlirToHloArgs, FunctionToHloArgs>& computation,
    const XLA_TpuMeshState* mesh_state,
    const std::vector<TensorShape>& arg_shapes,
    TpuProgramGroupInterface* tpu_program_group) {
  TF_ASSIGN_OR_RETURN(
      TpuCompilationRequestProto compilation_request,
      CreateTpuCompilationRequest(computation, metadata_, arg_shapes));

  return TpuProgramGroup::CompileAndBuild(compilation_request, mesh_state,
                                          tpu_program_group);
}
}  // namespace tpu
}  // namespace tensorflow
<commit_msg>Updating std::variant with absl::variant.<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

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.
==============================================================================*/
#include "tensorflow/core/tpu/kernels/tpu_compile_op_impl.h"

#include "tensorflow/compiler/xla/status.h"
#include "tensorflow/core/tpu/kernels/tpu_compile.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/tpu/kernels/tpu_mesh_state_c_api.h"
#include "tensorflow/core/tpu/kernels/tpu_program_group.h"
#include "tensorflow/core/tpu/kernels/tpu_program_group_interface.h"

namespace tensorflow {
namespace tpu {
Status TpuCompileOpKernelImpl::Compile(
    const absl::variant<MlirToHloArgs, FunctionToHloArgs>& computation,
    const XLA_TpuMeshState* mesh_state,
    const std::vector<TensorShape>& arg_shapes,
    TpuProgramGroupInterface* tpu_program_group) {
  TF_ASSIGN_OR_RETURN(
      TpuCompilationRequestProto compilation_request,
      CreateTpuCompilationRequest(computation, metadata_, arg_shapes));

  return TpuProgramGroup::CompileAndBuild(compilation_request, mesh_state,
                                          tpu_program_group);
}
}  // namespace tpu
}  // namespace tensorflow
<|endoftext|>


<commit_before>/*
    fastsim.cpp  -  Don Cross  -  1 July 2016.

    A simulation of mutually repulsive particles trapped on the surface of a sphere.
    This is also known as the Thompson Problem.
    See:  https://en.wikipedia.org/wiki/Thomson_problem
*/

#include <iostream>
#include <iomanip>
#include <fstream>
#include <vector>
#include <utility>
#include <cstdint>
#include <cmath>
#include <cstring>

namespace Electrons
{
    // Vector ---------------------------------------------------------------

    struct Vector
    {
        double  x;
        double  y;
        double  z;

        Vector(): x(0), y(0), z(0) {}
        Vector(double _x, double _y, double _z): x(_x), y(_y), z(_z) {}

        void Reset()
        {
            x = 0;
            y = 0;
            z = 0;
        }

        double MagSquared() const
        {
            return (x*x) + (y*y) + (z*z);
        }

        double Mag() const
        {
            return sqrt(MagSquared());
        }

        Vector UnitVector() const
        {
            double r = Mag();
            if (r < 1.0e-9)
            {
                throw "Cannot find unit vector for near-zero magnitude vector.";
            }
            return Vector(x/r, y/r, z/r);
        }

        Vector& operator += (const Vector& other)
        {
            x += other.x;
            y += other.y;
            z += other.z;
            return *this;
        }

        Vector& operator -= (const Vector& other)
        {
            x -= other.x;
            y -= other.y;
            z -= other.z;
            return *this;
        }

        static double Dot(const Vector& a, const Vector& b)
        {
            return (a.x*b.x) + (a.y*b.y) + (a.z*b.z);
        }
    };

    typedef std::vector<Vector> VectorList;

    inline Vector operator + (const Vector& a, const Vector& b)
    {
        return Vector(a.x + b.x, a.y + b.y, a.z + b.z);
    }

    inline Vector operator - (const Vector& a, const Vector& b)
    {
        return Vector(a.x - b.x, a.y - b.y, a.z - b.z);
    }

    inline Vector operator * (double s, const Vector& v)
    {
        return Vector(s*v.x, s*v.y, s*v.z);
    }

    // Particle ---------------------------------------------------------------

    struct Particle
    {
        Vector position;
        Vector force;

        Particle(): position(), force() {}
        Particle(const Vector& _position): position(_position), force() {}
        Particle(const Vector& _position, const Vector& _force): position(_position), force(_force) {}
    };

    typedef std::vector<Particle> ParticleList;

    // JSON output ------------------------------------------------------------

    void Print(std::ostream& output, double t)
    {
        using namespace std;
        output << setprecision(14) << fixed;
        if (t >= 0.0) output << " ";
        output << t;
    }

    void JsonIndent(std::ostream& output, int indent)
    {
        int spaces = 4*indent;
        for (int i=0; i < spaces; ++i)
        {
            output << " ";
        }
    }

    void JsonPrint(std::ostream& output, const Vector& v)
    {
        output << "{\"x\":";
        Print(output, v.x);
        output << ", \"y\":";
        Print(output, v.y);
        output << ", \"z\":";
        Print(output, v.z);
        output << "}";
    }

    void JsonPrint(std::ostream& output, const Particle& p)
    {
        output << "{\"position\":";
        JsonPrint(output, p.position);
        output << ", \"force\":";
        JsonPrint(output, p.force);
        output << "}";
    }

    void JsonPrint(std::ostream& output, const VectorList& list, int indent)
    {
        using namespace std;
        JsonIndent(output, indent);
        output << "[\n";
        bool first = true;
        for (const Vector& v : list)
        {
            JsonIndent(output, indent);
            output << (first ? "    " : ",   ");
            JsonPrint(output, v);
            output << "\n";
            first = false;
        }
        JsonIndent(output, indent);
        output << "]\n";
    }

    // Random generators ------------------------------------------------------------

    double Random(std::ifstream& infile)
    {
        using namespace std;

        uint64_t data;
        infile.read((char *)&data, sizeof(data));
        if (!infile) throw "Read failure generating random number.";

        // Convert the 64-bit integer to double-precision floating point.
        // Divide by maximum possible value to get a number in the closed range [0, +1].
        double x = static_cast<double>(data) / static_cast<double>(UINT64_MAX);

        // Convert to the closed range [-1, +1].
        return 1.0 - (2.0 * x);
    }

    Vector RandomSpherePoint(std::ifstream& infile)
    {
        // Algorithm for picking a random point on a sphere.
        // Avoids any clustering of points.
        // See equations (9), (10), (11) in:
        // http://mathworld.wolfram.com/SpherePointPicking.html
        while (true)
        {
            double a = Random(infile);
            double b = Random(infile);
            double mag = (a*a) + (b*b);
            if (mag < 1.0)
            {
                double root = 2.0 * sqrt(1.0 - mag);
                return Vector(a*root, b*root, 1.0 - (2.0*mag));
            }
        }
    }

    // Simulation ---------------------------------------------------------------

    class Simulation
    {
    private:
        ParticleList particles;
        int frame;

    public:
        Simulation(int numPoints)       // create an initial random state
            : frame(0)
        {
            using namespace std;

            if (numPoints < 0) throw "Number of points not allowed to be negative.";
            ifstream infile("/dev/urandom", ios::in | ios::binary);
            if (!infile) throw "Could not open /dev/urandom to obtain random numbers.";
            particles.reserve(static_cast<ParticleList::size_type>(numPoints));
            for (int i=0; i < numPoints; ++i)
            {
                particles.push_back(Particle(RandomSpherePoint(infile)));
            }
        }

        int ParticleCount() const
        {
            return static_cast<int>(particles.size());
        }

        int FrameNumber() const
        {
            return frame;
        }

        void JsonPrint(std::ostream& output, int indent) const
        {
            using namespace std;

            JsonIndent(output, indent);
            output << "{\"frame\": " << frame << ", \"particles\": [\n";
            bool first = true;
            for (const Particle& p : particles)
            {
                JsonIndent(output, indent);
                output << (first ? "    " : ",   ");
                Electrons::JsonPrint(output, p);
                output << "\n";
                first = false;
            }
            JsonIndent(output, indent);
            output << "]}\n";
        }

        bool AutoConverge()
        {
            using namespace std;

            // Theory: the simulation converges if total potential energy always decreases.

            // Create an auxiliary particle list to hold candidate next frames.
            const ParticleList::size_type n = particles.size();
            ParticleList nextlist;
            nextlist.reserve(n);
            for (ParticleList::size_type i = 0; i < n; ++i)
            {
                nextlist.push_back(Particle());
            }

            // Start out with a very large dt. Make it smaller until potential energy decreases.
            double dt = 4.0;    // known to optimally converge the n=2 case

            const double DeltaEnergyTolerance = 1.0e-10;
            CalcTangentialForces(particles);
            double energy = PotentialEnergy(particles);
            while (true)    // frame loop: each iteration updates the particles' positions
            {
                ++frame;
                //cout << "frame=" << frame << setprecision(12) << ", energy=" << energy << endl;
                double nextenergy;
                while (true)    // dt adjustment loop: adjust dt as needed for potential energy to decrease
                {
                    UpdatePositions(particles, nextlist, dt);
                    CalcTangentialForces(nextlist);
                    nextenergy = PotentialEnergy(nextlist);

                    if (fabs(nextenergy - energy) < DeltaEnergyTolerance)
                        return true;    // the simulation has settled down enough!

                    if (nextenergy < energy)
                        break;

                    dt *= 0.99;

#if 0
                    cout << "Decreased dt=" << setprecision(6) << fixed << dt << setprecision(12) <<
                        ", energy=" << energy <<
                        ", nextenergy=" << nextenergy <<
                        ", dE=" << scientific << (nextenergy - energy) <<
                        endl;
#endif

                    if (dt < 1.0e-20)
                        return false;   // dt has decreased so much that we are probably stuck!
                }

                swap(particles, nextlist);
                energy = nextenergy;
            }
        }

    private:
        static double CalcTangentialForces(ParticleList& particles)
        {
            // Reset each particle's force to a zero vector.
            for (Particle& p : particles)
            {
                p.force.Reset();
            }

            // Compute force between each unique pair of particles.
            const ParticleList::size_type numParticles = particles.size();
            for (ParticleList::size_type i=0; i < numParticles-1; ++i)
            {
                for (ParticleList::size_type j=i+1; j < numParticles; ++j)
                {
                    AddForces(particles[i], particles[j]);
                }
            }

            // The particles can only move along the surface of the sphere,
            // so eliminate the radial component of all forces,
            // leaving tangential forces.
            double maxForceMag = 0.0;
            for (Particle& p : particles)
            {
                // Calculate radial component using dot product and subtract
                // to get tangential component.
                p.force -= Vector::Dot(p.force, p.position) * p.position;
                double forceMag = p.force.Mag();
                if (forceMag > maxForceMag)
                {
                    maxForceMag = forceMag;
                }
            }

            return maxForceMag;
        }

        static void UpdatePositions(ParticleList& inlist, ParticleList& outlist, double dt)
        {
            const ParticleList::size_type numParticles = inlist.size();
            if (numParticles != outlist.size())
            {
                throw "Inconsistent particle list sizes.";
            }

            for (ParticleList::size_type i=0; i < numParticles; ++i)
            {
                Particle& p = inlist[i];
                Particle& q = outlist[i];
                q.position = ((dt * p.force) + p.position).UnitVector();
            }
        }

        static void AddForces(Particle& a, Particle& b)
        {
            // Force of electrically charged particles:
            // F = k*q1*q2/r^2.
            // We simplify the problem as:  F = 1/r^2.
            // Force is along the direction of the line passing through both.
            Vector dp = a.position - b.position;
            double forcemag = 1.0 / dp.MagSquared();
            Vector force = forcemag * dp.UnitVector();
            a.force += force;
            b.force -= force;
        }

        static double PotentialEnergy(const ParticleList& particles)
        {
            // Potential energy is proportional to the sum of reciprocals of
            // distances between all unique pairs of particles.

            double energy = 0.0;
            const ParticleList::size_type numParticles = particles.size();
            for (ParticleList::size_type i=0; i < numParticles-1; ++i)
            {
                for (ParticleList::size_type j=i+1; j < numParticles; ++j)
                {
                    energy += 1.0 / (particles[i].position - particles[j].position).Mag();
                }
            }
            return energy;
        }
    };
}

//======================================================================================

const int MaxParticles = 1000;

void PrintUsage()
{
    std::cout <<
        "\n"
        "USAGE:\n"
        "\n"
        "fastsim auto N\n"
        "    Simulate N particles using automatic convergence algorithm.\n"
        "\n";
}

int ScanNumParticles(const char *text)
{
    int n = atoi(text);
    if (n<1 || n>MaxParticles)
    {
        throw "Invalid number of particles.";
    }
    return n;
}

//======================================================================================

int main(int argc, const char *argv[])
{
    using namespace std;
    using namespace Electrons;
    try
    {
        if (argc > 1)
        {
            const char *verb = argv[1];

            if (!strcmp(verb, "auto") && (argc == 3))
            {
                int n = ScanNumParticles(argv[2]);
                Simulation sim(n);
                if (sim.AutoConverge())
                {
                    sim.JsonPrint(cout, 0);
                    return 0;
                }
            }
        }

        PrintUsage();
        return 2;
    }
    catch (const char *message)
    {
        cerr << "EXCEPTION: " << message << endl;
        return 3;
    }
}

<commit_msg>Detect convergence based on dP=dE/dt, so small values of dt don't fool us.<commit_after>/*
    fastsim.cpp  -  Don Cross  -  1 July 2016.

    A simulation of mutually repulsive particles trapped on the surface of a sphere.
    This is also known as the Thompson Problem.
    See:  https://en.wikipedia.org/wiki/Thomson_problem
*/

#include <iostream>
#include <iomanip>
#include <fstream>
#include <vector>
#include <utility>
#include <cstdint>
#include <cmath>
#include <cstring>

namespace Electrons
{
    // Vector ---------------------------------------------------------------

    struct Vector
    {
        double  x;
        double  y;
        double  z;

        Vector(): x(0), y(0), z(0) {}
        Vector(double _x, double _y, double _z): x(_x), y(_y), z(_z) {}

        void Reset()
        {
            x = 0;
            y = 0;
            z = 0;
        }

        double MagSquared() const
        {
            return (x*x) + (y*y) + (z*z);
        }

        double Mag() const
        {
            return sqrt(MagSquared());
        }

        Vector UnitVector() const
        {
            double r = Mag();
            if (r < 1.0e-9)
            {
                throw "Cannot find unit vector for near-zero magnitude vector.";
            }
            return Vector(x/r, y/r, z/r);
        }

        Vector& operator += (const Vector& other)
        {
            x += other.x;
            y += other.y;
            z += other.z;
            return *this;
        }

        Vector& operator -= (const Vector& other)
        {
            x -= other.x;
            y -= other.y;
            z -= other.z;
            return *this;
        }

        static double Dot(const Vector& a, const Vector& b)
        {
            return (a.x*b.x) + (a.y*b.y) + (a.z*b.z);
        }
    };

    typedef std::vector<Vector> VectorList;

    inline Vector operator + (const Vector& a, const Vector& b)
    {
        return Vector(a.x + b.x, a.y + b.y, a.z + b.z);
    }

    inline Vector operator - (const Vector& a, const Vector& b)
    {
        return Vector(a.x - b.x, a.y - b.y, a.z - b.z);
    }

    inline Vector operator * (double s, const Vector& v)
    {
        return Vector(s*v.x, s*v.y, s*v.z);
    }

    // Particle ---------------------------------------------------------------

    struct Particle
    {
        Vector position;
        Vector force;

        Particle(): position(), force() {}
        Particle(const Vector& _position): position(_position), force() {}
        Particle(const Vector& _position, const Vector& _force): position(_position), force(_force) {}
    };

    typedef std::vector<Particle> ParticleList;

    // JSON output ------------------------------------------------------------

    void Print(std::ostream& output, double t)
    {
        using namespace std;
        output << setprecision(14) << fixed;
        if (t >= 0.0) output << " ";
        output << t;
    }

    void JsonIndent(std::ostream& output, int indent)
    {
        int spaces = 4*indent;
        for (int i=0; i < spaces; ++i)
        {
            output << " ";
        }
    }

    void JsonPrint(std::ostream& output, const Vector& v)
    {
        output << "{\"x\":";
        Print(output, v.x);
        output << ", \"y\":";
        Print(output, v.y);
        output << ", \"z\":";
        Print(output, v.z);
        output << "}";
    }

    void JsonPrint(std::ostream& output, const Particle& p)
    {
        output << "{\"position\":";
        JsonPrint(output, p.position);
        output << ", \"force\":";
        JsonPrint(output, p.force);
        output << "}";
    }

    void JsonPrint(std::ostream& output, const VectorList& list, int indent)
    {
        using namespace std;
        JsonIndent(output, indent);
        output << "[\n";
        bool first = true;
        for (const Vector& v : list)
        {
            JsonIndent(output, indent);
            output << (first ? "    " : ",   ");
            JsonPrint(output, v);
            output << "\n";
            first = false;
        }
        JsonIndent(output, indent);
        output << "]\n";
    }

    // Random generators ------------------------------------------------------------

    double Random(std::ifstream& infile)
    {
        using namespace std;

        uint64_t data;
        infile.read((char *)&data, sizeof(data));
        if (!infile) throw "Read failure generating random number.";

        // Convert the 64-bit integer to double-precision floating point.
        // Divide by maximum possible value to get a number in the closed range [0, +1].
        double x = static_cast<double>(data) / static_cast<double>(UINT64_MAX);

        // Convert to the closed range [-1, +1].
        return 1.0 - (2.0 * x);
    }

    Vector RandomSpherePoint(std::ifstream& infile)
    {
        // Algorithm for picking a random point on a sphere.
        // Avoids any clustering of points.
        // See equations (9), (10), (11) in:
        // http://mathworld.wolfram.com/SpherePointPicking.html
        while (true)
        {
            double a = Random(infile);
            double b = Random(infile);
            double mag = (a*a) + (b*b);
            if (mag < 1.0)
            {
                double root = 2.0 * sqrt(1.0 - mag);
                return Vector(a*root, b*root, 1.0 - (2.0*mag));
            }
        }
    }

    // Simulation ---------------------------------------------------------------

    class Simulation
    {
    private:
        ParticleList particles;
        int frame;

    public:
        Simulation(int numPoints)       // create an initial random state
            : frame(0)
        {
            using namespace std;

            if (numPoints < 0) throw "Number of points not allowed to be negative.";
            ifstream infile("/dev/urandom", ios::in | ios::binary);
            if (!infile) throw "Could not open /dev/urandom to obtain random numbers.";
            particles.reserve(static_cast<ParticleList::size_type>(numPoints));
            for (int i=0; i < numPoints; ++i)
            {
                particles.push_back(Particle(RandomSpherePoint(infile)));
            }
        }

        int ParticleCount() const
        {
            return static_cast<int>(particles.size());
        }

        int FrameNumber() const
        {
            return frame;
        }

        void JsonPrint(std::ostream& output, int indent) const
        {
            using namespace std;

            JsonIndent(output, indent);
            output << "{\"frame\": " << frame << ", \"particles\": [\n";
            bool first = true;
            for (const Particle& p : particles)
            {
                JsonIndent(output, indent);
                output << (first ? "    " : ",   ");
                Electrons::JsonPrint(output, p);
                output << "\n";
                first = false;
            }
            JsonIndent(output, indent);
            output << "]}\n";
        }

        bool AutoConverge()
        {
            using namespace std;

            // Theory: the simulation converges if total potential energy always decreases.

            // Create an auxiliary particle list to hold candidate next frames.
            const ParticleList::size_type n = particles.size();
            ParticleList nextlist;
            nextlist.reserve(n);
            for (ParticleList::size_type i = 0; i < n; ++i)
            {
                nextlist.push_back(Particle());
            }

            // Start out with a very large dt. Make it smaller until potential energy decreases.
            double dt = 4.0;    // known to optimally converge the n=2 case

            const double DeltaPowerTolerance = 1.0e-30;
            CalcTangentialForces(particles);
            double energy = PotentialEnergy(particles);
            while (true)    // frame loop: each iteration updates the particles' positions
            {
                ++frame;
                //cout << "frame=" << frame << setprecision(12) << ", energy=" << energy << ", dt=" << dt << endl;
                double nextenergy;
                while (true)    // dt adjustment loop: adjust dt as needed for potential energy to decrease
                {
                    UpdatePositions(particles, nextlist, dt);
                    CalcTangentialForces(nextlist);
                    nextenergy = PotentialEnergy(nextlist);

                    // We want to detect convergence based on potential energy settling down.
                    // But because dt can change, this alone could cause dE to appear very small.
                    // So calculate the instantaneous power dP = dE/dt as a metric of convergence.
                    double dP = (nextenergy - energy) / dt;
                    //cout << "dP = " << scientific << setprecision(15) << dP << endl;

                    if (fabs(dP) < DeltaPowerTolerance)
                    {
                        //cout << "final dt = " << dt << endl;
                        swap(particles, nextlist);  // get that last little bit of refinement
                        return true;    // the simulation has settled down enough!
                    }

                    if (nextenergy < energy)
                        break;

                    dt *= 0.99;

#if 0
                    cout << "Decreased dt=" << setprecision(6) << fixed << dt << setprecision(12) <<
                        ", energy=" << energy <<
                        ", nextenergy=" << nextenergy <<
                        ", dE=" << scientific << (nextenergy - energy) <<
                        endl;
#endif

                    if (dt < 1.0e-20)
                        return false;   // dt has decreased so much that we are probably stuck!
                }

                swap(particles, nextlist);
                energy = nextenergy;
            }
        }

    private:
        static double CalcTangentialForces(ParticleList& particles)
        {
            // Reset each particle's force to a zero vector.
            for (Particle& p : particles)
            {
                p.force.Reset();
            }

            // Compute force between each unique pair of particles.
            const ParticleList::size_type numParticles = particles.size();
            for (ParticleList::size_type i=0; i < numParticles-1; ++i)
            {
                for (ParticleList::size_type j=i+1; j < numParticles; ++j)
                {
                    AddForces(particles[i], particles[j]);
                }
            }

            // The particles can only move along the surface of the sphere,
            // so eliminate the radial component of all forces,
            // leaving tangential forces.
            double maxForceMag = 0.0;
            for (Particle& p : particles)
            {
                // Calculate radial component using dot product and subtract
                // to get tangential component.
                p.force -= Vector::Dot(p.force, p.position) * p.position;
                double forceMag = p.force.Mag();
                if (forceMag > maxForceMag)
                {
                    maxForceMag = forceMag;
                }
            }

            return maxForceMag;
        }

        static void UpdatePositions(ParticleList& inlist, ParticleList& outlist, double dt)
        {
            const ParticleList::size_type numParticles = inlist.size();
            if (numParticles != outlist.size())
            {
                throw "Inconsistent particle list sizes.";
            }

            for (ParticleList::size_type i=0; i < numParticles; ++i)
            {
                Particle& p = inlist[i];
                Particle& q = outlist[i];
                q.position = ((dt * p.force) + p.position).UnitVector();
            }
        }

        static void AddForces(Particle& a, Particle& b)
        {
            // Force of electrically charged particles:
            // F = k*q1*q2/r^2.
            // We simplify the problem as:  F = 1/r^2.
            // Force is along the direction of the line passing through both.
            Vector dp = a.position - b.position;
            double forcemag = 1.0 / dp.MagSquared();
            Vector force = forcemag * dp.UnitVector();
            a.force += force;
            b.force -= force;
        }

        static double PotentialEnergy(const ParticleList& particles)
        {
            // Potential energy is proportional to the sum of reciprocals of
            // distances between all unique pairs of particles.

            double energy = 0.0;
            const ParticleList::size_type numParticles = particles.size();
            for (ParticleList::size_type i=0; i < numParticles-1; ++i)
            {
                for (ParticleList::size_type j=i+1; j < numParticles; ++j)
                {
                    energy += 1.0 / (particles[i].position - particles[j].position).Mag();
                }
            }
            return energy;
        }
    };
}

//======================================================================================

const int MaxParticles = 1000;

void PrintUsage()
{
    std::cout <<
        "\n"
        "USAGE:\n"
        "\n"
        "fastsim auto N\n"
        "    Simulate N particles using automatic convergence algorithm.\n"
        "\n";
}

int ScanNumParticles(const char *text)
{
    int n = atoi(text);
    if (n<1 || n>MaxParticles)
    {
        throw "Invalid number of particles.";
    }
    return n;
}

//======================================================================================

int main(int argc, const char *argv[])
{
    using namespace std;
    using namespace Electrons;
    try
    {
        if (argc > 1)
        {
            const char *verb = argv[1];

            if (!strcmp(verb, "auto") && (argc == 3))
            {
                int n = ScanNumParticles(argv[2]);
                Simulation sim(n);
                if (sim.AutoConverge())
                {
                    sim.JsonPrint(cout, 0);
                    return 0;
                }
            }
        }

        PrintUsage();
        return 2;
    }
    catch (const char *message)
    {
        cerr << "EXCEPTION: " << message << endl;
        return 3;
    }
}

<|endoftext|>