<commit_before>//=======================================================================
// Copyright (c) 2015 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#include<iostream>
#include<cerrno>

#include<wiringPi.h>

#include<stdio.h>
#include<stdlib.h>
#include<stdint.h>
#include<sys/socket.h>
#include<sys/un.h>
#include<sys/types.h>
#include<unistd.h>
#include<signal.h>

namespace {

const std::size_t UNIX_PATH_MAX = 108;
const std::size_t gpio_pin = 24;
const std::size_t max_timings = 85;
const std::size_t buffer_size = 4096;

// Configuration (this should be in a configuration file)
const char* server_socket_path = "/tmp/asgard_socket";
const char* client_socket_path = "/tmp/asgard_dht11_socket";
const std::size_t delay_ms = 20000;

//Buffers
char write_buffer[buffer_size + 1];
char receive_buffer[buffer_size + 1];

// The socket file descriptor
int socket_fd;

// The socket addresses
struct sockaddr_un client_address;
struct sockaddr_un server_address;

// The remote IDs
int source_id = -1;
int temperature_sensor_id = -1;
int humidity_sensor_id = -1;

void stop(){
    std::cout << "asgard:dht11: stop the driver" << std::endl;

    // Unregister the temperature sensor, if necessary
    if(temperature_sensor_id >= 0){
        auto nbytes = snprintf(write_buffer, buffer_size, "UNREG_SENSOR %d %d", source_id, temperature_sensor_id);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));
    }

    // Unregister the humidity sensor, if necessary
    if(temperature_sensor_id >= 0){
        auto nbytes = snprintf(write_buffer, buffer_size, "UNREG_SENSOR %d %d", source_id, humidity_sensor_id);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));
    }

    // Unregister the source, if necessary
    if(source_id >= 0){
        auto nbytes = snprintf(write_buffer, buffer_size, "UNREG_SOURCE %d", source_id);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));
    }

    // Unlink the client socket
    unlink(client_socket_path);

    // Close the socket
    close(socket_fd);
}

void terminate(int){
    stop();

    std::exit(0);
}

void read_data(){
    uint8_t laststate   = HIGH;

    int dht11_dat[5] = { 0, 0, 0, 0, 0 };

    /* pull pin down for 18 milliseconds */
    pinMode(gpio_pin, OUTPUT);
    digitalWrite(gpio_pin, LOW);
    delay(18);

    /* then pull it up for 40 microseconds */
    digitalWrite(gpio_pin, HIGH);
    delayMicroseconds(40);

    /* prepare to read the pin */
    pinMode(gpio_pin, INPUT);

    /* detect change and read data */
    uint8_t j = 0;
    for(uint8_t i = 0; i < max_timings; i++){
        uint8_t counter = 0;
        while (digitalRead(gpio_pin) == laststate){
            counter++;
            delayMicroseconds(1);
            if (counter == 255){
                break;
            }
        }

        laststate = digitalRead(gpio_pin);

        if (counter == 255){
            break;
        }

        /* ignore first 3 transitions */
        if (i >= 4 && i % 2 == 0){
            /* shove each bit into the storage bytes */
            dht11_dat[j / 8] <<= 1;
            if (counter > 16){
                dht11_dat[j / 8] |= 1;
            }
            j++;
        }
    }

    /*
     * check we read 40 bits (8bit x 5 ) + verify checksum in the last byte
     */
    if (j >= 40 && dht11_dat[4] == ((dht11_dat[0] + dht11_dat[1] + dht11_dat[2] + dht11_dat[3]) & 0xFF)){
        //Send the humidity to the server
        auto nbytes = snprintf(write_buffer, buffer_size, "DATA %d %d %d", source_id, humidity_sensor_id, dht11_dat[0]);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));

        //Send the temperature to the server
        nbytes = snprintf(write_buffer, buffer_size, "DATA %d %d %d", source_id, temperature_sensor_id, dht11_dat[2]);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));
    }
}

bool revoke_root(){
    if (getuid() == 0) {
        if (setgid(1000) != 0){
            std::cout << "asgard:dht11: setgid: Unable to drop group privileges: " << strerror(errno) << std::endl;
            return false;
        }

        if (setuid(1000) != 0){
            std::cout << "asgard:dht11: setgid: Unable to drop user privileges: " << strerror(errno) << std::endl;
            return false;
        }
    }

    if (setuid(0) != -1){
        std::cout << "asgard:dht11: managed to regain root privileges, exiting..." << std::endl;
        return false;
    }

    return true;
}

} //end of namespace

int main(){
    //Run the wiringPi setup (as root)
    wiringPiSetup();

    //Drop root privileges and run as pi:pi again
    if(!revoke_root()){
       std::cout << "asgard:dht11: unable to revoke root privileges, exiting..." << std::endl;
       return 1;
    }

    // Open the socket
    socket_fd = socket(AF_UNIX, SOCK_DGRAM, 0);
    if(socket_fd < 0){
        std::cerr << "asgard:dht11: socket() failed" << std::endl;
        return 1;
    }

    // Init the client address
    memset(&client_address, 0, sizeof(struct sockaddr_un));
    client_address.sun_family = AF_UNIX;
    snprintf(client_address.sun_path, UNIX_PATH_MAX, client_socket_path);

    // Unlink the client socket
    unlink(client_socket_path);

    // Bind to client socket
    if(bind(socket_fd, (const struct sockaddr *) &client_address, sizeof(struct sockaddr_un)) < 0){
        std::cerr << "asgard:dht11: bind() failed" << std::endl;
        return 1;
    }

    //Register signals for "proper" shutdown
    signal(SIGTERM, terminate);
    signal(SIGINT, terminate);

    // Init the server address
    memset(&server_address, 0, sizeof(struct sockaddr_un));
    server_address.sun_family = AF_UNIX;
    snprintf(server_address.sun_path, UNIX_PATH_MAX, server_socket_path);

    socklen_t address_length = sizeof(struct sockaddr_un);

    // Register the source
    auto nbytes = snprintf(write_buffer, buffer_size, "REG_SOURCE dht11");
    sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));

    auto bytes_received = recvfrom(socket_fd, receive_buffer, buffer_size, 0, (struct sockaddr *) &(server_address), &address_length);
    receive_buffer[bytes_received] = '\0';

    source_id = atoi(receive_buffer);

    std::cout << "asgard:dht11: remote source: " << source_id << std::endl;

    // Register the temperature sensor
    nbytes = snprintf(write_buffer, buffer_size, "REG_SENSOR %d %s %s", source_id, "TEMPERATURE", "local");
    sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));

    bytes_received = recvfrom(socket_fd, receive_buffer, buffer_size, 0, (struct sockaddr *) &(server_address), &address_length);
    receive_buffer[bytes_received] = '\0';

    temperature_sensor_id = atoi(receive_buffer);

    std::cout << "asgard:random: remote temperature sensor: " << temperature_sensor_id << std::endl;

    // Register the humidity sensor
    nbytes = snprintf(write_buffer, buffer_size, "REG_SENSOR %d %s %s", source_id, "HUMIDITY", "local");
    sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));

    bytes_received = recvfrom(socket_fd, receive_buffer, buffer_size, 0, (struct sockaddr *) &(server_address), &address_length);
    receive_buffer[bytes_received] = '\0';

    humidity_sensor_id = atoi(receive_buffer);

    std::cout << "asgard:random: remote humidity sensor: " << humidity_sensor_id << std::endl;

    //Read data continuously
    while (1){
        read_data();
        delay(delay_ms);
    }

    stop();

    return 0;
}
<commit_msg>Cleanup messages<commit_after>//=======================================================================
// Copyright (c) 2015 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#include<iostream>
#include<cerrno>

#include<wiringPi.h>

#include<stdio.h>
#include<stdlib.h>
#include<stdint.h>
#include<sys/socket.h>
#include<sys/un.h>
#include<sys/types.h>
#include<unistd.h>
#include<signal.h>

namespace {

const std::size_t UNIX_PATH_MAX = 108;
const std::size_t gpio_pin = 24;
const std::size_t max_timings = 85;
const std::size_t buffer_size = 4096;

// Configuration (this should be in a configuration file)
const char* server_socket_path = "/tmp/asgard_socket";
const char* client_socket_path = "/tmp/asgard_dht11_socket";
const std::size_t delay_ms = 20000;

//Buffers
char write_buffer[buffer_size + 1];
char receive_buffer[buffer_size + 1];

// The socket file descriptor
int socket_fd;

// The socket addresses
struct sockaddr_un client_address;
struct sockaddr_un server_address;

// The remote IDs
int source_id = -1;
int temperature_sensor_id = -1;
int humidity_sensor_id = -1;

void stop(){
    std::cout << "asgard:dht11: stop the driver" << std::endl;

    // Unregister the temperature sensor, if necessary
    if(temperature_sensor_id >= 0){
        auto nbytes = snprintf(write_buffer, buffer_size, "UNREG_SENSOR %d %d", source_id, temperature_sensor_id);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));
    }

    // Unregister the humidity sensor, if necessary
    if(temperature_sensor_id >= 0){
        auto nbytes = snprintf(write_buffer, buffer_size, "UNREG_SENSOR %d %d", source_id, humidity_sensor_id);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));
    }

    // Unregister the source, if necessary
    if(source_id >= 0){
        auto nbytes = snprintf(write_buffer, buffer_size, "UNREG_SOURCE %d", source_id);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));
    }

    // Unlink the client socket
    unlink(client_socket_path);

    // Close the socket
    close(socket_fd);
}

void terminate(int){
    stop();

    std::exit(0);
}

void read_data(){
    uint8_t laststate   = HIGH;

    int dht11_dat[5] = { 0, 0, 0, 0, 0 };

    /* pull pin down for 18 milliseconds */
    pinMode(gpio_pin, OUTPUT);
    digitalWrite(gpio_pin, LOW);
    delay(18);

    /* then pull it up for 40 microseconds */
    digitalWrite(gpio_pin, HIGH);
    delayMicroseconds(40);

    /* prepare to read the pin */
    pinMode(gpio_pin, INPUT);

    /* detect change and read data */
    uint8_t j = 0;
    for(uint8_t i = 0; i < max_timings; i++){
        uint8_t counter = 0;
        while (digitalRead(gpio_pin) == laststate){
            counter++;
            delayMicroseconds(1);
            if (counter == 255){
                break;
            }
        }

        laststate = digitalRead(gpio_pin);

        if (counter == 255){
            break;
        }

        /* ignore first 3 transitions */
        if (i >= 4 && i % 2 == 0){
            /* shove each bit into the storage bytes */
            dht11_dat[j / 8] <<= 1;
            if (counter > 16){
                dht11_dat[j / 8] |= 1;
            }
            j++;
        }
    }

    /*
     * check we read 40 bits (8bit x 5 ) + verify checksum in the last byte
     */
    if (j >= 40 && dht11_dat[4] == ((dht11_dat[0] + dht11_dat[1] + dht11_dat[2] + dht11_dat[3]) & 0xFF)){
        //Send the humidity to the server
        auto nbytes = snprintf(write_buffer, buffer_size, "DATA %d %d %d", source_id, humidity_sensor_id, dht11_dat[0]);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));

        //Send the temperature to the server
        nbytes = snprintf(write_buffer, buffer_size, "DATA %d %d %d", source_id, temperature_sensor_id, dht11_dat[2]);
        sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));
    }
}

bool revoke_root(){
    if (getuid() == 0) {
        if (setgid(1000) != 0){
            std::cout << "asgard:dht11: setgid: Unable to drop group privileges: " << strerror(errno) << std::endl;
            return false;
        }

        if (setuid(1000) != 0){
            std::cout << "asgard:dht11: setgid: Unable to drop user privileges: " << strerror(errno) << std::endl;
            return false;
        }
    }

    if (setuid(0) != -1){
        std::cout << "asgard:dht11: managed to regain root privileges, exiting..." << std::endl;
        return false;
    }

    return true;
}

} //end of namespace

