<commit_before>#include <iostream>
#include <fstream>
#include <memory>

#include "HexagonOffload.h"
#include "Closure.h"
#include "InjectHostDevBufferCopies.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "LLVM_Output.h"
#include "LLVM_Headers.h"
#include "Param.h"
#include "RemoveTrivialForLoops.h"
#include "Substitute.h"

namespace Halide {
namespace Internal {

using std::string;
using std::vector;

namespace {

// Replace the parameter objects of loads/stores with a new parameter
// object.
class ReplaceParams : public IRMutator {
    const std::map<std::string, Parameter> &replacements;

    using IRMutator::visit;

    void visit(const Load *op) {
        auto i = replacements.find(op->name);
        if (i != replacements.end()) {
            expr = Load::make(op->type, op->name, mutate(op->index), op->image,
                              i->second, mutate(op->predicate));
        } else {
            IRMutator::visit(op);
        }
    }

    void visit(const Store *op) {
        auto i = replacements.find(op->name);
        if (i != replacements.end()) {
            stmt = Store::make(op->name, mutate(op->value), mutate(op->index),
                               i->second, mutate(op->predicate));
        } else {
            IRMutator::visit(op);
        }
    }

public:
    ReplaceParams(const std::map<std::string, Parameter> &replacements)
        : replacements(replacements) {}
};

Stmt replace_params(Stmt s, const std::map<std::string, Parameter> &replacements) {
    return ReplaceParams(replacements).mutate(s);
}

// Wrap the stmt in a call to power_hvx_on, calling power_hvx_off
// as a destructor if successful.
Stmt power_hvx_on(Stmt stmt) {
    Expr power_on = Call::make(Int(32), "halide_hexagon_power_hvx_on", {}, Call::Extern);
    string power_on_result_name = unique_name("power_on_result");
    Expr power_on_result_var = Variable::make(Int(32), power_on_result_name);
    Stmt check_power_on = LetStmt::make(power_on_result_name, power_on,
                                        AssertStmt::make(EQ::make(power_on_result_var, 0), power_on_result_var));

    Expr dummy_obj = reinterpret(Handle(), cast<uint64_t>(1));
    Expr power_off = Call::make(Int(32), Call::register_destructor,
                                {Expr("halide_hexagon_power_hvx_off_as_destructor"), dummy_obj}, Call::Intrinsic);

    stmt = Block::make(Evaluate::make(power_off), stmt);
    stmt = Block::make(check_power_on, stmt);
    return stmt;
}

class InjectHexagonRpc : public IRMutator {
    std::map<std::string, Expr> state_vars;

    Module device_code;

    /** Alignment info for Int(32) variables in scope, so we don't
     * lose the information when creating Hexagon kernels. */
    Scope<ModulusRemainder> alignment_info;

    Expr state_var(const std::string& name, Type type) {
        Expr& var = state_vars[name];
        if (!var.defined()) {
            auto storage = Buffer<void *>::make_scalar(name + "_buf");
            storage() = nullptr;
            var = Load::make(type_of<void*>(), storage.name(), 0, storage, Parameter(), const_true());
        }
        return var;
    }

    Expr state_var_ptr(const std::string& name, Type type) {
        Expr var = state_var(name, type);
        return Call::make(Handle(), Call::address_of, {var}, Call::Intrinsic);
    }

    Expr module_state() {
        return state_var("hexagon_module_state", type_of<void*>());
    }

    Expr module_state_ptr() {
        return state_var_ptr("hexagon_module_state", type_of<void*>());
    }

    // Create a Buffer containing the given buffer/size, and return an
    // expression for a pointer to the first element.
    Expr buffer_ptr(const uint8_t* buffer, size_t size, const char* name) {
        Buffer<uint8_t> code((int)size, name);
        memcpy(code.data(), buffer, (int)size);
        Expr ptr_0 = Load::make(type_of<uint8_t>(), name, 0, code, Parameter(), const_true());
        return Call::make(Handle(), Call::address_of, {ptr_0}, Call::Intrinsic);
    }

    using IRMutator::visit;

    void visit(const For *loop) {
        if (loop->device_api == DeviceAPI::Hexagon) {
            // Unrolling or loop partitioning might generate multiple
            // loops with the same name, so we need to make them unique.
            std::string hex_name = unique_name("hex_" + loop->name);

            // After moving this to Hexagon, it doesn't need to be
            // marked Hexagon anymore.
            Stmt body = For::make(loop->name, loop->min, loop->extent, loop->for_type,
                                  DeviceAPI::None, loop->body);
            body = remove_trivial_for_loops(body);

            // Build a closure for the device code.
            // TODO: Should this move the body of the loop to Hexagon,
            // or the loop itself? Currently, this moves the loop itself.
            Closure c(body);

            // Make an argument list, and generate a function in the
            // device_code module. The hexagon runtime code expects
            // the arguments to appear in the order of (input buffers,
            // output buffers, input scalars).  Scalars must be last
            // for the scalar arguments to shadow the symbols of the
            // buffer that get generated by CodeGen_LLVM.
            std::vector<LoweredArgument> input_buffers, output_buffers;
            std::map<std::string, Parameter> replacement_params;
            for (const auto& i : c.buffers) {
                if (i.second.write) {
                    Argument::Kind kind = Argument::OutputBuffer;
                    output_buffers.push_back(LoweredArgument(i.first, kind, i.second.type, i.second.dimensions));
                } else {
                    Argument::Kind kind = Argument::InputBuffer;
                    input_buffers.push_back(LoweredArgument(i.first, kind, i.second.type, i.second.dimensions));
                }

                // Build a parameter to replace.
                Parameter p(i.second.type, true, i.second.dimensions);
                // Assert that buffers are aligned to one HVX vector.
                const int alignment = 128;
                p.set_host_alignment(alignment);
                // The other parameter constraints are already
                // accounted for by the closure grabbing those
                // arguments, so we only need to provide the host
                // alignment.
                replacement_params[i.first] = p;

                // Add an assert to the body that validates the
                // alignment of the buffer.
                if (!device_code.target().has_feature(Target::NoAsserts)) {
                    Expr host_ptr = reinterpret<uint64_t>(Variable::make(Handle(), i.first + ".host"));
                    Expr error = Call::make(Int(32), "halide_error_unaligned_host_ptr",
                                            {i.first, alignment}, Call::Extern);
                    body = Block::make(AssertStmt::make(host_ptr % alignment == 0, error), body);
                }
            }
            body = replace_params(body, replacement_params);

            std::vector<LoweredArgument> args;
            args.insert(args.end(), input_buffers.begin(), input_buffers.end());
            args.insert(args.end(), output_buffers.begin(), output_buffers.end());
            for (const auto& i : c.vars) {
                LoweredArgument arg(i.first, Argument::InputScalar, i.second, 0);
                if (alignment_info.contains(i.first)) {
                    arg.alignment = alignment_info.get(i.first);
                }
                args.push_back(arg);
            }
            device_code.append(LoweredFunc(hex_name, args, body, LoweredFunc::External));

            // Generate a call to hexagon_device_run.
            std::vector<Expr> arg_sizes;
            std::vector<Expr> arg_ptrs;
            std::vector<Expr> arg_flags;

            for (const auto& i : c.buffers) {
                arg_sizes.push_back(Expr((uint64_t) sizeof(buffer_t*)));
                arg_ptrs.push_back(Variable::make(type_of<buffer_t *>(), i.first + ".buffer"));
                // In the flags parameter, bit 0 set indicates the
                // buffer is read, bit 1 set indicates the buffer is
                // written. If neither are set, the argument is a scalar.
                int flags = 0;
                if (i.second.read) flags |= 0x1;
                if (i.second.write) flags |= 0x2;
                arg_flags.push_back(flags);
            }
            for (const auto& i : c.vars) {
                Expr arg = Variable::make(i.second, i.first);
                Expr arg_ptr = Call::make(type_of<void *>(), Call::make_struct, {arg}, Call::Intrinsic);
                arg_sizes.push_back(Expr((uint64_t) i.second.bytes()));
                arg_ptrs.push_back(arg_ptr);
                arg_flags.push_back(0x0);
            }

            // The argument list is terminated with an argument of size 0.
            arg_sizes.push_back(Expr((uint64_t) 0));

            std::string pipeline_name = hex_name + "_argv";
            std::vector<Expr> params;
            params.push_back(module_state());
            params.push_back(pipeline_name);
            params.push_back(state_var_ptr(hex_name, type_of<int>()));
            params.push_back(Call::make(type_of<size_t*>(), Call::make_struct, arg_sizes, Call::Intrinsic));
            params.push_back(Call::make(type_of<void**>(), Call::make_struct, arg_ptrs, Call::Intrinsic));
            params.push_back(Call::make(type_of<int*>(), Call::make_struct, arg_flags, Call::Intrinsic));

            stmt = call_extern_and_assert("halide_hexagon_run", params);

        } else {
            IRMutator::visit(loop);
        }
    }

    void visit(const Let *op) {
        if (op->value.type() == Int(32)) {
            alignment_info.push(op->name, modulus_remainder(op->value, alignment_info));
        }

        IRMutator::visit(op);

        if (op->value.type() == Int(32)) {
            alignment_info.pop(op->name);
        }
    }

    void visit(const LetStmt *op) {
        if (op->value.type() == Int(32)) {
            alignment_info.push(op->name, modulus_remainder(op->value, alignment_info));
        }

        IRMutator::visit(op);

        if (op->value.type() == Int(32)) {
            alignment_info.pop(op->name);
        }
    }

public:
    InjectHexagonRpc(const Target &target) : device_code("hexagon", target) {}

    Stmt inject(Stmt s) {
        s = mutate(s);

        // Skip if there are no device kernels.
        if (device_code.functions().empty()) {
            return s;
        }

        // If we got here, it means the pipeline runs at least one
        // Hexagon kernel. To reduce overhead of individual
        // sub-pipelines running on Hexagon, we can power on HVX once
        // for the duration of this pipeline.
        s = power_hvx_on(s);

        // Compile the device code
        debug(1) << "Hexagon device code module: " << device_code << "\n";

        llvm::LLVMContext context;
        std::unique_ptr<llvm::Module> llvm_module(compile_module_to_llvm_module(device_code, context));

        llvm::SmallVector<char, 4096> object;
        llvm::raw_svector_ostream object_stream(object);
        compile_llvm_module_to_object(*llvm_module, object_stream);

        if (debug::debug_level >= 2) {
            debug(2) << "Hexagon device code assembly: " << "\n";
            llvm::SmallString<4096> assembly;
            llvm::raw_svector_ostream assembly_stream(assembly);
            compile_llvm_module_to_assembly(*llvm_module, assembly_stream);
            debug(2) << assembly.c_str() << "\n";
        }

        // Wrap the statement in calls to halide_initialize_kernels.
        size_t code_size = object.size();
        Expr code_ptr = buffer_ptr(reinterpret_cast<uint8_t*>(&object[0]), code_size, "hexagon_code");
        Stmt init_kernels = call_extern_and_assert("halide_hexagon_initialize_kernels",
                                                   {module_state_ptr(), code_ptr, Expr((uint64_t) code_size)});
        s = Block::make(init_kernels, s);

        return s;
    }
};

}

Stmt inject_hexagon_rpc(Stmt s, const Target &host_target) {
    // Make a new target for the device module.
    Target target(Target::NoOS, Target::Hexagon, 32);

    // These feature flags are propagated from the host target to the
    // device module.
    //
    // TODO: We'd like Target::Debug to be in this list too, but trunk
    // llvm currently disagrees with hexagon clang as to what
    // constitutes valid debug info.
    static const Target::Feature shared_features[] = {
        Target::Profile,
        Target::NoAsserts,
        Target::HVX_64,
        Target::HVX_128,
        Target::HVX_v62
    };
    for (Target::Feature i : shared_features) {
        if (host_target.has_feature(i)) {
            target = target.with_feature(i);
        }
    }

    InjectHexagonRpc injector(target);
    s = injector.inject(s);
    return s;
}

}  // namespace Internal
}  // namespace Halide
<commit_msg>Restructure For visitor to reduce the depth of a large block.<commit_after>#include <iostream>
#include <fstream>
#include <memory>

#include "HexagonOffload.h"
#include "Closure.h"
#include "InjectHostDevBufferCopies.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "LLVM_Output.h"
#include "LLVM_Headers.h"
#include "Param.h"
#include "RemoveTrivialForLoops.h"
#include "Substitute.h"

namespace Halide {
namespace Internal {

using std::string;
using std::vector;

namespace {

// Replace the parameter objects of loads/stores with a new parameter
// object.
class ReplaceParams : public IRMutator {
    const std::map<std::string, Parameter> &replacements;

    using IRMutator::visit;

    void visit(const Load *op) {
        auto i = replacements.find(op->name);
        if (i != replacements.end()) {
            expr = Load::make(op->type, op->name, mutate(op->index), op->image,
                              i->second, mutate(op->predicate));
        } else {
            IRMutator::visit(op);
        }
    }

