<commit_before>//
// This file is a part of pomerol - a scientific ED code for obtaining
// properties of a Hubbard model on a finite-size lattice
//
// Copyright (C) 2010-2012 Andrey Antipov <Andrey.E.Antipov@gmail.com>
// Copyright (C) 2010-2012 Igor Krivenko <Igor.S.Krivenko@gmail.com>
//
// pomerol 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.
//
// pomerol 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 pomerol.  If not, see <http://www.gnu.org/licenses/>.

/** \file src/TwoParticleGF.cpp
** \brief Two-particle Green's function in the Matsubara representation.
**
** \author Igor Krivenko (Igor.S.Krivenko@gmail.com)
** \author Andrey Antipov (Andrey.E.Antipov@gmail.com)
*/
#include "TwoParticleGF.h"
#include <boost/serialization/complex.hpp>
#include <boost/serialization/vector.hpp>

namespace Pomerol{

TwoParticleGF::TwoParticleGF(const StatesClassification& S, const Hamiltonian& H,
                const AnnihilationOperator& C1, const AnnihilationOperator& C2,
                const CreationOperator& CX3, const CreationOperator& CX4,
                const DensityMatrix& DM) :
    Thermal(DM.beta), ComputableObject(Constructed),
    S(S), H(H), C1(C1), C2(C2), CX3(CX3), CX4(CX4), DM(DM),
    parts(0), Vanishing(true)
{}

TwoParticleGF::~TwoParticleGF()
{
      for(std::vector<TwoParticleGFPart*>::iterator iter = parts.begin(); iter != parts.end(); iter++)
          delete *iter;
}

BlockNumber TwoParticleGF::getLeftIndex(size_t PermutationNumber, size_t OperatorPosition, BlockNumber RightIndex) const
{
    switch(permutations3[PermutationNumber].perm[OperatorPosition]){
        case 0: return C1.getLeftIndex(RightIndex);
        case 1: return C2.getLeftIndex(RightIndex);
        case 2: return CX3.getLeftIndex(RightIndex);
        default: return ERROR_BLOCK_NUMBER;
    }
}

BlockNumber TwoParticleGF::getRightIndex(size_t PermutationNumber, size_t OperatorPosition, BlockNumber LeftIndex) const
{
    switch(permutations3[PermutationNumber].perm[OperatorPosition]){
        case 0: return C1.getRightIndex(LeftIndex);
        case 1: return C2.getRightIndex(LeftIndex);
        case 2: return CX3.getRightIndex(LeftIndex);
        default: return ERROR_BLOCK_NUMBER;
    }
}

const FieldOperatorPart& TwoParticleGF::OperatorPartAtPosition(size_t PermutationNumber, size_t OperatorPosition, BlockNumber LeftIndex) const
{
    switch(permutations3[PermutationNumber].perm[OperatorPosition]){
        case 0: return C1.getPartFromLeftIndex(LeftIndex);
        case 1: return C2.getPartFromLeftIndex(LeftIndex);
        case 2: return CX3.getPartFromLeftIndex(LeftIndex);
        default: assert(0);
    }
}

void TwoParticleGF::prepare(void)
{
    if(Status>=Prepared) return;

    // Find out non-trivial blocks of CX4.
    FieldOperator::BlocksBimap const& CX4NontrivialBlocks = CX4.getBlockMapping();
    for(FieldOperator::BlocksBimap::right_const_iterator outer_iter = CX4NontrivialBlocks.right.begin();
        outer_iter != CX4NontrivialBlocks.right.end(); outer_iter++){ // Iterate over the outermost index.
            for(size_t p=0; p<6; ++p){ // Choose a permutation
                  BlockNumber LeftIndices[4];
                  LeftIndices[0] = outer_iter->first;
                  LeftIndices[3] = outer_iter->second;
                  LeftIndices[2] = getLeftIndex(p,2,LeftIndices[3]);
                  LeftIndices[1] = getRightIndex(p,0,LeftIndices[0]);
                  // < LeftIndices[0] | O_1 | LeftIndices[1] >
                  // < LeftIndices[1] | O_2 | getRightIndex(p,1,LeftIndices[1]) >
                  // < LeftIndices[2]| O_3 | LeftIndices[3] >
                  // < LeftIndices[3] | CX4 | LeftIndices[0] >
                  // Select a relevant 'world stripe' (sequence of blocks).
                  if(getRightIndex(p,1,LeftIndices[1]) == LeftIndices[2] && LeftIndices[1].isCorrect() && LeftIndices[2].isCorrect()){
                      // DEBUG
                      /*DEBUG("new part: "  << S.getBlockInfo(LeftIndices[0]) << " "
                                          << S.getBlockInfo(LeftIndices[1]) << " "
                                          << S.getBlockInfo(LeftIndices[2]) << " "
                                          << S.getBlockInfo(LeftIndices[3]) << " "
                      <<"BlockNumbers part: "  << LeftIndices[0] << " " << LeftIndices[1] << " " << LeftIndices[2] << " " << LeftIndices[3]);
                      */
                      parts.push_back(new TwoParticleGFPart(
                            OperatorPartAtPosition(p,0,LeftIndices[0]),
                            OperatorPartAtPosition(p,1,LeftIndices[1]),
                            OperatorPartAtPosition(p,2,LeftIndices[2]),
                            (CreationOperatorPart&)CX4.getPartFromLeftIndex(LeftIndices[3]),
                            H.getPart(LeftIndices[0]), H.getPart(LeftIndices[1]), H.getPart(LeftIndices[2]), H.getPart(LeftIndices[3]),
                            DM.getPart(LeftIndices[0]), DM.getPart(LeftIndices[1]), DM.getPart(LeftIndices[2]), DM.getPart(LeftIndices[3]),
                      permutations3[p]));

                      (*parts.rbegin())->KroneckerSymbolTolerance = KroneckerSymbolTolerance;
                      (*parts.rbegin())->ReduceResonanceTolerance = ReduceResonanceTolerance;
                      (*parts.rbegin())->CoefficientTolerance = CoefficientTolerance;
                      (*parts.rbegin())->ReduceInvocationThreshold = ReduceInvocationThreshold;
                      (*parts.rbegin())->MultiTermCoefficientTolerance = MultiTermCoefficientTolerance;
                      }
            }
    }
    if ( parts.size() > 0 ) {
        Vanishing = false;
        INFO("TwoParticleGF(" << getIndex(0) << getIndex(1) << getIndex(2) << getIndex(3) << "): " << parts.size() << " parts will be calculated");
        }
    Status = Prepared;
}

bool TwoParticleGF::isVanishing(void) const
{
    return Vanishing;
}

void TwoParticleGF::compute(const boost::mpi::communicator & comm)
{
    comm.barrier();
    int comm_size = comm.size();
    int comm_rank = comm.rank();
    if (Status >= Computed) return;
    if (Status < Prepared) { throw (exStatusMismatch()); };
    if (comm_rank==0) { INFO("Calculating using " << comm_size << " procs."); };

    if(!Vanishing){
        for (size_t p = 0; p<parts.size(); p++) {
            if (comm_rank == p%comm_size) {
                std::cout << "["<<p+1<<"/"<<parts.size()<< "] Proc " << comm.rank() << " : " << std::flush;
                parts[p]->compute();
                boost::mpi::broadcast(comm, parts[p]->NonResonantTerms, comm_rank);
                boost::mpi::broadcast(comm, parts[p]->ResonantTerms, comm_rank);
                }
            else {
                boost::mpi::broadcast(comm, parts[p]->NonResonantTerms, p%comm_size);
                boost::mpi::broadcast(comm, parts[p]->ResonantTerms, p%comm_size);
                parts[p]->Status = TwoParticleGFPart::Computed;
                 };
                };
        };
    Status = Computed;
}

// size_t TwoParticleGF::getNumResonantTerms() const
// {
//     size_t num = 0;
//     for(std::vector<TwoParticleGFPart*>::const_iterator iter = parts.begin(); iter != parts.end(); iter++)
//         num += (*iter)->getNumResonantTerms();
//     return num;
// }
//
// size_t TwoParticleGF::getNumNonResonantTerms() const
// {
//     size_t num = 0;
//     for(std::vector<TwoParticleGFPart*>::const_iterator iter = parts.begin(); iter != parts.end(); iter++)
//         num += (*iter)->getNumNonResonantTerms();
//     return num;
// }

ParticleIndex TwoParticleGF::getIndex(size_t Position) const
{
    switch(Position){
        case 0: return C1.getIndex();
        case 1: return C2.getIndex();
        case 2: return CX3.getIndex();
        case 3: return CX4.getIndex();
        default: assert(0);
    }
}

unsigned short TwoParticleGF::getPermutationNumber ( const Permutation3& in )
{
    for (unsigned short i=0; i<6; ++i) if (in == permutations3[i]) return i;
    ERROR("TwoParticleGF: Permutation " << in << " not found in all permutations3");
    return 0;
}

} // end of namespace Pomerol

<commit_msg>Separate calculation and broadcasting to gain anything from mpi<commit_after>//
// This file is a part of pomerol - a scientific ED code for obtaining
// properties of a Hubbard model on a finite-size lattice
//
// Copyright (C) 2010-2012 Andrey Antipov <Andrey.E.Antipov@gmail.com>
// Copyright (C) 2010-2012 Igor Krivenko <Igor.S.Krivenko@gmail.com>
//
// pomerol 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.
//
// pomerol 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 pomerol.  If not, see <http://www.gnu.org/licenses/>.

/** \file src/TwoParticleGF.cpp
** \brief Two-particle Green's function in the Matsubara representation.
**
** \author Igor Krivenko (Igor.S.Krivenko@gmail.com)
** \author Andrey Antipov (Andrey.E.Antipov@gmail.com)
*/
#include "TwoParticleGF.h"
#include <boost/serialization/complex.hpp>
#include <boost/serialization/vector.hpp>