int main(){
    //Run the wiringPi setup (as root)
    wiringPiSetup();

    //Drop root privileges and run as pi:pi again
    if(!revoke_root()){
       std::cout << "asgard:dht11: unable to revoke root privileges, exiting..." << std::endl;
       return 1;
    }

    // Open the socket
    socket_fd = socket(AF_UNIX, SOCK_DGRAM, 0);
    if(socket_fd < 0){
        std::cerr << "asgard:dht11: socket() failed" << std::endl;
        return 1;
    }

    // Init the client address
    memset(&client_address, 0, sizeof(struct sockaddr_un));
    client_address.sun_family = AF_UNIX;
    snprintf(client_address.sun_path, UNIX_PATH_MAX, client_socket_path);

    // Unlink the client socket
    unlink(client_socket_path);

    // Bind to client socket
    if(bind(socket_fd, (const struct sockaddr *) &client_address, sizeof(struct sockaddr_un)) < 0){
        std::cerr << "asgard:dht11: bind() failed" << std::endl;
        return 1;
    }

    //Register signals for "proper" shutdown
    signal(SIGTERM, terminate);
    signal(SIGINT, terminate);

    // Init the server address
    memset(&server_address, 0, sizeof(struct sockaddr_un));
    server_address.sun_family = AF_UNIX;
    snprintf(server_address.sun_path, UNIX_PATH_MAX, server_socket_path);

    socklen_t address_length = sizeof(struct sockaddr_un);

    // Register the source
    auto nbytes = snprintf(write_buffer, buffer_size, "REG_SOURCE dht11");
    sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));

    auto bytes_received = recvfrom(socket_fd, receive_buffer, buffer_size, 0, (struct sockaddr *) &(server_address), &address_length);
    receive_buffer[bytes_received] = '\0';

    source_id = atoi(receive_buffer);

    std::cout << "asgard:dht11: remote source: " << source_id << std::endl;

    // Register the temperature sensor
    nbytes = snprintf(write_buffer, buffer_size, "REG_SENSOR %d %s %s", source_id, "TEMPERATURE", "local");
    sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));

    bytes_received = recvfrom(socket_fd, receive_buffer, buffer_size, 0, (struct sockaddr *) &(server_address), &address_length);
    receive_buffer[bytes_received] = '\0';

    temperature_sensor_id = atoi(receive_buffer);

    std::cout << "asgard:dht11: remote temperature sensor: " << temperature_sensor_id << std::endl;

    // Register the humidity sensor
    nbytes = snprintf(write_buffer, buffer_size, "REG_SENSOR %d %s %s", source_id, "HUMIDITY", "local");
    sendto(socket_fd, write_buffer, nbytes, 0, (struct sockaddr *) &server_address, sizeof(struct sockaddr_un));

    bytes_received = recvfrom(socket_fd, receive_buffer, buffer_size, 0, (struct sockaddr *) &(server_address), &address_length);
    receive_buffer[bytes_received] = '\0';

    humidity_sensor_id = atoi(receive_buffer);

    std::cout << "asgard:dht11: remote humidity sensor: " << humidity_sensor_id << std::endl;

    //Read data continuously
    while (1){
        read_data();
        delay(delay_ms);
    }

    stop();

    return 0;
}
<|endoftext|>


<commit_before>// SimEngine.cc -- a simulation engine
// Computing using SI units.
// This case is a hard-coded simulation until we figure out how to
// do otherwise.

#include <stdlib.h>
#include <dlfcn.h>
#include <iostream>
#include <cstdio>
#include <string>

#include "models/Atmos.h"
#include "models/FlightVehicle.h"
#include "models/SimpleMotion.h"
#include "utils/Writer.h"
#include "SimEngine.h"

#define LOGSUFFIX       ".log"
#ifndef DLSFX
#error DLSFX is not defined.
#endif

using namespace std;

// n is file handle number
void
SimEngine::setHandle(const int n, FILE* f) {
    switch (n) {
      case fnOut:
        fOut = f; break;
      case fnLog:
        fLog = f; break;
      default:
        fprintf(stderr, "setHandle failed\n");
        fflush(stderr);
        exit(1);
        break;
    }
}

Scenario*
SimEngine::loadScenario(const char* scenario_name) {
    //  void *scen_dyn;     // dynamic loading address
    //  Scenario *scen;
    bool load_fail = false;

    string module_name = string(scenario_name) + "." + DLSFX;

    scen_dyn = dlopen(module_name.c_str(), RTLD_LAZY);
    if (!scen_dyn) {
        cerr << "Cannot load scenario: " << dlerror() << '\n';
        cerr.flush();
        exit(EXIT_FAILURE);
    }

    // reset errors
    const char* dlsym_error;
    dlerror();

    // load the symbols
    dlscen_table.create =
            reinterpret_cast<scen_create_t*> (dlsym(scen_dyn, "create"));
    dlsym_error = dlerror();
    if (dlsym_error) {
        cerr << "Cannot load Scenario symbol create: " << dlsym_error << endl;
        load_fail = true;
    }

    dlscen_table.destroy =
        reinterpret_cast<scen_destroy_t*> (dlsym(scen_dyn, "destroy"));
    dlsym_error = dlerror();
    if (dlsym_error) {
        cerr << "Cannot load Scenario symbol destroy: " << dlsym_error << endl;
        load_fail = true;
    }

    if (load_fail) {
        exit(EXIT_FAILURE);
    }
    scen = dlscen_table.create();

    return scen;
}

void
SimEngine::unloadScenario() {
    dlscen_table.destroy(scen);
    dlclose(scen_dyn);
    cerr << "unloaded" << endl;
}

void
SimEngine::setScenario(const char* scenario_name) {
    scenario = scenario_name;
}

unsigned int
SimEngine::run() {

    Writer* writer = new Writer();
    writer->putmsg((const float)0.0, "veh1", "prolog", "start of new log");

    string logname = string(scenario) + LOGSUFFIX;
    fLog = fopen(logname.c_str(), "w");
    setHandle(SimEngine::fnLog, fLog);
    setHandle(SimEngine::fnOut, stdout);

    fprintf(fLog, "output to log\n");
    fprintf(stdout, "output to stdout\n");

    // Load the scenario module
    scen = loadScenario(scenario);
    cerr << scen->getName() << endl;

    // allow log to have CSV-like entries
    // headers.
    // step#, vehid, pos[0:2], world velo[0:2], veh velo[0:2], veh accel[0:2]
    SimpleMotionState* sms;
    scen->initcond();
    // INITIAL WRITE-OUT GOES HERE
    // initial step
    writer->initLine(0.0, "veh1", "traj");
    writer->putCart3Head("x", "0", "1", "2");
    writer->putCart3Head("dx", "0", "1", "2");
    writer->putCart3Head("ddx", "0", "1", "2");
    writer->finiLine();
#if 0
    writer->initLine(0.0, "veh1", "SMLD");
    writer->finiLine();
#endif /*0*/

    while (scen->moresteps()) {
        scen->step();
        sms = scen->getSimpleMotion();

        writer->initLine(sms->t, "veh1", "traj");
        writer->putCart3fLine(sms->x);
        writer->putCart3fLine(sms->dx);
        writer->putCart3fLine(sms->ddx);
        writer->finiLine();
#if 0
        writer->initLine(sms->t, "veh1", "SMLD");
        writer->finiLine();
#endif /*0*/

        // INCREMENTAL WRITE-OUT GOES HERE
        // flight angles
        // atmospheric conditions
        // summary
    }

    unloadScenario();

    return 0;
}

<commit_msg>Put scenario name in sim.log<commit_after>// SimEngine.cc -- a simulation engine
// Computing using SI units.
// This case is a hard-coded simulation until we figure out how to
// do otherwise.

#include <stdlib.h>
#include <dlfcn.h>
#include <iostream>
#include <cstdio>
#include <string>

#include "models/Atmos.h"
#include "models/FlightVehicle.h"
#include "models/SimpleMotion.h"
#include "utils/Writer.h"
#include "SimEngine.h"

#define LOGSUFFIX       ".log"
#ifndef DLSFX
#error DLSFX is not defined.
#endif

using namespace std;

// n is file handle number
void
SimEngine::setHandle(const int n, FILE* f) {
    switch (n) {
      case fnOut:
        fOut = f; break;
      case fnLog:
        fLog = f; break;
      default:
        fprintf(stderr, "setHandle failed\n");
        fflush(stderr);
        exit(1);
        break;
    }
}

Scenario*
SimEngine::loadScenario(const char* scenario_name) {
    //  void *scen_dyn;     // dynamic loading address
    //  Scenario *scen;
    bool load_fail = false;

    string module_name = string(scenario_name) + "." + DLSFX;

    scen_dyn = dlopen(module_name.c_str(), RTLD_LAZY);
    if (!scen_dyn) {
        cerr << "Cannot load scenario: " << dlerror() << '\n';
        cerr.flush();
        exit(EXIT_FAILURE);
    }

    // reset errors
    const char* dlsym_error;
    dlerror();

    // load the symbols
    dlscen_table.create =
            reinterpret_cast<scen_create_t*> (dlsym(scen_dyn, "create"));
    dlsym_error = dlerror();
    if (dlsym_error) {
        cerr << "Cannot load Scenario symbol create: " << dlsym_error << endl;
        load_fail = true;
    }

    dlscen_table.destroy =
        reinterpret_cast<scen_destroy_t*> (dlsym(scen_dyn, "destroy"));
    dlsym_error = dlerror();
    if (dlsym_error) {
        cerr << "Cannot load Scenario symbol destroy: " << dlsym_error << endl;
        load_fail = true;
    }

    if (load_fail) {
        exit(EXIT_FAILURE);
    }
    scen = dlscen_table.create();

    return scen;
}

void
SimEngine::unloadScenario() {
    dlscen_table.destroy(scen);
    dlclose(scen_dyn);
    cerr << "unloaded" << endl;
}

void
SimEngine::setScenario(const char* scenario_name) {
    scenario = scenario_name;
}

unsigned int
SimEngine::run() {

    Writer* writer = new Writer();
    writer->putmsg((const float)0.0, scenario, "prolog", "start of new log");

    string logname = string(scenario) + LOGSUFFIX;
    fLog = fopen(logname.c_str(), "w");
    setHandle(SimEngine::fnLog, fLog);
    setHandle(SimEngine::fnOut, stdout);

    fprintf(fLog, "output to log\n");
    fprintf(stdout, "output to stdout\n");

    // Load the scenario module
    scen = loadScenario(scenario);
    cerr << scen->getName() << endl;

    // allow log to have CSV-like entries
    // headers.
    // step#, vehid, pos[0:2], world velo[0:2], veh velo[0:2], veh accel[0:2]
    SimpleMotionState* sms;
    scen->initcond();
    // INITIAL WRITE-OUT GOES HERE
    // initial step
    writer->initLine(0.0, scenario, "traj");
    writer->putCart3Head("x", "0", "1", "2");
    writer->putCart3Head("dx", "0", "1", "2");
    writer->putCart3Head("ddx", "0", "1", "2");
    writer->finiLine();

    while (scen->moresteps()) {
        scen->step();
        sms = scen->getSimpleMotion();

        writer->initLine(sms->t, scenario, "traj");
        writer->putCart3fLine(sms->x);
        writer->putCart3fLine(sms->dx);
        writer->putCart3fLine(sms->ddx);
        writer->finiLine();

        // INCREMENTAL WRITE-OUT GOES HERE
        // flight angles
        // atmospheric conditions
        // summary
    }

    unloadScenario();

    return 0;
}

<|endoftext|>


<commit_before>#include "app.hpp"
#include "gui.hpp"
#include "util/request.hpp"
#include "../ext/osdialog/osdialog.h"
#include <string.h>
#include <thread>