    void visit(const Store *op) {
        auto i = replacements.find(op->name);
        if (i != replacements.end()) {
            stmt = Store::make(op->name, mutate(op->value), mutate(op->index),
                               i->second, mutate(op->predicate));
        } else {
            IRMutator::visit(op);
        }
    }

public:
    ReplaceParams(const std::map<std::string, Parameter> &replacements)
        : replacements(replacements) {}
};

Stmt replace_params(Stmt s, const std::map<std::string, Parameter> &replacements) {
    return ReplaceParams(replacements).mutate(s);
}

// Wrap the stmt in a call to power_hvx_on, calling power_hvx_off
// as a destructor if successful.
Stmt power_hvx_on(Stmt stmt) {
    Expr power_on = Call::make(Int(32), "halide_hexagon_power_hvx_on", {}, Call::Extern);
    string power_on_result_name = unique_name("power_on_result");
    Expr power_on_result_var = Variable::make(Int(32), power_on_result_name);
    Stmt check_power_on = LetStmt::make(power_on_result_name, power_on,
                                        AssertStmt::make(EQ::make(power_on_result_var, 0), power_on_result_var));

    Expr dummy_obj = reinterpret(Handle(), cast<uint64_t>(1));
    Expr power_off = Call::make(Int(32), Call::register_destructor,
                                {Expr("halide_hexagon_power_hvx_off_as_destructor"), dummy_obj}, Call::Intrinsic);

    stmt = Block::make(Evaluate::make(power_off), stmt);
    stmt = Block::make(check_power_on, stmt);
    return stmt;
}

class InjectHexagonRpc : public IRMutator {
    std::map<std::string, Expr> state_vars;

    Module device_code;

    /** Alignment info for Int(32) variables in scope, so we don't
     * lose the information when creating Hexagon kernels. */
    Scope<ModulusRemainder> alignment_info;

    Expr state_var(const std::string& name, Type type) {
        Expr& var = state_vars[name];
        if (!var.defined()) {
            auto storage = Buffer<void *>::make_scalar(name + "_buf");
            storage() = nullptr;
            var = Load::make(type_of<void*>(), storage.name(), 0, storage, Parameter(), const_true());
        }
        return var;
    }

    Expr state_var_ptr(const std::string& name, Type type) {
        Expr var = state_var(name, type);
        return Call::make(Handle(), Call::address_of, {var}, Call::Intrinsic);
    }

    Expr module_state() {
        return state_var("hexagon_module_state", type_of<void*>());
    }

    Expr module_state_ptr() {
        return state_var_ptr("hexagon_module_state", type_of<void*>());
    }

    // Create a Buffer containing the given buffer/size, and return an
    // expression for a pointer to the first element.
    Expr buffer_ptr(const uint8_t* buffer, size_t size, const char* name) {
        Buffer<uint8_t> code((int)size, name);
        memcpy(code.data(), buffer, (int)size);
        Expr ptr_0 = Load::make(type_of<uint8_t>(), name, 0, code, Parameter(), const_true());
        return Call::make(Handle(), Call::address_of, {ptr_0}, Call::Intrinsic);
    }

    using IRMutator::visit;

    void visit(const For *loop) {
        if (loop->device_api != DeviceAPI::Hexagon) {
            IRMutator::visit(loop);
            return;
        }

        // Unrolling or loop partitioning might generate multiple
        // loops with the same name, so we need to make them unique.
        std::string hex_name = unique_name("hex_" + loop->name);

        // After moving this to Hexagon, it doesn't need to be
        // marked Hexagon anymore.
        Stmt body = For::make(loop->name, loop->min, loop->extent, loop->for_type,
                              DeviceAPI::None, loop->body);
        body = remove_trivial_for_loops(body);

        // Build a closure for the device code.
        // TODO: Should this move the body of the loop to Hexagon,
        // or the loop itself? Currently, this moves the loop itself.
        Closure c(body);

        // Make an argument list, and generate a function in the
        // device_code module. The hexagon runtime code expects
        // the arguments to appear in the order of (input buffers,
        // output buffers, input scalars).  Scalars must be last
        // for the scalar arguments to shadow the symbols of the
        // buffer that get generated by CodeGen_LLVM.
        std::vector<LoweredArgument> input_buffers, output_buffers;
        std::map<std::string, Parameter> replacement_params;
        for (const auto& i : c.buffers) {
            if (i.second.write) {
                Argument::Kind kind = Argument::OutputBuffer;
                output_buffers.push_back(LoweredArgument(i.first, kind, i.second.type, i.second.dimensions));
            } else {
                Argument::Kind kind = Argument::InputBuffer;
                input_buffers.push_back(LoweredArgument(i.first, kind, i.second.type, i.second.dimensions));
            }

            // Build a parameter to replace.
            Parameter p(i.second.type, true, i.second.dimensions);
            // Assert that buffers are aligned to one HVX vector.
            const int alignment = 128;
            p.set_host_alignment(alignment);
            // The other parameter constraints are already
            // accounted for by the closure grabbing those
            // arguments, so we only need to provide the host
            // alignment.
            replacement_params[i.first] = p;

            // Add an assert to the body that validates the
            // alignment of the buffer.
            if (!device_code.target().has_feature(Target::NoAsserts)) {
                Expr host_ptr = reinterpret<uint64_t>(Variable::make(Handle(), i.first + ".host"));
                Expr error = Call::make(Int(32), "halide_error_unaligned_host_ptr",
                                        {i.first, alignment}, Call::Extern);
                body = Block::make(AssertStmt::make(host_ptr % alignment == 0, error), body);
            }
        }
        body = replace_params(body, replacement_params);

        std::vector<LoweredArgument> args;
        args.insert(args.end(), input_buffers.begin(), input_buffers.end());
        args.insert(args.end(), output_buffers.begin(), output_buffers.end());
        for (const auto& i : c.vars) {
            LoweredArgument arg(i.first, Argument::InputScalar, i.second, 0);
            if (alignment_info.contains(i.first)) {
                arg.alignment = alignment_info.get(i.first);
            }
            args.push_back(arg);
        }
        device_code.append(LoweredFunc(hex_name, args, body, LoweredFunc::External));

        // Generate a call to hexagon_device_run.
        std::vector<Expr> arg_sizes;
        std::vector<Expr> arg_ptrs;
        std::vector<Expr> arg_flags;

        for (const auto& i : c.buffers) {
            arg_sizes.push_back(Expr((uint64_t) sizeof(buffer_t*)));
            arg_ptrs.push_back(Variable::make(type_of<buffer_t *>(), i.first + ".buffer"));
            // In the flags parameter, bit 0 set indicates the
            // buffer is read, bit 1 set indicates the buffer is
            // written. If neither are set, the argument is a scalar.
            int flags = 0;
            if (i.second.read) flags |= 0x1;
            if (i.second.write) flags |= 0x2;
            arg_flags.push_back(flags);
        }
        for (const auto& i : c.vars) {
            Expr arg = Variable::make(i.second, i.first);
            Expr arg_ptr = Call::make(type_of<void *>(), Call::make_struct, {arg}, Call::Intrinsic);
            arg_sizes.push_back(Expr((uint64_t) i.second.bytes()));
            arg_ptrs.push_back(arg_ptr);
            arg_flags.push_back(0x0);
        }

        // The argument list is terminated with an argument of size 0.
        arg_sizes.push_back(Expr((uint64_t) 0));

        std::string pipeline_name = hex_name + "_argv";
        std::vector<Expr> params;
        params.push_back(module_state());
        params.push_back(pipeline_name);
        params.push_back(state_var_ptr(hex_name, type_of<int>()));
        params.push_back(Call::make(type_of<size_t*>(), Call::make_struct, arg_sizes, Call::Intrinsic));
        params.push_back(Call::make(type_of<void**>(), Call::make_struct, arg_ptrs, Call::Intrinsic));
        params.push_back(Call::make(type_of<int*>(), Call::make_struct, arg_flags, Call::Intrinsic));

        stmt = call_extern_and_assert("halide_hexagon_run", params);
    }

    void visit(const Let *op) {
        if (op->value.type() == Int(32)) {
            alignment_info.push(op->name, modulus_remainder(op->value, alignment_info));
        }

        IRMutator::visit(op);

        if (op->value.type() == Int(32)) {
            alignment_info.pop(op->name);
        }
    }

    void visit(const LetStmt *op) {
        if (op->value.type() == Int(32)) {
            alignment_info.push(op->name, modulus_remainder(op->value, alignment_info));
        }

        IRMutator::visit(op);

        if (op->value.type() == Int(32)) {
            alignment_info.pop(op->name);
        }
    }

public:
    InjectHexagonRpc(const Target &target) : device_code("hexagon", target) {}

    Stmt inject(Stmt s) {
        s = mutate(s);

        // Skip if there are no device kernels.
        if (device_code.functions().empty()) {
            return s;
        }

        // If we got here, it means the pipeline runs at least one
        // Hexagon kernel. To reduce overhead of individual
        // sub-pipelines running on Hexagon, we can power on HVX once
        // for the duration of this pipeline.
        s = power_hvx_on(s);

        // Compile the device code
        debug(1) << "Hexagon device code module: " << device_code << "\n";

        llvm::LLVMContext context;
        std::unique_ptr<llvm::Module> llvm_module(compile_module_to_llvm_module(device_code, context));

        llvm::SmallVector<char, 4096> object;
        llvm::raw_svector_ostream object_stream(object);
        compile_llvm_module_to_object(*llvm_module, object_stream);

        if (debug::debug_level >= 2) {
            debug(2) << "Hexagon device code assembly: " << "\n";
            llvm::SmallString<4096> assembly;
            llvm::raw_svector_ostream assembly_stream(assembly);
            compile_llvm_module_to_assembly(*llvm_module, assembly_stream);
            debug(2) << assembly.c_str() << "\n";
        }

        // Wrap the statement in calls to halide_initialize_kernels.
        size_t code_size = object.size();
        Expr code_ptr = buffer_ptr(reinterpret_cast<uint8_t*>(&object[0]), code_size, "hexagon_code");
        Stmt init_kernels = call_extern_and_assert("halide_hexagon_initialize_kernels",
                                                   {module_state_ptr(), code_ptr, Expr((uint64_t) code_size)});
        s = Block::make(init_kernels, s);

        return s;
    }
};

}

Stmt inject_hexagon_rpc(Stmt s, const Target &host_target) {
    // Make a new target for the device module.
    Target target(Target::NoOS, Target::Hexagon, 32);

    // These feature flags are propagated from the host target to the
    // device module.
    //
    // TODO: We'd like Target::Debug to be in this list too, but trunk
    // llvm currently disagrees with hexagon clang as to what
    // constitutes valid debug info.
    static const Target::Feature shared_features[] = {
        Target::Profile,
        Target::NoAsserts,
        Target::HVX_64,
        Target::HVX_128,
        Target::HVX_v62
    };
    for (Target::Feature i : shared_features) {
        if (host_target.has_feature(i)) {
            target = target.with_feature(i);
        }
    }

    InjectHexagonRpc injector(target);
    s = injector.inject(s);
    return s;
}

}  // namespace Internal
}  // namespace Halide
<|endoftext|>


<commit_before>/******************************************************************************
 *  Copyright (c) 2015 Jamis Hoo
 *  Distributed under the MIT license
 *  (See accompanying file LICENSE or copy at http://opensource.org/licenses/MIT)
 *
 *  Project:
 *  Filename: sm4_ecb.cc
 *  Version: 1.0
 *  Author: Jamis Hoo
 *  E-mail: hoojamis@gmail.com
 *  Date: May 16, 2015
 *  Time: 17:01:33
 *  Description: SM4(128 bit) Electronic Codebook Mode(ECB)
 *****************************************************************************/
#include <cstdio>
#include <iostream>
#include <fstream>
#include <cinttypes>
#include <vector>

inline uint32_t left_rotate(const uint32_t x, const size_t i) {
    return x << i | (x >> (sizeof(uint32_t) * 8 - i));
}