namespace Pomerol{

TwoParticleGF::TwoParticleGF(const StatesClassification& S, const Hamiltonian& H,
                const AnnihilationOperator& C1, const AnnihilationOperator& C2,
                const CreationOperator& CX3, const CreationOperator& CX4,
                const DensityMatrix& DM) :
    Thermal(DM.beta), ComputableObject(Constructed),
    S(S), H(H), C1(C1), C2(C2), CX3(CX3), CX4(CX4), DM(DM),
    parts(0), Vanishing(true)
{}

TwoParticleGF::~TwoParticleGF()
{
      for(std::vector<TwoParticleGFPart*>::iterator iter = parts.begin(); iter != parts.end(); iter++)
          delete *iter;
}

BlockNumber TwoParticleGF::getLeftIndex(size_t PermutationNumber, size_t OperatorPosition, BlockNumber RightIndex) const
{
    switch(permutations3[PermutationNumber].perm[OperatorPosition]){
        case 0: return C1.getLeftIndex(RightIndex);
        case 1: return C2.getLeftIndex(RightIndex);
        case 2: return CX3.getLeftIndex(RightIndex);
        default: return ERROR_BLOCK_NUMBER;
    }
}

BlockNumber TwoParticleGF::getRightIndex(size_t PermutationNumber, size_t OperatorPosition, BlockNumber LeftIndex) const
{
    switch(permutations3[PermutationNumber].perm[OperatorPosition]){
        case 0: return C1.getRightIndex(LeftIndex);
        case 1: return C2.getRightIndex(LeftIndex);
        case 2: return CX3.getRightIndex(LeftIndex);
        default: return ERROR_BLOCK_NUMBER;
    }
}

const FieldOperatorPart& TwoParticleGF::OperatorPartAtPosition(size_t PermutationNumber, size_t OperatorPosition, BlockNumber LeftIndex) const
{
    switch(permutations3[PermutationNumber].perm[OperatorPosition]){
        case 0: return C1.getPartFromLeftIndex(LeftIndex);
        case 1: return C2.getPartFromLeftIndex(LeftIndex);
        case 2: return CX3.getPartFromLeftIndex(LeftIndex);
        default: assert(0);
    }
}

void TwoParticleGF::prepare(void)
{
    if(Status>=Prepared) return;

    // Find out non-trivial blocks of CX4.
    FieldOperator::BlocksBimap const& CX4NontrivialBlocks = CX4.getBlockMapping();
    for(FieldOperator::BlocksBimap::right_const_iterator outer_iter = CX4NontrivialBlocks.right.begin();
        outer_iter != CX4NontrivialBlocks.right.end(); outer_iter++){ // Iterate over the outermost index.
            for(size_t p=0; p<6; ++p){ // Choose a permutation
                  BlockNumber LeftIndices[4];
                  LeftIndices[0] = outer_iter->first;
                  LeftIndices[3] = outer_iter->second;
                  LeftIndices[2] = getLeftIndex(p,2,LeftIndices[3]);
                  LeftIndices[1] = getRightIndex(p,0,LeftIndices[0]);
                  // < LeftIndices[0] | O_1 | LeftIndices[1] >
                  // < LeftIndices[1] | O_2 | getRightIndex(p,1,LeftIndices[1]) >
                  // < LeftIndices[2]| O_3 | LeftIndices[3] >
                  // < LeftIndices[3] | CX4 | LeftIndices[0] >
                  // Select a relevant 'world stripe' (sequence of blocks).
                  if(getRightIndex(p,1,LeftIndices[1]) == LeftIndices[2] && LeftIndices[1].isCorrect() && LeftIndices[2].isCorrect()){
                      // DEBUG
                      /*DEBUG("new part: "  << S.getBlockInfo(LeftIndices[0]) << " "
                                          << S.getBlockInfo(LeftIndices[1]) << " "
                                          << S.getBlockInfo(LeftIndices[2]) << " "
                                          << S.getBlockInfo(LeftIndices[3]) << " "
                      <<"BlockNumbers part: "  << LeftIndices[0] << " " << LeftIndices[1] << " " << LeftIndices[2] << " " << LeftIndices[3]);
                      */
                      parts.push_back(new TwoParticleGFPart(
                            OperatorPartAtPosition(p,0,LeftIndices[0]),
                            OperatorPartAtPosition(p,1,LeftIndices[1]),
                            OperatorPartAtPosition(p,2,LeftIndices[2]),
                            (CreationOperatorPart&)CX4.getPartFromLeftIndex(LeftIndices[3]),
                            H.getPart(LeftIndices[0]), H.getPart(LeftIndices[1]), H.getPart(LeftIndices[2]), H.getPart(LeftIndices[3]),
                            DM.getPart(LeftIndices[0]), DM.getPart(LeftIndices[1]), DM.getPart(LeftIndices[2]), DM.getPart(LeftIndices[3]),
                      permutations3[p]));

                      (*parts.rbegin())->KroneckerSymbolTolerance = KroneckerSymbolTolerance;
                      (*parts.rbegin())->ReduceResonanceTolerance = ReduceResonanceTolerance;
                      (*parts.rbegin())->CoefficientTolerance = CoefficientTolerance;
                      (*parts.rbegin())->ReduceInvocationThreshold = ReduceInvocationThreshold;
                      (*parts.rbegin())->MultiTermCoefficientTolerance = MultiTermCoefficientTolerance;
                      }
            }
    }
    if ( parts.size() > 0 ) {
        Vanishing = false;
        INFO("TwoParticleGF(" << getIndex(0) << getIndex(1) << getIndex(2) << getIndex(3) << "): " << parts.size() << " parts will be calculated");
        }
    Status = Prepared;
}

bool TwoParticleGF::isVanishing(void) const
{
    return Vanishing;
}

void TwoParticleGF::compute(const boost::mpi::communicator & comm)
{
    comm.barrier();
    int comm_size = comm.size();
    int comm_rank = comm.rank();
    if (Status >= Computed) return;
    if (Status < Prepared) { throw (exStatusMismatch()); };

    std::map<size_t,int> jobs_dispatcher;

    if(!Vanishing){
        if (comm_rank==0) { INFO("Calculating using " << comm_size << " procs."); };
        for (size_t p = 0; p<parts.size(); p++) {
            jobs_dispatcher[p] = p%comm_size;
            if (comm_rank == p%comm_size) {
                std::cout << "["<<p+1<<"/"<<parts.size()<< "] Proc " << comm.rank() << " : " << std::flush;
                parts[p]->compute();
            };
        };

        for (size_t p = 0; p<parts.size(); p++) {
            if (comm_rank == jobs_dispatcher[p]) {
                boost::mpi::broadcast(comm, parts[p]->NonResonantTerms, comm_rank);
                boost::mpi::broadcast(comm, parts[p]->ResonantTerms, comm_rank);
                }
            else {
                boost::mpi::broadcast(comm, parts[p]->NonResonantTerms, jobs_dispatcher[p]);
                boost::mpi::broadcast(comm, parts[p]->ResonantTerms, jobs_dispatcher[p]);
                parts[p]->Status = TwoParticleGFPart::Computed;
                 };
                };
        };
    Status = Computed;
}

// size_t TwoParticleGF::getNumResonantTerms() const
// {
//     size_t num = 0;
//     for(std::vector<TwoParticleGFPart*>::const_iterator iter = parts.begin(); iter != parts.end(); iter++)
//         num += (*iter)->getNumResonantTerms();
//     return num;
// }
//
// size_t TwoParticleGF::getNumNonResonantTerms() const
// {
//     size_t num = 0;
//     for(std::vector<TwoParticleGFPart*>::const_iterator iter = parts.begin(); iter != parts.end(); iter++)
//         num += (*iter)->getNumNonResonantTerms();
//     return num;
// }

ParticleIndex TwoParticleGF::getIndex(size_t Position) const
{
    switch(Position){
        case 0: return C1.getIndex();
        case 1: return C2.getIndex();
        case 2: return CX3.getIndex();
        case 3: return CX4.getIndex();
        default: assert(0);
    }
}

unsigned short TwoParticleGF::getPermutationNumber ( const Permutation3& in )
{
    for (unsigned short i=0; i<6; ++i) if (in == permutations3[i]) return i;
    ERROR("TwoParticleGF: Permutation " << in << " not found in all permutations3");
    return 0;
}

} // end of namespace Pomerol

<|endoftext|>


<commit_before>#include <icvar.h>
#include "main.h"
#include "cvar.h"
#include "sound.h"

/*
 * Define ConVars
 *
 * Note:
 *
 *  Notable constructors of ConVar
 *
        ConVar( const char *pName, const char *pDefaultValue, int flags = 0);

        ConVar( const char *pName, const char *pDefaultValue, int flags, const char *pHelpString );

        ConVar( const char *pName, const char *pDefaultValue, int flags, const char *pHelpString,
            bool bMin, float fMin, bool bMax, float fMax );

        ConVar( const char *pName, const char *pDefaultValue, int flags, const char *pHelpString,
            FnChangeCallback_t callback );

        ConVar( const char *pName, const char *pDefaultValue, int flags,
            const char *pHelpString, bool bMin, float fMin, bool bMax, float fMax,
            FnChangeCallback_t callback );

 *
 *  FCVAR_* flags can be found in tier1/iconvar.h
 *  FCVAR_ARCHIVE will allow cvars to be archived
 *
 *  Definition of FnChangeCallback_t:
 *
        typedef void ( *FnChangeCallback_t )( IConVar *var, const char *pOldValue, float flOldValue );
 *
 */

static ConVar prec_version("prec_version", OPENPREC_VERSION, FCVAR_NONE,
        "The current version of open-prec");

ConVar prec_next_demoname("prec_next_demoname", "", FCVAR_ARCHIVE,
        "Name of the next demo (unimplemented)");

ConVar prec_notify("prec_notify", "1", FCVAR_ARCHIVE,
        "Where to log notifications (started/stopped recording, bookmarks, etc.)\n"
        "\t0 - Log to console\n"
        "\t1 - Log to team chat\n"
        "\t2 - Display on HUD (unimplemented)");

enum struct PrecNotify {
    Console,
    Chat,
    Hud
};

ConVar prec_mode("prec_mode", "2", FCVAR_ARCHIVE,
        "\t0 - Turn off addon\n"
        "\t1 - Record only custom named demos\n"
        "\t2 - Record on servers with mp_tournament 1\n"
        "\t3 - Always record");

ConVar prec_min_streak("prec_min_streak", "4", FCVAR_ARCHIVE,
        "Minimum killstreak to log (unimplemented)",
        true, 2.0, false, 0);

ConVar prec_kill_delay("prec_kill_delay", "15", FCVAR_ARCHIVE,
        "Max delay between kills in a killstreak (unimplemented)",
        true, 5.0, false, 0);

ConVar prec_dir("prec_dir", "/home/matt/", FCVAR_ARCHIVE,
        "Directory to store demos in (unimplemented)");

ConVar prec_sound("prec_sound", "1", FCVAR_ARCHIVE,
        "Play sounds when starting/stopping recording\n"
        "\t0 - Off\n"
        "\t1 - On");

ConVar prec_screens("prec_screens", "0", FCVAR_ARCHIVE,
        "Take screenshots of score and status (unimplemented)\n"
        "\t0 - Off"
        "\t1 - On");

ConVar prec_log("prec_log", "1", FCVAR_ARCHIVE,
        "\t0 - Don't log killstreaks or bookmarks (unimplemented)\n"
        "\t1 - Log killstreaks and bookmarks to killstreaks.txt (unimplemented)\n"
        "\t2 - Log separately for each demo (unimplemented)\n"
        "\t3 - Log separately for each map (unimplemented)");

ConVar prec_delete_useless_demo("prec_delete_useless_demo", "0", FCVAR_ARCHIVE,
        "Delete demo files without killstreaks or bookmarks (unimplemented)\n"
        "\t0 - Off \n"
        "\t1 - On");

ConVar prec_stv_status("prec_stv_status", "0", FCVAR_ARCHIVE,
        "Show SourceTV statistics after match (unimplemented)\n"
        "\t0 - Off\n"
        "\t1 - On");

ConVar prec_tag("prec_tag", "", FCVAR_ARCHIVE,
        "Tag for demo file name (unimplemented)");

static const int g_numCvars = 13;
static ConVar *g_cvarList[] = {
    &prec_version,
    &prec_next_demoname,
    &prec_notify,
    &prec_mode,
    &prec_min_streak,
    &prec_kill_delay,
    &prec_dir,
    &prec_sound,
    &prec_screens,
    &prec_log,
    &prec_delete_useless_demo,
    &prec_stv_status,
    &prec_tag
};

/*
 * Define ConCommands
 *
 * Note:
 *
 *  CON_COMMAND(name, ...) expands to the following
 *
 *      static ConCommand name##_command(...) // Variable
 *      static void name(...) { ... } // Function
 *
 *  First argument passed can safely be assumed to be garbage
 *  Probably a legacy thing, because no built in convars read
 *  the first arg anyways
 */

static const int g_numCommands = 5;
extern ConCommand *g_commandList[g_numCommands];

CON_COMMAND(prec_about, "Display useful information about the plugin") {
    Msg("%s\n%s\n%s\n",
        OPENPREC_NAME " is an open-source demo recording tool for TF2",
        "Version: " OPENPREC_VERSION, "Author: " OPENPREC_AUTH
    );
    play_sound(Sound::About);
}

CON_COMMAND(prec_info, "List commands and cvars") {
    ConCommand *help = g_pCVar->FindCommand("help");
    const char *argv[2];
    argv[0] = "help";
    for (int i = 0; i < g_numCommands; ++i) {
        argv[1] = g_commandList[i]->GetName();
        CCommand args(2, argv);
        help->Dispatch(args);
    }
    for (int i = 0; i < g_numCvars; ++i) {
        argv[1] = g_cvarList[i]->GetName();
        CCommand args(2, argv);
        help->Dispatch(args);
    }
}

#define RECORD_NOTIFICATION "**RECORDING STARTED**"
CON_COMMAND_EXTERN(prec_record, prec_record, "Record a demo") {
    ConCommand *record = g_pCVar->FindCommand("record");
    int mp_tournament = g_pCVar->FindVar("mp_tournament")->GetInt();
    int mode = prec_mode.GetInt();
    const char *argv[2];
    const char *nextdemoname = prec_next_demoname.GetString();
    const char *basedir;
    const char *tag;
    const char *bluteam;
    const char *redteam;
    char mapname[256];
    argv[0] = "record";
    switch (static_cast<PrecMode>(mode)) {
        case PrecMode::Off: return;
        case PrecMode::Named: {
            if (strcmp(nextdemoname, "") == 0) return;
        } break;
        case PrecMode::Tournament: {
            if (mp_tournament != 1) return;
        } break;
        case PrecMode::Always: break;
        default: {
            ConNotify("Invalid recording mode, doing nothing");
            return;
        }
    }
    basedir = prec_dir.GetString();
    tag = prec_tag.GetString();
    bluteam = g_pCVar->FindVar("mp_tournament_blueteamname")->GetString();
    redteam = g_pCVar->FindVar("mp_tournament_redteamname")->GetString();
    g_pEngineClient->GetChapterName(mapname, sizeof(mapname));
    if (strcmp(nextdemoname, "") == 0) {
        // TODO: Custom name formatting
        nextdemoname = "foobar_demo";
    }
    argv[1] = nextdemoname;
    record->Dispatch(CCommand(2, argv));
    if (g_pEngineClient->IsRecordingDemo()) {
        int notify = prec_notify.GetInt();
        ConNotify(RECORD_NOTIFICATION);
        play_sound(Sound::Recording);
        if (mp_tournament == 1) {
            switch (static_cast<PrecNotify>(notify)) {
                default:
                case PrecNotify::Console:
                    break;
                case PrecNotify::Chat: {
                    g_pEngineClient->ClientCmd("say_team [" OPENPREC_NAME "]: " RECORD_NOTIFICATION);
                } break;
                case PrecNotify::Hud: {
                    ConNotifyf(OPENPREC_NOFEATURE_MSGF, "Hud notifications");
                } break;
            }
        }
    }
}