namespace rack {

static std::string newVersion = "";

static void checkVersion() {
	json_t *resJ = requestJson(METHOD_GET, gApiHost + "/version", NULL);

	if (resJ) {
		json_t *versionJ = json_object_get(resJ, "version");
		if (versionJ) {
			const char *version = json_string_value(versionJ);
			if (version && strlen(version) > 0 && version != gApplicationVersion) {
				newVersion = version;
			}
		}
		json_decref(resJ);
	}
}

RackScene::RackScene() {
	scrollWidget = new RackScrollWidget();
	{
		zoomWidget = new ZoomWidget();
		{
			assert(!gRackWidget);
			gRackWidget = new RackWidget();
			zoomWidget->addChild(gRackWidget);
		}
		scrollWidget->container->addChild(zoomWidget);
	}
	addChild(scrollWidget);

	gToolbar = new Toolbar();
	addChild(gToolbar);
	scrollWidget->box.pos.y = gToolbar->box.size.y;

	// Check for new version
	if (!gApplicationVersion.empty()) {
		std::thread versionThread(checkVersion);
		versionThread.detach();
	}
}

void RackScene::step() {
	// Resize owned descendants
	gToolbar->box.size.x = box.size.x;
	scrollWidget->box.size = box.size.minus(scrollWidget->box.pos);

	// Resize to be a bit larger than the ScrollWidget viewport
	gRackWidget->box.size = scrollWidget->box.size
		.minus(scrollWidget->container->box.pos)
		.plus(Vec(500, 500))
		.div(zoomWidget->zoom);

	Scene::step();

	zoomWidget->box.size = gRackWidget->box.size.mult(zoomWidget->zoom);

	// Version popup message
	if (!newVersion.empty()) {
		std::string versionMessage = stringf("Rack %s is available.\n\nYou have Rack %s.\n\nWould you like to download the new version on the website?", newVersion.c_str(), gApplicationVersion.c_str());
		if (osdialog_message(OSDIALOG_INFO, OSDIALOG_YES_NO, versionMessage.c_str())) {
			std::thread t(openBrowser, "https://vcvrack.com/");
			t.detach();
			guiClose();
		}
		newVersion = "";
	}
}

void RackScene::draw(NVGcontext *vg) {
	Scene::draw(vg);
}

void RackScene::onHoverKey(EventHoverKey &e) {
	switch (e.key) {
		case GLFW_KEY_N:
			if (guiIsModPressed() && !guiIsShiftPressed()) {
				gRackWidget->reset();
				e.consumed = true;
				return;
			}
			break;
		case GLFW_KEY_Q:
			if (guiIsModPressed() && !guiIsShiftPressed()) {
				guiClose();
				e.consumed = true;
				return;
			}
			break;
		case GLFW_KEY_O:
			if (guiIsModPressed() && !guiIsShiftPressed()) {
				gRackWidget->openDialog();
				e.consumed = true;
				return;
			}
			break;
		case GLFW_KEY_S:
			if (guiIsModPressed() && !guiIsShiftPressed()) {
				gRackWidget->saveDialog();
				e.consumed = true;
				return;
			}
			if (guiIsModPressed() && guiIsShiftPressed()) {
				gRackWidget->saveAsDialog();
				e.consumed = true;
				return;
			}
			break;
	}

	Widget::onHoverKey(e);
}

void RackScene::onPathDrop(EventPathDrop &e) {
	if (e.paths.size() >= 1) {
		const std::string& firstPath = e.paths.front();
		if (extractExtension(firstPath) == "vcv") {
			gRackWidget->loadPatch(firstPath);
			e.consumed = true;
		}
	}
}

} // namespace rack
<commit_msg>Pass onPathDrop to RackScene children<commit_after>#include "app.hpp"
#include "gui.hpp"
#include "util/request.hpp"
#include "../ext/osdialog/osdialog.h"
#include <string.h>
#include <thread>

namespace rack {

static std::string newVersion = "";

static void checkVersion() {
	json_t *resJ = requestJson(METHOD_GET, gApiHost + "/version", NULL);

	if (resJ) {
		json_t *versionJ = json_object_get(resJ, "version");
		if (versionJ) {
			const char *version = json_string_value(versionJ);
			if (version && strlen(version) > 0 && version != gApplicationVersion) {
				newVersion = version;
			}
		}
		json_decref(resJ);
	}
}

RackScene::RackScene() {
	scrollWidget = new RackScrollWidget();
	{
		zoomWidget = new ZoomWidget();
		{
			assert(!gRackWidget);
			gRackWidget = new RackWidget();
			zoomWidget->addChild(gRackWidget);
		}
		scrollWidget->container->addChild(zoomWidget);
	}
	addChild(scrollWidget);

	gToolbar = new Toolbar();
	addChild(gToolbar);
	scrollWidget->box.pos.y = gToolbar->box.size.y;

	// Check for new version
	if (!gApplicationVersion.empty()) {
		std::thread versionThread(checkVersion);
		versionThread.detach();
	}
}

void RackScene::step() {
	// Resize owned descendants
	gToolbar->box.size.x = box.size.x;
	scrollWidget->box.size = box.size.minus(scrollWidget->box.pos);

	// Resize to be a bit larger than the ScrollWidget viewport
	gRackWidget->box.size = scrollWidget->box.size
		.minus(scrollWidget->container->box.pos)
		.plus(Vec(500, 500))
		.div(zoomWidget->zoom);

	Scene::step();

	zoomWidget->box.size = gRackWidget->box.size.mult(zoomWidget->zoom);

	// Version popup message
	if (!newVersion.empty()) {
		std::string versionMessage = stringf("Rack %s is available.\n\nYou have Rack %s.\n\nWould you like to download the new version on the website?", newVersion.c_str(), gApplicationVersion.c_str());
		if (osdialog_message(OSDIALOG_INFO, OSDIALOG_YES_NO, versionMessage.c_str())) {
			std::thread t(openBrowser, "https://vcvrack.com/");
			t.detach();
			guiClose();
		}
		newVersion = "";
	}
}

void RackScene::draw(NVGcontext *vg) {
	Scene::draw(vg);
}

void RackScene::onHoverKey(EventHoverKey &e) {
	switch (e.key) {
		case GLFW_KEY_N:
			if (guiIsModPressed() && !guiIsShiftPressed()) {
				gRackWidget->reset();
				e.consumed = true;
				return;
			}
			break;
		case GLFW_KEY_Q:
			if (guiIsModPressed() && !guiIsShiftPressed()) {
				guiClose();
				e.consumed = true;
				return;
			}
			break;
		case GLFW_KEY_O:
			if (guiIsModPressed() && !guiIsShiftPressed()) {
				gRackWidget->openDialog();
				e.consumed = true;
				return;
			}
			break;
		case GLFW_KEY_S:
			if (guiIsModPressed() && !guiIsShiftPressed()) {
				gRackWidget->saveDialog();
				e.consumed = true;
				return;
			}
			if (guiIsModPressed() && guiIsShiftPressed()) {
				gRackWidget->saveAsDialog();
				e.consumed = true;
				return;
			}
			break;
	}

	Widget::onHoverKey(e);
}

void RackScene::onPathDrop(EventPathDrop &e) {
	if (e.paths.size() >= 1) {
		const std::string& firstPath = e.paths.front();
		if (extractExtension(firstPath) == "vcv") {
			gRackWidget->loadPatch(firstPath);
			e.consumed = true;
		}
	}

	if (!e.consumed)
		Scene::onPathDrop(e);
}

} // namespace rack
<|endoftext|>


<commit_before>#include <QStringList>
#include <QString>

#include "plow/plow.h"

#include "JobTree.h"

namespace Plow { namespace Gui {

JobTree::JobTree(QWidget *parent) : QWidget(parent)
{
    layout = new QVBoxLayout(this);

    QStringList header;
    header << "Job" << "Cores" << "Max" << "Waiting";

    treeWidget = new QTreeWidget;
    treeWidget->setHeaderLabels(header);
    treeWidget->setColumnCount(4);
    treeWidget->setColumnWidth(0, 300);

    layout->addWidget(treeWidget);

}

void JobTree::updateJobs()
{

    std::vector<JobT> jobs;
    JobFilterT filter;

    std::vector<JobState::type> states;
    states.push_back(JobState::RUNNING);
    filter.states = states;

    getJobs(jobs, filter);

    for (std::vector<JobT>::iterator i = jobs.begin();
                                     i != jobs.end();
                                     ++i)
    {
        QStringList data;
        data << QString::fromStdString(i->name)
             << QString::number(i->runningCoreCount)
             << QString::number(i->maxCores)
             << QString::number(i->waitingTaskCount);

        QTreeWidgetItem *item = new QTreeWidgetItem((QTreeWidget*)0, data);
        treeWidget->insertTopLevelItem(0, item);
    }
}

} // Gui
} // Plow<commit_msg>Updated for latest thrift file.<commit_after>#include <QStringList>
#include <QString>

#include "plow/plow.h"

#include "JobTree.h"

namespace Plow { namespace Gui {

JobTree::JobTree(QWidget *parent) : QWidget(parent)
{
    layout = new QVBoxLayout(this);

    QStringList header;
    header << "Job" << "Cores" << "Max" << "Waiting";

    treeWidget = new QTreeWidget;
    treeWidget->setHeaderLabels(header);
    treeWidget->setColumnCount(4);
    treeWidget->setColumnWidth(0, 300);

    layout->addWidget(treeWidget);

}

void JobTree::updateJobs()
{

    std::vector<JobT> jobs;
    JobFilterT filter;

    std::vector<JobState::type> states;
    states.push_back(JobState::RUNNING);
    filter.states = states;

    getJobs(jobs, filter);

    for (std::vector<JobT>::iterator i = jobs.begin();
                                     i != jobs.end();
                                     ++i)
    {
        QStringList data;
        data << QString::fromStdString(i->name)
             << QString::number(i->runCores)
             << QString::number(i->maxCores)
             << QString::number(i->totals.waitingTaskCount);

        QTreeWidgetItem *item = new QTreeWidgetItem((QTreeWidget*)0, data);
        treeWidget->insertTopLevelItem(0, item);
    }
}

} // Gui
} // Plow<|endoftext|>