inline uint32_t tauTransformation(const uint32_t x) {
    constexpr uint8_t Sbox[256] = {
        0xd6, 0x90, 0xe9, 0xfe, 0xcc, 0xe1, 0x3d, 0xb7, 0x16, 0xb6, 0x14, 0xc2, 0x28, 0xfb, 0x2c, 0x05,
        0x2b, 0x67, 0x9a, 0x76, 0x2a, 0xbe, 0x04, 0xc3, 0xaa, 0x44, 0x13, 0x26, 0x49, 0x86, 0x06, 0x99,
        0x9c, 0x42, 0x50, 0xf4, 0x91, 0xef, 0x98, 0x7a, 0x33, 0x54, 0x0b, 0x43, 0xed, 0xcf, 0xac, 0x62,
        0xe4, 0xb3, 0x1c, 0xa9, 0xc9, 0x08, 0xe8, 0x95, 0x80, 0xdf, 0x94, 0xfa, 0x75, 0x8f, 0x3f, 0xa6,
        0x47, 0x07, 0xa7, 0xfc, 0xf3, 0x73, 0x17, 0xba, 0x83, 0x59, 0x3c, 0x19, 0xe6, 0x85, 0x4f, 0xa8,
        0x68, 0x6b, 0x81, 0xb2, 0x71, 0x64, 0xda, 0x8b, 0xf8, 0xeb, 0x0f, 0x4b, 0x70, 0x56, 0x9d, 0x35,
        0x1e, 0x24, 0x0e, 0x5e, 0x63, 0x58, 0xd1, 0xa2, 0x25, 0x22, 0x7c, 0x3b, 0x01, 0x21, 0x78, 0x87,
        0xd4, 0x00, 0x46, 0x57, 0x9f, 0xd3, 0x27, 0x52, 0x4c, 0x36, 0x02, 0xe7, 0xa0, 0xc4, 0xc8, 0x9e,
        0xea, 0xbf, 0x8a, 0xd2, 0x40, 0xc7, 0x38, 0xb5, 0xa3, 0xf7, 0xf2, 0xce, 0xf9, 0x61, 0x15, 0xa1,
        0xe0, 0xae, 0x5d, 0xa4, 0x9b, 0x34, 0x1a, 0x55, 0xad, 0x93, 0x32, 0x30, 0xf5, 0x8c, 0xb1, 0xe3,
        0x1d, 0xf6, 0xe2, 0x2e, 0x82, 0x66, 0xca, 0x60, 0xc0, 0x29, 0x23, 0xab, 0x0d, 0x53, 0x4e, 0x6f,
        0xd5, 0xdb, 0x37, 0x45, 0xde, 0xfd, 0x8e, 0x2f, 0x03, 0xff, 0x6a, 0x72, 0x6d, 0x6c, 0x5b, 0x51,
        0x8d, 0x1b, 0xaf, 0x92, 0xbb, 0xdd, 0xbc, 0x7f, 0x11, 0xd9, 0x5c, 0x41, 0x1f, 0x10, 0x5a, 0xd8,
        0x0a, 0xc1, 0x31, 0x88, 0xa5, 0xcd, 0x7b, 0xbd, 0x2d, 0x74, 0xd0, 0x12, 0xb8, 0xe5, 0xb4, 0xb0,
        0x89, 0x69, 0x97, 0x4a, 0x0c, 0x96, 0x77, 0x7e, 0x65, 0xb9, 0xf1, 0x09, 0xc5, 0x6e, 0xc6, 0x84,
        0x18, 0xf0, 0x7d, 0xec, 0x3a, 0xdc, 0x4d, 0x20, 0x79, 0xee, 0x5f, 0x3e, 0xd7, 0xcb, 0x39, 0x48
    };
    uint32_t val;
    val = Sbox[x >>  0 & 0xff] | Sbox[x >>  8 & 0xff] |
          Sbox[x >> 16 & 0xff] | Sbox[x >> 24 & 0xff];
    return val;
}

inline uint32_t LTransformation(const uint32_t x) {
    uint32_t big_endian = (x >>  0 & 0xff) << 24 |
                          (x >>  8 & 0xff) << 16 |
                          (x >> 16 & 0xff) <<  8 |
                          (x >> 24 & 0xff) <<  0;
    uint32_t val = big_endian ^ left_rotate(big_endian,  2) ^
                                left_rotate(big_endian, 10) ^
                                left_rotate(big_endian, 18) ^
                                left_rotate(big_endian, 24);
    uint32_t small_endian = (big_endian >>  0 & 0xff) << 24 |
                            (big_endian >>  8 & 0xff) << 16 |
                            (big_endian >> 16 & 0xff) <<  8 |
                            (big_endian >> 24 & 0xff) <<  0;
    return small_endian;
}

inline uint32_t TTransformation(const uint32_t x) {
    return LTransformation(tauTransformation(x));
}

void keyExpansion(const uint8_t key[], uint8_t keys[]) {

}

void sm4_ecb(const void* plain, const size_t length, const void* key, void* cipher) {
    uint8_t* plain_ = (uint8_t*)(plain);
    uint8_t* cipher_ = (uint8_t*)(cipher);

    uint8_t keys[128];
    keyExpansion((uint8_t*)(key), keys);

}

int main(int argc, char** argv) {
    if (argc == 1) return 0;

    std::ifstream fin(argv[1]);

    fin.seekg(0, std::ios::end);
    std::string buffer;
    size_t len = fin.tellg();
    buffer.reserve(len);
    fin.seekg(0, std::ios::beg);

    if (len % 16) {
        printf("Length of plain text should be multiple of 16 bytes. \n");
        return 0;
    }

    buffer.assign((std::istreambuf_iterator<char>(fin)),
                   std::istreambuf_iterator<char>());

    fin.close();

    // 128 bit key size
    unsigned char key[16] = { 0 };

    if (argc == 3) {
        fin.open(argv[2]);
        fin.seekg(0, std::ios::beg);
        char buffer[2];
        for (size_t i = 0; i < 16; ++i) {
            fin.read(buffer, 2);
            key[i] = std::stoi(buffer, 0, 16);
        }
        fin.close();
    }

    std::vector<char> cipher(buffer.length(), 0);
    aes_ecb(buffer.data(), buffer.length(), key, &cipher[0]);

    for (size_t i = 0; i < buffer.length(); ++i)
        printf("%02x", int(cipher[i]) & 0xff);
    printf("\n");

}

<commit_msg>sm4 ecb finished<commit_after>/******************************************************************************
 *  Copyright (c) 2015 Jamis Hoo
 *  Distributed under the MIT license
 *  (See accompanying file LICENSE or copy at http://opensource.org/licenses/MIT)
 *
 *  Project:
 *  Filename: sm4_ecb.cc
 *  Version: 1.0
 *  Author: Jamis Hoo
 *  E-mail: hoojamis@gmail.com
 *  Date: May 16, 2015
 *  Time: 17:01:33
 *  Description: SM4(128 bit) Electronic Codebook Mode(ECB)
 *****************************************************************************/
#include <cstdio>
#include <iostream>
#include <fstream>
#include <cinttypes>
#include <vector>

inline uint32_t left_rotate(const uint32_t x, const size_t i) {
    return x << i | (x >> (sizeof(uint32_t) * 8 - i));
}

inline uint32_t tauTransformation(const uint32_t x) {
    constexpr uint8_t Sbox[256] = {
        0xd6, 0x90, 0xe9, 0xfe, 0xcc, 0xe1, 0x3d, 0xb7, 0x16, 0xb6, 0x14, 0xc2, 0x28, 0xfb, 0x2c, 0x05,
        0x2b, 0x67, 0x9a, 0x76, 0x2a, 0xbe, 0x04, 0xc3, 0xaa, 0x44, 0x13, 0x26, 0x49, 0x86, 0x06, 0x99,
        0x9c, 0x42, 0x50, 0xf4, 0x91, 0xef, 0x98, 0x7a, 0x33, 0x54, 0x0b, 0x43, 0xed, 0xcf, 0xac, 0x62,
        0xe4, 0xb3, 0x1c, 0xa9, 0xc9, 0x08, 0xe8, 0x95, 0x80, 0xdf, 0x94, 0xfa, 0x75, 0x8f, 0x3f, 0xa6,
        0x47, 0x07, 0xa7, 0xfc, 0xf3, 0x73, 0x17, 0xba, 0x83, 0x59, 0x3c, 0x19, 0xe6, 0x85, 0x4f, 0xa8,
        0x68, 0x6b, 0x81, 0xb2, 0x71, 0x64, 0xda, 0x8b, 0xf8, 0xeb, 0x0f, 0x4b, 0x70, 0x56, 0x9d, 0x35,
        0x1e, 0x24, 0x0e, 0x5e, 0x63, 0x58, 0xd1, 0xa2, 0x25, 0x22, 0x7c, 0x3b, 0x01, 0x21, 0x78, 0x87,
        0xd4, 0x00, 0x46, 0x57, 0x9f, 0xd3, 0x27, 0x52, 0x4c, 0x36, 0x02, 0xe7, 0xa0, 0xc4, 0xc8, 0x9e,
        0xea, 0xbf, 0x8a, 0xd2, 0x40, 0xc7, 0x38, 0xb5, 0xa3, 0xf7, 0xf2, 0xce, 0xf9, 0x61, 0x15, 0xa1,
        0xe0, 0xae, 0x5d, 0xa4, 0x9b, 0x34, 0x1a, 0x55, 0xad, 0x93, 0x32, 0x30, 0xf5, 0x8c, 0xb1, 0xe3,
        0x1d, 0xf6, 0xe2, 0x2e, 0x82, 0x66, 0xca, 0x60, 0xc0, 0x29, 0x23, 0xab, 0x0d, 0x53, 0x4e, 0x6f,
        0xd5, 0xdb, 0x37, 0x45, 0xde, 0xfd, 0x8e, 0x2f, 0x03, 0xff, 0x6a, 0x72, 0x6d, 0x6c, 0x5b, 0x51,
        0x8d, 0x1b, 0xaf, 0x92, 0xbb, 0xdd, 0xbc, 0x7f, 0x11, 0xd9, 0x5c, 0x41, 0x1f, 0x10, 0x5a, 0xd8,
        0x0a, 0xc1, 0x31, 0x88, 0xa5, 0xcd, 0x7b, 0xbd, 0x2d, 0x74, 0xd0, 0x12, 0xb8, 0xe5, 0xb4, 0xb0,
        0x89, 0x69, 0x97, 0x4a, 0x0c, 0x96, 0x77, 0x7e, 0x65, 0xb9, 0xf1, 0x09, 0xc5, 0x6e, 0xc6, 0x84,
        0x18, 0xf0, 0x7d, 0xec, 0x3a, 0xdc, 0x4d, 0x20, 0x79, 0xee, 0x5f, 0x3e, 0xd7, 0xcb, 0x39, 0x48
    };
    uint32_t val;
    val = Sbox[x >>  0 & 0xff] <<  0 | Sbox[x >>  8 & 0xff] << 8 |
          Sbox[x >> 16 & 0xff] << 16 | Sbox[x >> 24 & 0xff] << 24;
    return val;
}

inline uint32_t endianConvert(const uint32_t x) {
    return (x << 24 & 0xff000000) | (x <<  8 & 0x00ff0000) |
           (x >>  8 & 0x0000ff00) | (x >> 24 & 0x000000ff);
}

inline uint32_t LTransformation(const uint32_t x) {
    return x                  ^
           left_rotate(x,  2) ^ left_rotate(x, 10) ^
           left_rotate(x, 18) ^ left_rotate(x, 24);
}

inline uint32_t L1Transformation(const uint32_t x) {
    return x ^ left_rotate(x, 13) ^ left_rotate(x, 23);
}

inline uint32_t TTransformation(const uint32_t x) {
    return LTransformation(tauTransformation(x));
}

inline uint32_t FFunction(const uint32_t x[], const uint32_t rkey) {
    return x[0] ^ TTransformation(x[1] ^ x[2] ^ x[3] ^ rkey);
}

// key: 16 bytes
// keys: 32 * 4 bytes
void keyExpansion(const uint8_t key[], uint32_t keys[]) {
    uint32_t mk[4];
    mk[0] = key[ 0] << 24 | key[ 1] << 16 | key[ 2] << 8 | key[ 3] << 0;
    mk[1] = key[ 4] << 24 | key[ 5] << 16 | key[ 6] << 8 | key[ 7] << 0;
    mk[2] = key[ 8] << 24 | key[ 9] << 16 | key[10] << 8 | key[11] << 0;
    mk[3] = key[12] << 24 | key[13] << 16 | key[14] << 8 | key[15] << 0;
    constexpr uint32_t FK[4] = { 0xA3B1BAC6, 0x56AA3350, 0x677D9197, 0xB27022DC };
    constexpr uint32_t CK[32] = {
        0x00070e15, 0x1c232a31, 0x383f464d, 0x545b6269,
        0x70777e85, 0x8c939aa1, 0xa8afb6bd, 0xc4cbd2d9,
        0xe0e7eef5, 0xfc030a11, 0x181f262d, 0x343b4249,
        0x50575e65, 0x6c737a81, 0x888f969d, 0xa4abb2b9,
        0xc0c7ced5, 0xdce3eaf1, 0xf8ff060d, 0x141b2229,
        0x30373e45, 0x4c535a61, 0x686f767d, 0x848b9299,
        0xa0a7aeb5, 0xbcc3cad1, 0xd8dfe6ed, 0xf4fb0209,
        0x10171e25, 0x2c333a41, 0x484f565d, 0x646b7279
    };
    uint32_t k[36];
    k[0] = mk[0] ^ FK[0], k[1] = mk[1] ^ FK[1],
    k[2] = mk[2] ^ FK[2], k[3] = mk[3] ^ FK[3];

    for (size_t i = 0; i < 32; ++i)
        keys[i] = k[i + 4] = k[i] ^ L1Transformation(tauTransformation(k[i + 1] ^ k[i + 2] ^ k[i + 3] ^ CK[i]));
}

void sm4Iteration(const uint32_t plain[], const uint32_t keys[], uint32_t cipher[]) {
    uint32_t x[36];
    x[0] = endianConvert(plain[0]);
    x[1] = endianConvert(plain[1]);
    x[2] = endianConvert(plain[2]);
    x[3] = endianConvert(plain[3]);

    for (size_t i = 0; i < 32; ++i)
        x[i + 4] = FFunction(x + i, keys[i]);

    cipher[0] = endianConvert(x[35]);
    cipher[1] = endianConvert(x[34]);
    cipher[2] = endianConvert(x[33]);
    cipher[3] = endianConvert(x[32]);
}

void sm4_ecb(const void* plain, const size_t length, const void* key, void* cipher) {
    const uint32_t* plain_ = (const uint32_t*)(plain);
    uint32_t* cipher_ = (uint32_t*)(cipher);

    uint32_t keys[32];
    keyExpansion((const uint8_t*)(key), keys);

    for (size_t i = 0; i < length / 16; ++i) {
        sm4Iteration(plain_ + 4 * i, keys, cipher_ + 4 * i);
    }
}