CON_COMMAND(prec_mark, "Make a bookmark at the current tick") {
}

CON_COMMAND(prec_delete_demo, "Delete previous demo") {
}

ConCommand *g_commandList[g_numCommands] = {
    &prec_about_command,
    &prec_info_command,
    &prec_record_command,
    &prec_mark_command,
    &prec_delete_demo_command
};

bool register_cvars() {
    for (int i = 0; i < g_numCvars; ++i) {
        g_pCVar->RegisterConCommand(g_cvarList[i]);
    }
    return true;
}

void unregister_cvars() {
    for (int i = 0; i < g_numCvars; ++i) {
        g_pCVar->UnregisterConCommand(g_cvarList[i]);
    }
}

bool register_concommands() {
    for (int i = 0; i < g_numCommands; ++i) {
        g_pCVar->RegisterConCommand(g_commandList[i]);
    }
    return true;
}

void unregister_concommands() {
    for (int i = 0; i < g_numCommands; ++i) {
        g_pCVar->UnregisterConCommand(g_commandList[i]);
    }
}
<commit_msg>Don't archive next demoname<commit_after>#include <icvar.h>
#include "main.h"
#include "cvar.h"
#include "sound.h"

/*
 * Define ConVars
 *
 * Note:
 *
 *  Notable constructors of ConVar
 *
        ConVar( const char *pName, const char *pDefaultValue, int flags = 0);

        ConVar( const char *pName, const char *pDefaultValue, int flags, const char *pHelpString );

        ConVar( const char *pName, const char *pDefaultValue, int flags, const char *pHelpString,
            bool bMin, float fMin, bool bMax, float fMax );

        ConVar( const char *pName, const char *pDefaultValue, int flags, const char *pHelpString,
            FnChangeCallback_t callback );

        ConVar( const char *pName, const char *pDefaultValue, int flags,
            const char *pHelpString, bool bMin, float fMin, bool bMax, float fMax,
            FnChangeCallback_t callback );

 *
 *  FCVAR_* flags can be found in tier1/iconvar.h
 *  FCVAR_ARCHIVE will allow cvars to be archived
 *
 *  Definition of FnChangeCallback_t:
 *
        typedef void ( *FnChangeCallback_t )( IConVar *var, const char *pOldValue, float flOldValue );
 *
 */

static ConVar prec_version("prec_version", OPENPREC_VERSION, FCVAR_NONE,
        "The current version of open-prec");

ConVar prec_next_demoname("prec_next_demoname", "", FCVAR_NONE,
        "Name of the next demo (unimplemented)");

ConVar prec_notify("prec_notify", "1", FCVAR_ARCHIVE,
        "Where to log notifications (started/stopped recording, bookmarks, etc.)\n"
        "\t0 - Log to console\n"
        "\t1 - Log to team chat\n"
        "\t2 - Display on HUD (unimplemented)");

enum struct PrecNotify {
    Console,
    Chat,
    Hud
};

ConVar prec_mode("prec_mode", "2", FCVAR_ARCHIVE,
        "\t0 - Turn off addon\n"
        "\t1 - Record only custom named demos\n"
        "\t2 - Record on servers with mp_tournament 1\n"
        "\t3 - Always record");

ConVar prec_min_streak("prec_min_streak", "4", FCVAR_ARCHIVE,
        "Minimum killstreak to log (unimplemented)",
        true, 2.0, false, 0);

ConVar prec_kill_delay("prec_kill_delay", "15", FCVAR_ARCHIVE,
        "Max delay between kills in a killstreak (unimplemented)",
        true, 5.0, false, 0);

ConVar prec_dir("prec_dir", "/home/matt/", FCVAR_ARCHIVE,
        "Directory to store demos in (unimplemented)");

ConVar prec_sound("prec_sound", "1", FCVAR_ARCHIVE,
        "Play sounds when starting/stopping recording\n"
        "\t0 - Off\n"
        "\t1 - On");

ConVar prec_screens("prec_screens", "0", FCVAR_ARCHIVE,
        "Take screenshots of score and status (unimplemented)\n"
        "\t0 - Off"
        "\t1 - On");

ConVar prec_log("prec_log", "1", FCVAR_ARCHIVE,
        "\t0 - Don't log killstreaks or bookmarks (unimplemented)\n"
        "\t1 - Log killstreaks and bookmarks to killstreaks.txt (unimplemented)\n"
        "\t2 - Log separately for each demo (unimplemented)\n"
        "\t3 - Log separately for each map (unimplemented)");

ConVar prec_delete_useless_demo("prec_delete_useless_demo", "0", FCVAR_ARCHIVE,
        "Delete demo files without killstreaks or bookmarks (unimplemented)\n"
        "\t0 - Off \n"
        "\t1 - On");

ConVar prec_stv_status("prec_stv_status", "0", FCVAR_ARCHIVE,
        "Show SourceTV statistics after match (unimplemented)\n"
        "\t0 - Off\n"
        "\t1 - On");

ConVar prec_tag("prec_tag", "", FCVAR_ARCHIVE,
        "Tag for demo file name (unimplemented)");

static const int g_numCvars = 13;
static ConVar *g_cvarList[] = {
    &prec_version,
    &prec_next_demoname,
    &prec_notify,
    &prec_mode,
    &prec_min_streak,
    &prec_kill_delay,
    &prec_dir,
    &prec_sound,
    &prec_screens,
    &prec_log,
    &prec_delete_useless_demo,
    &prec_stv_status,
    &prec_tag
};

/*
 * Define ConCommands
 *
 * Note:
 *
 *  CON_COMMAND(name, ...) expands to the following
 *
 *      static ConCommand name##_command(...) // Variable
 *      static void name(...) { ... } // Function
 *
 *  First argument passed can safely be assumed to be garbage
 *  Probably a legacy thing, because no built in convars read
 *  the first arg anyways
 */

static const int g_numCommands = 5;
extern ConCommand *g_commandList[g_numCommands];

CON_COMMAND(prec_about, "Display useful information about the plugin") {
    Msg("%s\n%s\n%s\n",
        OPENPREC_NAME " is an open-source demo recording tool for TF2",
        "Version: " OPENPREC_VERSION, "Author: " OPENPREC_AUTH
    );
    play_sound(Sound::About);
}

CON_COMMAND(prec_info, "List commands and cvars") {
    ConCommand *help = g_pCVar->FindCommand("help");
    const char *argv[2];
    argv[0] = "help";
    for (int i = 0; i < g_numCommands; ++i) {
        argv[1] = g_commandList[i]->GetName();
        CCommand args(2, argv);
        help->Dispatch(args);
    }
    for (int i = 0; i < g_numCvars; ++i) {
        argv[1] = g_cvarList[i]->GetName();
        CCommand args(2, argv);
        help->Dispatch(args);
    }
}

#define RECORD_NOTIFICATION "**RECORDING STARTED**"
CON_COMMAND_EXTERN(prec_record, prec_record, "Record a demo") {
    ConCommand *record = g_pCVar->FindCommand("record");
    int mp_tournament = g_pCVar->FindVar("mp_tournament")->GetInt();
    int mode = prec_mode.GetInt();
    const char *argv[2];
    const char *nextdemoname = prec_next_demoname.GetString();
    const char *basedir;
    const char *tag;
    const char *bluteam;
    const char *redteam;
    char mapname[256];
    argv[0] = "record";
    switch (static_cast<PrecMode>(mode)) {
        case PrecMode::Off: return;
        case PrecMode::Named: {
            if (strcmp(nextdemoname, "") == 0) return;
        } break;
        case PrecMode::Tournament: {
            if (mp_tournament != 1) return;
        } break;
        case PrecMode::Always: break;
        default: {
            ConNotify("Invalid recording mode, doing nothing");
            return;
        }
    }
    basedir = prec_dir.GetString();
    tag = prec_tag.GetString();
    bluteam = g_pCVar->FindVar("mp_tournament_blueteamname")->GetString();
    redteam = g_pCVar->FindVar("mp_tournament_redteamname")->GetString();
    g_pEngineClient->GetChapterName(mapname, sizeof(mapname));
    if (strcmp(nextdemoname, "") == 0) {
        // TODO: Custom name formatting
        nextdemoname = "foobar_demo";
    }
    argv[1] = nextdemoname;
    record->Dispatch(CCommand(2, argv));
    if (g_pEngineClient->IsRecordingDemo()) {
        int notify = prec_notify.GetInt();
        ConNotify(RECORD_NOTIFICATION);
        play_sound(Sound::Recording);
        if (mp_tournament == 1) {
            switch (static_cast<PrecNotify>(notify)) {
                default:
                case PrecNotify::Console:
                    break;
                case PrecNotify::Chat: {
                    g_pEngineClient->ClientCmd("say_team [" OPENPREC_NAME "]: " RECORD_NOTIFICATION);
                } break;
                case PrecNotify::Hud: {
                    ConNotifyf(OPENPREC_NOFEATURE_MSGF, "Hud notifications");
                } break;
            }
        }
    }
}

CON_COMMAND(prec_mark, "Make a bookmark at the current tick") {
}

CON_COMMAND(prec_delete_demo, "Delete previous demo") {
}

ConCommand *g_commandList[g_numCommands] = {
    &prec_about_command,
    &prec_info_command,
    &prec_record_command,
    &prec_mark_command,
    &prec_delete_demo_command
};

bool register_cvars() {
    for (int i = 0; i < g_numCvars; ++i) {
        g_pCVar->RegisterConCommand(g_cvarList[i]);
    }
    return true;
}

void unregister_cvars() {
    for (int i = 0; i < g_numCvars; ++i) {
        g_pCVar->UnregisterConCommand(g_cvarList[i]);
    }
}

bool register_concommands() {
    for (int i = 0; i < g_numCommands; ++i) {
        g_pCVar->RegisterConCommand(g_commandList[i]);
    }
    return true;
}

void unregister_concommands() {
    for (int i = 0; i < g_numCommands; ++i) {
        g_pCVar->UnregisterConCommand(g_commandList[i]);
    }
}
<|endoftext|>


<commit_before>
#include <SFML/Graphics.hpp>
#include <SFML/Window.hpp>
#include <SFML/System.hpp>

namespace demo
{

template <typename T, int W, int H>
class frameb
{
public:
	T& xy(int x, int y) { return _data[y * W + x]; }
	T& ofs(int o) { return _data[o]; }
	const T& xy(int x, int y) const { return _data[y * W + x]; }
	const T& ofs(int o) const { return _data[o]; }
	T* data() { return _data; }
private:
	T _data[W * H];
};

template <typename T, int W, int H, T (*F)(int, int)>
class procfn
{
public:
	T xy(int x, int y) const { return F(x, y); }
	T ofs(int o) const { return F(o % W, o / W); }
};

template <typename T, int W, int H, class P, T (*F)(int, int, const P&)>
class procst
{
public:
	P _params;
	T xy(int x, int y) const { return F(x, y, _params); }
	T ofs(int o) const { return F(o % W, o / W, _params); }
};

template <typename T, int W, int H, class I>
class buffer : public I
{
public:
	typedef buffer<T, W, H, I> tBuffer;

	using I::xy;
	using I::ofs;

	template <class U>
	tBuffer& copyXY(const U& src)
	{
		for (int y = 0, o = 0; y < H; ++y)
			for (int x = 0; x < W; ++x, ++o)
				ofs(o) = src.xy(x, y);
		return *this;
	}

	template <class U>
	tBuffer& copyOfs(const U& src)
	{
		for (int o = 0; o < W * H; ++o)
			ofs(o) = src.ofs(o);
		return *this;
	}

	template <typename T0, class I0, typename F>
	tBuffer& transformXY(const buffer<T0, W, H, I0>& src0, const F& func)
	{
		for (int y = 0, o = 0; y < H; ++y)
			for (int x = 0; x < W; ++x, ++o)
				ofs(o) = func(src0.xy(x, y));
		return *this;
	}

	template <typename T0, class I0, typename F>
	tBuffer& transformOfs(const buffer<T0, W, H, I0>& src0, const F& func)
	{
		for (int o = 0; o < W * H; ++o)
			ofs(o) = func(src0.ofs(o));
		return *this;
	}

	template <typename T0, class I0, typename T1, class I1, typename F>
	tBuffer& transformXY(const buffer<T0, W, H, I0>& src0, const buffer<T1, W, H, I1>& src1, const F& func)
	{
		for (int y = 0, o = 0; y < H; ++y)
			for (int x = 0; x < W; ++x, ++o)
				ofs(o) = func(src0.xy(x, y), src1.xy(x, y));
		return *this;
	}