<commit_before>/*
 * Copyright (C) 2009 Toni Gundogdu.
 *
 * This file is part of cclive.
 *
 * cclive is free software: you can redistribute it and/or modify it under the
 * terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * cclive is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <string>

#include "quvi.h"
#include "except.h"

RuntimeException::RuntimeException()
    : rc(CCLIVE_OK), error("")
{
}

RuntimeException::RuntimeException(const ReturnCode& rc)
    : rc(rc), error("")
{
}

RuntimeException::RuntimeException(
    const ReturnCode& rc,
    const std::string& error)
    : rc(rc), error(error)
{
}

RuntimeException::~RuntimeException() {
}

const std::string
RuntimeException::what() const {
    static const char
    errorStrings[_CCLIVE_MAX_RETURNCODES][48] = {
        "no error",
        "(reserved)", // gengetopt uses this (1)
        "invalid option argument",
        "curl_easy_init returned null",
        "file already fully retrieved; nothing to do",
        "system call failed",
        "no support",
        "network error",
        "fetch failed",
        "parse failed",
        "internal error",
    };

    ReturnCode _rc = rc;
    if (_rc >= _CCLIVE_MAX_RETURNCODES)
        _rc = CCLIVE_INTERNAL;

    std::string msg = errorStrings[_rc];

    if (error.length() > 0) {
        if (msg.length() > 0)
            msg += ": ";

        msg += error;

        if (_rc == CCLIVE_INTERNAL)
            msg += ": " + static_cast<int>(rc);
    }
    return msg;
}

const ReturnCode&
RuntimeException::getReturnCode() const {
    return rc;
}

QuviException::QuviException(const std::string& error)
    : RuntimeException(CCLIVE_QUVI, error),
      httpcode(0),
      curlcode(0)
{
    const quvi_t quvi = quvimgr.handle();
    assert(quvi != 0);
    quvi_getinfo(quvi, QUVII_HTTPCODE, &httpcode);
    quvi_getinfo(quvi, QUVII_CURLCODE, &curlcode);
}

QuviException::QuviException(
    const std::string& err,
    const long& httpcode)
    : RuntimeException(CCLIVE_FETCH, err),
      httpcode(httpcode),
      curlcode(0)
{
}

const long&
QuviException::getHttpCode() const {
    return httpcode;
}

const long&
QuviException::getCurlCode() const {
    return curlcode;
}

NoSupportException::NoSupportException(const std::string& err)
    : RuntimeException(CCLIVE_NOSUPPORT, err)
{
}

ParseException::ParseException(const std::string& err)
    : RuntimeException(CCLIVE_PARSE, err)
{
}

NothingToDoException::NothingToDoException()
    : RuntimeException(CCLIVE_NOTHINGTODO)
{
}

<commit_msg>tweak static error strings.<commit_after>/*
 * Copyright (C) 2009 Toni Gundogdu.
 *
 * This file is part of cclive.
 *
 * cclive is free software: you can redistribute it and/or modify it under the
 * terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * cclive is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <string>

#include "quvi.h"
#include "except.h"

RuntimeException::RuntimeException()
    : rc(CCLIVE_OK), error("")
{
}

RuntimeException::RuntimeException(const ReturnCode& rc)
    : rc(rc), error("")
{
}

RuntimeException::RuntimeException(
    const ReturnCode& rc,
    const std::string& error)
    : rc(rc), error(error)
{
}

RuntimeException::~RuntimeException() {
}

const std::string
RuntimeException::what() const {
    static const char
    errorStrings[_CCLIVE_MAX_RETURNCODES][48] = {
        "no error",
        "(reserved)", // gengetopt uses this (1)
        "invalid option argument",
        "curl_easy_init returned null",
        "file already fully retrieved; nothing to do",
        "system",
        "no support",
        "network",
        "fetch",
        "parse",
        "internal",
    };

    ReturnCode _rc = rc;
    if (_rc >= _CCLIVE_MAX_RETURNCODES)
        _rc = CCLIVE_INTERNAL;

    std::string msg = errorStrings[_rc];

    if (error.length() > 0) {
        if (msg.length() > 0)
            msg += ": ";

        msg += error;

        if (_rc == CCLIVE_INTERNAL)
            msg += ": " + static_cast<int>(rc);
    }
    return msg;
}

const ReturnCode&
RuntimeException::getReturnCode() const {
    return rc;
}

QuviException::QuviException(const std::string& error)
    : RuntimeException(CCLIVE_QUVI, error),
      httpcode(0),
      curlcode(0)
{
    const quvi_t quvi = quvimgr.handle();
    assert(quvi != 0);
    quvi_getinfo(quvi, QUVII_HTTPCODE, &httpcode);
    quvi_getinfo(quvi, QUVII_CURLCODE, &curlcode);
}

QuviException::QuviException(
    const std::string& err,
    const long& httpcode)
    : RuntimeException(CCLIVE_FETCH, err),
      httpcode(httpcode),
      curlcode(0)
{
}

const long&
QuviException::getHttpCode() const {
    return httpcode;
}

const long&
QuviException::getCurlCode() const {
    return curlcode;
}

NoSupportException::NoSupportException(const std::string& err)
    : RuntimeException(CCLIVE_NOSUPPORT, err)
{
}

ParseException::ParseException(const std::string& err)
    : RuntimeException(CCLIVE_PARSE, err)
{
}

NothingToDoException::NothingToDoException()
    : RuntimeException(CCLIVE_NOTHINGTODO)
{
}

<|endoftext|>


<commit_before>/* Copyright 2018 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/kernels/data/optimize_dataset_op.h"

// On mobile we do not provide optimize dataset op because not all of its
// dependencies are available there. The op is replaced with a no-op.
#if !defined(IS_MOBILE_PLATFORM)
#include <map>

#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/data/dataset_utils.h"
#include "tensorflow/core/kernels/data/rewrite_utils.h"
#include "tensorflow/core/lib/random/random.h"
#include "tensorflow/core/platform/host_info.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"

namespace tensorflow {
namespace data {

/* static */ constexpr const char* const OptimizeDatasetOp::kDatasetType;
/* static */ constexpr const char* const OptimizeDatasetOp::kInputDataset;
/* static */ constexpr const char* const OptimizeDatasetOp::kOptimizations;
/* static */ constexpr const char* const
    OptimizeDatasetOp::kOptimizationsEnabled;
/* static */ constexpr const char* const
    OptimizeDatasetOp::kOptimizationsDisabled;
/* static */ constexpr const char* const
    OptimizeDatasetOp::kOptimizationsDefault;
/* static */ constexpr const char* const OptimizeDatasetOp::kOutputTypes;
/* static */ constexpr const char* const OptimizeDatasetOp::kOutputShapes;
/* static */ constexpr const char* const
    OptimizeDatasetOp::kOptimizationConfigs;
/* static */ constexpr const char* const OptimizeDatasetOp::kOptimizeDatasetV1;
/* static */ constexpr const char* const OptimizeDatasetOp::kOptimizeDatasetV2;

constexpr char kOptimizerName[] = "tf_data_meta_optimizer";
constexpr char kOptimizers[] = "optimizers";
constexpr char kOptimizerConfigs[] = "optimizer_configs";

OptimizeDatasetOp::OptimizeDatasetOp(OpKernelConstruction* ctx)
    : UnaryDatasetOpKernel(ctx) {
  auto& op_name = ctx->def().op();
  if (op_name == kOptimizeDatasetV1) {
    op_version_ = 1;
  } else if (op_name == kOptimizeDatasetV2) {
    op_version_ = 2;
  }
  OP_REQUIRES_OK(ctx,
                 ctx->GetAttr(kOptimizationConfigs, &optimization_configs_));
}

void OptimizeDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
                                    DatasetBase** output) {
  std::vector<tstring> optimizations;
  if (op_version_ == 1) {
    OP_REQUIRES_OK(
        ctx, ParseVectorArgument<tstring>(ctx, kOptimizations, &optimizations));
  } else if (op_version_ == 2) {
    std::vector<tstring> optimizations_enabled, optimizations_disabled,
        optimizations_default;
    OP_REQUIRES_OK(ctx, ParseVectorArgument<tstring>(ctx, kOptimizationsEnabled,
                                                     &optimizations_enabled));
    OP_REQUIRES_OK(ctx,
                   ParseVectorArgument<tstring>(ctx, kOptimizationsDisabled,
                                                &optimizations_disabled));
    OP_REQUIRES_OK(ctx, ParseVectorArgument<tstring>(ctx, kOptimizationsDefault,
                                                     &optimizations_default));

    string job_name = port::JobName();
    // The map that stores the live experiment names and for how much percentage
    // of the Borg jobs, the experiments will be randomly turned on.
    // clang-format off
    absl::flat_hash_map<string, uint64> live_experiments = {
        {"enable_gradient_descent", 0},
        {"map_parallelization", 5}
    };
    // clang-format on
    auto hash_func = [](const string& str) { return Hash64(str); };
    optimizations = SelectOptimizations(
        job_name, live_experiments, optimizations_enabled,
        optimizations_disabled, optimizations_default, hash_func);

    // Log and record the live experiments that will be applied.
    if (!job_name.empty() && !live_experiments.empty()) {
      VLOG(1) << "The input pipeline is subject to tf.data experiment. "
                 "Please see `go/tf-data-experiments` for more details.";

      for (auto& pair : live_experiments) {
        string experiment = pair.first;
        if (std::find(optimizations.begin(), optimizations.end(), experiment) !=
            optimizations.end()) {
          VLOG(1) << "The live experiment \"" << experiment << "\" is applied.";
          metrics::RecordTFDataExperiment(experiment);
        }
      }
    }
  }

  // The vector stores the graduated experiment names which will be turned on
  // for all input pipelines.
  // clang-format off
  std::vector<string> graduated_experiments = {"disable_intra_op_parallelism"};
  // clang-format on

  // Add the graduated experiments to the optimization list and log them.
  for (auto& experiment : graduated_experiments) {
    if (std::find(optimizations.begin(), optimizations.end(), experiment) ==
        optimizations.end()) {
      optimizations.push_back(experiment);
    }
    VLOG(1) << "The graduated experiment \"" << experiment << "\" is applied.";
  }

  // If there are no optimizations to be applied, directly return the input.
  if (optimizations.empty()) {
    *output = input;
    input->Ref();
    return;
  }

  auto config_factory = [this, &optimizations]() {
    return CreateConfig(optimizations, optimization_configs_);
  };
  Status s = RewriteDataset(ctx, input, std::move(config_factory),
                            /*record_fingerprint=*/true, output);
  if (errors::IsDeadlineExceeded(s)) {
    // Ignore DeadlineExceeded as it implies that the attempted rewrite took too
    // long which should not prevent further computation.
    LOG(WARNING) << s.ToString();

    *output = input;
    input->Ref();
    return;
  }
  OP_REQUIRES_OK(ctx, s);
}

RewriterConfig OptimizeDatasetOp::CreateConfig(
    std::vector<tstring> optimizations,
    std::vector<string> optimizations_configs) {
  RewriterConfig rewriter_config;
  rewriter_config.add_optimizers(kOptimizerName);
  rewriter_config.set_meta_optimizer_iterations(RewriterConfig::ONE);
  rewriter_config.set_fail_on_optimizer_errors(true);
  auto custom_optimizer = rewriter_config.add_custom_optimizers();
  custom_optimizer->set_name(kOptimizerName);
  auto* custom_optimizations_list =
      (*custom_optimizer->mutable_parameter_map())[kOptimizers].mutable_list();
  for (const auto& opt : optimizations) {
    custom_optimizations_list->add_s(opt.data(), opt.size());
  }
  auto* config_list =
      (*custom_optimizer->mutable_parameter_map())[kOptimizerConfigs]
          .mutable_list();
  for (const auto& config : optimizations_configs) {
    config_list->add_s(config.data(), config.size());
  }
  return rewriter_config;
}

namespace {
REGISTER_KERNEL_BUILDER(Name("OptimizeDataset").Device(DEVICE_CPU),
                        OptimizeDatasetOp);
REGISTER_KERNEL_BUILDER(Name("OptimizeDatasetV2").Device(DEVICE_CPU),
                        OptimizeDatasetOp);
}  // namespace
}  // namespace data
}  // namespace tensorflow
#else  // !IS_MOBILE_PLATFORM
namespace tensorflow {
namespace data {

OptimizeDatasetOp::OptimizeDatasetOp(OpKernelConstruction* ctx)
    : UnaryDatasetOpKernel(ctx) {}

void OptimizeDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
                                    DatasetBase** output) {
  input->Ref();
  *output = input;
}

namespace {
REGISTER_KERNEL_BUILDER(Name("OptimizeDataset").Device(DEVICE_CPU),
                        OptimizeDatasetOp);
REGISTER_KERNEL_BUILDER(Name("OptimizeDatasetV2").Device(DEVICE_CPU),
                        OptimizeDatasetOp);
}  // namespace
}  // namespace data
}  // namespace tensorflow
#endif  // !IS_MOBILE_PLATFORM
<commit_msg>[tf.data] Rolls out the optimization `map_parallelization` as experiment to 20% of Borg jobs.<commit_after>/* Copyright 2018 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/kernels/data/optimize_dataset_op.h"

// On mobile we do not provide optimize dataset op because not all of its
// dependencies are available there. The op is replaced with a no-op.
#if !defined(IS_MOBILE_PLATFORM)
#include <map>

#include "tensorflow/core/framework/partial_tensor_shape.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/data/dataset_utils.h"
#include "tensorflow/core/kernels/data/rewrite_utils.h"
#include "tensorflow/core/lib/random/random.h"
#include "tensorflow/core/platform/host_info.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"

namespace tensorflow {
namespace data {

/* static */ constexpr const char* const OptimizeDatasetOp::kDatasetType;
/* static */ constexpr const char* const OptimizeDatasetOp::kInputDataset;
/* static */ constexpr const char* const OptimizeDatasetOp::kOptimizations;
/* static */ constexpr const char* const
    OptimizeDatasetOp::kOptimizationsEnabled;
/* static */ constexpr const char* const
    OptimizeDatasetOp::kOptimizationsDisabled;
/* static */ constexpr const char* const
    OptimizeDatasetOp::kOptimizationsDefault;
/* static */ constexpr const char* const OptimizeDatasetOp::kOutputTypes;
/* static */ constexpr const char* const OptimizeDatasetOp::kOutputShapes;
/* static */ constexpr const char* const
    OptimizeDatasetOp::kOptimizationConfigs;
/* static */ constexpr const char* const OptimizeDatasetOp::kOptimizeDatasetV1;
/* static */ constexpr const char* const OptimizeDatasetOp::kOptimizeDatasetV2;

constexpr char kOptimizerName[] = "tf_data_meta_optimizer";
constexpr char kOptimizers[] = "optimizers";
constexpr char kOptimizerConfigs[] = "optimizer_configs";