int main(int argc, char** argv) {
    if (argc == 1) return 0;

    std::ifstream fin(argv[1]);

    fin.seekg(0, std::ios::end);
    std::string buffer;
    size_t len = fin.tellg();
    buffer.reserve(len);
    fin.seekg(0, std::ios::beg);

    if (len % 16) {
        printf("Length of plain text should be multiple of 16 bytes. \n");
        return 0;
    }

    buffer.assign((std::istreambuf_iterator<char>(fin)),
                   std::istreambuf_iterator<char>());

    fin.close();

    // 128 bit key size
    unsigned char key[16] = { 0 };

    if (argc == 3) {
        fin.open(argv[2]);
        fin.seekg(0, std::ios::beg);
        char buffer[3] = { 0 };
        for (size_t i = 0; i < 16; ++i) {
            fin.read(buffer, 2);
            key[i] = std::stoi(buffer, 0, 16);
        }
        fin.close();
    }

    std::vector<char> cipher(buffer.length(), 0);
    sm4_ecb(buffer.data(), buffer.length(), key, &cipher[0]);

    for (size_t i = 0; i < buffer.length(); ++i)
        printf("%02x", int(cipher[i]) & 0xff);
    printf("\n");

}

<|endoftext|>


<commit_before>/*
 * This file is part of the Camera Streaming Daemon
 *
 * Copyright (C) 2017  Intel Corporation. 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 <arpa/inet.h>
#include <assert.h>
#include <fcntl.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#include "log.h"
#include "socket.h"

#define DEFAULT_BUF_SIZE 1024

Socket::Socket()
    : _read_cb([](const struct buffer &buf, const struct sockaddr_in &sockaddr) {})
    , _write_buf{0, nullptr}
{
}

Socket::~Socket()
{
    free(_write_buf.data);
}

int Socket::write(const struct buffer &buf, const struct sockaddr_in &sockaddr)
{
    int r;
    r = _do_write(buf, sockaddr);
    if (r >= 0)
        return r;

    if (_write_buf.data) {
        delete[] _write_buf.data;
        log_warning("There was another packet to be sent. Dropping that packet.");
    }

    _write_buf = {buf.len, new uint8_t[buf.len]};
    memcpy(_write_buf.data, buf.data, buf.len);
    sockaddr_buf = sockaddr;
    monitor_write(true);

    return 0;
}

void Socket::set_read_callback(
    std::function<void(const struct buffer &buf, const struct sockaddr_in &sockaddr)> cb)
{
    _read_cb = cb;
}

bool Socket::_can_read()
{
    struct buffer read_buf {
        DEFAULT_BUF_SIZE, new uint8_t[DEFAULT_BUF_SIZE]
    };
    struct sockaddr_in sockaddr = {};
    int r = _do_read(read_buf, sockaddr);

    if (r == -EAGAIN) {
        log_debug("Read packet failed. Trying again.");
        return true;
    }
    if (r < 0) {
        log_debug("Read failed. Droping packet.");
        return false;
    }
    if (r > 0) {
        read_buf.len = (unsigned int)r;
        _read_cb(read_buf, sockaddr);
    }
    delete[] read_buf.data;
    return true;
}

bool Socket::_can_write()
{
    int r;

    if (!_write_buf.data)
        return false;

    r = _do_write(_write_buf, sockaddr_buf);
    if (r == -EAGAIN) {
        log_debug("Write package failed. Trying again.");
        return true;
    } else if (r < 0) {
        log_error("Write package failed. Droping packet.");
    }

    delete[] _write_buf.data;
    _write_buf = {0, nullptr};
    return false;
}

UDPSocket::UDPSocket()
{
}

UDPSocket::~UDPSocket()
{
    close();
}

void UDPSocket::close()
{
    if (_fd >= 0)
        ::close(_fd);
    _fd = -1;
}

int UDPSocket::open(bool broadcast)
{
    const int broadcast_val = 1;

    _fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (_fd == -1) {
        log_error("Could not create socket (%m)");
        return -1;
    }

    if (broadcast) {
        if (setsockopt(_fd, SOL_SOCKET, SO_BROADCAST, &broadcast_val, sizeof(broadcast_val))) {
            log_error("Error enabling broadcast in socket (%m)");
            goto fail;
        }
    }

    if (fcntl(_fd, F_SETFL, O_NONBLOCK | FASYNC) < 0) {
        log_error("Error setting socket _fd as non-blocking (%m)");
        goto fail;
    }

    log_info("Open UDP [%d]", _fd);

    monitor_read(true);
    return _fd;

fail:
    if (_fd >= 0) {
        ::close(_fd);
        _fd = -1;
    }
    return -1;
}

int UDPSocket::bind(const char *ip, unsigned long port)
{
    assert(ip);

    struct sockaddr_in sockaddr = {};

    sockaddr.sin_family = AF_INET;
    sockaddr.sin_addr.s_addr = inet_addr(ip);
    sockaddr.sin_port = htons(port);

    if (_fd < 0) {
        log_error("Trying to bind an invalid _fd");
        return -EINVAL;
    }

    if (::bind(_fd, (struct sockaddr *)&sockaddr, sizeof(sockaddr)) < 0) {
        log_error("Error binding socket (%m)");
        return -1;
    }

    log_info("Bind UDP [%d] %s:%lu", _fd, ip, port);
    return 0;
}

int UDPSocket::_do_write(const struct buffer &buf, const struct sockaddr_in &sockaddr)
{
    if (_fd < 0) {
        log_error("Trying to write to an invalid _fd");
        return -EINVAL;
    }

    if (!sockaddr.sin_port) {
        log_debug("No one ever connected to %d. No one to write for", _fd);
        return 0;
    }

    ssize_t r = ::sendto(_fd, buf.data, buf.len, 0, (struct sockaddr *)&sockaddr, sizeof(sockaddr));
    if (r == -1) {
        if (errno != EAGAIN && errno != ECONNREFUSED && errno != ENETUNREACH)
            log_error("Error sending udp packet (%m)");
        return -errno;
    };

    /* Incomplete packet, we warn and discard the rest */
    if (r != (ssize_t)buf.len) {
        log_debug("Discarding packet, incomplete write %zd but len=%u", r, buf.len);
    }

    log_debug("UDP: [%d] wrote %zd bytes", _fd, r);

    return r;
}

int UDPSocket::_do_read(const struct buffer &buf, struct sockaddr_in &sockaddr)
{
    socklen_t addrlen = sizeof(sockaddr);
    ssize_t r = ::recvfrom(_fd, buf.data, buf.len, 0, (struct sockaddr *)&sockaddr, &addrlen);
    if (r == -1) {
        if (errno != EAGAIN)
            log_error("Error reading udp packet (%m)");
        return -errno;
    }

    return r;
}
<commit_msg>Socket: Fix memory leak on errors in _can_read method<commit_after>/*
 * This file is part of the Camera Streaming Daemon
 *
 * Copyright (C) 2017  Intel Corporation. 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 <arpa/inet.h>
#include <assert.h>
#include <fcntl.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#include "log.h"
#include "socket.h"

#define DEFAULT_BUF_SIZE 1024

Socket::Socket()
    : _read_cb([](const struct buffer &buf, const struct sockaddr_in &sockaddr) {})
    , _write_buf{0, nullptr}
{
}

Socket::~Socket()
{
    free(_write_buf.data);
}

int Socket::write(const struct buffer &buf, const struct sockaddr_in &sockaddr)
{
    int r;
    r = _do_write(buf, sockaddr);
    if (r >= 0)
        return r;

    if (_write_buf.data) {
        delete[] _write_buf.data;
        log_warning("There was another packet to be sent. Dropping that packet.");
    }

    _write_buf = {buf.len, new uint8_t[buf.len]};
    memcpy(_write_buf.data, buf.data, buf.len);
    sockaddr_buf = sockaddr;
    monitor_write(true);

    return 0;
}

void Socket::set_read_callback(
    std::function<void(const struct buffer &buf, const struct sockaddr_in &sockaddr)> cb)
{
    _read_cb = cb;
}

bool Socket::_can_read()
{
    struct buffer read_buf {
        DEFAULT_BUF_SIZE, new uint8_t[DEFAULT_BUF_SIZE]
    };
    struct sockaddr_in sockaddr = {};
    int r = _do_read(read_buf, sockaddr);

    if (r == -EAGAIN) {
        log_debug("Read packet failed. Trying again.");
        goto end;
    }
    if (r < 0) {
        log_debug("Read failed. Droping packet.");
        delete[] read_buf.data;
        return false;
    }
    if (r > 0) {
        read_buf.len = (unsigned int)r;
        _read_cb(read_buf, sockaddr);
    }

end:
    delete[] read_buf.data;
    return true;
}

bool Socket::_can_write()
{
    int r;

    if (!_write_buf.data)
        return false;

    r = _do_write(_write_buf, sockaddr_buf);
    if (r == -EAGAIN) {
        log_debug("Write package failed. Trying again.");
        return true;
    } else if (r < 0) {
        log_error("Write package failed. Droping packet.");
    }

    delete[] _write_buf.data;
    _write_buf = {0, nullptr};
    return false;
}

UDPSocket::UDPSocket()
{
}

UDPSocket::~UDPSocket()
{
    close();
}

void UDPSocket::close()
{
    if (_fd >= 0)
        ::close(_fd);
    _fd = -1;
}

int UDPSocket::open(bool broadcast)
{
    const int broadcast_val = 1;

    _fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (_fd == -1) {
        log_error("Could not create socket (%m)");
        return -1;
    }

    if (broadcast) {
        if (setsockopt(_fd, SOL_SOCKET, SO_BROADCAST, &broadcast_val, sizeof(broadcast_val))) {
            log_error("Error enabling broadcast in socket (%m)");
            goto fail;
        }
    }

    if (fcntl(_fd, F_SETFL, O_NONBLOCK | FASYNC) < 0) {
        log_error("Error setting socket _fd as non-blocking (%m)");
        goto fail;
    }

    log_info("Open UDP [%d]", _fd);

    monitor_read(true);
    return _fd;

fail:
    if (_fd >= 0) {
        ::close(_fd);
        _fd = -1;
    }
    return -1;
}

int UDPSocket::bind(const char *ip, unsigned long port)
{
    assert(ip);

    struct sockaddr_in sockaddr = {};

    sockaddr.sin_family = AF_INET;
    sockaddr.sin_addr.s_addr = inet_addr(ip);
    sockaddr.sin_port = htons(port);

    if (_fd < 0) {
        log_error("Trying to bind an invalid _fd");
        return -EINVAL;
    }

    if (::bind(_fd, (struct sockaddr *)&sockaddr, sizeof(sockaddr)) < 0) {
        log_error("Error binding socket (%m)");
        return -1;
    }

    log_info("Bind UDP [%d] %s:%lu", _fd, ip, port);
    return 0;
}

int UDPSocket::_do_write(const struct buffer &buf, const struct sockaddr_in &sockaddr)
{
    if (_fd < 0) {
        log_error("Trying to write to an invalid _fd");
        return -EINVAL;
    }

    if (!sockaddr.sin_port) {
        log_debug("No one ever connected to %d. No one to write for", _fd);
        return 0;
    }

    ssize_t r = ::sendto(_fd, buf.data, buf.len, 0, (struct sockaddr *)&sockaddr, sizeof(sockaddr));
    if (r == -1) {
        if (errno != EAGAIN && errno != ECONNREFUSED && errno != ENETUNREACH)
            log_error("Error sending udp packet (%m)");
        return -errno;
    };

    /* Incomplete packet, we warn and discard the rest */
    if (r != (ssize_t)buf.len) {
        log_debug("Discarding packet, incomplete write %zd but len=%u", r, buf.len);
    }

    log_debug("UDP: [%d] wrote %zd bytes", _fd, r);

    return r;
}

int UDPSocket::_do_read(const struct buffer &buf, struct sockaddr_in &sockaddr)
{
    socklen_t addrlen = sizeof(sockaddr);
    ssize_t r = ::recvfrom(_fd, buf.data, buf.len, 0, (struct sockaddr *)&sockaddr, &addrlen);
    if (r == -1) {
        if (errno != EAGAIN)
            log_error("Error reading udp packet (%m)");
        return -errno;
    }

    return r;
}
<|endoftext|>


<commit_before>#include "solver.h"
#include "loss.h"
#include "logger.h"
#include "timer.h"

#include <cstring>
#include <stdio.h>
#include <iostream>
#include <algorithm>

using namespace BlitzL1;

using std::vector;
using std::cout;
using std::endl;
using std::min;

const int MAX_BACKTRACK = 20;
const value_t L_INCREASE_RATIO = 1.25;

void Solver::update_intercept(value_t &intercept,
                              const Loss *loss_function,
                              const Dataset *data) {
  if (!use_intercept)
    return;

  // 1-d newton method:
  vector<value_t> H;
  for (int itr = 0; itr < 10; ++itr) {
    value_t grad = sum_vector(theta);
    loss_function->compute_H(H, theta, aux_dual, data);
    value_t hess = sum_vector(H);
    value_t delta = -grad / hess;
    intercept += delta;
    loss_function->apply_intercept_update(delta, theta, aux_dual, data);

    if (std::abs(grad) <= 1e-14)
      break;
  }
}