	template <typename T0, class I0, typename T1, class I1, typename F>
	tBuffer& transformOfs(const buffer<T0, W, H, I0>& src0, const buffer<T1, W, H, I1>& src1, const F& func)
	{
		for (int o = 0; o < W * H; ++o)
			ofs(o) = func(src0.ofs(o), src1.ofs(o));
		return *this;
	}
};

template <int W, int H>
class demowin
{
public:
	typedef buffer<sf::Uint32, W, H, demo::frameb<sf::Uint32, W, H>> tBackBuffer;

	demowin()
	{
	}

	template <class T>
	buffer<T, W, H, demo::frameb<T, W, H>>* createBuffer()
	{
		return new buffer<T, W, H, demo::frameb<T, W, H>>();
	}

	template <class F>
	void run(const F& f)
	{
		auto bgFb = new buffer<sf::Uint32, W, H, demo::frameb<sf::Uint32, W, H>>();

		sf::RenderWindow win(sf::VideoMode(W, H), "toto");
		win.setFramerateLimit(60);
		win.setVerticalSyncEnabled(true);

		sf::Texture bg;
		bg.create(W, H);
		sf::Sprite bgSp(bg);

		sf::View view = win.getDefaultView();
		sf::Vector2u winSize;

		int frame = 0;
		while (win.isOpen())
		{
			sf::Event e;
			while (win.pollEvent(e))
			{
				switch (e.type)
				{
				case sf::Event::Closed:
					win.close();
					break;
				case sf::Event::Resized:
					winSize = sf::Vector2u(e.size.width, e.size.height);
					view.reset(sf::FloatRect(0.0f, 0.0f, e.size.width, e.size.height));
					win.setView(view);
					win.setSize(winSize);
					break;
				}
			}

			if (!f(*bgFb, frame))
				win.close();
			bg.update((sf::Uint8*)&bgFb->ofs(0));

			float s0 = winSize.x / float(W);
			float s1 = winSize.y / float(H);
			float s = std::min(s0, s1);
			float x = 0.5f * (winSize.x - s * W);
			float y = 0.5f * (winSize.y - s * H);
			bgSp.setScale(s, s);
			bgSp.setPosition(x, y);

			win.clear(sf::Color::Black);
			win.draw(bgSp);
			win.display();

			++frame;
		}
	}
};

inline sf::Uint32 rgba(sf::Uint8 a, sf::Uint8 r, sf::Uint8 g, sf::Uint8 b)
{
	return (a << 24) | (r << 16) | (g << 8) | b;
}

inline sf::Uint32 modulate(sf::Uint32 a, sf::Uint32 b)
{
	typedef sf::Uint8 tColor[4];
	tColor& c = *reinterpret_cast<tColor*>(&a);
	tColor& d = *reinterpret_cast<tColor*>(&b);
	tColor r = {
		sf::Uint8(c[0] * d[0] / 255u),
		sf::Uint8(c[1] * d[1] / 255u),
		sf::Uint8(c[2] * d[2] / 255u),
		sf::Uint8(c[3] * d[3] / 255u)
	};
	return *reinterpret_cast<sf::Uint32*>(&r);
}

template <class T, size_t N, class F>
std::array<T, N> makePal(const F& f)
{
	std::array<T, N> r;
	for (size_t i = 0; i < N; ++i)
		r[i] = f(i);
	return r;
}

}

const int ScrWidth = 320;
const int ScrHeight = 200;

// ---------------------------------------------------------------------------------------
// circles
// ---------------------------------------------------------------------------------------

struct ccparams
{
	int x0, y0, x1, y1;
};

inline int dist(int x0, int y0, int x1, int y1)
{
	int dx = x1 - x0;
	int dy = y1 - y0;
	return (dx * dx + dy * dy) / 512;
}

static inline sf::Uint8 computeCC(int x, int y, const ccparams& p)
{
	int u = dist(p.x0, p.y0, x, y);
	int v = dist(p.x1, p.y1, x, y);
	return u ^ v;
}

// ---------------------------------------------------------------------------------------
// plasma
// ---------------------------------------------------------------------------------------

inline float term(float x)
{
	return (1.0f + x) / 2.0f;
}

static inline sf::Uint8 computePlasma(int x, int y, const int& frame)
{
	const float coeff = 3.0f;
	const float cm = 1.1234f;
	const float ff = 2.0f * frame;

	const float fy = y;
	const float v0 = term(sin((ff - fy) / (coeff * 10.0f)));
	const float v1 = term(sin((2.0f * ff + cm * fy) / (coeff * 15.0f)));

	const float fx = x;
	const float fz = x + y;
	const float fw = x - y;
	const float v2 = term(sin((ff + fx) / (coeff * 10.0f)));
	const float v3 = term(sin((2.0f * ff - cm * fx) / (coeff * 15.0f)));
	const float v5 = term(sin((ff + fw) / (coeff * 10.0f)));
	const float v4 = term(sin((2.0f * ff + cm * fz) / (coeff * 15.0f)));
	const float vf = (v0 + v1 + v2 + v3 + v4 + v5) / 6.f;
	const int idx = 255.0f * vf;

	return idx;
}

// ---------------------------------------------------------------------------------------
// fire
// ---------------------------------------------------------------------------------------

// ---------------------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------------------

sf::Uint8 lerp8(int v, int a, int b)
{
	return sf::Uint8((255u * (v - a)) / (b - a));
}

int main(int, char**)
{
	typedef demo::demowin<ScrWidth, ScrHeight> tWin320x200;
	tWin320x200 win;

	auto fb8 = win.createBuffer<sf::Uint8>();
	auto cc = new demo::buffer<sf::Uint8, ScrWidth, ScrHeight, demo::procst<sf::Uint8, ScrWidth, ScrHeight, ccparams, computeCC>>();
	auto plasma = new demo::buffer<sf::Uint8, ScrWidth, ScrHeight, demo::procst<sf::Uint8, ScrWidth, ScrHeight, int, computePlasma>>();

	auto pal0 = demo::makePal<sf::Uint32, 256>([] (size_t i) { return demo::rgba(255, i, 0, i); });
	auto pal1 = demo::makePal<sf::Uint32, 256>([] (size_t i) { return demo::rgba(255, 255 - i, 0, i); });
	auto pal2 = demo::makePal<sf::Uint32, 256>([] (size_t i) {
		if (i < 64)  return demo::rgba(255, 0            , 0           , 4 * i);
		if (i < 128) return demo::rgba(255, 0            , 4 * (i - 64), 255);
		if (i < 192) return demo::rgba(255, 4 * (i - 128), 255         , 255);
		return              demo::rgba(255, 255          , 255         , 255);
	});

	win.run([&] (tWin320x200::tBackBuffer& bgFb, int frame) {
		const int fxDuration = 1000;
		if (frame < 1 * fxDuration)
		{
			sf::Uint8* d = fb8->data();
			if (frame % 8 == 0)
			{
				for (int j = 0; j < ScrWidth;)
				{
					sf::Uint8 r = 192 + 16 * (rand() % 4);
					for (int i = 0; i < 10; ++i, ++j)
					{
						d[(ScrHeight - 1) * ScrWidth + j] = r;
					}
				}
			}
			for (int j = 0; j < 2; ++j)
			{
				for (int i = 0; i < (ScrHeight - 1) * ScrWidth; ++i)
				{
					d[i] = (2 * d[i] + 1 * d[i + ScrWidth - 1] + 4 * d[i + ScrWidth] + 1 * d[i + ScrWidth + 1]) / 8;
					if (d[i] > 0)
						--d[i];
				}
			}
			bgFb.transformOfs(*fb8, [&pal2] (sf::Uint8 l) { return pal2[l]; });
		}
		else if (frame < 2 * fxDuration)
		{
			plasma->_params = frame;
			if (frame < 300)
				bgFb.transformOfs(fb8->copyXY(*plasma), [&pal1, frame] (sf::Uint8 l) { return pal1[l * lerp8(frame, 0 ,300) / 255]; });
			else
				bgFb.transformOfs(fb8->copyXY(*plasma), [&pal1] (sf::Uint8 l) { return pal1[l]; });
		}
		else if (frame < 3 * fxDuration)
		{
			cc->_params = {
				int(160 + 150 * sinf(0.03f * frame)), 100,
				160, int(100 + 90 * sinf(0.04f * frame))
			};
			bgFb.transformOfs(fb8->copyXY(*cc), [&pal0] (sf::Uint8 l) { return pal0[l]; });
		}
		else
		{
			return false;
		}
		return true;
	});

	return 0;
}

<commit_msg>improve fire FX<commit_after>
#include <SFML/Graphics.hpp>
#include <SFML/Window.hpp>
#include <SFML/System.hpp>

namespace demo
{

template <typename T, int W, int H>
class frameb
{
public:
	T& xy(int x, int y) { return _data[y * W + x]; }
	T& ofs(int o) { return _data[o]; }
	const T& xy(int x, int y) const { return _data[y * W + x]; }
	const T& ofs(int o) const { return _data[o]; }
	T* data() { return _data; }
private:
	T _data[W * (H + 1)];
};

template <typename T, int W, int H, T (*F)(int, int)>
class procfn
{
public:
	T xy(int x, int y) const { return F(x, y); }
	T ofs(int o) const { return F(o % W, o / W); }
};

template <typename T, int W, int H, class P, T (*F)(int, int, const P&)>
class procst
{
public:
	P _params;
	T xy(int x, int y) const { return F(x, y, _params); }
	T ofs(int o) const { return F(o % W, o / W, _params); }
};

template <typename T, int W, int H, class I>
class buffer : public I
{
public:
	typedef buffer<T, W, H, I> tBuffer;

	using I::xy;
	using I::ofs;

	template <class U>
	tBuffer& copyXY(const U& src)
	{
		for (int y = 0, o = 0; y < H; ++y)
			for (int x = 0; x < W; ++x, ++o)
				ofs(o) = src.xy(x, y);
		return *this;
	}

	template <class U>
	tBuffer& copyOfs(const U& src)
	{
		for (int o = 0; o < W * H; ++o)
			ofs(o) = src.ofs(o);
		return *this;
	}

	template <typename T0, class I0, typename F>
	tBuffer& transformXY(const buffer<T0, W, H, I0>& src0, const F& func)
	{
		for (int y = 0, o = 0; y < H; ++y)
			for (int x = 0; x < W; ++x, ++o)
				ofs(o) = func(src0.xy(x, y));
		return *this;
	}

	template <typename T0, class I0, typename F>
	tBuffer& transformOfs(const buffer<T0, W, H, I0>& src0, const F& func)
	{
		for (int o = 0; o < W * H; ++o)
			ofs(o) = func(src0.ofs(o));
		return *this;
	}

	template <typename T0, class I0, typename T1, class I1, typename F>
	tBuffer& transformXY(const buffer<T0, W, H, I0>& src0, const buffer<T1, W, H, I1>& src1, const F& func)
	{
		for (int y = 0, o = 0; y < H; ++y)
			for (int x = 0; x < W; ++x, ++o)
				ofs(o) = func(src0.xy(x, y), src1.xy(x, y));
		return *this;
	}

	template <typename T0, class I0, typename T1, class I1, typename F>
	tBuffer& transformOfs(const buffer<T0, W, H, I0>& src0, const buffer<T1, W, H, I1>& src1, const F& func)
	{
		for (int o = 0; o < W * H; ++o)
			ofs(o) = func(src0.ofs(o), src1.ofs(o));
		return *this;
	}
};

template <int W, int H>
class demowin
{
public:
	typedef buffer<sf::Uint32, W, H, demo::frameb<sf::Uint32, W, H>> tBackBuffer;

	demowin()
	{
	}

	template <class T>
	buffer<T, W, H, demo::frameb<T, W, H>>* createBuffer()
	{
		return new buffer<T, W, H, demo::frameb<T, W, H>>();
	}