OptimizeDatasetOp::OptimizeDatasetOp(OpKernelConstruction* ctx)
    : UnaryDatasetOpKernel(ctx) {
  auto& op_name = ctx->def().op();
  if (op_name == kOptimizeDatasetV1) {
    op_version_ = 1;
  } else if (op_name == kOptimizeDatasetV2) {
    op_version_ = 2;
  }
  OP_REQUIRES_OK(ctx,
                 ctx->GetAttr(kOptimizationConfigs, &optimization_configs_));
}

void OptimizeDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
                                    DatasetBase** output) {
  std::vector<tstring> optimizations;
  if (op_version_ == 1) {
    OP_REQUIRES_OK(
        ctx, ParseVectorArgument<tstring>(ctx, kOptimizations, &optimizations));
  } else if (op_version_ == 2) {
    std::vector<tstring> optimizations_enabled, optimizations_disabled,
        optimizations_default;
    OP_REQUIRES_OK(ctx, ParseVectorArgument<tstring>(ctx, kOptimizationsEnabled,
                                                     &optimizations_enabled));
    OP_REQUIRES_OK(ctx,
                   ParseVectorArgument<tstring>(ctx, kOptimizationsDisabled,
                                                &optimizations_disabled));
    OP_REQUIRES_OK(ctx, ParseVectorArgument<tstring>(ctx, kOptimizationsDefault,
                                                     &optimizations_default));

    string job_name = port::JobName();
    // The map that stores the live experiment names and for how much percentage
    // of the Borg jobs, the experiments will be randomly turned on.
    // clang-format off
    absl::flat_hash_map<string, uint64> live_experiments = {
        {"enable_gradient_descent", 0},
        {"map_parallelization", 20}
    };
    // clang-format on
    auto hash_func = [](const string& str) { return Hash64(str); };
    optimizations = SelectOptimizations(
        job_name, live_experiments, optimizations_enabled,
        optimizations_disabled, optimizations_default, hash_func);

    // Log and record the live experiments that will be applied.
    if (!job_name.empty() && !live_experiments.empty()) {
      VLOG(1) << "The input pipeline is subject to tf.data experiment. "
                 "Please see `go/tf-data-experiments` for more details.";

      for (auto& pair : live_experiments) {
        string experiment = pair.first;
        if (std::find(optimizations.begin(), optimizations.end(), experiment) !=
            optimizations.end()) {
          VLOG(1) << "The live experiment \"" << experiment << "\" is applied.";
          metrics::RecordTFDataExperiment(experiment);
        }
      }
    }
  }

  // The vector stores the graduated experiment names which will be turned on
  // for all input pipelines.
  // clang-format off
  std::vector<string> graduated_experiments = {"disable_intra_op_parallelism"};
  // clang-format on

  // Add the graduated experiments to the optimization list and log them.
  for (auto& experiment : graduated_experiments) {
    if (std::find(optimizations.begin(), optimizations.end(), experiment) ==
        optimizations.end()) {
      optimizations.push_back(experiment);
    }
    VLOG(1) << "The graduated experiment \"" << experiment << "\" is applied.";
  }

  // If there are no optimizations to be applied, directly return the input.
  if (optimizations.empty()) {
    *output = input;
    input->Ref();
    return;
  }

  auto config_factory = [this, &optimizations]() {
    return CreateConfig(optimizations, optimization_configs_);
  };
  Status s = RewriteDataset(ctx, input, std::move(config_factory),
                            /*record_fingerprint=*/true, output);
  if (errors::IsDeadlineExceeded(s)) {
    // Ignore DeadlineExceeded as it implies that the attempted rewrite took too
    // long which should not prevent further computation.
    LOG(WARNING) << s.ToString();

    *output = input;
    input->Ref();
    return;
  }
  OP_REQUIRES_OK(ctx, s);
}

RewriterConfig OptimizeDatasetOp::CreateConfig(
    std::vector<tstring> optimizations,
    std::vector<string> optimizations_configs) {
  RewriterConfig rewriter_config;
  rewriter_config.add_optimizers(kOptimizerName);
  rewriter_config.set_meta_optimizer_iterations(RewriterConfig::ONE);
  rewriter_config.set_fail_on_optimizer_errors(true);
  auto custom_optimizer = rewriter_config.add_custom_optimizers();
  custom_optimizer->set_name(kOptimizerName);
  auto* custom_optimizations_list =
      (*custom_optimizer->mutable_parameter_map())[kOptimizers].mutable_list();
  for (const auto& opt : optimizations) {
    custom_optimizations_list->add_s(opt.data(), opt.size());
  }
  auto* config_list =
      (*custom_optimizer->mutable_parameter_map())[kOptimizerConfigs]
          .mutable_list();
  for (const auto& config : optimizations_configs) {
    config_list->add_s(config.data(), config.size());
  }
  return rewriter_config;
}

namespace {
REGISTER_KERNEL_BUILDER(Name("OptimizeDataset").Device(DEVICE_CPU),
                        OptimizeDatasetOp);
REGISTER_KERNEL_BUILDER(Name("OptimizeDatasetV2").Device(DEVICE_CPU),
                        OptimizeDatasetOp);
}  // namespace
}  // namespace data
}  // namespace tensorflow
#else  // !IS_MOBILE_PLATFORM
namespace tensorflow {
namespace data {

OptimizeDatasetOp::OptimizeDatasetOp(OpKernelConstruction* ctx)
    : UnaryDatasetOpKernel(ctx) {}

void OptimizeDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
                                    DatasetBase** output) {
  input->Ref();
  *output = input;
}

namespace {
REGISTER_KERNEL_BUILDER(Name("OptimizeDataset").Device(DEVICE_CPU),
                        OptimizeDatasetOp);
REGISTER_KERNEL_BUILDER(Name("OptimizeDatasetV2").Device(DEVICE_CPU),
                        OptimizeDatasetOp);
}  // namespace
}  // namespace data
}  // namespace tensorflow
#endif  // !IS_MOBILE_PLATFORM
<|endoftext|>


<commit_before>/* Copyright 2019 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/lite/delegates/gpu/gl/kernels/softmax.h"

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/memory/memory.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/common/status.h"
#include "tensorflow/lite/delegates/gpu/common/types.h"
#include "tensorflow/lite/delegates/gpu/common/util.h"
#include "tensorflow/lite/delegates/gpu/gl/variable.h"

namespace tflite {
namespace gpu {
namespace gl {
namespace {

float4 GetMask(int num_channels) {
  float4 mask(0.0f);
  const int remainder = num_channels % 4 == 0 ? 4 : num_channels % 4;
  for (int i = 0; i < remainder; ++i) mask[i] = 1.0f;
  return mask;
}

class Softmax : public NodeShader {
 public:
  absl::Status GenerateCode(const GenerationContext& ctx,
                            GeneratedCode* generated_code) const final {
    const auto* input = ctx.graph->FindInputs(ctx.node->id)[0];
    const auto* output = ctx.graph->FindOutputs(ctx.node->id)[0];
    const auto& attr = absl::any_cast<const SoftmaxAttributes&>(
        ctx.node->operation.attributes);
    if (input->tensor.shape != output->tensor.shape) {
      return absl::InvalidArgumentError(
          "Input and output shapes do not match.");
    }
    if (attr.axis != Axis::CHANNELS) {
      return absl::UnimplementedError(
          "Softmax is only supported for channels axis.");
    }
    return input->tensor.shape.h == 1 && input->tensor.shape.w == 1
               ? GenerateCodeFor1x1(ctx, generated_code)
               : GenerateCodeGeneral(ctx, generated_code);
  }

 private:
  absl::Status GenerateCodeFor1x1(const GenerationContext& ctx,
                                  GeneratedCode* generated_code) const {
    const auto* output = ctx.graph->FindOutputs(ctx.node->id)[0];
    const int depth = IntegralDivideRoundUp(output->tensor.shape.c, 4);
    std::vector<Variable> shared_variables = {
        {"partial_sum", std::vector<float4>(8)},
    };
    std::vector<Variable> uniform_parameters = {
        {"depth", depth},
        {"depth_div_32", IntegralDivideRoundUp(depth, 32)},
        {"mask", GetMask(output->tensor.shape.c)},
    };
    std::string source_code = R"(
  highp vec4 kOnes = vec4(1.0);
  highp float sum = 0.0;
  int offset = 0;
  int s = 0;
  int tid = int(gl_LocalInvocationID.x);
  do {
    int z = offset + tid;
    if (z < $depth$) {
      highp vec4 mask_temp = z == $depth$ - 1 ? $mask$ : kOnes;
      highp vec4 src = $input_data_0[0, 0, z]$;
      sum += dot(mask_temp, exp(src));
      offset += 32;
    }
    s++;
  } while (s < $depth_div_32$);

  partial_sum[tid / 4][tid % 4] = sum;

  memoryBarrierShared();
  barrier();

  if (tid == 0) {
    sum = dot(kOnes, partial_sum[0]);
    sum += dot(kOnes, partial_sum[1]);
    sum += dot(kOnes, partial_sum[2]);
    sum += dot(kOnes, partial_sum[3]);
    sum += dot(kOnes, partial_sum[4]);
    sum += dot(kOnes, partial_sum[5]);
    sum += dot(kOnes, partial_sum[6]);
    sum += dot(kOnes, partial_sum[7]);
    partial_sum[0][0] = 1.0 / sum;
  }

  memoryBarrierShared();
  barrier();

  sum = partial_sum[0][0];

  offset = 0;
  s = 0;
  do {
    int z = offset + tid;
    if (z < $depth$) {
      highp vec4 src = $input_data_0[0, 0, z]$;
      highp vec4 temp = exp(src) * sum;
      $output_data_0[0, 0, z]$ = temp;
      offset += 32;
    }
    s++;
  } while (s < $depth_div_32$);
)";
    *generated_code = {
        /*parameters=*/std::move(uniform_parameters),
        /*objects=*/{},
        /*shared_variables=*/std::move(shared_variables),
        /*workload=*/uint3(32, 1, 1),
        /*workgroup=*/uint3(32, 1, 1),
        /*source_code=*/std::move(source_code),
        /*input=*/IOStructure::ONLY_DEFINITIONS,
        /*output=*/IOStructure::ONLY_DEFINITIONS,
    };
    return absl::OkStatus();
  }

  absl::Status GenerateCodeGeneral(const GenerationContext& ctx,
                                   GeneratedCode* generated_code) const {
    const auto* output = ctx.graph->FindOutputs(ctx.node->id)[0];
    std::vector<Variable> parameters = {
        {"src_depth", IntegralDivideRoundUp(output->tensor.shape.c, 4)},
        {"mask", GetMask(output->tensor.shape.c)},
    };

    std::string source_code = R"(
  highp vec4 kOnes = vec4(1.0);
  highp float sum = 0.0;
  for (int d = 0; d < $src_depth$ - 1; ++d) {
    highp vec4 src = $input_data_0[gid.x, gid.y, d]$;
    sum += dot(kOnes, exp(src));
  }
  {
    int d = $src_depth$ - 1;
    highp vec4 src = $input_data_0[gid.x, gid.y, d]$;
    sum += dot($mask$, exp(src));
  }
  for (int d = 0; d < $src_depth$; ++d) {
    highp vec4 src = $input_data_0[gid.x, gid.y, d]$;
    highp vec4 temp_sum = exp(src) / sum;
    $output_data_0[gid.x, gid.y, d] = temp_sum$;
  }
)";
    *generated_code = {
        /*parameters=*/std::move(parameters),
        /*objects=*/{},
        /*shared_variables=*/{},
        /*workload=*/uint3(output->tensor.shape.w, output->tensor.shape.h, 1),
        /*workgroup=*/uint3(),
        /*source_code=*/std::move(source_code),
        /*input=*/IOStructure::ONLY_DEFINITIONS,
        /*output=*/IOStructure::ONLY_DEFINITIONS,
    };
    return absl::OkStatus();
  }
};

}  // namespace