Loss* get_loss_function(const char* loss_type) {
  if (strcmp(loss_type, "logistic") == 0)
    return new LogisticLoss();
  else if (strcmp(loss_type, "squared") == 0)
    return new SquaredLoss();
  else
    throw loss_type;
}

value_t Solver::compute_lambda_max(Dataset *data, const char* loss_type) {

  Loss* loss_function = get_loss_function(loss_type);
  index_t d = data->get_d();
  vector<value_t> x(d, 0.0);
  value_t intercept = 0.0;
  loss_function->compute_dual_points(
                    theta, aux_dual, &x[0], intercept, data);
  update_intercept(intercept, loss_function, data);

  value_t lambda_max = 0.0;
  for (index_t j = 0; j < d; ++j) {
    value_t ip = data->get_column(j)->inner_product(theta);
    value_t val = std::abs(ip);
    if (val > lambda_max)
      lambda_max = val;
  }

  return lambda_max;
}

void Solver::run_prox_newton_iteration(value_t *x,
                                       value_t &intercept,
                                       value_t lambda,
                                       const Loss *loss_function,
                                       const Dataset *data) {

  // Initialize iteration:
  index_t n = data->get_n();
  vector<value_t> Delta_x(working_set_size, 0.0);
  vector<value_t> A_Delta_x(n, 0.0);
  value_t Delta_intercept = 0.0;

  // Set up gradient and hessian values:
  vector<value_t> H;
  loss_function->compute_H(H, theta, aux_dual, data);
  vector<value_t> grad_cache(working_set_size, 0.0);
  vector<value_t> hess_cache(working_set_size, 0.0);
  for (index_t ind = 0; ind < working_set_size; ++ind) {
    index_t j = prioritized_features[ind];
    grad_cache[ind] = data->get_column(j)->inner_product(theta);
    hess_cache[ind] = data->get_column(j)->h_norm_sq(H);
  }

  // Cache values for updating intercept:
  value_t sum_theta = sum_vector(theta);
  value_t sum_H = sum_vector(H);

  // Approximately solve for newton direction:
  for (int cd_itr = 0; cd_itr < 10; ++cd_itr) {

    for (index_t ind = 0; ind < working_set_size; ++ind) {
      index_t j = prioritized_features[ind];
      const Column *col = data->get_column(j);

      // Compute hessian and gradient for coordinate j:
      value_t hess = hess_cache[ind];
      if (hess == 0.0)
        continue;

      value_t grad = grad_cache[ind]
                   + col->weighted_inner_product(A_Delta_x, H);

      // Apply coordinate descent update:
      value_t old_value = x[j] + Delta_x[ind];
      value_t proposal = old_value - grad / hess;
      value_t new_value = soft_threshold(proposal, lambda / hess);
      value_t diff = new_value - old_value;

      if (diff == 0.0)
        continue;

      Delta_x[ind] = new_value - x[j];
      col->add_multiple(A_Delta_x, diff);
    }

    // Update intercept:
    if (use_intercept) {
      value_t ip = inner_product(H, A_Delta_x);
      value_t diff = -(ip + sum_theta)/sum_H;
      Delta_intercept += diff;
      add_scaler(A_Delta_x, diff);
    }
  }

  // Apply update:
  for (index_t ind = 0; ind < working_set_size; ++ind) {
    index_t j = prioritized_features[ind];
    x[j] += Delta_x[ind];
  }
  intercept += Delta_intercept;
  loss_function->compute_dual_points(
                    theta, aux_dual, x, intercept, data);
}

value_t Solver::priority_norm_j(index_t j, const Dataset* data) {
  if (use_intercept)
    return data->get_column(j)->l2_norm_centered();
  else
    return data->get_column(j)->l2_norm();
}

value_t Solver::compute_alpha(const Dataset* data, value_t lambda, value_t theta_scale) {
  value_t best_alpha = 1.0;
  index_t d = data->get_d();
  for (index_t j = 0; j < d; ++j) {
    value_t norm = priority_norm_j(j, data);
    if (norm <= 0.0)
      continue;

    value_t l = ATphi[j];
    value_t r = theta_scale * ATtheta[j];
    if (std::abs(r) <= lambda)
      continue;

    value_t alpha;
    if (r >= 0)
      alpha = (lambda - l)/(r - l);
    else
      alpha = (-lambda - l)/(r - l);

    if (alpha < best_alpha)
      best_alpha = alpha;
  }
  return best_alpha;
}

void Solver::update_phi(value_t alpha, value_t theta_scale) {
  for (size_t j = 0; j < ATphi.size(); ++j)
    ATphi[j] = (1 - alpha) * ATphi[j] + alpha * theta_scale * ATtheta[j];
  for (size_t i = 0; i < phi.size(); ++i)
    phi[i] = (1 - alpha) * phi[i] + alpha * theta_scale * theta[i];
}

void Solver::prioritize_features(const Dataset *data, value_t lambda) {
  index_t d = data->get_d();
  feature_priorities.resize(d);
  for (index_t j = 0; j < d; ++j) {
    value_t AjTphi = ATphi[j];
    value_t norm = priority_norm_j(j, data);
    value_t priority_value = (lambda - std::abs(AjTphi)) / norm;
    feature_priorities[j] = priority_value;
  }
  IndirectComparator cmp(feature_priorities);
  sort(
    prioritized_features.begin(),
    prioritized_features.end(),
    cmp);
}

void Solver::solve(const Dataset *data,
                   value_t lambda,
                   const char *loss_type,
                   value_t* x,
                   value_t &intercept,
                   char *solution_status,
                   value_t &primal_obj,
                   value_t &duality_gap,
                   int &num_iterations,
                   const char *log_directory) {

  Timer timer;
  Logger logger(log_directory);

  num_iterations = 0;
  Loss* loss_function = get_loss_function(loss_type);
  index_t n = data->get_n();
  index_t d = data->get_d();
  working_set_size = 0;
  prioritized_features.resize(d);
  for (index_t j = 0; j < d; ++j)
    prioritized_features[j] = j;

  // Initialize dual points:
  loss_function->compute_dual_points(
                    theta, aux_dual, x, intercept, data);
  phi.assign(n, 0.0);
  ATphi.assign(d, 0.0);
  ATtheta.assign(d, 0.0);

  // Initialize objective values:
  value_t primal_loss = loss_function->primal_loss(theta, aux_dual);
  primal_obj = primal_loss + lambda * l1_norm(x, d);

  // Main Blitz loop:
  while (++num_iterations) {
    value_t primal_obj_last = primal_obj;

    // Update intercept:
    update_intercept(intercept, loss_function, data);

    // Update ATtheta:
    for (index_t j = 0; j < d; ++j) {
      ATtheta[j] = data->get_column(j)->inner_product(theta);
    }

    // Compute theta scale:
    value_t max_grad_working_set = 0.0;
    for (index_t ind = 0; ind < working_set_size; ++ind) {
      index_t j = prioritized_features[ind];
      value_t AjTtheta = data->get_column(j)->inner_product(theta);
      if (std::abs(AjTtheta) > max_grad_working_set)
        max_grad_working_set = std::abs(AjTtheta);
    }
    value_t theta_scale = 1.0;
    if (max_grad_working_set > lambda)
      theta_scale = lambda / max_grad_working_set;

    // Update phi:
    value_t alpha = compute_alpha(data, lambda, theta_scale);
    update_phi(alpha, theta_scale);
    value_t dual_obj = loss_function->dual_obj(phi, data);

    // Prioritize features:
    prioritize_features(data, lambda);

    // Determine working set size:
    working_set_size = 2 * l0_norm(x, d);
    if (working_set_size < 100)
      working_set_size = 100;
    if (working_set_size > d)
      working_set_size = d;

    // Solve subproblem:
    for (int prox_newt_itr = 0; prox_newt_itr < 5; prox_newt_itr++){
      run_prox_newton_iteration(x, intercept, lambda, loss_function, data);
      update_intercept(intercept, loss_function, data);
    }

    primal_loss = loss_function->primal_loss(theta, aux_dual);
    primal_obj = primal_loss + lambda * l1_norm(x, d);
    duality_gap = primal_obj - dual_obj;

    // Print/record results:
    double elapsed_time = timer.elapsed_time();
    timer.pause_timing();
    if (verbose)
      cout << "Time: " << elapsed_time
           << " Objective: " << primal_obj
           << " alpha: " << alpha
           << " |C|: " << working_set_size
           << " Duality gap: " << duality_gap << endl;

    logger.log_point(elapsed_time, primal_obj);
    timer.continue_timing();

    // Test stopping conditions:
    if ((duality_gap / std::abs(dual_obj) < tolerance) &&
        (elapsed_time > min_time)) {
      sprintf(solution_status, "reached stopping tolerance");
      break;
    }

    if (primal_obj >= primal_obj_last) {
      sprintf(solution_status, "reached machine precision");
      break;
    }

    if (elapsed_time > max_time) {
      sprintf(solution_status, "reached time limit");
      break;
    }
  }
}

<commit_msg>removed unused constants<commit_after>#include "solver.h"
#include "loss.h"
#include "logger.h"
#include "timer.h"

#include <cstring>
#include <stdio.h>
#include <iostream>
#include <algorithm>

using namespace BlitzL1;

using std::vector;
using std::cout;
using std::endl;
using std::min;

void Solver::update_intercept(value_t &intercept,
                              const Loss *loss_function,
                              const Dataset *data) {
  if (!use_intercept)
    return;

  // 1-d newton method:
  vector<value_t> H;
  for (int itr = 0; itr < 10; ++itr) {
    value_t grad = sum_vector(theta);
    loss_function->compute_H(H, theta, aux_dual, data);
    value_t hess = sum_vector(H);
    value_t delta = -grad / hess;
    intercept += delta;
    loss_function->apply_intercept_update(delta, theta, aux_dual, data);

    if (std::abs(grad) <= 1e-14)
      break;
  }
}

Loss* get_loss_function(const char* loss_type) {
  if (strcmp(loss_type, "logistic") == 0)
    return new LogisticLoss();
  else if (strcmp(loss_type, "squared") == 0)
    return new SquaredLoss();
  else
    throw loss_type;
}

value_t Solver::compute_lambda_max(Dataset *data, const char* loss_type) {

  Loss* loss_function = get_loss_function(loss_type);
  index_t d = data->get_d();
  vector<value_t> x(d, 0.0);
  value_t intercept = 0.0;
  loss_function->compute_dual_points(
                    theta, aux_dual, &x[0], intercept, data);
  update_intercept(intercept, loss_function, data);

  value_t lambda_max = 0.0;
  for (index_t j = 0; j < d; ++j) {
    value_t ip = data->get_column(j)->inner_product(theta);
    value_t val = std::abs(ip);
    if (val > lambda_max)
      lambda_max = val;
  }

  return lambda_max;
}

void Solver::run_prox_newton_iteration(value_t *x,
                                       value_t &intercept,
                                       value_t lambda,
                                       const Loss *loss_function,
                                       const Dataset *data) {

  // Initialize iteration:
  index_t n = data->get_n();
  vector<value_t> Delta_x(working_set_size, 0.0);
  vector<value_t> A_Delta_x(n, 0.0);
  value_t Delta_intercept = 0.0;

  // Set up gradient and hessian values:
  vector<value_t> H;
  loss_function->compute_H(H, theta, aux_dual, data);
  vector<value_t> grad_cache(working_set_size, 0.0);
  vector<value_t> hess_cache(working_set_size, 0.0);
  for (index_t ind = 0; ind < working_set_size; ++ind) {
    index_t j = prioritized_features[ind];
    grad_cache[ind] = data->get_column(j)->inner_product(theta);
    hess_cache[ind] = data->get_column(j)->h_norm_sq(H);
  }

  // Cache values for updating intercept:
  value_t sum_theta = sum_vector(theta);
  value_t sum_H = sum_vector(H);

  // Approximately solve for newton direction:
  for (int cd_itr = 0; cd_itr < 10; ++cd_itr) {

    for (index_t ind = 0; ind < working_set_size; ++ind) {
      index_t j = prioritized_features[ind];
      const Column *col = data->get_column(j);

      // Compute hessian and gradient for coordinate j:
      value_t hess = hess_cache[ind];
      if (hess == 0.0)
        continue;

      value_t grad = grad_cache[ind]
                   + col->weighted_inner_product(A_Delta_x, H);

      // Apply coordinate descent update:
      value_t old_value = x[j] + Delta_x[ind];
      value_t proposal = old_value - grad / hess;
      value_t new_value = soft_threshold(proposal, lambda / hess);
      value_t diff = new_value - old_value;

      if (diff == 0.0)
        continue;

      Delta_x[ind] = new_value - x[j];
      col->add_multiple(A_Delta_x, diff);
    }

    // Update intercept:
    if (use_intercept) {
      value_t ip = inner_product(H, A_Delta_x);
      value_t diff = -(ip + sum_theta)/sum_H;
      Delta_intercept += diff;
      add_scaler(A_Delta_x, diff);
    }
  }

  // Apply update:
  for (index_t ind = 0; ind < working_set_size; ++ind) {
    index_t j = prioritized_features[ind];
    x[j] += Delta_x[ind];
  }
  intercept += Delta_intercept;
  loss_function->compute_dual_points(
                    theta, aux_dual, x, intercept, data);
}