	template <class F>
	void run(const F& f)
	{
		auto bgFb = new buffer<sf::Uint32, W, H, demo::frameb<sf::Uint32, W, H>>();

		sf::RenderWindow win(sf::VideoMode(W, H), "toto");
		win.setFramerateLimit(60);
		win.setVerticalSyncEnabled(true);

		sf::Texture bg;
		bg.create(W, H);
		sf::Sprite bgSp(bg);

		sf::View view = win.getDefaultView();
		sf::Vector2u winSize;

		int frame = 0;
		while (win.isOpen())
		{
			sf::Event e;
			while (win.pollEvent(e))
			{
				switch (e.type)
				{
				case sf::Event::Closed:
					win.close();
					break;
				case sf::Event::Resized:
					winSize = sf::Vector2u(e.size.width, e.size.height);
					view.reset(sf::FloatRect(0.0f, 0.0f, e.size.width, e.size.height));
					win.setView(view);
					win.setSize(winSize);
					break;
				}
			}

			if (!f(*bgFb, frame))
				win.close();
			bg.update((sf::Uint8*)&bgFb->ofs(0));

			float s0 = winSize.x / float(W);
			float s1 = winSize.y / float(H);
			float s = std::min(s0, s1);
			float x = 0.5f * (winSize.x - s * W);
			float y = 0.5f * (winSize.y - s * H);
			bgSp.setScale(s, s);
			bgSp.setPosition(x, y);

			win.clear(sf::Color::Black);
			win.draw(bgSp);
			win.display();

			++frame;
		}
	}
};

inline sf::Uint32 rgba(sf::Uint8 a, sf::Uint8 r, sf::Uint8 g, sf::Uint8 b)
{
	return (a << 24) | (r << 16) | (g << 8) | b;
}

inline sf::Uint32 modulate(sf::Uint32 a, sf::Uint32 b)
{
	typedef sf::Uint8 tColor[4];
	tColor& c = *reinterpret_cast<tColor*>(&a);
	tColor& d = *reinterpret_cast<tColor*>(&b);
	tColor r = {
		sf::Uint8(c[0] * d[0] / 255u),
		sf::Uint8(c[1] * d[1] / 255u),
		sf::Uint8(c[2] * d[2] / 255u),
		sf::Uint8(c[3] * d[3] / 255u)
	};
	return *reinterpret_cast<sf::Uint32*>(&r);
}

template <class T, size_t N, class F>
std::array<T, N> makePal(const F& f)
{
	std::array<T, N> r;
	for (size_t i = 0; i < N; ++i)
		r[i] = f(i);
	return r;
}

}

const int ScrWidth = 320;
const int ScrHeight = 200;

// ---------------------------------------------------------------------------------------
// circles
// ---------------------------------------------------------------------------------------

struct ccparams
{
	int x0, y0, x1, y1;
};

inline int dist(int x0, int y0, int x1, int y1)
{
	int dx = x1 - x0;
	int dy = y1 - y0;
	return (dx * dx + dy * dy) / 512;
}

static inline sf::Uint8 computeCC(int x, int y, const ccparams& p)
{
	int u = dist(p.x0, p.y0, x, y);
	int v = dist(p.x1, p.y1, x, y);
	return u ^ v;
}

// ---------------------------------------------------------------------------------------
// plasma
// ---------------------------------------------------------------------------------------

inline float term(float x)
{
	return (1.0f + x) / 2.0f;
}

static inline sf::Uint8 computePlasma(int x, int y, const int& frame)
{
	const float coeff = 3.0f;
	const float cm = 1.1234f;
	const float ff = 2.0f * frame;

	const float fy = y;
	const float v0 = term(sin((ff - fy) / (coeff * 10.0f)));
	const float v1 = term(sin((2.0f * ff + cm * fy) / (coeff * 15.0f)));

	const float fx = x;
	const float fz = x + y;
	const float fw = x - y;
	const float v2 = term(sin((ff + fx) / (coeff * 10.0f)));
	const float v3 = term(sin((2.0f * ff - cm * fx) / (coeff * 15.0f)));
	const float v5 = term(sin((ff + fw) / (coeff * 10.0f)));
	const float v4 = term(sin((2.0f * ff + cm * fz) / (coeff * 15.0f)));
	const float vf = (v0 + v1 + v2 + v3 + v4 + v5) / 6.f;
	const int idx = 255.0f * vf;

	return idx;
}

// ---------------------------------------------------------------------------------------
// fire
// ---------------------------------------------------------------------------------------
inline void setFire(sf::Uint16* d, int frame)
{
	if (frame % 4 == 0)
	{
		for (int j = 0; j < ScrWidth;)
		{
			sf::Uint8 r = 192 + 63 * (rand() % 2);
			for (int i = 0; i < 10; ++i, ++j)
			{
				d[(ScrHeight - 1) * ScrWidth + j] = (r << 8) + rand() % 8;
			}
		}
	}
	for (int j = 0; j < 2; ++j)
	{
		for (int i = 0; i < (ScrHeight - 1) * ScrWidth; ++i)
		{
			d[i] = (2 * d[i] + 1 * d[i + ScrWidth - 1] + 3 * d[i + ScrWidth] + 2 * d[i + ScrWidth + 1]) / 8;
			if (d[i] > 255 /*&& d[i] < (200 << 8)*/)
				d[i] -= 256;
		}
	}
}

// ---------------------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------------------

sf::Uint8 lerp8(int v, int a, int b)
{
	return sf::Uint8((255u * (v - a)) / (b - a));
}

int main(int, char**)
{
	typedef demo::demowin<ScrWidth, ScrHeight> tWin320x200;
	tWin320x200 win;

	auto fb8 = win.createBuffer<sf::Uint8>();
	auto fb16 = win.createBuffer<sf::Uint16>();
	auto cc = new demo::buffer<sf::Uint8, ScrWidth, ScrHeight, demo::procst<sf::Uint8, ScrWidth, ScrHeight, ccparams, computeCC>>();
	auto plasma = new demo::buffer<sf::Uint8, ScrWidth, ScrHeight, demo::procst<sf::Uint8, ScrWidth, ScrHeight, int, computePlasma>>();

	auto pal0 = demo::makePal<sf::Uint32, 256>([] (size_t i) { return demo::rgba(255, i, 0, i); });
	auto pal1 = demo::makePal<sf::Uint32, 256>([] (size_t i) { return demo::rgba(255, 255 - i, 0, i); });
	auto pal2 = demo::makePal<sf::Uint32, 256>([] (size_t i) {
		if (i < 64 ) return demo::rgba(255, 0            , 0           , 4 * i);
		if (i < 128) return demo::rgba(255, 0            , 4 * (i - 64), 255  );
		if (i < 192) return demo::rgba(255, 4 * (i - 128), 255         , 255  );
		return              demo::rgba(255, 255          , 255         , 255  );
	});

	win.run([&] (tWin320x200::tBackBuffer& bgFb, int frame) {
		const int fxDuration = 1000;
		if (frame < 1 * fxDuration)
		{
			// use 16bpp buffer to increase quality
			setFire(fb16->data(), frame);
			bgFb.transformOfs(*fb16, [&pal2] (sf::Uint16 l) { return pal2[l >> 8]; });
		}
		else if (frame < 2 * fxDuration)
		{
			plasma->_params = frame;
			if (frame < 300)
				bgFb.transformOfs(fb8->copyXY(*plasma), [&pal1, frame] (sf::Uint8 l) { return pal1[l * lerp8(frame, 0 ,300) / 255]; });
			else
				bgFb.transformOfs(fb8->copyXY(*plasma), [&pal1] (sf::Uint8 l) { return pal1[l]; });
		}
		else if (frame < 3 * fxDuration)
		{
			cc->_params = {
				int(160 + 150 * sinf(0.03f * frame)), 100,
				160, int(100 + 90 * sinf(0.04f * frame))
			};
			bgFb.transformOfs(fb8->copyXY(*cc), [&pal0] (sf::Uint8 l) { return pal0[l]; });
		}
		else
		{
			return false;
		}
		return true;
	});

	return 0;
}

<|endoftext|>


<commit_before>//
//  main.cpp
//  loadcaffe
//
//  Created by Sergey Zagoruyko on 28/11/14.
//  Copyright (c) 2014 Sergey Zagoruyko. All rights reserved.
//

#include <iostream>
#include <fstream>
#include <fcntl.h>
#include <unistd.h>

#include <TH/TH.h>

#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <google/protobuf/text_format.h>

#include "build/caffe.pb.h"

using google::protobuf::io::FileInputStream;
using google::protobuf::io::FileOutputStream;
using google::protobuf::io::ZeroCopyInputStream;
using google::protobuf::io::CodedInputStream;
using google::protobuf::Message;

extern "C" {
void loadBinary(void** handle, const char* prototxt_name, const char* binary_name);
void convertProtoToLua(void** handle, const char* lua_name, const char* cuda_package);
void loadModule(const void** handle, const char* name, THFloatTensor* weight, THFloatTensor* bias);
void destroyBinary(void** handle);
}

bool ReadProtoFromTextFile(const char* filename, Message* proto) {
    int fd = open(filename, O_RDONLY);
    if(fd < 0)
      return false;

    FileInputStream* input = new FileInputStream(fd);
    bool success = google::protobuf::TextFormat::Parse(input, proto);
    delete input;
    close(fd);
    return success;
}

bool ReadProtoFromBinaryFile(const char* filename, Message* proto) {
    int fd = open(filename, O_RDONLY);
    if(fd < 0)
      return false;

    ZeroCopyInputStream* raw_input = new FileInputStream(fd);
    CodedInputStream* coded_input = new CodedInputStream(raw_input);
    coded_input->SetTotalBytesLimit(1073741824, 536870912);

    bool success = proto->ParseFromCodedStream(coded_input);

    delete coded_input;
    delete raw_input;
    close(fd);
    return success;
}

void convertProtoToLua(void** handle, const char* lua_name, const char* cuda_package)
{
  const caffe::NetParameter netparam = *(const caffe::NetParameter*)handle[1];

  std::ofstream ofs (lua_name);

  ofs << "require '" << cuda_package << "'\n";
  ofs << "require 'cunn'\n";
  ofs << "model = {}\n";
  if(std::string(cuda_package)=="ccn2")
    ofs<< "table.insert(model, {'torch_transpose_dwhb', nn.Transpose({1,4},{1,3},{1,2})})\n";
  else if(std::string(cuda_package)=="nn")
    ofs<< "require 'inn'\n";

  int num_output = netparam.input_dim_size();
  for (int i=0; i<netparam.layers_size(); ++i)
  {
    std::vector<std::pair<std::string, std::string>> lines;
    auto& layer = netparam.layers(i);
    switch(layer.type())
    {
      case caffe::LayerParameter::CONVOLUTION:
      {
	auto &param = layer.convolution_param();
	int groups = param.group() == 0 ? 1 : param.group();
	int nInputPlane = layer.blobs(0).channels()*groups;
        int nOutputPlane = layer.blobs(0).num();
	//int nOutputPlane = param.num_output();
	num_output = nOutputPlane;
	int kW = param.kernel_w();
	int kH = param.kernel_h();
	int dW = param.stride_w();
	int dH = param.stride_h();
	if(kW==0 || kH==0)
	{
	  kW = param.kernel_size();
	  kH = kW;
	}
	if(dW==0 || dH==0)
	{
	  dW = param.stride();
	  dH = dW;
	}
	int pad_w = param.pad_w();
	int pad_h = param.pad_h();
        if(pad_w==0 || pad_h==0)
        {
          pad_w = param.pad();
          pad_h = pad_w;
        }
	if(std::string(cuda_package) == "ccn2")
	{
          if(kW != kH || dW != dH || pad_w != pad_h)
          {
            std::cout << "ccn2 only supports square images!\n";
            break;
          }
	  char buf[1024];
	  sprintf(buf, "ccn2.SpatialConvolution(%d, %d, %d, %d, %d, %d)",
	      nInputPlane, nOutputPlane, kW, dW, pad_w, groups);
	  lines.emplace_back(layer.name(), buf);
	}
	else
	{
	  char buf[1024];
	  const char* mm_or_not = std::string(cuda_package)=="nn" ? "MM" : "";
	  sprintf(buf, "%s.SpatialConvolution%s(%d, %d, %d, %d, %d, %d, %d, %d)",
	      cuda_package, mm_or_not, nInputPlane, nOutputPlane, kW, kH, dW, dH, pad_w, pad_h);
	  lines.emplace_back(layer.name(), buf);
	}
	break;
      }
      case caffe::LayerParameter::POOLING:
      {
	auto &param = layer.pooling_param();
	std::string ptype = param.pool() == caffe::PoolingParameter::MAX ? "Max" : "Avg";
	int kW = param.kernel_w();
	int kH = param.kernel_h();
	int dW = param.stride_w();
	int dH = param.stride_h();
	if(kW==0 || kH==0)
	{
	  kW = param.kernel_size();
	  kH = kW;
	}
	if(dW==0 || dH==0)
	{
	  dW = param.stride();
	  dH = dW;
	}

	if(std::string(cuda_package) == "ccn2")
	{
	  char buf[1024];
	  sprintf(buf, "ccn2.Spatial%sPooling(%d, %d)", ptype.c_str(), kW, dW);
	  lines.emplace_back(layer.name(), buf);
	}
	else if(std::string(cuda_package) == "cudnn")
	{
	  char buf[1024];
	  sprintf(buf, "%s.Spatial%sPooling(%d, %d, %d, %d):ceil()", cuda_package, ptype=="Avg" ? "Average" : "Max", kW, kH, dW, dH);
	  lines.emplace_back(layer.name(), buf);
	}
        else if(std::string(cuda_package) == "nn")
        {
          char buf[1024];
	  sprintf(buf, "inn.Spatial%sPooling(%d, %d, %d, %d)", ptype=="Avg" ? "Average" : "Max", kW, kH, dW, dH);
	  lines.emplace_back(layer.name(), buf);
        }
	break;
      }
      case caffe::LayerParameter::RELU:
      {
	lines.emplace_back(layer.name(), "nn.ReLU()");
	break;
      }
      case caffe::LayerParameter::LRN:
      {
	if(std::string(cuda_package) == "ccn2")
	{
	  auto &param = layer.lrn_param();
	  int local_size = param.local_size();
	  float alpha = param.alpha();
	  float beta = param.beta();
	  char buf[1024];
	  sprintf(buf, "ccn2.SpatialCrossResponseNormalization(%d, %.6f, %.4f)", local_size, alpha, beta);
	  lines.emplace_back(layer.name(), buf);
	}
	break;
      }
      case caffe::LayerParameter::INNER_PRODUCT:
      {
	auto &param = layer.inner_product_param();
	int nInputPlane = layer.blobs(0).width();
	int nOutputPlane = param.num_output();
	char buf[1024];
	sprintf(buf, "nn.Linear(%d, %d)", nInputPlane, nOutputPlane);
	if(num_output != nInputPlane)
	{
	  if(std::string(cuda_package) == "ccn2")
	    lines.emplace_back("torch_transpose_bdwh", "nn.Transpose({4,1},{4,2},{4,3})");
	  lines.emplace_back("torch_view", "nn.View(-1):setNumInputDims(3)");
	}
	lines.emplace_back(layer.name(), buf);
	num_output = nOutputPlane;
	break;
      }
      case caffe::LayerParameter::DROPOUT:
      {
	char buf[1024];
	sprintf(buf, "nn.Dropout(%f)", netparam.layers(i).dropout_param().dropout_ratio());
	lines.emplace_back(layer.name(), buf);
	break;
      }
      case caffe::LayerParameter::SOFTMAX_LOSS:
      {
	lines.emplace_back(layer.name(), "nn.SoftMax()");
	break;
      }
      case caffe::LayerParameter::SOFTMAX:
      {
	lines.emplace_back(layer.name(), "nn.SoftMax()");
	break;
      }
      default:
      {
	std::cout << "MODULE " << netparam.layers(i).name() << " UNDEFINED\n";
	break;
      }
    }

    if(!lines.empty())
      for(auto& it: lines)
	ofs << "table.insert(model, {'" << it.first << "', " << it.second << "})\n";
    else
    {
      ofs << "-- module '" << layer.name() << "' not found\n";
      std::cout << "module '" << layer.name() << "' not found\n";
    }
  }
}