std::unique_ptr<NodeShader> NewSoftmaxNodeShader() {
  return absl::make_unique<Softmax>();
}

}  // namespace gl
}  // namespace gpu
}  // namespace tflite
<commit_msg>Fix output pattern in 1x1 SOFTMAX case.<commit_after>/* Copyright 2019 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/lite/delegates/gpu/gl/kernels/softmax.h"

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/memory/memory.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/common/status.h"
#include "tensorflow/lite/delegates/gpu/common/types.h"
#include "tensorflow/lite/delegates/gpu/common/util.h"
#include "tensorflow/lite/delegates/gpu/gl/variable.h"

namespace tflite {
namespace gpu {
namespace gl {
namespace {

float4 GetMask(int num_channels) {
  float4 mask(0.0f);
  const int remainder = num_channels % 4 == 0 ? 4 : num_channels % 4;
  for (int i = 0; i < remainder; ++i) mask[i] = 1.0f;
  return mask;
}

class Softmax : public NodeShader {
 public:
  absl::Status GenerateCode(const GenerationContext& ctx,
                            GeneratedCode* generated_code) const final {
    const auto* input = ctx.graph->FindInputs(ctx.node->id)[0];
    const auto* output = ctx.graph->FindOutputs(ctx.node->id)[0];
    const auto& attr = absl::any_cast<const SoftmaxAttributes&>(
        ctx.node->operation.attributes);
    if (input->tensor.shape != output->tensor.shape) {
      return absl::InvalidArgumentError(
          "Input and output shapes do not match.");
    }
    if (attr.axis != Axis::CHANNELS) {
      return absl::UnimplementedError(
          "Softmax is only supported for channels axis.");
    }
    return input->tensor.shape.h == 1 && input->tensor.shape.w == 1
               ? GenerateCodeFor1x1(ctx, generated_code)
               : GenerateCodeGeneral(ctx, generated_code);
  }

 private:
  absl::Status GenerateCodeFor1x1(const GenerationContext& ctx,
                                  GeneratedCode* generated_code) const {
    const auto* output = ctx.graph->FindOutputs(ctx.node->id)[0];
    const int depth = IntegralDivideRoundUp(output->tensor.shape.c, 4);
    std::vector<Variable> shared_variables = {
        {"partial_sum", std::vector<float4>(8)},
    };
    std::vector<Variable> uniform_parameters = {
        {"depth", depth},
        {"depth_div_32", IntegralDivideRoundUp(depth, 32)},
        {"mask", GetMask(output->tensor.shape.c)},
    };
    std::string source_code = R"(
  highp vec4 kOnes = vec4(1.0);
  highp float sum = 0.0;
  int offset = 0;
  int s = 0;
  int tid = int(gl_LocalInvocationID.x);
  do {
    int z = offset + tid;
    if (z < $depth$) {
      highp vec4 mask_temp = z == $depth$ - 1 ? $mask$ : kOnes;
      highp vec4 src = $input_data_0[0, 0, z]$;
      sum += dot(mask_temp, exp(src));
      offset += 32;
    }
    s++;
  } while (s < $depth_div_32$);

  partial_sum[tid / 4][tid % 4] = sum;

  memoryBarrierShared();
  barrier();

  if (tid == 0) {
    sum = dot(kOnes, partial_sum[0]);
    sum += dot(kOnes, partial_sum[1]);
    sum += dot(kOnes, partial_sum[2]);
    sum += dot(kOnes, partial_sum[3]);
    sum += dot(kOnes, partial_sum[4]);
    sum += dot(kOnes, partial_sum[5]);
    sum += dot(kOnes, partial_sum[6]);
    sum += dot(kOnes, partial_sum[7]);
    partial_sum[0][0] = 1.0 / sum;
  }

  memoryBarrierShared();
  barrier();

  sum = partial_sum[0][0];

  offset = 0;
  s = 0;
  do {
    int z = offset + tid;
    if (z < $depth$) {
      highp vec4 src = $input_data_0[0, 0, z]$;
      highp vec4 temp = exp(src) * sum;
      $output_data_0[0, 0, z] = temp$;
      offset += 32;
    }
    s++;
  } while (s < $depth_div_32$);
)";
    *generated_code = {
        /*parameters=*/std::move(uniform_parameters),
        /*objects=*/{},
        /*shared_variables=*/std::move(shared_variables),
        /*workload=*/uint3(32, 1, 1),
        /*workgroup=*/uint3(32, 1, 1),
        /*source_code=*/std::move(source_code),
        /*input=*/IOStructure::ONLY_DEFINITIONS,
        /*output=*/IOStructure::ONLY_DEFINITIONS,
    };
    return absl::OkStatus();
  }

  absl::Status GenerateCodeGeneral(const GenerationContext& ctx,
                                   GeneratedCode* generated_code) const {
    const auto* output = ctx.graph->FindOutputs(ctx.node->id)[0];
    std::vector<Variable> parameters = {
        {"src_depth", IntegralDivideRoundUp(output->tensor.shape.c, 4)},
        {"mask", GetMask(output->tensor.shape.c)},
    };

    std::string source_code = R"(
  highp vec4 kOnes = vec4(1.0);
  highp float sum = 0.0;
  for (int d = 0; d < $src_depth$ - 1; ++d) {
    highp vec4 src = $input_data_0[gid.x, gid.y, d]$;
    sum += dot(kOnes, exp(src));
  }
  {
    int d = $src_depth$ - 1;
    highp vec4 src = $input_data_0[gid.x, gid.y, d]$;
    sum += dot($mask$, exp(src));
  }
  for (int d = 0; d < $src_depth$; ++d) {
    highp vec4 src = $input_data_0[gid.x, gid.y, d]$;
    highp vec4 temp_sum = exp(src) / sum;
    $output_data_0[gid.x, gid.y, d] = temp_sum$;
  }
)";
    *generated_code = {
        /*parameters=*/std::move(parameters),
        /*objects=*/{},
        /*shared_variables=*/{},
        /*workload=*/uint3(output->tensor.shape.w, output->tensor.shape.h, 1),
        /*workgroup=*/uint3(),
        /*source_code=*/std::move(source_code),
        /*input=*/IOStructure::ONLY_DEFINITIONS,
        /*output=*/IOStructure::ONLY_DEFINITIONS,
    };
    return absl::OkStatus();
  }
};

}  // namespace

std::unique_ptr<NodeShader> NewSoftmaxNodeShader() {
  return absl::make_unique<Softmax>();
}

}  // namespace gl
}  // namespace gpu
}  // namespace tflite
<|endoftext|>