value_t Solver::priority_norm_j(index_t j, const Dataset* data) {
  if (use_intercept)
    return data->get_column(j)->l2_norm_centered();
  else
    return data->get_column(j)->l2_norm();
}

value_t Solver::compute_alpha(const Dataset* data, value_t lambda, value_t theta_scale) {
  value_t best_alpha = 1.0;
  index_t d = data->get_d();
  for (index_t j = 0; j < d; ++j) {
    value_t norm = priority_norm_j(j, data);
    if (norm <= 0.0)
      continue;

    value_t l = ATphi[j];
    value_t r = theta_scale * ATtheta[j];
    if (std::abs(r) <= lambda)
      continue;

    value_t alpha;
    if (r >= 0)
      alpha = (lambda - l)/(r - l);
    else
      alpha = (-lambda - l)/(r - l);

    if (alpha < best_alpha)
      best_alpha = alpha;
  }
  return best_alpha;
}

void Solver::update_phi(value_t alpha, value_t theta_scale) {
  for (size_t j = 0; j < ATphi.size(); ++j)
    ATphi[j] = (1 - alpha) * ATphi[j] + alpha * theta_scale * ATtheta[j];
  for (size_t i = 0; i < phi.size(); ++i)
    phi[i] = (1 - alpha) * phi[i] + alpha * theta_scale * theta[i];
}

void Solver::prioritize_features(const Dataset *data, value_t lambda) {
  index_t d = data->get_d();
  feature_priorities.resize(d);
  for (index_t j = 0; j < d; ++j) {
    value_t AjTphi = ATphi[j];
    value_t norm = priority_norm_j(j, data);
    value_t priority_value = (lambda - std::abs(AjTphi)) / norm;
    feature_priorities[j] = priority_value;
  }
  IndirectComparator cmp(feature_priorities);
  sort(
    prioritized_features.begin(),
    prioritized_features.end(),
    cmp);
}

void Solver::solve(const Dataset *data,
                   value_t lambda,
                   const char *loss_type,
                   value_t* x,
                   value_t &intercept,
                   char *solution_status,
                   value_t &primal_obj,
                   value_t &duality_gap,
                   int &num_iterations,
                   const char *log_directory) {

  Timer timer;
  Logger logger(log_directory);

  num_iterations = 0;
  Loss* loss_function = get_loss_function(loss_type);
  index_t n = data->get_n();
  index_t d = data->get_d();
  working_set_size = 0;
  prioritized_features.resize(d);
  for (index_t j = 0; j < d; ++j)
    prioritized_features[j] = j;

  // Initialize dual points:
  loss_function->compute_dual_points(
                    theta, aux_dual, x, intercept, data);
  phi.assign(n, 0.0);
  ATphi.assign(d, 0.0);
  ATtheta.assign(d, 0.0);

  // Initialize objective values:
  value_t primal_loss = loss_function->primal_loss(theta, aux_dual);
  primal_obj = primal_loss + lambda * l1_norm(x, d);

  // Main Blitz loop:
  while (++num_iterations) {
    value_t primal_obj_last = primal_obj;

    // Update intercept:
    update_intercept(intercept, loss_function, data);

    // Update ATtheta:
    for (index_t j = 0; j < d; ++j) {
      ATtheta[j] = data->get_column(j)->inner_product(theta);
    }

    // Compute theta scale:
    value_t max_grad_working_set = 0.0;
    for (index_t ind = 0; ind < working_set_size; ++ind) {
      index_t j = prioritized_features[ind];
      value_t AjTtheta = data->get_column(j)->inner_product(theta);
      if (std::abs(AjTtheta) > max_grad_working_set)
        max_grad_working_set = std::abs(AjTtheta);
    }
    value_t theta_scale = 1.0;
    if (max_grad_working_set > lambda)
      theta_scale = lambda / max_grad_working_set;

    // Update phi:
    value_t alpha = compute_alpha(data, lambda, theta_scale);
    update_phi(alpha, theta_scale);
    value_t dual_obj = loss_function->dual_obj(phi, data);

    // Prioritize features:
    prioritize_features(data, lambda);

    // Determine working set size:
    working_set_size = 2 * l0_norm(x, d);
    if (working_set_size < 100)
      working_set_size = 100;
    if (working_set_size > d)
      working_set_size = d;

    // Solve subproblem:
    for (int prox_newt_itr = 0; prox_newt_itr < 5; prox_newt_itr++){
      run_prox_newton_iteration(x, intercept, lambda, loss_function, data);
      update_intercept(intercept, loss_function, data);
    }

    primal_loss = loss_function->primal_loss(theta, aux_dual);
    primal_obj = primal_loss + lambda * l1_norm(x, d);
    duality_gap = primal_obj - dual_obj;

    // Print/record results:
    double elapsed_time = timer.elapsed_time();
    timer.pause_timing();
    if (verbose)
      cout << "Time: " << elapsed_time
           << " Objective: " << primal_obj
           << " alpha: " << alpha
           << " |C|: " << working_set_size
           << " Duality gap: " << duality_gap << endl;

    logger.log_point(elapsed_time, primal_obj);
    timer.continue_timing();

    // Test stopping conditions:
    if ((duality_gap / std::abs(dual_obj) < tolerance) &&
        (elapsed_time > min_time)) {
      sprintf(solution_status, "reached stopping tolerance");
      break;
    }

    if (primal_obj >= primal_obj_last) {
      sprintf(solution_status, "reached machine precision");
      break;
    }

    if (elapsed_time > max_time) {
      sprintf(solution_status, "reached time limit");
      break;
    }
  }
}

<|endoftext|>


<commit_before>#include "eval_state.h"

#include <fstream>
#include <iostream>
#include <random>
#include <string>
#include <unordered_map>
#include <vector>
#include <cassert>
#include <cstdlib>

#include "player.h"
#include "zs_wrap.h"

std::unique_ptr<std::ofstream> dout;
std::unordered_map<int, double> memoization;

// Computes the probability that the first player will win
double eval_state(const struct game_state &s) {
  int state_rep = game_state_rep(s);

  if (memoization.count(state_rep)) {
    // already calculated; return memoized result
    return memoization[state_rep];
  }

  // Check if game is already over (that is, whether either player has
  // negative HP)
  if ((s.p1.hp <= 0) || (s.p2.hp <= 0)) {
    // the player with less HP loses
    double result;
    if (s.p1.hp < s.p2.hp) {
      result = 0;
    } else if (s.p1.hp > s.p2.hp) {
      result = 1;
    } else {
      result = 0.5;
    }
    return result;
  }

  // Autoload if neither player has a bullet
  if (s.p1.bullets == 0 && s.p2.bullets == 0) {
    double res = eval_state(act(s, -1, -1));
    return res;
  }

  // Otherwise we need to do the recursive step; we call eval_state on
  // every position which can be reached from the current position and then
  // call zs_solve on it

  // we have s.p1.bullets + 2 columns and s.p2.bullets + 2 rows
  double *payoffMatrix = reinterpret_cast<double *>(malloc((s.p1.bullets + 2) *
                                                           (s.p2.bullets + 2) *
                                                           sizeof(double)));
  for (int p2Move = -1; p2Move <= s.p2.bullets; p2Move++) {
    int offset = (p2Move + 1) * (s.p1.bullets + 2);
    for (int p1Move = -1; p1Move <= s.p1.bullets; p1Move++) {
      payoffMatrix[offset + (p1Move + 1)] = 1 +
        eval_state(act(s, p1Move, p2Move));
      // adding offset of 1 to avoid some rounding problems
    }
  }

  double *results = zs_solve(s.p2.bullets + 2,
                             s.p1.bullets + 2,
                             payoffMatrix);
  free(payoffMatrix);
  if (results == NULL) {
    std::cerr << "Unsolved position?!" << std::endl;
    // It'd be weird, because zero-sum games are always solvable
    exit(-1);
  }
  results[0] -= 1;

  if (results[0] < 0.00000001) {
          results[0] = 0;
  }
  if (results[0] > 1) {
          results[0] = 1;
  }
  double value = results[0];
  free(results);
  if (dout) {
    dout->write((const char *)&state_rep, sizeof(state_rep));
    dout->write((const char *)&value, sizeof(value));
  }
  memoization[state_rep] = value;
  return value;
}

namespace {
  std::random_device rd;
  std::default_random_engine generator(rd());
  std::uniform_real_distribution<double> distribution(0.0, 1.0);
}

// Prints a move generated according to spe distribution
void get_strat(const struct game_state &s) {
  using std::cout;
  using std::endl;
  // Autoload
  if (s.p1.bullets == 0 && s.p2.bullets == 0) {
    cout << "Autoload! ";
    get_strat(act(s, -1, -1));
    return;
  }

  // Check if game is already over (that is, whether either player has
  // negative HP)
  if ((s.p1.hp <= 0) || (s.p2.hp <= 0)) {
    cout << "Reload (Game is already over)" << endl;
    return;
  }

  double *payoffMatrix = reinterpret_cast<double *>(malloc((s.p1.bullets + 2) *
                                                           (s.p2.bullets + 2) *
                                                           sizeof(double)));
  for (int p2Move = -1; p2Move <= s.p2.bullets; p2Move++) {
    int offset = (p2Move + 1) * (s.p1.bullets + 2);
    for (int p1Move = -1; p1Move <= s.p1.bullets; p1Move++) {
      payoffMatrix[offset + (p1Move + 1)] = 1 +
        eval_state(act(s, p1Move, p2Move));
    }
  }

  double *results = zs_solve(s.p2.bullets + 2, s.p1.bullets + 2, payoffMatrix);
  // print out actual strategy
  /*
  cout << "Mixed-strategy subgame perfect equilibrium:" << endl;
  cout << "Probability to win: " << results[0] - 1 << endl;
  cout << "Reload with probability " << results[1] << endl;
  cout << "Block with probability " << results[2] << endl;
  for (int i = 1; i <= s.p1.bullets; ++i) {
    cout << "Fire " << i << " with probability " << results[2+i] << endl;
  }
  */
  double move = distribution(generator);
  for (int i = 0; i < s.p1.bullets + 2; ++i) {
    move -= results[1 + i];
    if (move <= 0 || i == s.p1.bullets + 1) {
      // choose the action at i-1
      if (i == 0) {
        cout << "Reload" << endl;
      } else if (i == 1) {
        cout << "Block" << endl;
      } else {
        cout << "Fire " << i-1 << endl;
      }
      break;
    }
  }
  free(payoffMatrix);
  free(results);
}

// Prints doubletime strat at a given position
void get_doubletime_strat(const struct game_state &s) {
  using std::cout;
  using std::endl;
  // Check whether the position is actually memoized
  if (!(memoization.count(game_state_rep(s)) || memoization.count(game_state_rev_rep(s)))) {
    cout << "Unreachable state" << endl;
    return;
  }

  double *payoffMatrix = reinterpret_cast<double *>(malloc((s.p1.bullets + 2) *
                                                           (s.p2.bullets + 2) *
                                                           sizeof(double)));
  for (int p2Move = -1; p2Move <= s.p2.bullets; p2Move++) {
    int offset = (p2Move + 1) * (s.p1.bullets + 2);
    for (int p1Move = -1; p1Move <= s.p1.bullets; p1Move++) {
      payoffMatrix[offset + (p1Move + 1)] = 1 +
        eval_state(act(s, p1Move, p2Move));
    }
  }

  double *results = zs_solve(s.p2.bullets + 2, s.p1.bullets + 2, payoffMatrix);
  if (results[0] - 1 < 0.000001) {
    cout << "Note: You are 0% to win under optimal play" << endl;
  }

  double move = distribution(generator);
  for (int i = 0; i < s.p1.bullets + 2; ++i) {
    move -= results[1 + i];
    if (move <= 0 || i == s.p1.bullets + 1) {
      cout << "\tAction: ";
      // choose the action at i-1
      if (i == 0) {
        cout << "Reload" << endl;
      } else if (i == 1) {
        cout << "Block" << endl;
      } else {
        cout << "Fire " << i-1 << endl;
      }
      cout << "\tDoubletimes" << endl;
      cout << "\tReload: ";
      get_strat(act(s, i-1, -1));
      cout << "\tBlock:  ";
      get_strat(act(s, i-1, 0));
      for (int j = 1; j <= s.p2.bullets; ++j) {
        cout << "\tFire " << j << ": ";
        get_strat(act(s, i-1, j));
      }
      break;
    }
  }
  free(payoffMatrix);
  free(results);
}
<commit_msg>Update eval_state.cc<commit_after>#include "eval_state.h"

#include <fstream>
#include <iostream>
#include <random>
#include <string>
#include <unordered_map>
#include <vector>
#include <cassert>
#include <cstdlib>

#include "player.h"
#include "zs_wrap.h"

std::unique_ptr<std::ofstream> dout;
std::unordered_map<int, double> memoization;