void loadBinary(void** handle, const char* prototxt_name, const char* binary_name)
{
  caffe::NetParameter* netparam = (caffe::NetParameter*)handle;
  netparam = new caffe::NetParameter();
  ReadProtoFromTextFile(prototxt_name, netparam);
  bool success = ReadProtoFromBinaryFile(binary_name, netparam);
  if(success)
  {
    std::cout << "Successfully loaded " << binary_name << std::endl;
    handle[1] = netparam;
  }
  else
    std::cout << "Couldn't load " << binary_name << std::endl;
}

void destroyBinary(void** handle)
{
  const caffe::NetParameter* netparam2 = (const caffe::NetParameter*)handle[1];
  delete netparam2;
}

void loadModule(const void** handle, const char* name, THFloatTensor* weight, THFloatTensor* bias)
{
  if(handle == NULL)
  {
    std::cout << "network not loaded!\n";
    return;
  }

  const caffe::NetParameter* netparam = (const caffe::NetParameter*)handle[1];

  int n = netparam->layers_size();
  for(int i=0; i<n; ++i)
  {
    auto &layer = netparam->layers(i);
    if(std::string(name) == layer.name())
    {
      int nInputPlane = layer.blobs(0).channels();
      int nOutputPlane = layer.blobs(0).num();
      int kW = layer.blobs(0).width();
      int kH = layer.blobs(0).height();
      printf("%s: %d %d %d %d\n", name, nOutputPlane, nInputPlane, kW, kH);

      THFloatTensor_resize4d(weight, nOutputPlane, nInputPlane, kW, kH);
      memcpy(THFloatTensor_data(weight), layer.blobs(0).data().data(), sizeof(float)*nOutputPlane*nInputPlane*kW*kH);

      THFloatTensor_resize1d(bias, layer.blobs(1).data_size());
      memcpy(THFloatTensor_data(bias), layer.blobs(1).data().data(), sizeof(float)*layer.blobs(1).data_size());
    }
  }
}
<commit_msg>typo fix<commit_after>//
//  main.cpp
//  loadcaffe
//
//  Created by Sergey Zagoruyko on 28/11/14.
//  Copyright (c) 2014 Sergey Zagoruyko. All rights reserved.
//

#include <iostream>
#include <fstream>
#include <fcntl.h>
#include <unistd.h>

#include <TH/TH.h>

#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include <google/protobuf/text_format.h>

#include "build/caffe.pb.h"

using google::protobuf::io::FileInputStream;
using google::protobuf::io::FileOutputStream;
using google::protobuf::io::ZeroCopyInputStream;
using google::protobuf::io::CodedInputStream;
using google::protobuf::Message;

extern "C" {
void loadBinary(void** handle, const char* prototxt_name, const char* binary_name);
void convertProtoToLua(void** handle, const char* lua_name, const char* cuda_package);
void loadModule(const void** handle, const char* name, THFloatTensor* weight, THFloatTensor* bias);
void destroyBinary(void** handle);
}

bool ReadProtoFromTextFile(const char* filename, Message* proto) {
    int fd = open(filename, O_RDONLY);
    if(fd < 0)
      return false;

    FileInputStream* input = new FileInputStream(fd);
    bool success = google::protobuf::TextFormat::Parse(input, proto);
    delete input;
    close(fd);
    return success;
}

bool ReadProtoFromBinaryFile(const char* filename, Message* proto) {
    int fd = open(filename, O_RDONLY);
    if(fd < 0)
      return false;

    ZeroCopyInputStream* raw_input = new FileInputStream(fd);
    CodedInputStream* coded_input = new CodedInputStream(raw_input);
    coded_input->SetTotalBytesLimit(1073741824, 536870912);

    bool success = proto->ParseFromCodedStream(coded_input);

    delete coded_input;
    delete raw_input;
    close(fd);
    return success;
}

void convertProtoToLua(void** handle, const char* lua_name, const char* cuda_package)
{
  const caffe::NetParameter netparam = *(const caffe::NetParameter*)handle[1];

  std::ofstream ofs (lua_name);

  ofs << "require '" << cuda_package << "'\n";
  ofs << "require 'cunn'\n";
  ofs << "model = {}\n";
  if(std::string(cuda_package)=="ccn2")
    ofs<< "table.insert(model, {'torch_transpose_dwhb', nn.Transpose({1,4},{1,3},{1,2})})\n";
  else if(std::string(cuda_package)=="nn")
    ofs<< "require 'inn'\n";

  int num_output = netparam.input_dim_size();
  for (int i=0; i<netparam.layers_size(); ++i)
  {
    std::vector<std::pair<std::string, std::string>> lines;
    auto& layer = netparam.layers(i);
    switch(layer.type())
    {
      case caffe::LayerParameter::CONVOLUTION:
      {
	auto &param = layer.convolution_param();
	int groups = param.group() == 0 ? 1 : param.group();
	int nInputPlane = layer.blobs(0).channels()*groups;
        int nOutputPlane = layer.blobs(0).num();
	//int nOutputPlane = param.num_output();
	num_output = nOutputPlane;
	int kW = param.kernel_w();
	int kH = param.kernel_h();
	int dW = param.stride_w();
	int dH = param.stride_h();
	if(kW==0 || kH==0)
	{
	  kW = param.kernel_size();
	  kH = kW;
	}
	if(dW==0 || dH==0)
	{
	  dW = param.stride();
	  dH = dW;
	}
	int pad_w = param.pad_w();
	int pad_h = param.pad_h();
        if(pad_w==0 || pad_h==0)
        {
          pad_w = param.pad();
          pad_h = pad_w;
        }
	if(std::string(cuda_package) == "ccn2")
	{
          if(kW != kH || dW != dH || pad_w != pad_h)
          {
            std::cout << "ccn2 only supports square images!\n";
            break;
          }
	  char buf[1024];
	  sprintf(buf, "ccn2.SpatialConvolution(%d, %d, %d, %d, %d, %d)",
	      nInputPlane, nOutputPlane, kW, dW, pad_w, groups);
	  lines.emplace_back(layer.name(), buf);
	}
	else
	{
	  char buf[1024];
	  const char* mm_or_not = std::string(cuda_package)=="nn" ? "MM" : "";
	  sprintf(buf, "%s.SpatialConvolution%s(%d, %d, %d, %d, %d, %d, %d, %d)",
	      cuda_package, mm_or_not, nInputPlane, nOutputPlane, kW, kH, dW, dH, pad_w, pad_h);
	  lines.emplace_back(layer.name(), buf);
	}
	break;
      }
      case caffe::LayerParameter::POOLING:
      {
	auto &param = layer.pooling_param();
	std::string ptype = param.pool() == caffe::PoolingParameter::MAX ? "Max" : "Avg";
	int kW = param.kernel_w();
	int kH = param.kernel_h();
	int dW = param.stride_w();
	int dH = param.stride_h();
	if(kW==0 || kH==0)
	{
	  kW = param.kernel_size();
	  kH = kW;
	}
	if(dW==0 || dH==0)
	{
	  dW = param.stride();
	  dH = dW;
	}

	if(std::string(cuda_package) == "ccn2")
	{
	  char buf[1024];
	  sprintf(buf, "ccn2.Spatial%sPooling(%d, %d)", ptype.c_str(), kW, dW);
	  lines.emplace_back(layer.name(), buf);
	}
	else if(std::string(cuda_package) == "cudnn")
	{
	  char buf[1024];
	  sprintf(buf, "%s.Spatial%sPooling(%d, %d, %d, %d):ceil()", cuda_package, ptype=="Avg" ? "Average" : "Max", kW, kH, dW, dH);
	  lines.emplace_back(layer.name(), buf);
	}
        else if(std::string(cuda_package) == "nn")
        {
          char buf[1024];
	  sprintf(buf, "inn.Spatial%sPooling(%d, %d, %d, %d)", ptype=="Avg" ? "Average" : "Max", kW, kH, dW, dH);
	  lines.emplace_back(layer.name(), buf);
        }
	break;
      }
      case caffe::LayerParameter::RELU:
      {
	lines.emplace_back(layer.name(), "nn.ReLU()");
	break;
      }
      case caffe::LayerParameter::LRN:
      {
	if(std::string(cuda_package) == "ccn2")
	{
	  auto &param = layer.lrn_param();
	  int local_size = param.local_size();
	  float alpha = param.alpha();
	  float beta = param.beta();
	  char buf[1024];
	  sprintf(buf, "ccn2.SpatialCrossResponseNormalization(%d, %.6f, %.4f)", local_size, alpha, beta);
	  lines.emplace_back(layer.name(), buf);
	}
	break;
      }
      case caffe::LayerParameter::INNER_PRODUCT:
      {
	auto &param = layer.inner_product_param();
	int nInputPlane = layer.blobs(0).width();
	int nOutputPlane = param.num_output();
	char buf[1024];
	sprintf(buf, "nn.Linear(%d, %d)", nInputPlane, nOutputPlane);
	if(num_output != nInputPlane)
	{
	  if(std::string(cuda_package) == "ccn2")
	    lines.emplace_back("torch_transpose_bdwh", "nn.Transpose({4,1},{4,2},{4,3})");
	  lines.emplace_back("torch_view", "nn.View(-1):setNumInputDims(3)");
	}
	lines.emplace_back(layer.name(), buf);
	num_output = nOutputPlane;
	break;
      }
      case caffe::LayerParameter::DROPOUT:
      {
	char buf[1024];
	sprintf(buf, "nn.Dropout(%f)", netparam.layers(i).dropout_param().dropout_ratio());
	lines.emplace_back(layer.name(), buf);
	break;
      }
      case caffe::LayerParameter::SOFTMAX_LOSS:
      {
	lines.emplace_back(layer.name(), "nn.SoftMax()");
	break;
      }
      case caffe::LayerParameter::SOFTMAX:
      {
	lines.emplace_back(layer.name(), "nn.SoftMax()");
	break;
      }
      default:
      {
	std::cout << "MODULE " << netparam.layers(i).name() << " UNDEFINED\n";
	break;
      }
    }

    if(!lines.empty())
      for(auto& it: lines)
	ofs << "table.insert(model, {'" << it.first << "', " << it.second << "})\n";
    else
    {
      ofs << "-- module '" << layer.name() << "' not found\n";
      std::cout << "module '" << layer.name() << "' not found\n";
    }
  }
}