<commit_before>/* Copyright 2019 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/lite/delegates/gpu/metal/compute_task.h"

#include <Availability.h>

#include <map>
#include <string>
#include <tuple>

#include "absl/strings/match.h"
#include "absl/strings/substitute.h"
#include "tensorflow/lite/delegates/gpu/common/kernel_info.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/common/status.h"
#include "tensorflow/lite/delegates/gpu/common/types.h"
#include "tensorflow/lite/delegates/gpu/common/util.h"
#include "tensorflow/lite/delegates/gpu/metal/common.h"

namespace tflite {
namespace gpu {
namespace metal {
namespace {
int3 GetWorkGroupsCount(int grid_dimension, const int3& grid_size,
                        const int3& work_group_size,
                        const int3& work_group_launch_order) {
  int3 work_groups_count;
  if (grid_dimension == 1) {
    work_groups_count.x = DivideRoundUp(grid_size.x, work_group_size.x);
    work_groups_count.y = 1;
    work_groups_count.z = 1;
  } else if (grid_dimension == 2) {
    int3 wgs;
    wgs.x = DivideRoundUp(grid_size.x, work_group_size.x);
    wgs.y = DivideRoundUp(grid_size.y, work_group_size.y);
    work_groups_count.x = wgs[work_group_launch_order[0]];
    work_groups_count.y = wgs[work_group_launch_order[1]];
    work_groups_count.z = 1;
  } else {  // grid_dimension == 3
    int3 wgs;
    wgs.x = DivideRoundUp(grid_size.x, work_group_size.x);
    wgs.y = DivideRoundUp(grid_size.y, work_group_size.y);
    wgs.z = DivideRoundUp(grid_size.z, work_group_size.z);
    work_groups_count.x = wgs[work_group_launch_order[0]];
    work_groups_count.y = wgs[work_group_launch_order[1]];
    work_groups_count.z = wgs[work_group_launch_order[2]];
  }
  return work_groups_count;
}
}  // namespace

void ComputeTask::Init(std::unique_ptr<GPUOperation>&& operation) {
  operation_ = std::move(operation);
}

const OperationDef& ComputeTask::GetDefinition() const {
  return operation_->definition_;
}

bool ComputeTask::IsLinkable() const { return operation_->IsLinkable(); }

absl::Status ComputeTask::AddTask(ComputeTask* task) {
  return operation_->AddOperation(task->operation_.get());
}

absl::Status ComputeTask::Compile(MetalDevice* device) {
  operation_->AssembleCode(device->GetInfo());
  const std::map<std::string, std::string> linkables = {
      {operation_->dst_tensors_names_[0], operation_->elementwise_code_}};
  RETURN_IF_ERROR(metal_args_.Init(linkables, device, &operation_->args_,
                                   &operation_->code_));

  operation_->args_.ReleaseCPURepresentation();

  return CompileProgram(device, operation_->definition_.precision,
                        operation_->code_);
}

absl::Status ComputeTask::CompileProgram(MetalDevice* device,
                                         CalculationsPrecision precision,
                                         const std::string& kernel_code) {
  NSString* barrier;
  // simdgroup_barrier is supported since Metal shading language version 2.0
  if (device->IsLanguageVersion2orHigher()) {
    barrier = @"simdgroup_barrier";
  } else {
    barrier = @"threadgroup_barrier";
  }
  NSString* storageType;
  NSString* accumulatorType;
  NSString* toAccumulatorType4 = @"";
  if (precision == CalculationsPrecision::F32) {
    storageType = @"float";
    accumulatorType = @"float";
  } else {
    // FP16
    storageType = @"half";
    if (precision == CalculationsPrecision::F32_F16) {
      accumulatorType = @"float";
      toAccumulatorType4 = @"float4";
    } else {
      accumulatorType = @"half";
    }
  }
  NSDictionary<NSString*, NSString*>* macros = @{
    @"float16" : @"float4x4",
    @"half16" : @"half4x4",
    @"float8" : @"float2x4",
    @"half8" : @"half2x4",
    @"FLT16_0123(V)" : @"V[0]",
    @"FLT16_4567(V)" : @"V[1]",
    @"FLT16_89ab(V)" : @"V[2]",
    @"FLT16_cdef(V)" : @"V[3]",
    @"FLT" : storageType,
    @"FLT2" : [NSString stringWithFormat:@"%@2", storageType],
    @"FLT3" : [NSString stringWithFormat:@"%@3", storageType],
    @"FLT4" : [NSString stringWithFormat:@"%@4", storageType],
    @"ACCUM_FLT" : accumulatorType,
    @"ACCUM_FLT2" : [NSString stringWithFormat:@"%@2", accumulatorType],
    @"ACCUM_FLT3" : [NSString stringWithFormat:@"%@3", accumulatorType],
    @"ACCUM_FLT4" : [NSString stringWithFormat:@"%@4", accumulatorType],
    @"INIT_ACCUM_FLT4(value)" :
        [NSString stringWithFormat:@"%@4(value)", accumulatorType],
    @"TO_ACCUM_TYPE" : toAccumulatorType4,
    @"TO_ACCUM_FLT" : accumulatorType,
    @"TO_ACCUM_FLT2" : [NSString stringWithFormat:@"%@2", accumulatorType],
    @"TO_ACCUM_FLT3" : [NSString stringWithFormat:@"%@3", accumulatorType],
    @"TO_ACCUM_FLT4" : [NSString stringWithFormat:@"%@4", accumulatorType],
    @"TO_FLT4" : [NSString stringWithFormat:@"%@4", storageType],
    @"SIMDGROUP_BARRIER" : barrier,
    @"SIMD_LOCAL_MEM_BARRIER" : barrier,
    @"MAIN_FUNCTION" : @"\"kernel void ComputeFunction\"",
    @"GLOBAL_ID_0" : @"static_cast<int>(reserved_gid.x)",
    @"GLOBAL_ID_1" : @"static_cast<int>(reserved_gid.y)",
    @"GLOBAL_ID_2" : @"static_cast<int>(reserved_gid.z)",
    @"LOCAL_ID_0" : @"static_cast<int>(reserved_lid.x)",
    @"LOCAL_ID_1" : @"static_cast<int>(reserved_lid.y)",
    @"LOCAL_ID_2" : @"static_cast<int>(reserved_lid.z)",
    @"GROUP_ID_0" : @"static_cast<int>(reserved_group_id.x)",
    @"GROUP_ID_1" : @"static_cast<int>(reserved_group_id.y)",
    @"GROUP_ID_2" : @"static_cast<int>(reserved_group_id.z)",
    @"GROUP_SIZE_0" : @"static_cast<int>(reserved_group_size.x)",
    @"GROUP_SIZE_1" : @"static_cast<int>(reserved_group_size.y)",
    @"GROUP_SIZE_2" : @"static_cast<int>(reserved_group_size.z)",
    @"SUB_GROUP_LOCAL_ID" : @"static_cast<int>(reserved_simd_id)",
    @"\"SUB_GROUP_BROADCAST(V, ID)\"" : @"\"simd_broadcast(V, ID)\"",
    @"__local" : @"threadgroup",
    @"__global" : @"device",
    @"__constant" : @"constant",
    @"LOCAL_MEM_BARRIER" : @"threadgroup_barrier(mem_flags::mem_threadgroup)",
    @"INIT_FLT(value)" : [NSString stringWithFormat:@"%@(value)", storageType],
    @"INIT_FLT4(value)" :
        [NSString stringWithFormat:@"%@4(value)", storageType],
    @"\"INIT_FLT4v4(v0, v1, v2, v3)\"" :
        [NSString stringWithFormat:@"\"%@4(v0, v1, v2, v3)\"", storageType],
    @"INIT_FLOAT(value)" : @"float(value)",
    @"INIT_FLOAT2(value)" : @"float2(value)",
    @"\"INIT_FLOAT2v2(v0, v1)\"" : @"\"float2(v0, v1)\"",
    @"INIT_FLOAT3(value)" : @"float3(value)",
    @"\"INIT_FLOAT3v3(v0, v1, v2)\"" : @"\"float3(v0, v1, v2)\"",
    @"INIT_FLOAT4(value)" : @"float4(value)",
    @"\"INIT_FLOAT4v4(v0, v1, v2, v3)\"" : @"\"float4(v0, v1, v2, v3)\"",
    @"INIT_INT(value)" : @"int(value)",
    @"\"INIT_INT2v2(v0, v1)\"" : @"\"int2(v0, v1)\"",
    @"\"INIT_INT4v4(v0, v1, v2, v3)\"" : @"\"int4(v0, v1, v2, v3)\"",
    @"CONVERT_TO_INT4(value)" : @"int4(value)",
    @"\"SELECT_BY_INDEX_FROM_FLT4(value, index)\"" : @"\"(value)[index]\"",
  };

  NSString* code =
      [NSString stringWithCString:kernel_code.c_str()
                         encoding:[NSString defaultCStringEncoding]];
  id<MTLComputePipelineState> program;
  RETURN_IF_ERROR(CreateComputeProgram(device->device(), code,
                                       @"ComputeFunction", macros, &program));
  if (!program) {
    return absl::InternalError("Unknown shader compilation error");
  }
  program_ = program;
  return absl::OkStatus();
}

absl::Status ComputeTask::UpdateParams() {
  for (int i = 0; i < operation_->src_tensors_names_.size(); ++i) {
    const auto* metal_spatial_tensor =
        dynamic_cast<const MetalSpatialTensor*>(operation_->src_[i]);
    if (!metal_spatial_tensor) {
      return absl::InvalidArgumentError("Expected MetalSpatialTensor.");
    }
    RETURN_IF_ERROR(metal_args_.SetObjectRef(operation_->src_tensors_names_[i],
                                             *metal_spatial_tensor));
  }
  for (int i = 0; i < operation_->dst_tensors_names_.size(); ++i) {
    const auto* metal_spatial_tensor =
        dynamic_cast<const MetalSpatialTensor*>(operation_->dst_[i]);
    if (!metal_spatial_tensor) {
      return absl::InvalidArgumentError("Expected MetalSpatialTensor.");
    }
    RETURN_IF_ERROR(metal_args_.SetObjectRef(operation_->dst_tensors_names_[i],
                                             *metal_spatial_tensor));
  }
  RETURN_IF_ERROR(operation_->BindArguments(&metal_args_));
  operation_->grid_size_ = operation_->GetGridSize();
  operation_->work_groups_count_ = GetWorkGroupsCount(
      operation_->grid_dimension_, operation_->grid_size_,
      operation_->work_group_size_, operation_->work_group_launch_order_);
  return absl::OkStatus();
}

void ComputeTask::Encode(id<MTLComputeCommandEncoder> encoder) {
  [encoder setComputePipelineState:program_];
  metal_args_.Encode(encoder, 0);
  MTLSize groupsCount, groupsSize;
  groupsCount.width = operation_->work_groups_count_.x;
  groupsCount.height = operation_->work_groups_count_.y;
  groupsCount.depth = operation_->work_groups_count_.z;
  groupsSize.width = operation_->work_group_size_.x;
  groupsSize.height = operation_->work_group_size_.y;
  groupsSize.depth = operation_->work_group_size_.z;
  [encoder dispatchThreadgroups:groupsCount threadsPerThreadgroup:groupsSize];
}

void ComputeTask::SetSrcTensor(MetalSpatialTensor* tensor, int index) {
  operation_->SetSrc(tensor, index);
  auto status =
      metal_args_.SetObjectRef(operation_->src_tensors_names_[index], *tensor);
}

void ComputeTask::SetDstTensor(MetalSpatialTensor* tensor, int index) {
  operation_->SetDst(tensor, index);
  auto status =
      metal_args_.SetObjectRef(operation_->dst_tensors_names_[index], *tensor);
}

absl::Status ComputeTask::Tune(TuningType tuning_type, MetalDevice* device) {
  std::vector<int3> possible_work_groups;
  KernelInfo kernel_info;
  kernel_info.max_work_group_size = [program_ maxTotalThreadsPerThreadgroup];
  kernel_info.private_memory_size = 0;
  operation_->GetPossibleKernelWorkGroups(tuning_type, device->GetInfo(),
                                          kernel_info, &possible_work_groups);
  if (possible_work_groups.empty()) {
    return absl::NotFoundError(
        "Can not found work_group size to launch kernel");
  }
  operation_->work_group_size_ = possible_work_groups[0];
  operation_->work_groups_count_ = GetWorkGroupsCount(
      operation_->grid_dimension_, operation_->grid_size_,
      operation_->work_group_size_, operation_->work_group_launch_order_);
  return absl::OkStatus();
}

}  // namespace metal
}  // namespace gpu
}  // namespace tflite
<commit_msg>Remove macros already defined in `metal_stdlib` on iOS15.<commit_after>/* Copyright 2019 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/lite/delegates/gpu/metal/compute_task.h"

#include <Availability.h>

#include <map>
#include <string>
#include <tuple>

#include "absl/strings/match.h"
#include "absl/strings/substitute.h"
#include "tensorflow/lite/delegates/gpu/common/kernel_info.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/common/status.h"
#include "tensorflow/lite/delegates/gpu/common/types.h"
#include "tensorflow/lite/delegates/gpu/common/util.h"
#include "tensorflow/lite/delegates/gpu/metal/common.h"

namespace tflite {
namespace gpu {
namespace metal {
namespace {
int3 GetWorkGroupsCount(int grid_dimension, const int3& grid_size,
                        const int3& work_group_size,
                        const int3& work_group_launch_order) {
  int3 work_groups_count;
  if (grid_dimension == 1) {
    work_groups_count.x = DivideRoundUp(grid_size.x, work_group_size.x);
    work_groups_count.y = 1;
    work_groups_count.z = 1;
  } else if (grid_dimension == 2) {
    int3 wgs;
    wgs.x = DivideRoundUp(grid_size.x, work_group_size.x);
    wgs.y = DivideRoundUp(grid_size.y, work_group_size.y);
    work_groups_count.x = wgs[work_group_launch_order[0]];
    work_groups_count.y = wgs[work_group_launch_order[1]];
    work_groups_count.z = 1;
  } else {  // grid_dimension == 3
    int3 wgs;
    wgs.x = DivideRoundUp(grid_size.x, work_group_size.x);
    wgs.y = DivideRoundUp(grid_size.y, work_group_size.y);
    wgs.z = DivideRoundUp(grid_size.z, work_group_size.z);
    work_groups_count.x = wgs[work_group_launch_order[0]];
    work_groups_count.y = wgs[work_group_launch_order[1]];
    work_groups_count.z = wgs[work_group_launch_order[2]];
  }
  return work_groups_count;
}
}  // namespace

void ComputeTask::Init(std::unique_ptr<GPUOperation>&& operation) {
  operation_ = std::move(operation);
}

const OperationDef& ComputeTask::GetDefinition() const {
  return operation_->definition_;
}

bool ComputeTask::IsLinkable() const { return operation_->IsLinkable(); }

absl::Status ComputeTask::AddTask(ComputeTask* task) {
  return operation_->AddOperation(task->operation_.get());
}

absl::Status ComputeTask::Compile(MetalDevice* device) {
  operation_->AssembleCode(device->GetInfo());
  const std::map<std::string, std::string> linkables = {
      {operation_->dst_tensors_names_[0], operation_->elementwise_code_}};
  RETURN_IF_ERROR(metal_args_.Init(linkables, device, &operation_->args_,
                                   &operation_->code_));

  operation_->args_.ReleaseCPURepresentation();

  return CompileProgram(device, operation_->definition_.precision,
                        operation_->code_);
}

absl::Status ComputeTask::CompileProgram(MetalDevice* device,
                                         CalculationsPrecision precision,
                                         const std::string& kernel_code) {
  NSString* barrier;
  // simdgroup_barrier is supported since Metal shading language version 2.0
  if (device->IsLanguageVersion2orHigher()) {
    barrier = @"simdgroup_barrier";
  } else {
    barrier = @"threadgroup_barrier";
  }
  NSString* storageType;
  NSString* accumulatorType;
  NSString* toAccumulatorType4 = @"";
  if (precision == CalculationsPrecision::F32) {
    storageType = @"float";
    accumulatorType = @"float";
  } else {
    // FP16
    storageType = @"half";
    if (precision == CalculationsPrecision::F32_F16) {
      accumulatorType = @"float";
      toAccumulatorType4 = @"float4";
    } else {
      accumulatorType = @"half";
    }
  }

  NSMutableDictionary<NSString*, NSString*>* mutableMacros =
      [[NSMutableDictionary alloc] init];
  if (@available(iOS 15.0, *)) {
    // float16, float8 and half8  are already defined
    // in metal_stdlib on iOS15.
  } else {
    [mutableMacros addEntriesFromDictionary:@{
      @"float16" : @"float4x4",
      @"float8" : @"float2x4",
      @"half8" : @"half2x4",
    }];
  }
  [mutableMacros addEntriesFromDictionary:@{
    @"half16" : @"half4x4",
    @"FLT16_0123(V)" : @"V[0]",
    @"FLT16_4567(V)" : @"V[1]",
    @"FLT16_89ab(V)" : @"V[2]",
    @"FLT16_cdef(V)" : @"V[3]",
    @"FLT" : storageType,
    @"FLT2" : [NSString stringWithFormat:@"%@2", storageType],
    @"FLT3" : [NSString stringWithFormat:@"%@3", storageType],
    @"FLT4" : [NSString stringWithFormat:@"%@4", storageType],
    @"ACCUM_FLT" : accumulatorType,
    @"ACCUM_FLT2" : [NSString stringWithFormat:@"%@2", accumulatorType],
    @"ACCUM_FLT3" : [NSString stringWithFormat:@"%@3", accumulatorType],
    @"ACCUM_FLT4" : [NSString stringWithFormat:@"%@4", accumulatorType],
    @"INIT_ACCUM_FLT4(value)" :
        [NSString stringWithFormat:@"%@4(value)", accumulatorType],
    @"TO_ACCUM_TYPE" : toAccumulatorType4,
    @"TO_ACCUM_FLT" : accumulatorType,
    @"TO_ACCUM_FLT2" : [NSString stringWithFormat:@"%@2", accumulatorType],
    @"TO_ACCUM_FLT3" : [NSString stringWithFormat:@"%@3", accumulatorType],
    @"TO_ACCUM_FLT4" : [NSString stringWithFormat:@"%@4", accumulatorType],
    @"TO_FLT4" : [NSString stringWithFormat:@"%@4", storageType],
    @"SIMDGROUP_BARRIER" : barrier,
    @"SIMD_LOCAL_MEM_BARRIER" : barrier,
    @"MAIN_FUNCTION" : @"\"kernel void ComputeFunction\"",
    @"GLOBAL_ID_0" : @"static_cast<int>(reserved_gid.x)",
    @"GLOBAL_ID_1" : @"static_cast<int>(reserved_gid.y)",
    @"GLOBAL_ID_2" : @"static_cast<int>(reserved_gid.z)",
    @"LOCAL_ID_0" : @"static_cast<int>(reserved_lid.x)",
    @"LOCAL_ID_1" : @"static_cast<int>(reserved_lid.y)",
    @"LOCAL_ID_2" : @"static_cast<int>(reserved_lid.z)",
    @"GROUP_ID_0" : @"static_cast<int>(reserved_group_id.x)",
    @"GROUP_ID_1" : @"static_cast<int>(reserved_group_id.y)",
    @"GROUP_ID_2" : @"static_cast<int>(reserved_group_id.z)",
    @"GROUP_SIZE_0" : @"static_cast<int>(reserved_group_size.x)",
    @"GROUP_SIZE_1" : @"static_cast<int>(reserved_group_size.y)",
    @"GROUP_SIZE_2" : @"static_cast<int>(reserved_group_size.z)",
    @"SUB_GROUP_LOCAL_ID" : @"static_cast<int>(reserved_simd_id)",
    @"\"SUB_GROUP_BROADCAST(V, ID)\"" : @"\"simd_broadcast(V, ID)\"",
    @"__local" : @"threadgroup",
    @"__global" : @"device",
    @"__constant" : @"constant",
    @"LOCAL_MEM_BARRIER" : @"threadgroup_barrier(mem_flags::mem_threadgroup)",
    @"INIT_FLT(value)" : [NSString stringWithFormat:@"%@(value)", storageType],
    @"INIT_FLT4(value)" :
        [NSString stringWithFormat:@"%@4(value)", storageType],
    @"\"INIT_FLT4v4(v0, v1, v2, v3)\"" :
        [NSString stringWithFormat:@"\"%@4(v0, v1, v2, v3)\"", storageType],
    @"INIT_FLOAT(value)" : @"float(value)",
    @"INIT_FLOAT2(value)" : @"float2(value)",
    @"\"INIT_FLOAT2v2(v0, v1)\"" : @"\"float2(v0, v1)\"",
    @"INIT_FLOAT3(value)" : @"float3(value)",
    @"\"INIT_FLOAT3v3(v0, v1, v2)\"" : @"\"float3(v0, v1, v2)\"",
    @"INIT_FLOAT4(value)" : @"float4(value)",
    @"\"INIT_FLOAT4v4(v0, v1, v2, v3)\"" : @"\"float4(v0, v1, v2, v3)\"",
    @"INIT_INT(value)" : @"int(value)",
    @"\"INIT_INT2v2(v0, v1)\"" : @"\"int2(v0, v1)\"",
    @"\"INIT_INT4v4(v0, v1, v2, v3)\"" : @"\"int4(v0, v1, v2, v3)\"",
    @"CONVERT_TO_INT4(value)" : @"int4(value)",
    @"\"SELECT_BY_INDEX_FROM_FLT4(value, index)\"" : @"\"(value)[index]\"",
  }];
  NSDictionary<NSString*, NSString*>* macros = [mutableMacros copy];

  NSString* code =
      [NSString stringWithCString:kernel_code.c_str()
                         encoding:[NSString defaultCStringEncoding]];
  id<MTLComputePipelineState> program;
  RETURN_IF_ERROR(CreateComputeProgram(device->device(), code,
                                       @"ComputeFunction", macros, &program));
  if (!program) {
    return absl::InternalError("Unknown shader compilation error");
  }
  program_ = program;
  return absl::OkStatus();
}

absl::Status ComputeTask::UpdateParams() {
  for (int i = 0; i < operation_->src_tensors_names_.size(); ++i) {
    const auto* metal_spatial_tensor =
        dynamic_cast<const MetalSpatialTensor*>(operation_->src_[i]);
    if (!metal_spatial_tensor) {
      return absl::InvalidArgumentError("Expected MetalSpatialTensor.");
    }
    RETURN_IF_ERROR(metal_args_.SetObjectRef(operation_->src_tensors_names_[i],
                                             *metal_spatial_tensor));
  }
  for (int i = 0; i < operation_->dst_tensors_names_.size(); ++i) {
    const auto* metal_spatial_tensor =
        dynamic_cast<const MetalSpatialTensor*>(operation_->dst_[i]);
    if (!metal_spatial_tensor) {
      return absl::InvalidArgumentError("Expected MetalSpatialTensor.");
    }
    RETURN_IF_ERROR(metal_args_.SetObjectRef(operation_->dst_tensors_names_[i],
                                             *metal_spatial_tensor));
  }
  RETURN_IF_ERROR(operation_->BindArguments(&metal_args_));
  operation_->grid_size_ = operation_->GetGridSize();
  operation_->work_groups_count_ = GetWorkGroupsCount(
      operation_->grid_dimension_, operation_->grid_size_,
      operation_->work_group_size_, operation_->work_group_launch_order_);
  return absl::OkStatus();
}

void ComputeTask::Encode(id<MTLComputeCommandEncoder> encoder) {
  [encoder setComputePipelineState:program_];
  metal_args_.Encode(encoder, 0);
  MTLSize groupsCount, groupsSize;
  groupsCount.width = operation_->work_groups_count_.x;
  groupsCount.height = operation_->work_groups_count_.y;
  groupsCount.depth = operation_->work_groups_count_.z;
  groupsSize.width = operation_->work_group_size_.x;
  groupsSize.height = operation_->work_group_size_.y;
  groupsSize.depth = operation_->work_group_size_.z;
  [encoder dispatchThreadgroups:groupsCount threadsPerThreadgroup:groupsSize];
}

void ComputeTask::SetSrcTensor(MetalSpatialTensor* tensor, int index) {
  operation_->SetSrc(tensor, index);
  auto status =
      metal_args_.SetObjectRef(operation_->src_tensors_names_[index], *tensor);
}

void ComputeTask::SetDstTensor(MetalSpatialTensor* tensor, int index) {
  operation_->SetDst(tensor, index);
  auto status =
      metal_args_.SetObjectRef(operation_->dst_tensors_names_[index], *tensor);
}

absl::Status ComputeTask::Tune(TuningType tuning_type, MetalDevice* device) {
  std::vector<int3> possible_work_groups;
  KernelInfo kernel_info;
  kernel_info.max_work_group_size = [program_ maxTotalThreadsPerThreadgroup];
  kernel_info.private_memory_size = 0;
  operation_->GetPossibleKernelWorkGroups(tuning_type, device->GetInfo(),
                                          kernel_info, &possible_work_groups);
  if (possible_work_groups.empty()) {
    return absl::NotFoundError(
        "Can not found work_group size to launch kernel");
  }
  operation_->work_group_size_ = possible_work_groups[0];
  operation_->work_groups_count_ = GetWorkGroupsCount(
      operation_->grid_dimension_, operation_->grid_size_,
      operation_->work_group_size_, operation_->work_group_launch_order_);
  return absl::OkStatus();
}

}  // namespace metal
}  // namespace gpu
}  // namespace tflite
<|endoftext|>


<commit_before>// Copyright 2010-2012 RethinkDB, all rights reserved.
#ifndef ARCH_SPINLOCK_HPP_
#define ARCH_SPINLOCK_HPP_

#include <pthread.h>
#include <string.h>

#include "errors.hpp"

// TODO(OSX) Reconsider how to detect spinlock feature.
#if __APPLE__
#define PTHREAD_HAS_SPINLOCK 0
#else
#define PTHREAD_HAS_SPINLOCK 1
#endif

// TODO: we should use regular mutexes on single core CPU
// instead of spinlocks

class spinlock_t {
public:
    friend class spinlock_acq_t;

    spinlock_t() {
#if PTHREAD_HAS_SPINLOCK
        pthread_spin_init(&l, PTHREAD_PROCESS_PRIVATE);
#else
        pthread_mutex_init(&l, NULL);
#endif
    }
    ~spinlock_t() {
#if PTHREAD_HAS_SPINLOCK
        pthread_spin_destroy(&l);
#else
        pthread_mutex_destroy(&l);
#endif
    }

private:
    void lock() {
#if PTHREAD_HAS_SPINLOCK
        int res = pthread_spin_lock(&l);
#else
        int res = pthread_mutex_lock(&l);
#endif
        guarantee_err(res == 0, "could not lock spin lock");
    }
    void unlock() {
#if PTHREAD_HAS_SPINLOCK
        int res = pthread_spin_unlock(&l);
#else
        int res = pthread_mutex_unlock(&l);
#endif
        guarantee_err(res == 0, "could not unlock spin lock");
    }

#if PTHREAD_HAS_SPINLOCK
    pthread_spinlock_t l;
#else
    pthread_mutex_t l;
#endif

    DISABLE_COPYING(spinlock_t);
};

class spinlock_acq_t {
public:
    explicit spinlock_acq_t(spinlock_t *the_lock) : the_lock_(the_lock) {
        the_lock_->lock();
    }
    ~spinlock_acq_t() {
        the_lock_->unlock();
    }

private:
    spinlock_t *the_lock_;

    DISABLE_COPYING(spinlock_acq_t);
};

#endif /* ARCH_SPINLOCK_HPP_ */
<commit_msg>Made spinlock_t check return codes and handle them properly with guarantee_xerr in place of guarantee_err.<commit_after>// Copyright 2010-2012 RethinkDB, all rights reserved.
#ifndef ARCH_SPINLOCK_HPP_
#define ARCH_SPINLOCK_HPP_