// Computes the probability that the first player will win
double eval_state(const struct game_state &s) {
  int state_rep = game_state_rep(s);

  if (memoization.count(state_rep)) {
    // already calculated; return memoized result
    return memoization[state_rep];
  }

  // Check if game is already over (that is, whether either player has
  // negative HP)
  if ((s.p1.hp <= 0) || (s.p2.hp <= 0)) {
    // the player with less HP loses
    double result;
    if (s.p1.hp < s.p2.hp) {
      result = 0;
    } else if (s.p1.hp > s.p2.hp) {
      result = 1;
    } else {
      result = 0.5;
    }
    return result;
  }

  // Autoload if neither player has a bullet
  if (s.p1.bullets == 0 && s.p2.bullets == 0) {
    double res = eval_state(act(s, -1, -1));
    return res;
  }

  // Otherwise we need to do the recursive step; we call eval_state on
  // every position which can be reached from the current position and then
  // call zs_solve on it

  // we have s.p1.bullets + 2 columns and s.p2.bullets + 2 rows
  double *payoffMatrix = reinterpret_cast<double *>(malloc((s.p1.bullets + 2) *
                                                           (s.p2.bullets + 2) *
                                                           sizeof(double)));
  for (int p2Move = -1; p2Move <= s.p2.bullets; p2Move++) {
    int offset = (p2Move + 1) * (s.p1.bullets + 2);
    for (int p1Move = -1; p1Move <= s.p1.bullets; p1Move++) {
      payoffMatrix[offset + (p1Move + 1)] = 1 +
        eval_state(act(s, p1Move, p2Move));
      // adding offset of 1 to avoid some rounding problems
    }
  }

  double *results = zs_solve(s.p2.bullets + 2,
                             s.p1.bullets + 2,
                             payoffMatrix);
  free(payoffMatrix);
  if (results == NULL) {
    std::cerr << "Unsolved position?!" << std::endl;
    // It'd be weird, because zero-sum games are always solvable
    exit(-1);
  }
  results[0] -= 1;

  if (results[0] < 0.00000001) {
          results[0] = 0;
  }
  if (results[0] > 1) {
          results[0] = 1;
  }
  double value = results[0];
  free(results);
  if (dout) {
    dout->write((const char *)&state_rep, sizeof(state_rep));
    dout->write((const char *)&value, sizeof(value));
  }
  memoization[state_rep] = value;
  return value;
}

namespace {
  std::random_device rd;
  std::default_random_engine generator(rd());
  std::uniform_real_distribution<double> distribution(0.0, 1.0);
}

// Prints a move generated according to spe distribution
void get_strat(const struct game_state &s) {
  using std::cout;
  using std::endl;
  // Autoload
  if (s.p1.bullets == 0 && s.p2.bullets == 0) {
    cout << "Autoload! ";
    get_strat(act(s, -1, -1));
    return;
  }

  // Check if game is already over (that is, whether either player has
  // negative HP)
  if ((s.p1.hp <= 0) || (s.p2.hp <= 0)) {
    cout << "Reload (Game is already over)" << endl;
    return;
  }

  double *payoffMatrix = reinterpret_cast<double *>(malloc((s.p1.bullets + 2) *
                                                           (s.p2.bullets + 2) *
                                                           sizeof(double)));
  for (int p2Move = -1; p2Move <= s.p2.bullets; p2Move++) {
    int offset = (p2Move + 1) * (s.p1.bullets + 2);
    for (int p1Move = -1; p1Move <= s.p1.bullets; p1Move++) {
      payoffMatrix[offset + (p1Move + 1)] = 1 +
        eval_state(act(s, p1Move, p2Move));
    }
  }

  double *results = zs_solve(s.p2.bullets + 2, s.p1.bullets + 2, payoffMatrix);
  // print out actual strategy
  /*
  cout << "Mixed-strategy subgame perfect equilibrium:" << endl;
  cout << "Probability to win: " << results[0] - 1 << endl;
  cout << "Reload with probability " << results[1] << endl;
  cout << "Block with probability " << results[2] << endl;
  for (int i = 1; i <= s.p1.bullets; ++i) {
    cout << "Fire " << i << " with probability " << results[2+i] << endl;
  }
  */
  double move = distribution(generator);
  for (int i = 0; i < s.p1.bullets + 2; ++i) {
    move -= results[1 + i];
    if (move <= 0 || i == s.p1.bullets + 1) {
      // choose the action at i-1
      if (i == 0) {
        cout << "Reload" << endl;
      } else if (i == 1) {
        cout << "Block" << endl;
      } else {
        cout << "Fire " << i-1 << endl;
      }
      break;
    }
  }
  free(payoffMatrix);
  free(results);
}

// Prints doubletime strat at a given position
void get_doubletime_strat(const struct game_state &s) {
  using std::cout;
  using std::endl;
  // Check whether the position is actually memoized
  if (!memoization.count(game_state_rep(s))) {
    cout << "Unreachable state" << endl;
    return;
  }

  double *payoffMatrix = reinterpret_cast<double *>(malloc((s.p1.bullets + 2) *
                                                           (s.p2.bullets + 2) *
                                                           sizeof(double)));
  for (int p2Move = -1; p2Move <= s.p2.bullets; p2Move++) {
    int offset = (p2Move + 1) * (s.p1.bullets + 2);
    for (int p1Move = -1; p1Move <= s.p1.bullets; p1Move++) {
      payoffMatrix[offset + (p1Move + 1)] = 1 +
        eval_state(act(s, p1Move, p2Move));
    }
  }

  double *results = zs_solve(s.p2.bullets + 2, s.p1.bullets + 2, payoffMatrix);
  if (results[0] - 1 < 0.000001) {
    cout << "Note: You are 0% to win under optimal play" << endl;
  }

  double move = distribution(generator);
  for (int i = 0; i < s.p1.bullets + 2; ++i) {
    move -= results[1 + i];
    if (move <= 0 || i == s.p1.bullets + 1) {
      cout << "\tAction: ";
      // choose the action at i-1
      if (i == 0) {
        cout << "Reload" << endl;
      } else if (i == 1) {
        cout << "Block" << endl;
      } else {
        cout << "Fire " << i-1 << endl;
      }
      cout << "\tDoubletimes" << endl;
      cout << "\tReload: ";
      get_strat(act(s, i-1, -1));
      cout << "\tBlock:  ";
      get_strat(act(s, i-1, 0));
      for (int j = 1; j <= s.p2.bullets; ++j) {
        cout << "\tFire " << j << ": ";
        get_strat(act(s, i-1, j));
      }
      break;
    }
  }
  free(payoffMatrix);
  free(results);
}
<|endoftext|>


<commit_before>// Copyright(c) 2015-present, Gabi Melman & spdlog contributors.
// Distributed under the MIT License (http://opensource.org/licenses/MIT)

#ifndef SPDLOG_COMPILED_LIB
#error Please define SPDLOG_COMPILED_LIB to compile this file.
#endif

#include "spdlog/spdlog-inl.h"
#include "spdlog/common-inl.h"
#include "spdlog/details/backtracer-inl.h"
#include "spdlog/details/registry-inl.h"
#include "spdlog/details/os-inl.h"
#include "spdlog/details/pattern_formatter-inl.h"
#include "spdlog/details/log_msg-inl.h"
#include "spdlog/details/log_msg_buffer-inl.h"
#include "spdlog/logger-inl.h"
#include "spdlog/sinks/sink-inl.h"

template spdlog::logger::logger(std::string name, sinks_init_list::iterator begin, sinks_init_list::iterator end);
<commit_msg>Update spdlog.cpp<commit_after>// Copyright(c) 2015-present, Gabi Melman & spdlog contributors.
// Distributed under the MIT License (http://opensource.org/licenses/MIT)

#ifndef SPDLOG_COMPILED_LIB
#error Please define SPDLOG_COMPILED_LIB to compile this file.
#endif

#include "spdlog/spdlog-inl.h"
#include "spdlog/common-inl.h"
#include "spdlog/details/backtracer-inl.h"
#include "spdlog/details/registry-inl.h"
#include "spdlog/details/os-inl.h"
#include "spdlog/details/pattern_formatter-inl.h"
#include "spdlog/details/log_msg-inl.h"
#include "spdlog/details/log_msg_buffer-inl.h"
#include "spdlog/logger-inl.h"
#include "spdlog/sinks/sink-inl.h"

// template instantiate logger with sinks init list
template spdlog::logger::logger(std::string name, sinks_init_list::iterator begin, sinks_init_list::iterator end);
<|endoftext|>


<commit_before>///
/// @file  IteratorHelper.cpp
///        Functions used to calculate the next start and stop
///        numbers for primesieve::iterator.
///
/// Copyright (C) 2022 Kim Walisch, <kim.walisch@gmail.com>
///
/// This file is distributed under the BSD License. See the COPYING
/// file in the top level directory.
///

#include <primesieve/config.hpp>
#include <primesieve/IteratorHelper.hpp>
#include <primesieve/PrimeGenerator.hpp>
#include <primesieve/pmath.hpp>

#include <stdint.h>
#include <algorithm>
#include <cmath>
#include <limits>

using namespace primesieve;

namespace {

uint64_t getNextDist(uint64_t start, uint64_t dist)
{
  uint64_t minDist = (uint64_t) std::sqrt(start);
  uint64_t maxDist = 1ull << 60;

  dist *= 4;
  dist = std::max(PrimeGenerator::maxCachedPrime(), dist);
  return inBetween(minDist, dist, maxDist);
}

uint64_t getPrevDist(uint64_t stop, uint64_t dist)
{
  double x = std::max(10.0, (double) stop);
  uint64_t logx = (uint64_t) std::log(x);
  uint64_t minDist = (config::MIN_CACHE_ITERATOR / sizeof(uint64_t)) * logx;
  uint64_t maxDist = (config::MAX_CACHE_ITERATOR / sizeof(uint64_t)) * logx;
  uint64_t tinyDist = PrimeGenerator::maxCachedPrime() * 4;
  uint64_t defaultDist = (uint64_t) (std::sqrt(stop) * 2);

  dist *= 4;
  minDist = inBetween(tinyDist, dist, minDist);
  return inBetween(minDist, defaultDist, maxDist);
}

} // namespace

namespace primesieve {

void IteratorHelper::next(uint64_t* start,
                          uint64_t* stop,
                          uint64_t stopHint,
                          uint64_t* dist)
{
  *start = checkedAdd(*stop, 1);
  *dist = getNextDist(*start, *dist);

  if (stopHint >= *start &&
      stopHint < std::numeric_limits<uint64_t>::max())
  {
    // For primesieve::iterator it is advantageous to buffer
    // slightly more primes than the stopHint since the
    // stopHint is often not 100% accurate and the user
    // might iterate over a few primes > stopHint.
    *stop = checkedAdd(stopHint, maxPrimeGap(stopHint));
  }
  else
  {
    // In case the user has used the default stopHint=UINT64_MAX
    // we take a conservative approach and only buffer a
    // small number of primes. If the user uses more primes
    // than we have buffered, then we will increase the sieving
    // distance and buffer more primes (than last time).
    *stop = checkedAdd(*start, *dist);
  }
}

void IteratorHelper::prev(uint64_t* start,
                          uint64_t* stop,
                          uint64_t stopHint,
                          uint64_t* dist)
{
  *stop = checkedSub(*start, 1);
  *dist = getPrevDist(*stop, *dist);
  *start = checkedSub(*stop, *dist);

  if (stopHint >= *start &&
      stopHint <= *stop)
    *start = checkedSub(stopHint, maxPrimeGap(stopHint));
}

} // namespace
<commit_msg>Refactor<commit_after>///
/// @file  IteratorHelper.cpp
///        Functions used to calculate the next start and stop
///        numbers for primesieve::iterator.
///
/// Copyright (C) 2022 Kim Walisch, <kim.walisch@gmail.com>
///
/// This file is distributed under the BSD License. See the COPYING
/// file in the top level directory.
///

#include <primesieve/config.hpp>
#include <primesieve/IteratorHelper.hpp>
#include <primesieve/PrimeGenerator.hpp>
#include <primesieve/pmath.hpp>

#include <stdint.h>
#include <algorithm>
#include <cmath>
#include <limits>

using namespace primesieve;

namespace {

uint64_t getNextDist(uint64_t start, uint64_t dist)
{
  uint64_t minDist = (uint64_t) std::sqrt(start);
  uint64_t maxDist = 1ull << 60;

  dist *= 4;
  minDist = std::max(minDist, PrimeGenerator::maxCachedPrime());
  return inBetween(minDist, dist, maxDist);
}

uint64_t getPrevDist(uint64_t stop, uint64_t dist)
{
  double x = std::max(10.0, (double) stop);
  uint64_t logx = (uint64_t) std::log(x);
  uint64_t minDist = (config::MIN_CACHE_ITERATOR / sizeof(uint64_t)) * logx;
  uint64_t maxDist = (config::MAX_CACHE_ITERATOR / sizeof(uint64_t)) * logx;
  uint64_t tinyDist = PrimeGenerator::maxCachedPrime() * 4;
  uint64_t defaultDist = (uint64_t) (std::sqrt(stop) * 2);

  dist *= 4;
  minDist = inBetween(tinyDist, dist, minDist);
  return inBetween(minDist, defaultDist, maxDist);
}

} // namespace