void loadBinary(void** handle, const char* prototxt_name, const char* binary_name)
{
  caffe::NetParameter* netparam = new caffe::NetParameter();
  ReadProtoFromTextFile(prototxt_name, netparam);
  bool success = ReadProtoFromBinaryFile(binary_name, netparam);
  if(success)
  {
    std::cout << "Successfully loaded " << binary_name << std::endl;
    handle[1] = netparam;
  }
  else
    std::cout << "Couldn't load " << binary_name << std::endl;
}

void destroyBinary(void** handle)
{
  const caffe::NetParameter* netparam2 = (const caffe::NetParameter*)handle[1];
  delete netparam2;
}

void loadModule(const void** handle, const char* name, THFloatTensor* weight, THFloatTensor* bias)
{
  if(handle == NULL)
  {
    std::cout << "network not loaded!\n";
    return;
  }

  const caffe::NetParameter* netparam = (const caffe::NetParameter*)handle[1];

  int n = netparam->layers_size();
  for(int i=0; i<n; ++i)
  {
    auto &layer = netparam->layers(i);
    if(std::string(name) == layer.name())
    {
      int nInputPlane = layer.blobs(0).channels();
      int nOutputPlane = layer.blobs(0).num();
      int kW = layer.blobs(0).width();
      int kH = layer.blobs(0).height();
      printf("%s: %d %d %d %d\n", name, nOutputPlane, nInputPlane, kW, kH);

      THFloatTensor_resize4d(weight, nOutputPlane, nInputPlane, kW, kH);
      memcpy(THFloatTensor_data(weight), layer.blobs(0).data().data(), sizeof(float)*nOutputPlane*nInputPlane*kW*kH);

      THFloatTensor_resize1d(bias, layer.blobs(1).data_size());
      memcpy(THFloatTensor_data(bias), layer.blobs(1).data().data(), sizeof(float)*layer.blobs(1).data_size());
    }
  }
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 PaddlePaddle 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 "paddle/fluid/framework/io/shell.h"
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/platform/timer.h"

namespace paddle {
namespace framework {

std::shared_ptr<FILE> shell_fopen(const std::string& path,
                                  const std::string& mode) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return nullptr;
#else
  if (shell_verbose()) {
    LOG(INFO) << "Opening file[" << path << "] with mode[" << mode << "]";
  }
  FILE* fp;
  if (!(fp = fopen(path.c_str(), mode.c_str()))) {
    PADDLE_THROW(platform::errors::Unavailable(
        "Failed to open file, path[%s], mode[%s].", path, mode));
  }
  return {fp, [path](FILE* fp) {
            if (shell_verbose()) {
              LOG(INFO) << "Closing file[" << path << "]";
            }
            if (0 != fclose(fp)) {
              PADDLE_THROW(platform::errors::Unavailable(
                  "Failed to close file, path[%s].", path));
            }
          }};
#endif
}

// Close all open file descriptors
// The implementation is async signal safe
// Mostly copy from CPython code
static int close_open_fds_internal() {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return 0;
#else
  struct linux_dirent {
    long d_ino = 0;  // NOLINT
    off_t d_off;
    unsigned short d_reclen = 0;  // NOLINT
    char d_name[256];
  };

  int dir_fd = -1;
  if ((dir_fd = open("/proc/self/fd", O_RDONLY)) < 0) {
    PADDLE_THROW(platform::errors::Unavailable("Failed to open proc/self/fd."));
    return -1;
  }
  char buffer[sizeof(linux_dirent)];

  for (;;) {
    int bytes = 0;
    if ((bytes = syscall(SYS_getdents, dir_fd,
                         reinterpret_cast<linux_dirent*>(buffer),
                         sizeof(buffer))) < 0) {
      PADDLE_THROW(platform::errors::Unavailable(
          "System call failed via syscall function."));
      return -1;
    }

    if (bytes == 0) {
      break;
    }

    linux_dirent* entry = NULL;

    for (int offset = 0; offset < bytes; offset += entry->d_reclen) {
      entry = reinterpret_cast<linux_dirent*>(buffer + offset);
      int fd = 0;
      const char* s = entry->d_name;

      while (*s >= '0' && *s <= '9') {
        fd = fd * 10 + (*s - '0');
        s++;
      }

      if (s != entry->d_name && fd != dir_fd && fd >= 3) {
        close(fd);
      }
    }
  }

  close(dir_fd);
  return 0;
#endif
}

static int shell_popen_fork_internal(const char* real_cmd, bool do_read,
                                     int parent_end, int child_end,
                                     bool redirect_stderr = false) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return 0;
#else
  int child_pid = -1;
  // Too frequent calls to fork() makes openmpi very slow. Use vfork() instead.
  // But vfork() is very dangerous. Be careful.
  if ((child_pid = vfork()) < 0) {
    return -1;
  }

  // The following code is async signal safe (No memory allocation, no access to
  // global data, etc.)
  if (child_pid != 0) {
    return child_pid;
  }

  int child_std_end = do_read ? 1 : 0;
  close(parent_end);

  if (child_end != child_std_end) {
    PCHECK(dup2(child_end, child_std_end) == child_std_end);
    if (redirect_stderr && do_read) {
      PCHECK(dup2(child_end, 2) == 2);
    }
    close(child_end);
  }

  close_open_fds_internal();
  PCHECK(execl("/bin/bash", "bash", "-c", real_cmd, NULL) >= 0);
  // Note: just for compilation. the child don't run this line.
  _exit(0);
#endif
}

static int read_from_pipe(FILE* fp, std::string* output) {
  char buf[4096];
  while (1) {
    int n = fread(buf, 1, 4096, fp);
    if (n <= 0) {
      break;
    }

    output->append(buf, n);
  }

  if (!feof(fp)) {
    return -1;
  }

  return 0;
}

std::shared_ptr<FILE> shell_popen(const std::string& cmd,
                                  const std::string& mode, int* err_no,
                                  int* status, bool redirect_stderr) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return nullptr;
#else
  bool do_read = mode == "r";
  bool do_write = mode == "w";
  if (!(do_read || do_write)) {
    *err_no = -1;
    return NULL;
  }

  VLOG(3) << "Opening pipe[" << cmd << "] with mode[" << mode << "]";

  std::string real_cmd = "set -o pipefail; " + cmd;

  int pipe_fds[2];
  if (pipe(pipe_fds) != 0) {
    *err_no = -1;
    return NULL;
  }
  int parent_end = 0;
  int child_end = 0;

  if (do_read) {
    parent_end = pipe_fds[0];
    child_end = pipe_fds[1];
  } else if (do_write) {
    parent_end = pipe_fds[1];
    child_end = pipe_fds[0];
  }

  sighandler_t old_handler;
  old_handler = signal(SIGCHLD, SIG_DFL);

  fcntl(parent_end, F_SETFD, FD_CLOEXEC);

  int child_pid = shell_popen_fork_internal(
      real_cmd.c_str(), do_read, parent_end, child_end, redirect_stderr);

  close(child_end);

  FILE* fp = NULL;
  if ((fp = fdopen(parent_end, mode.c_str())) == NULL) {
    *err_no = -1;
    signal(SIGCHLD, old_handler);
    return NULL;
  }

  return {fp, [cmd, child_pid, old_handler, err_no, status](FILE* fp) {
            VLOG(3) << "Closing pipe[" << cmd << "]";
            if (fclose(fp)) {
              *err_no = -1;
            }

            int wstatus = -1;
            // don't do this before parent read data from child pipe
            // or when get the large data, it will hang!
            waitpid(child_pid, &wstatus, 0);

            if (status) {
              *status = wstatus;
            }

            if (WIFEXITED(wstatus) || wstatus == (128 + SIGPIPE) * 256) {
            } else {
              PADDLE_ENFORCE_NE(
                  errno, ECHILD,
                  platform::errors::Fatal("Must not be ECHILD errno here!"));
              *err_no = -1;
            }

            signal(SIGCHLD, old_handler);
          }};
#endif
}

static int shell_p2open_fork_internal(const char* real_cmd, int pipein_fds[2],
                                      int pipeout_fds[2]) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return 0;
#else
  int child_pid = -1;
  if ((child_pid = fork()) < 0) {
    return -1;
  }

  if (child_pid != 0) {
    return child_pid;
  }

  close(pipein_fds[0]);
  close(pipeout_fds[1]);

  if (pipein_fds[1] != 1) {
    if (dup2(pipein_fds[1], 1) != 1) {
      return -1;
    }
    close(pipein_fds[1]);
  }

  if (pipeout_fds[0] != 0) {
    if (dup2(pipeout_fds[0], 0) != 0) {
      return -1;
    }
    close(pipeout_fds[0]);
  }

  close_open_fds_internal();
  if (execl("/bin/sh", "sh", "-c", real_cmd, NULL) < 0) {
    return -1;
  }
  exit(127);
#endif
}

std::pair<std::shared_ptr<FILE>, std::shared_ptr<FILE>> shell_p2open(
    const std::string& cmd) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return {};
#else
  if (shell_verbose()) {
    LOG(INFO) << "Opening bidirectional pipe[" << cmd << "]";
  }

  std::string real_cmd = "set -o pipefail; " + cmd;

  int pipein_fds[2];
  int pipeout_fds[2];
  if (pipe(pipein_fds) != 0) {
    return {NULL, NULL};
  }
  if (pipe(pipeout_fds) != 0) {
    return {NULL, NULL};
  }

  int child_pid =
      shell_p2open_fork_internal(real_cmd.c_str(), pipein_fds, pipeout_fds);

  close(pipein_fds[1]);
  close(pipeout_fds[0]);
  fcntl(pipein_fds[0], F_SETFD, FD_CLOEXEC);
  fcntl(pipeout_fds[1], F_SETFD, FD_CLOEXEC);

  std::shared_ptr<int> child_life = {
      NULL, [child_pid, cmd](void*) {
        if (shell_verbose()) {
          LOG(INFO) << "Closing bidirectional pipe[" << cmd << "]";
        }

        int wstatus, ret;

        do {
          PCHECK((ret = waitpid(child_pid, &wstatus, 0)) >= 0 ||
                 (ret == -1 && errno == EINTR));
        } while (ret == -1 && errno == EINTR);

        PCHECK(wstatus == 0 || wstatus == (128 + SIGPIPE) * 256 ||
               (wstatus == -1 && errno == ECHILD))
            << "status[" << wstatus << "], cmd[" << cmd << "]";

        if (wstatus == -1 && errno == ECHILD) {
          // temporarily remove this warning
          // LOG(WARNING) << "errno is ECHILD";
        }
      }};

  FILE* in_fp;
  PCHECK((in_fp = fdopen(pipein_fds[0], "r")) != NULL);
  FILE* out_fp;
  PCHECK((out_fp = fdopen(pipeout_fds[1], "w")) != NULL);
  return {{in_fp, [child_life](FILE* fp) { PCHECK(fclose(fp) == 0); }},
          {out_fp, [child_life](FILE* fp) { PCHECK(fclose(fp) == 0); }}};
#endif
}

#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
#else
static int _get_err_no(int err_no, int status) {
  if (err_no == 0) {
    if (WIFEXITED(status)) {
      return WEXITSTATUS(status);
    }
    return -1;
  }

  return err_no;
}
#endif

static int _shell_execute_cmd(const std::string& cmd, std::string* output,
                              int time_out, int sleep_inter,
                              bool redirect_stderr = false) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  PADDLE_THROW(platform::errors::Unimplemented(
      "This function(shell_get_command_output) is not implemented under _WIN32 "
      "or __APPLE__."));
#else
  int err_no = 0;
  int status = 0;
  int cmd_status = 0;
  platform::Timer timer;
  do {
    VLOG(3) << "exec cmd:[" << cmd << "]";

    err_no = 0;
    status = 0;
    *output = "";
    auto pipe = shell_popen(cmd, "r", &err_no, &status, redirect_stderr);

    if (err_no == 0) {
      // read file
      err_no = read_from_pipe(&*pipe, output);
      if (err_no) {
        LOG(WARNING) << "status[" << status << "], cmd[" << cmd << "]"
                     << ", err_no[" << err_no << "]";
      }
    }

    // close file and etc.
    pipe = nullptr;
    if (err_no) {
      LOG(WARNING) << "status[" << status << "], cmd[" << cmd << "]"
                   << ", err_no[" << err_no << "]";
    }

    cmd_status = _get_err_no(err_no, status);
    // cmd run ok!
    if (cmd_status == 0) {
      return cmd_status;
    }

    // time out
    timer.Pause();
    if ((time_out > 0 && timer.ElapsedMS() >= time_out) || time_out == 0) {
      break;
    }
    timer.Resume();

    if (sleep_inter > 0) {
      usleep(sleep_inter * 1000);
    }
  } while (cmd_status);

  // log when check timeout!
  if (time_out != 0) {
    *output += string::Sprintf(
        " _shell_execute_cmd execute cmd:%s ElapsedMS:%d, err_no:%d status:%d",
        cmd, timer.ElapsedMS(), err_no, cmd_status);
    LOG(WARNING) << *output;
  }

  return cmd_status;