#include <pthread.h>
#include <string.h>

#include "errors.hpp"

// TODO(OSX) Reconsider how to detect spinlock feature.
#if __APPLE__
#define PTHREAD_HAS_SPINLOCK 0
#else
#define PTHREAD_HAS_SPINLOCK 1
#endif

// TODO: we should use regular mutexes on single core CPU
// instead of spinlocks

class spinlock_t {
public:
    friend class spinlock_acq_t;

    spinlock_t() {
#if PTHREAD_HAS_SPINLOCK
        // TODO(OSX) check that pthread_spin_* functions used in this file return error codes.
        int res = pthread_spin_init(&l, PTHREAD_PROCESS_PRIVATE);
#else
        int res = pthread_mutex_init(&l, NULL);
#endif
        guarantee_xerr(res == 0, res, "could not initialize spinlock");
    }
    ~spinlock_t() {
#if PTHREAD_HAS_SPINLOCK
        int res = pthread_spin_destroy(&l);
#else
        int res = pthread_mutex_destroy(&l);
#endif
        guarantee_xerr(res == 0, res, "could not destroy spinlock");
    }

private:
    void lock() {
#if PTHREAD_HAS_SPINLOCK
        int res = pthread_spin_lock(&l);
#else
        int res = pthread_mutex_lock(&l);
#endif
        guarantee_xerr(res == 0, res, "could not lock spin lock");
    }
    void unlock() {
#if PTHREAD_HAS_SPINLOCK
        int res = pthread_spin_unlock(&l);
#else
        int res = pthread_mutex_unlock(&l);
#endif
        guarantee_xerr(res == 0, res, "could not unlock spin lock");
    }

#if PTHREAD_HAS_SPINLOCK
    pthread_spinlock_t l;
#else
    pthread_mutex_t l;
#endif

    DISABLE_COPYING(spinlock_t);
};

class spinlock_acq_t {
public:
    explicit spinlock_acq_t(spinlock_t *the_lock) : the_lock_(the_lock) {
        the_lock_->lock();
    }
    ~spinlock_acq_t() {
        the_lock_->unlock();
    }

private:
    spinlock_t *the_lock_;

    DISABLE_COPYING(spinlock_acq_t);
};

#endif /* ARCH_SPINLOCK_HPP_ */
<|endoftext|>