namespace primesieve {

void IteratorHelper::next(uint64_t* start,
                          uint64_t* stop,
                          uint64_t stopHint,
                          uint64_t* dist)
{
  *start = checkedAdd(*stop, 1);
  *dist = getNextDist(*start, *dist);

  if (stopHint >= *start &&
      stopHint < std::numeric_limits<uint64_t>::max())
  {
    // For primesieve::iterator it is advantageous to buffer
    // slightly more primes than the stopHint since the
    // stopHint is often not 100% accurate and the user
    // might iterate over a few primes > stopHint.
    *stop = checkedAdd(stopHint, maxPrimeGap(stopHint));
  }
  else
  {
    // In case the user has used the default stopHint=UINT64_MAX
    // we take a conservative approach and only buffer a
    // small number of primes. If the user uses more primes
    // than we have buffered, then we will increase the sieving
    // distance and buffer more primes (than last time).
    *stop = checkedAdd(*start, *dist);
  }
}

void IteratorHelper::prev(uint64_t* start,
                          uint64_t* stop,
                          uint64_t stopHint,
                          uint64_t* dist)
{
  *stop = checkedSub(*start, 1);
  *dist = getPrevDist(*stop, *dist);
  *start = checkedSub(*stop, *dist);

  if (stopHint >= *start &&
      stopHint <= *stop)
    *start = checkedSub(stopHint, maxPrimeGap(stopHint));
}

} // namespace
<|endoftext|>


<commit_before>#pragma once
#include "../include/mozart/any.hpp"
#include "./core.hpp"
namespace cov_basic {
	class arglist final {
		template<typename _Tp>static int count_types(int index=1)
		{
			return index;
		}
		template<typename _Tp,typename _fT,typename...ArgTypes>
		static int count_types(int index=1)
		{
			return count_types<_fT,ArgTypes...>(++index);
		}
		template<typename _Tp>static std::string get_type(int expect,int index)
		{
			if(expect==index)
				return typeid(_Tp).name();
			else
				throw internal_error("Wrong index of arglist.");
		}
		template<typename _Tp,typename _fT,typename...ArgTypes>
		static std::string get_type(int expect,int index)
		{
			if(expect==index)
				return typeid(_Tp).name();
			else
				return get_type<_fT,ArgTypes...>(expect,++index);
		}
		template<typename _Tp>static int check_types(int index,const std::deque<cov::any>& args)
		{
			if(index<=args.size()&&args.at(index-1).type()==typeid(_Tp))
				return -1;
			else
				return index;
		}
		template<typename _Tp,typename _fT,typename...ArgTypes>
		static int check_types(int index,const std::deque<cov::any>& args)
		{
			if(index<=args.size()&&args.at(index-1).type()==typeid(_Tp))
				return check_types<_fT,ArgTypes...>(++index,args);
			else
				return index;
		}
	public:
		template<typename...ArgTypes>
		static void check(const std::deque<cov::any>& args)
		{
			if(count_types<ArgTypes...>()==args.size()) {
				int result=check_types<ArgTypes...>(1,args);
				if(result!=-1)
					throw syntax_error("Invalid Argument.At "+std::to_string(result)+".Expected "+get_type<ArgTypes...>(result,1));
			}
			else
				throw syntax_error("Wrong size of the arguments.Expected "+std::to_string(count_types<ArgTypes...>()));
		}
	};
}
<commit_msg>更新arglist<commit_after>#pragma once
#include "../include/mozart/any.hpp"
#include "./core.hpp"
namespace cov_basic {
	class arglist final {
		template<typename _Tp>static int count_types(int index=1)
		{
			return index;
		}
		template<typename _Tp,typename _fT,typename...ArgTypes>
		static int count_types(int index=1)
		{
			return count_types<_fT,ArgTypes...>(++index);
		}
		template<typename _Tp>static std::string get_type(int expect,int index)
		{
			if(expect==index)
				return typeid(_Tp).name();
			else
				throw internal_error("Wrong index of arglist.");
		}
		template<typename _Tp,typename _fT,typename...ArgTypes>
		static std::string get_type(int expect,int index)
		{
			if(expect==index)
				return typeid(_Tp).name();
			else
				return get_type<_fT,ArgTypes...>(expect,++index);
		}
		template<typename _Tp>static int check_types(int index,const std::deque<cov::any>& args)
		{
			if(index<=args.size()&&(typeid(_Tp)==typeid(cov::any)||args.at(index-1).type()==typeid(_Tp)))
				return -1;
			else
				return index;
		}
		template<typename _Tp,typename _fT,typename...ArgTypes>
		static int check_types(int index,const std::deque<cov::any>& args)
		{
			if(index<=args.size()&&(typeid(_Tp)==typeid(cov::any)||args.at(index-1).type()==typeid(_Tp)))
				return check_types<_fT,ArgTypes...>(++index,args);
			else
				return index;
		}
	public:
		template<typename...ArgTypes>
		static void check(const std::deque<cov::any>& args)
		{
			if(count_types<ArgTypes...>()==args.size()) {
				int result=check_types<ArgTypes...>(1,args);
				if(result!=-1)
					throw syntax_error("Invalid Argument.At "+std::to_string(result)+".Expected "+get_type<ArgTypes...>(result,1));
			}
			else
				throw syntax_error("Wrong size of the arguments.Expected "+std::to_string(count_types<ArgTypes...>()));
		}
	};
}
<|endoftext|>


<commit_before><commit_msg>MappingQC: fixed support for non-human genomes.<commit_after><|endoftext|>


<commit_before>#include "ROOT/TBufferMerger.hxx"

#include "TROOT.h" // For EnableThreadSafety

#include "TTree.h"

#include "benchmark/benchmark.h"

#include <random>
#include <string>
#include <sstream>
#include <sys/stat.h>

using namespace ROOT::Experimental;

static void BM_TBufferFile_CreateEmpty(benchmark::State &state)
{
   const char *filename = "empty.root";
   while (state.KeepRunning()) {
      TBufferMerger m(std::unique_ptr<TMemFile>(new TMemFile(filename, "RECREATE")));
   }
}
BENCHMARK(BM_TBufferFile_CreateEmpty);

TBufferMerger *Merger = nullptr;
static void BM_TBufferFile_GetFile(benchmark::State &state)
{
   ROOT::EnableThreadSafety();
   using namespace ROOT::Experimental;
   if (state.thread_index == 0) {
      // Setup code here.
      // FIXME: We should have a way to pass an externally constructed file or stream to
      // TFile*, this would allow us to create in-memory files and avoid killing disks
      // when we benchmark IO.
      Merger = new TBufferMerger(std::unique_ptr<TMemFile>(new TMemFile("virtual_file.root", "RECREATE")));
   }
   while (state.KeepRunning()) {
      // Run the test as normal.
      auto myFile = Merger->GetFile();
   }
   if (state.thread_index == 0) {
      // Teardown code here.
      delete Merger;
   }
}
BENCHMARK(BM_TBufferFile_GetFile)->Unit(benchmark::kMicrosecond);
BENCHMARK(BM_TBufferFile_GetFile)->Unit(benchmark::kMicrosecond)->UseRealTime()->ThreadPerCpu();
BENCHMARK(BM_TBufferFile_GetFile)->Unit(benchmark::kMicrosecond)->UseRealTime()->ThreadRange(1, 256);

/// Creates a TMemFile, fills a TTree with random numbers. The data is written if it exceeds 32MB.
inline void FillTreeWithRandomData(TBufferMerger &merger, size_t nEntriesPerWorker = 24 * 1024, int flush = 32 )
{
   thread_local std::default_random_engine g;
   std::uniform_real_distribution<double> dist(0.0, 1.0);

   auto f = merger.GetFile();
   auto t = new TTree("random", "random");
   t->ResetBit(kMustCleanup);
   t->SetAutoFlush(-(flush) * 1024 * 1024); // Flush at exceeding 32MB

   double rng;

   t->Branch("random", &rng);

   long entries = 0;
   for (size_t i = 0; i < nEntriesPerWorker; ++i) {
      rng = dist(g);
      t->Fill();
      ++entries;
      auto atflush = t->GetAutoFlush();
      if (entries == atflush) {
         entries = 0;
         f->Write();
      }
   }
   f->Write();
}

static void BM_TBufferFile_FillTreeWithRandomData(benchmark::State &state)
{
   ROOT::EnableThreadSafety();
   using namespace ROOT::Experimental;
   if (state.thread_index == 0) {
      // Setup code here.
      Merger = new TBufferMerger(std::unique_ptr<TMemFile>(new TMemFile("virtual_file.root", "RECREATE")));
   }
   int flush = state.range(0);
   while (state.KeepRunning())
      FillTreeWithRandomData(*Merger, flush);
   auto size = Merger->GetFile()->GetSize();
   std::stringstream ss;
   ss << size;
   state.SetLabel(ss.str());
   if (state.thread_index == 0) {
      // Teardown code here.
      delete Merger;
   }
}
BENCHMARK(BM_TBufferFile_FillTreeWithRandomData)->Unit(benchmark::kMicrosecond)->Arg(32);
BENCHMARK(BM_TBufferFile_FillTreeWithRandomData)->Unit(benchmark::kMicrosecond)->Arg(32)->UseRealTime()->ThreadPerCpu();
BENCHMARK(BM_TBufferFile_FillTreeWithRandomData)->Unit(benchmark::kMicrosecond)->Range(1, 32)->UseRealTime()->ThreadRange(1, 256);

// Define our main.
BENCHMARK_MAIN();
<commit_msg>Fixing detection of number of hardware threads to be tested for IO benchmark<commit_after>#include "ROOT/TBufferMerger.hxx"

#include "TROOT.h" // For EnableThreadSafety

#include "TTree.h"

#include "benchmark/benchmark.h"

#include <random>
#include <string>
#include <sstream>
#include <thread>
#include <sys/stat.h>

using namespace ROOT::Experimental;

int GetNumberHardwareThreads()
{
   return std::thread::hardware_concurrency();
}

static void BM_TBufferFile_CreateEmpty(benchmark::State &state)
{
   const char *filename = "empty.root";
   while (state.KeepRunning()) {
      TBufferMerger m(std::unique_ptr<TMemFile>(new TMemFile(filename, "RECREATE")));
   }
}
BENCHMARK(BM_TBufferFile_CreateEmpty);

TBufferMerger *Merger = nullptr;
static void BM_TBufferFile_GetFile(benchmark::State &state)
{
   ROOT::EnableThreadSafety();
   using namespace ROOT::Experimental;
   if (state.thread_index == 0) {
      // Setup code here.
      // FIXME: We should have a way to pass an externally constructed file or stream to
      // TFile*, this would allow us to create in-memory files and avoid killing disks
      // when we benchmark IO.
      Merger = new TBufferMerger(std::unique_ptr<TMemFile>(new TMemFile("virtual_file.root", "RECREATE")));
   }
   while (state.KeepRunning()) {
      // Run the test as normal.
      auto myFile = Merger->GetFile();
   }
   if (state.thread_index == 0) {
      // Teardown code here.
      delete Merger;
   }
}
BENCHMARK(BM_TBufferFile_GetFile)->Unit(benchmark::kMicrosecond);
BENCHMARK(BM_TBufferFile_GetFile)->Unit(benchmark::kMicrosecond)->UseRealTime()->ThreadPerCpu();
BENCHMARK(BM_TBufferFile_GetFile)->Unit(benchmark::kMicrosecond)->UseRealTime()->ThreadRange(1, GetNumberHardwareThreads()*4);

/// Creates a TMemFile, fills a TTree with random numbers. The data is written if it exceeds 32MB.
inline void FillTreeWithRandomData(TBufferMerger &merger, size_t nEntriesPerWorker = 24 * 1024, int flush = 32 )
{
   thread_local std::default_random_engine g;
   std::uniform_real_distribution<double> dist(0.0, 1.0);

   auto f = merger.GetFile();
   auto t = new TTree("random", "random");
   t->ResetBit(kMustCleanup);
   t->SetAutoFlush(-(flush) * 1024 * 1024); // Flush at exceeding 32MB

   double rng;

   t->Branch("random", &rng);

   long entries = 0;
   for (size_t i = 0; i < nEntriesPerWorker; ++i) {
      rng = dist(g);
      t->Fill();
      ++entries;
      auto atflush = t->GetAutoFlush();
      if (entries == atflush) {
         entries = 0;
         f->Write();
      }
   }
   f->Write();
}

static void BM_TBufferFile_FillTreeWithRandomData(benchmark::State &state)
{
   ROOT::EnableThreadSafety();
   using namespace ROOT::Experimental;
   if (state.thread_index == 0) {
      // Setup code here.
      Merger = new TBufferMerger(std::unique_ptr<TMemFile>(new TMemFile("virtual_file.root", "RECREATE")));
   }
   int flush = state.range(0);
   while (state.KeepRunning())
      FillTreeWithRandomData(*Merger, flush);
   auto size = Merger->GetFile()->GetSize();
   std::stringstream ss;
   ss << size;
   state.SetLabel(ss.str());
   if (state.thread_index == 0) {
      // Teardown code here.
      delete Merger;
   }
}
BENCHMARK(BM_TBufferFile_FillTreeWithRandomData)->Unit(benchmark::kMicrosecond)->Arg(32);
BENCHMARK(BM_TBufferFile_FillTreeWithRandomData)->Unit(benchmark::kMicrosecond)->Arg(32)->UseRealTime()->ThreadPerCpu();
BENCHMARK(BM_TBufferFile_FillTreeWithRandomData)->Unit(benchmark::kMicrosecond)->Range(1, 32)->UseRealTime()->ThreadRange(1, GetNumberHardwareThreads()*4);

// Define our main.
BENCHMARK_MAIN();
<|endoftext|>