#endif
}

std::string shell_get_command_output(const std::string& cmd, int time_out,
                                     int sleep_inter) {
  std::string output;
  _shell_execute_cmd(cmd, &output, time_out, sleep_inter);
  return output;
}

std::vector<std::string> shell_execute_cmd(const std::string& cmd, int time_out,
                                           int sleep_inter,
                                           bool redirect_stderr) {
  std::string output;
  int ret =
      _shell_execute_cmd(cmd, &output, time_out, sleep_inter, redirect_stderr);
  return std::vector<std::string>({string::Sprintf("%d", ret), output});
}

}  // end namespace framework
}  // end namespace paddle
<commit_msg>Fix a bug when running on an operating system without "bash." (#29131)<commit_after>// Copyright (c) 2019 PaddlePaddle 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 "paddle/fluid/framework/io/shell.h"
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/platform/timer.h"

namespace paddle {
namespace framework {

std::shared_ptr<FILE> shell_fopen(const std::string& path,
                                  const std::string& mode) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return nullptr;
#else
  if (shell_verbose()) {
    LOG(INFO) << "Opening file[" << path << "] with mode[" << mode << "]";
  }
  FILE* fp;
  if (!(fp = fopen(path.c_str(), mode.c_str()))) {
    PADDLE_THROW(platform::errors::Unavailable(
        "Failed to open file, path[%s], mode[%s].", path, mode));
  }
  return {fp, [path](FILE* fp) {
            if (shell_verbose()) {
              LOG(INFO) << "Closing file[" << path << "]";
            }
            if (0 != fclose(fp)) {
              PADDLE_THROW(platform::errors::Unavailable(
                  "Failed to close file, path[%s].", path));
            }
          }};
#endif
}

// Close all open file descriptors
// The implementation is async signal safe
// Mostly copy from CPython code
static int close_open_fds_internal() {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return 0;
#else
  struct linux_dirent {
    long d_ino = 0;  // NOLINT
    off_t d_off;
    unsigned short d_reclen = 0;  // NOLINT
    char d_name[256];
  };

  int dir_fd = -1;
  if ((dir_fd = open("/proc/self/fd", O_RDONLY)) < 0) {
    PADDLE_THROW(platform::errors::Unavailable("Failed to open proc/self/fd."));
    return -1;
  }
  char buffer[sizeof(linux_dirent)];

  for (;;) {
    int bytes = 0;
    if ((bytes = syscall(SYS_getdents, dir_fd,
                         reinterpret_cast<linux_dirent*>(buffer),
                         sizeof(buffer))) < 0) {
      PADDLE_THROW(platform::errors::Unavailable(
          "System call failed via syscall function."));
      return -1;
    }

    if (bytes == 0) {
      break;
    }

    linux_dirent* entry = NULL;

    for (int offset = 0; offset < bytes; offset += entry->d_reclen) {
      entry = reinterpret_cast<linux_dirent*>(buffer + offset);
      int fd = 0;
      const char* s = entry->d_name;

      while (*s >= '0' && *s <= '9') {
        fd = fd * 10 + (*s - '0');
        s++;
      }

      if (s != entry->d_name && fd != dir_fd && fd >= 3) {
        close(fd);
      }
    }
  }

  close(dir_fd);
  return 0;
#endif
}

static int shell_popen_fork_internal(const char* real_cmd, bool do_read,
                                     int parent_end, int child_end,
                                     bool redirect_stderr = false) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return 0;
#else
  int child_pid = -1;
  // Too frequent calls to fork() makes openmpi very slow. Use vfork() instead.
  // But vfork() is very dangerous. Be careful.
  if ((child_pid = vfork()) < 0) {
    return -1;
  }

  // The following code is async signal safe (No memory allocation, no access to
  // global data, etc.)
  if (child_pid != 0) {
    return child_pid;
  }

  int child_std_end = do_read ? 1 : 0;
  close(parent_end);

  if (child_end != child_std_end) {
    PCHECK(dup2(child_end, child_std_end) == child_std_end);
    if (redirect_stderr && do_read) {
      PCHECK(dup2(child_end, 2) == 2);
    }
    close(child_end);
  }

  close_open_fds_internal();

#if defined(PADDLE_WITH_MUSL)
  PCHECK(execl("/bin/sh", "sh", "-c", real_cmd, NULL) >= 0);
#else
  PCHECK(execl("/bin/bash", "bash", "-c", real_cmd, NULL) >= 0);
#endif
  // Note: just for compilation. the child don't run this line.
  _exit(0);
#endif
}

static int read_from_pipe(FILE* fp, std::string* output) {
  char buf[4096];
  while (1) {
    int n = fread(buf, 1, 4096, fp);
    if (n <= 0) {
      break;
    }

    output->append(buf, n);
  }

  if (!feof(fp)) {
    return -1;
  }

  return 0;
}

std::shared_ptr<FILE> shell_popen(const std::string& cmd,
                                  const std::string& mode, int* err_no,
                                  int* status, bool redirect_stderr) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return nullptr;
#else
  bool do_read = mode == "r";
  bool do_write = mode == "w";
  if (!(do_read || do_write)) {
    *err_no = -1;
    return NULL;
  }

  VLOG(3) << "Opening pipe[" << cmd << "] with mode[" << mode << "]";

  std::string real_cmd = "set -o pipefail; " + cmd;

  int pipe_fds[2];
  if (pipe(pipe_fds) != 0) {
    *err_no = -1;
    return NULL;
  }
  int parent_end = 0;
  int child_end = 0;

  if (do_read) {
    parent_end = pipe_fds[0];
    child_end = pipe_fds[1];
  } else if (do_write) {
    parent_end = pipe_fds[1];
    child_end = pipe_fds[0];
  }

  sighandler_t old_handler;
  old_handler = signal(SIGCHLD, SIG_DFL);

  fcntl(parent_end, F_SETFD, FD_CLOEXEC);

  int child_pid = shell_popen_fork_internal(
      real_cmd.c_str(), do_read, parent_end, child_end, redirect_stderr);

  close(child_end);

  FILE* fp = NULL;
  if ((fp = fdopen(parent_end, mode.c_str())) == NULL) {
    *err_no = -1;
    signal(SIGCHLD, old_handler);
    return NULL;
  }

  return {fp, [cmd, child_pid, old_handler, err_no, status](FILE* fp) {
            VLOG(3) << "Closing pipe[" << cmd << "]";
            if (fclose(fp)) {
              *err_no = -1;
            }

            int wstatus = -1;
            // don't do this before parent read data from child pipe
            // or when get the large data, it will hang!
            waitpid(child_pid, &wstatus, 0);

            if (status) {
              *status = wstatus;
            }

            if (WIFEXITED(wstatus) || wstatus == (128 + SIGPIPE) * 256) {
            } else {
              PADDLE_ENFORCE_NE(
                  errno, ECHILD,
                  platform::errors::Fatal("Must not be ECHILD errno here!"));
              *err_no = -1;
            }

            signal(SIGCHLD, old_handler);
          }};
#endif
}

static int shell_p2open_fork_internal(const char* real_cmd, int pipein_fds[2],
                                      int pipeout_fds[2]) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return 0;
#else
  int child_pid = -1;
  if ((child_pid = fork()) < 0) {
    return -1;
  }

  if (child_pid != 0) {
    return child_pid;
  }

  close(pipein_fds[0]);
  close(pipeout_fds[1]);

  if (pipein_fds[1] != 1) {
    if (dup2(pipein_fds[1], 1) != 1) {
      return -1;
    }
    close(pipein_fds[1]);
  }

  if (pipeout_fds[0] != 0) {
    if (dup2(pipeout_fds[0], 0) != 0) {
      return -1;
    }
    close(pipeout_fds[0]);
  }

  close_open_fds_internal();
  if (execl("/bin/sh", "sh", "-c", real_cmd, NULL) < 0) {
    return -1;
  }
  exit(127);
#endif
}

std::pair<std::shared_ptr<FILE>, std::shared_ptr<FILE>> shell_p2open(
    const std::string& cmd) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  return {};
#else
  if (shell_verbose()) {
    LOG(INFO) << "Opening bidirectional pipe[" << cmd << "]";
  }

  std::string real_cmd = "set -o pipefail; " + cmd;

  int pipein_fds[2];
  int pipeout_fds[2];
  if (pipe(pipein_fds) != 0) {
    return {NULL, NULL};
  }
  if (pipe(pipeout_fds) != 0) {
    return {NULL, NULL};
  }

  int child_pid =
      shell_p2open_fork_internal(real_cmd.c_str(), pipein_fds, pipeout_fds);

  close(pipein_fds[1]);
  close(pipeout_fds[0]);
  fcntl(pipein_fds[0], F_SETFD, FD_CLOEXEC);
  fcntl(pipeout_fds[1], F_SETFD, FD_CLOEXEC);

  std::shared_ptr<int> child_life = {
      NULL, [child_pid, cmd](void*) {
        if (shell_verbose()) {
          LOG(INFO) << "Closing bidirectional pipe[" << cmd << "]";
        }

        int wstatus, ret;

        do {
          PCHECK((ret = waitpid(child_pid, &wstatus, 0)) >= 0 ||
                 (ret == -1 && errno == EINTR));
        } while (ret == -1 && errno == EINTR);

        PCHECK(wstatus == 0 || wstatus == (128 + SIGPIPE) * 256 ||
               (wstatus == -1 && errno == ECHILD))
            << "status[" << wstatus << "], cmd[" << cmd << "]";

        if (wstatus == -1 && errno == ECHILD) {
          // temporarily remove this warning
          // LOG(WARNING) << "errno is ECHILD";
        }
      }};

  FILE* in_fp;
  PCHECK((in_fp = fdopen(pipein_fds[0], "r")) != NULL);
  FILE* out_fp;
  PCHECK((out_fp = fdopen(pipeout_fds[1], "w")) != NULL);
  return {{in_fp, [child_life](FILE* fp) { PCHECK(fclose(fp) == 0); }},
          {out_fp, [child_life](FILE* fp) { PCHECK(fclose(fp) == 0); }}};
#endif
}

#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
#else
static int _get_err_no(int err_no, int status) {
  if (err_no == 0) {
    if (WIFEXITED(status)) {
      return WEXITSTATUS(status);
    }
    return -1;
  }

  return err_no;
}
#endif

static int _shell_execute_cmd(const std::string& cmd, std::string* output,
                              int time_out, int sleep_inter,
                              bool redirect_stderr = false) {
#if defined(_WIN32) || defined(__APPLE__) || defined(PADDLE_ARM)
  PADDLE_THROW(platform::errors::Unimplemented(
      "This function(shell_get_command_output) is not implemented under _WIN32 "
      "or __APPLE__."));
#else
  int err_no = 0;
  int status = 0;
  int cmd_status = 0;
  platform::Timer timer;
  do {
    VLOG(3) << "exec cmd:[" << cmd << "]";

    err_no = 0;
    status = 0;
    *output = "";
    auto pipe = shell_popen(cmd, "r", &err_no, &status, redirect_stderr);

    if (err_no == 0) {
      // read file
      err_no = read_from_pipe(&*pipe, output);
      if (err_no) {
        LOG(WARNING) << "status[" << status << "], cmd[" << cmd << "]"
                     << ", err_no[" << err_no << "]";
      }
    }

    // close file and etc.
    pipe = nullptr;
    if (err_no) {
      LOG(WARNING) << "status[" << status << "], cmd[" << cmd << "]"
                   << ", err_no[" << err_no << "]";
    }

    cmd_status = _get_err_no(err_no, status);
    // cmd run ok!
    if (cmd_status == 0) {
      return cmd_status;
    }

    // time out
    timer.Pause();
    if ((time_out > 0 && timer.ElapsedMS() >= time_out) || time_out == 0) {
      break;
    }
    timer.Resume();

    if (sleep_inter > 0) {
      usleep(sleep_inter * 1000);
    }
  } while (cmd_status);

  // log when check timeout!
  if (time_out != 0) {
    *output += string::Sprintf(
        " _shell_execute_cmd execute cmd:%s ElapsedMS:%d, err_no:%d status:%d",
        cmd, timer.ElapsedMS(), err_no, cmd_status);
    LOG(WARNING) << *output;
  }

  return cmd_status;

#endif
}

std::string shell_get_command_output(const std::string& cmd, int time_out,
                                     int sleep_inter) {
  std::string output;
  _shell_execute_cmd(cmd, &output, time_out, sleep_inter);
  return output;
}

std::vector<std::string> shell_execute_cmd(const std::string& cmd, int time_out,
                                           int sleep_inter,
                                           bool redirect_stderr) {
  std::string output;
  int ret =
      _shell_execute_cmd(cmd, &output, time_out, sleep_inter, redirect_stderr);
  return std::vector<std::string>({string::Sprintf("%d", ret), output});
}

}  // end namespace framework
}  // end namespace paddle
<|endoftext|>