<commit_before>//
// Copyright (c) 2008-2019 the Urho3D project.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//

#include "../Precompiled.h"

#include "../Container/Allocator.h"
#include "../Core/Profiler.h"

#include "../DebugNew.h"

#if URHO3D_STATIC
URHO3D_API void* operator new[](size_t size, const char* pName, int flags, unsigned debugFlags, const char* file, int line)
{
    void* memory = malloc(size);
    if (memory == nullptr)
        throw std::bad_alloc();
    return memory;
}

URHO3D_API void* operator new[](size_t size, size_t alignment, size_t alignmentOffset, const char* pName, int flags, unsigned debugFlags, const char* file, int line)
{
    void* memory = aligned_alloc(alignment, size);
    if (memory == nullptr)
        throw std::bad_alloc();
    return memory;
}
URHO3D_API void operator delete[](void* memory)
{
    if (memory != nullptr)
        free(memory);
}
#endif

namespace Urho3D
{

AllocatorBlock* AllocatorReserveBlock(AllocatorBlock* allocator, unsigned nodeSize, unsigned capacity)
{
    URHO3D_PROFILE("AllocatorReserveBlock");

    if (!capacity)
        capacity = 1;

    auto* blockPtr = new unsigned char[sizeof(AllocatorBlock) + capacity * (sizeof(AllocatorNode) + nodeSize)];
    auto* newBlock = reinterpret_cast<AllocatorBlock*>(blockPtr);
    newBlock->nodeSize_ = nodeSize;
    newBlock->capacity_ = capacity;
    newBlock->free_ = nullptr;
    newBlock->next_ = nullptr;

    if (!allocator)
        allocator = newBlock;
    else
    {
        newBlock->next_ = allocator->next_;
        allocator->next_ = newBlock;
    }

    // Initialize the nodes. Free nodes are always chained to the first (parent) allocator
    unsigned char* nodePtr = blockPtr + sizeof(AllocatorBlock);
    auto* firstNewNode = reinterpret_cast<AllocatorNode*>(nodePtr);

    for (unsigned i = 0; i < capacity - 1; ++i)
    {
        auto* newNode = reinterpret_cast<AllocatorNode*>(nodePtr);
        newNode->next_ = reinterpret_cast<AllocatorNode*>(nodePtr + sizeof(AllocatorNode) + nodeSize);
        nodePtr += sizeof(AllocatorNode) + nodeSize;
    }
    // i == capacity - 1
    {
        auto* newNode = reinterpret_cast<AllocatorNode*>(nodePtr);
        newNode->next_ = nullptr;
    }

    allocator->free_ = firstNewNode;

    return newBlock;
}

AllocatorBlock* AllocatorInitialize(unsigned nodeSize, unsigned initialCapacity)
{
    URHO3D_PROFILE("AllocatorInitialize");

    AllocatorBlock* block = AllocatorReserveBlock(nullptr, nodeSize, initialCapacity);
    return block;
}

void AllocatorUninitialize(AllocatorBlock* allocator)
{
    URHO3D_PROFILE("AllocatorUninitialize");

    while (allocator)
    {
        AllocatorBlock* next = allocator->next_;
        delete[] reinterpret_cast<unsigned char*>(allocator);
        allocator = next;
    }
}

void* AllocatorReserve(AllocatorBlock* allocator)
{
    if (!allocator)
        return nullptr;

    URHO3D_PROFILE("AllocatorReserve");

    if (!allocator->free_)
    {
        // Free nodes have been exhausted. Allocate a new larger block
        unsigned newCapacity = (allocator->capacity_ + 1) >> 1u;
        AllocatorReserveBlock(allocator, allocator->nodeSize_, newCapacity);
        allocator->capacity_ += newCapacity;
    }

    // We should have new free node(s) chained
    AllocatorNode* freeNode = allocator->free_;
    void* ptr = (reinterpret_cast<unsigned char*>(freeNode)) + sizeof(AllocatorNode);
    allocator->free_ = freeNode->next_;
    freeNode->next_ = nullptr;

    return ptr;
}

void AllocatorFree(AllocatorBlock* allocator, void* ptr)
{
    if (!allocator || !ptr)
        return;

    URHO3D_PROFILE("AllocatorFree");

    auto* dataPtr = static_cast<unsigned char*>(ptr);
    auto* node = reinterpret_cast<AllocatorNode*>(dataPtr - sizeof(AllocatorNode));

    // Chain the node back to free nodes
    node->next_ = allocator->free_;
    allocator->free_ = node;
}

}
<commit_msg>Misc: Build fix for MacOS (missing include).<commit_after>//
// Copyright (c) 2008-2019 the Urho3D project.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//

#include "../Precompiled.h"

#include <stdlib.h> // For alligned_alloc on MacOS

#include "../Container/Allocator.h"
#include "../Core/Profiler.h"

#include "../DebugNew.h"

#if URHO3D_STATIC
URHO3D_API void* operator new[](size_t size, const char* pName, int flags, unsigned debugFlags, const char* file, int line)
{
    void* memory = malloc(size);
    if (memory == nullptr)
        throw std::bad_alloc();
    return memory;
}

URHO3D_API void* operator new[](size_t size, size_t alignment, size_t alignmentOffset, const char* pName, int flags, unsigned debugFlags, const char* file, int line)
{
    void* memory = aligned_alloc(alignment, size);
    if (memory == nullptr)
        throw std::bad_alloc();
    return memory;
}
URHO3D_API void operator delete[](void* memory)
{
    if (memory != nullptr)
        free(memory);
}
#endif

namespace Urho3D
{

AllocatorBlock* AllocatorReserveBlock(AllocatorBlock* allocator, unsigned nodeSize, unsigned capacity)
{
    URHO3D_PROFILE("AllocatorReserveBlock");

    if (!capacity)
        capacity = 1;

    auto* blockPtr = new unsigned char[sizeof(AllocatorBlock) + capacity * (sizeof(AllocatorNode) + nodeSize)];
    auto* newBlock = reinterpret_cast<AllocatorBlock*>(blockPtr);
    newBlock->nodeSize_ = nodeSize;
    newBlock->capacity_ = capacity;
    newBlock->free_ = nullptr;
    newBlock->next_ = nullptr;

    if (!allocator)
        allocator = newBlock;
    else
    {
        newBlock->next_ = allocator->next_;
        allocator->next_ = newBlock;
    }

    // Initialize the nodes. Free nodes are always chained to the first (parent) allocator
    unsigned char* nodePtr = blockPtr + sizeof(AllocatorBlock);
    auto* firstNewNode = reinterpret_cast<AllocatorNode*>(nodePtr);

    for (unsigned i = 0; i < capacity - 1; ++i)
    {
        auto* newNode = reinterpret_cast<AllocatorNode*>(nodePtr);
        newNode->next_ = reinterpret_cast<AllocatorNode*>(nodePtr + sizeof(AllocatorNode) + nodeSize);
        nodePtr += sizeof(AllocatorNode) + nodeSize;
    }
    // i == capacity - 1
    {
        auto* newNode = reinterpret_cast<AllocatorNode*>(nodePtr);
        newNode->next_ = nullptr;
    }

    allocator->free_ = firstNewNode;

    return newBlock;
}

AllocatorBlock* AllocatorInitialize(unsigned nodeSize, unsigned initialCapacity)
{
    URHO3D_PROFILE("AllocatorInitialize");

    AllocatorBlock* block = AllocatorReserveBlock(nullptr, nodeSize, initialCapacity);
    return block;
}

void AllocatorUninitialize(AllocatorBlock* allocator)
{
    URHO3D_PROFILE("AllocatorUninitialize");

    while (allocator)
    {
        AllocatorBlock* next = allocator->next_;
        delete[] reinterpret_cast<unsigned char*>(allocator);
        allocator = next;
    }
}

void* AllocatorReserve(AllocatorBlock* allocator)
{
    if (!allocator)
        return nullptr;

    URHO3D_PROFILE("AllocatorReserve");

    if (!allocator->free_)
    {
        // Free nodes have been exhausted. Allocate a new larger block
        unsigned newCapacity = (allocator->capacity_ + 1) >> 1u;
        AllocatorReserveBlock(allocator, allocator->nodeSize_, newCapacity);
        allocator->capacity_ += newCapacity;
    }

    // We should have new free node(s) chained
    AllocatorNode* freeNode = allocator->free_;
    void* ptr = (reinterpret_cast<unsigned char*>(freeNode)) + sizeof(AllocatorNode);
    allocator->free_ = freeNode->next_;
    freeNode->next_ = nullptr;

    return ptr;
}

void AllocatorFree(AllocatorBlock* allocator, void* ptr)
{
    if (!allocator || !ptr)
        return;

    URHO3D_PROFILE("AllocatorFree");

    auto* dataPtr = static_cast<unsigned char*>(ptr);
    auto* node = reinterpret_cast<AllocatorNode*>(dataPtr - sizeof(AllocatorNode));

    // Chain the node back to free nodes
    node->next_ = allocator->free_;
    allocator->free_ = node;
}

}
<|endoftext|>


<commit_before>#include <algorithm>
#include <map>
#include <string>
#include <limits>

#include "Prefetch.h"
#include "CodeGen_Internal.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "Bounds.h"
#include "Scope.h"
#include "Simplify.h"
#include "Substitute.h"
#include "Util.h"

namespace Halide {
namespace Internal {

using std::map;
using std::string;
using std::vector;
using std::stack;

// Prefetch debug levels
int dbg_prefetch  = 1;
int dbg_prefetch2 = 2;
int dbg_prefetch3 = 10;

class InjectPrefetch : public IRMutator {
public:
    InjectPrefetch(const map<string, Function> &e)
        : env(e) { }
private:
    const map<string, Function> &env;
    Scope<Interval> scope;
    unsigned long ptmp = 0;   // ID for all tmp vars in a prefetch op

private:
    using IRMutator::visit;

    // Strip down the tuple name, e.g. f.*.var into f
    string tuple_func(const string &name) {
        vector<string> v = split_string(name, ".");
        internal_assert(v.size() > 0);
        return v[0];
    }

    // Strip down the tuple name, e.g. f.*.var into var
    string tuple_var(const string &name) {
        vector<string> v = split_string(name, ".");
        internal_assert(v.size() > 0);
        return v[v.size()-1];
    }

    Function get_func(const string &name) {
        map<string, Function>::const_iterator iter = env.find(name);
        internal_assert(iter != env.end()) << "Realize node refers to function not in environment.\n";
        return iter->second;
    }

    void visit(const Let *op) {
        Interval in = bounds_of_expr_in_scope(op->value, scope);
        scope.push(op->name, in);
        IRMutator::visit(op);
        scope.pop(op->name);
    }
    void visit(const LetStmt *op) {
        Interval in = bounds_of_expr_in_scope(op->value, scope);
        scope.push(op->name, in);
        IRMutator::visit(op);
        scope.pop(op->name);
    }

    void visit(const For *op) {
        Stmt body = op->body;

        string func_name = tuple_func(op->name);
        string var_name  = tuple_var(op->name);
        vector<Prefetch> &prefetches = get_func(func_name).schedule().prefetches();

        if (prefetches.empty()) {
            debug(dbg_prefetch2) << "InjectPrefetch: " << op->name << " " << func_name << " " << var_name;
            debug(dbg_prefetch2) << " No prefetch\n";
        } else {
            debug(dbg_prefetch) << "InjectPrefetch: " << op->name << " " << func_name << " " << var_name;
            debug(dbg_prefetch) << " Found prefetch directive(s)\n";
        }

        for (const Prefetch &p : prefetches) {
            debug(dbg_prefetch) << "InjectPrefetch: check var:" << p.var
                               << " offset:" << p.offset << "\n";
            if (p.var == var_name) {
                debug(dbg_prefetch) << " prefetch on " << var_name << "\n";

                // Interval prein(op->name, op->name);
                // Add in prefetch offset
                Expr var = Variable::make(Int(32), op->name);
                Interval prein(var + p.offset, var + p.offset);
                scope.push(op->name, prein);

                map<string, Box> boxes;
                boxes = boxes_required(body, scope);

                int cnt = 0;            // Number of prefetch ops generated
                Stmt pstmts;            // Generated prefetch statements
                debug(dbg_prefetch) << "  boxes required:\n";
                for (auto &b : boxes) {
                    const string &varname = b.first;
                    Box &box = b.second;
                    int dims = box.size();
                    debug(dbg_prefetch) << "  var:" << varname << ":\n";
                    for (int i = 0; i < dims; i++) {
                        debug(dbg_prefetch) << "    ---\n";
                        debug(dbg_prefetch) << "    box[" << i << "].min: " << box[i].min << "\n";
                        debug(dbg_prefetch) << "    box[" << i << "].max: " << box[i].max << "\n";
                    }
                    debug(dbg_prefetch) << "    ---------\n";

                    string pstr = std::to_string(ptmp++);
                    string varname_prefetch_buf = varname + ".prefetch_" + pstr + "_buf";
                    Expr var_prefetch_buf = Variable::make(Int(32), varname_prefetch_buf);

                    // Make the names for accessing the buffer strides
                    vector<Expr> stride_var(dims);
                    for (int i = 0; i < dims; i++) {
                        string istr = std::to_string(i);
                        string stride_name = varname + ".stride." + istr;
                        stride_var[i] = Variable::make(Int(32), stride_name);
                    }

                    // Make the names for the prefetch box mins & maxes
                    vector<string> min_name(dims), max_name(dims);
                    for (int i = 0; i < dims; i++) {
                        string istr = std::to_string(i);
                        min_name[i] = varname + ".prefetch_" + pstr + "_min_" + istr;
                        max_name[i] = varname + ".prefetch_" + pstr + "_max_" + istr;
                    }
                    vector<Expr> min_var(dims), max_var(dims);
                    for (int i = 0; i < dims; i++) {
                        min_var[i] = Variable::make(Int(32), min_name[i]);
                        max_var[i] = Variable::make(Int(32), max_name[i]);
                    }

                    // Create a buffer_t object for this prefetch.
                    Type t = Int(16);  // FIXME (need type of box...or extract elem_size)
                    vector<Expr> args(dims*3 + 2);
                    Expr first_elem = Load::make(t, varname, 0, Buffer(), Parameter());
                    args[0] = Call::make(Handle(), Call::address_of, {first_elem}, Call::PureIntrinsic);
                    args[1] = make_zero(t);
                    for (int i = 0; i < dims; i++) {
                        args[3*i+2] = min_var[i];
                        args[3*i+3] = max_var[i] - min_var[i] + 1;
                        args[3*i+4] = stride_var[i];
                    }

                    // Inject the prefetch call
                    vector<Expr> args_prefetch(2);
                    args_prefetch[0] = dims;
                    args_prefetch[1] = var_prefetch_buf;
                    Stmt stmt_prefetch = Evaluate::make(Call::make(Int(32), Call::prefetch_buffer_t,
                                          args_prefetch, Call::Intrinsic));
                    if (cnt == 0) {
                        pstmts = stmt_prefetch;
                    } else {
                        pstmts = Block::make({stmt_prefetch, pstmts});
                    }
                    cnt++;

                    // Inject the create_buffer_t call
                    Expr prefetch_buf = Call::make(type_of<struct buffer_t *>(), Call::create_buffer_t,
                                          args, Call::Intrinsic);
                    pstmts = LetStmt::make(varname_prefetch_buf, prefetch_buf, pstmts);

                    // Inject bounds variable assignments
                    for (int i = dims-1; i >= 0; i--) {
                        pstmts = LetStmt::make(max_name[i], box[i].max, pstmts);
                        pstmts = LetStmt::make(min_name[i], box[i].min, pstmts);
                        // stride already defined by input buffer
                    }
                }

                if (cnt) {
#if 1
                    // Guard to not prefetch past the end of the iteration space
                    Expr pcond = likely((Variable::make(Int(32), op->name) + p.offset)
                                                        < (op->min + op->extent - 1));
                    Stmt pguard = IfThenElse::make(pcond, pstmts);
                    body = Block::make({pguard, body});

                    debug(dbg_prefetch3) << "    prefetch: (cnt:" << cnt << ")\n";
                    debug(dbg_prefetch3) << pguard << "\n";
#else
                    body = Block::make({pstmts, body});
#endif
                }

                scope.pop(op->name);
            }
        }

        body = mutate(body);
        stmt = For::make(op->name, op->min, op->extent, op->for_type, op->device_api, body);
    }

};

Stmt inject_prefetch(Stmt s, const std::map<std::string, Function> &env)
{
    size_t read;
    std::string lvl = get_env_variable("HL_DEBUG_PREFETCH", read);
    if (read) {
        int dbg_level = atoi(lvl.c_str());
        dbg_prefetch  -= dbg_level;
        dbg_prefetch2 -= dbg_level;
        dbg_prefetch3 -= dbg_level;
        if (dbg_prefetch < 0) {
            dbg_prefetch = 0;
        }
        if (dbg_prefetch2 < 0) {
            dbg_prefetch2 = 0;
        }
        if (dbg_prefetch3 < 0) {
            dbg_prefetch3 = 0;
        }
    }
    return InjectPrefetch(env).mutate(s);
}

}
}
<commit_msg>Get buffer type from the environment<commit_after>#include <algorithm>
#include <map>
#include <string>
#include <limits>

#include "Prefetch.h"
#include "CodeGen_Internal.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "Bounds.h"
#include "Scope.h"
#include "Simplify.h"
#include "Substitute.h"
#include "Util.h"

namespace Halide {
namespace Internal {

using std::map;
using std::string;
using std::vector;
using std::stack;

// Prefetch debug levels
int dbg_prefetch  = 1;
int dbg_prefetch2 = 2;
int dbg_prefetch3 = 10;

class InjectPrefetch : public IRMutator {
public:
    InjectPrefetch(const map<string, Function> &e)
        : env(e) { }
private:
    const map<string, Function> &env;
    Scope<Interval> scope;
    unsigned long ptmp = 0;   // ID for all tmp vars in a prefetch op

private:
    using IRMutator::visit;

    // Strip down the tuple name, e.g. f.*.var into f
    string tuple_func(const string &name) {
        vector<string> v = split_string(name, ".");
        internal_assert(v.size() > 0);
        return v[0];
    }

    // Strip down the tuple name, e.g. f.*.var into var
    string tuple_var(const string &name) {
        vector<string> v = split_string(name, ".");
        internal_assert(v.size() > 0);
        return v[v.size()-1];
    }

    // Lookup a function in the environment
    Function get_func(const string &name) {
        map<string, Function>::const_iterator iter = env.find(name);
        internal_assert(iter != env.end()) << "Realize node refers to function not in environment.\n";
        return iter->second;
    }

    // Determine the type of a named buffer
    Type get_type(string varname) {
        debug(dbg_prefetch2) << "getType(" << varname << ")\n";
        Type t = Int(8);
        map<string, Function>::const_iterator varit = env.find(varname);
        if (varit != env.end()) {
            Function varf = varit->second;
            debug(dbg_prefetch2) << "    found: " << varit->first << "\n";
            if (varf.outputs()) {
                vector<Type> varts = varf.output_types();
                t = varts[0];
                debug(dbg_prefetch2) << "    type: " << t << "\n";
            }
        } else {
            debug(dbg_prefetch2) << "    could not determine type\n";
        }
        return t;
    }

    void visit(const Let *op) {
        Interval in = bounds_of_expr_in_scope(op->value, scope);
        scope.push(op->name, in);
        IRMutator::visit(op);
        scope.pop(op->name);
    }
    void visit(const LetStmt *op) {
        Interval in = bounds_of_expr_in_scope(op->value, scope);
        scope.push(op->name, in);
        IRMutator::visit(op);
        scope.pop(op->name);
    }

    void visit(const For *op) {
        Stmt body = op->body;

        string func_name = tuple_func(op->name);
        string ivar_name = tuple_var(op->name);
        vector<Prefetch> &prefetches = get_func(func_name).schedule().prefetches();

        if (prefetches.empty()) {
            debug(dbg_prefetch2) << "InjectPrefetch: " << op->name << " " << func_name << " " << ivar_name;
            debug(dbg_prefetch2) << " No prefetch\n";
        } else {
            debug(dbg_prefetch) << "InjectPrefetch: " << op->name << " " << func_name << " " << ivar_name;
            debug(dbg_prefetch) << " Found prefetch directive(s)\n";
        }

        for (const Prefetch &p : prefetches) {
            debug(dbg_prefetch) << "InjectPrefetch: check ivar:" << p.var
                               << " offset:" << p.offset << "\n";
            if (p.var == ivar_name) {
                debug(dbg_prefetch) << " prefetch on " << ivar_name << "\n";

                // Interval prein(op->name, op->name);
                // Add in prefetch offset
                Expr var = Variable::make(Int(32), op->name);
                Interval prein(var + p.offset, var + p.offset);
                scope.push(op->name, prein);

                map<string, Box> boxes;
                boxes = boxes_required(body, scope);

                int cnt = 0;            // Number of prefetch ops generated
                Stmt pstmts;            // Generated prefetch statements
                debug(dbg_prefetch) << "  boxes required:\n";
                for (auto &b : boxes) {
                    const string &varname = b.first;
                    Box &box = b.second;
                    int dims = box.size();
                    Type t = get_type(varname);
                    debug(dbg_prefetch) << "  var: " << varname << " (" << t << "):\n";
                    for (int i = 0; i < dims; i++) {
                        debug(dbg_prefetch) << "    ---\n";
                        debug(dbg_prefetch) << "    box[" << i << "].min: " << box[i].min << "\n";
                        debug(dbg_prefetch) << "    box[" << i << "].max: " << box[i].max << "\n";
                    }
                    debug(dbg_prefetch) << "    ---------\n";

                    string pstr = std::to_string(ptmp++);
                    string varname_prefetch_buf = varname + ".prefetch_" + pstr + "_buf";
                    Expr var_prefetch_buf = Variable::make(Int(32), varname_prefetch_buf);

                    // Make the names for accessing the buffer strides
                    vector<Expr> stride_var(dims);
                    for (int i = 0; i < dims; i++) {
                        string istr = std::to_string(i);
                        string stride_name = varname + ".stride." + istr;
                        stride_var[i] = Variable::make(Int(32), stride_name);
                    }

                    // Make the names for the prefetch box mins & maxes
                    vector<string> min_name(dims), max_name(dims);
                    for (int i = 0; i < dims; i++) {
                        string istr = std::to_string(i);
                        min_name[i] = varname + ".prefetch_" + pstr + "_min_" + istr;
                        max_name[i] = varname + ".prefetch_" + pstr + "_max_" + istr;
                    }
                    vector<Expr> min_var(dims), max_var(dims);
                    for (int i = 0; i < dims; i++) {
                        min_var[i] = Variable::make(Int(32), min_name[i]);
                        max_var[i] = Variable::make(Int(32), max_name[i]);
                    }

                    // Create a buffer_t object for this prefetch.
                    vector<Expr> args(dims*3 + 2);
                    Expr first_elem = Load::make(t, varname, 0, Buffer(), Parameter());
                    args[0] = Call::make(Handle(), Call::address_of, {first_elem}, Call::PureIntrinsic);
                    args[1] = make_zero(t);
                    for (int i = 0; i < dims; i++) {
                        args[3*i+2] = min_var[i];
                        args[3*i+3] = max_var[i] - min_var[i] + 1;
                        args[3*i+4] = stride_var[i];
                    }

                    // Inject the prefetch call
                    vector<Expr> args_prefetch(2);
                    args_prefetch[0] = dims;
                    args_prefetch[1] = var_prefetch_buf;
                    Stmt stmt_prefetch = Evaluate::make(Call::make(Int(32), Call::prefetch_buffer_t,
                                          args_prefetch, Call::Intrinsic));
                    if (cnt == 0) {
                        pstmts = stmt_prefetch;
                    } else {
                        pstmts = Block::make({stmt_prefetch, pstmts});
                    }
                    cnt++;

                    // Inject the create_buffer_t call
                    Expr prefetch_buf = Call::make(type_of<struct buffer_t *>(), Call::create_buffer_t,
                                          args, Call::Intrinsic);
                    pstmts = LetStmt::make(varname_prefetch_buf, prefetch_buf, pstmts);

                    // Inject bounds variable assignments
                    for (int i = dims-1; i >= 0; i--) {
                        pstmts = LetStmt::make(max_name[i], box[i].max, pstmts);
                        pstmts = LetStmt::make(min_name[i], box[i].min, pstmts);
                        // stride already defined by input buffer
                    }
                }

                if (cnt) {
#if 1
                    // Guard to not prefetch past the end of the iteration space
                    Expr pcond = likely((Variable::make(Int(32), op->name) + p.offset)
                                                        < (op->min + op->extent - 1));
                    Stmt pguard = IfThenElse::make(pcond, pstmts);
                    body = Block::make({pguard, body});

                    debug(dbg_prefetch3) << "    prefetch: (cnt:" << cnt << ")\n";
                    debug(dbg_prefetch3) << pguard << "\n";
#else
                    body = Block::make({pstmts, body});
#endif
                }

                scope.pop(op->name);
            }
        }

        body = mutate(body);
        stmt = For::make(op->name, op->min, op->extent, op->for_type, op->device_api, body);
    }

};

Stmt inject_prefetch(Stmt s, const std::map<std::string, Function> &env)
{
    size_t read;
    std::string lvl = get_env_variable("HL_DEBUG_PREFETCH", read);
    if (read) {
        int dbg_level = atoi(lvl.c_str());
        dbg_prefetch  -= dbg_level;
        dbg_prefetch2 -= dbg_level;
        dbg_prefetch3 -= dbg_level;
        if (dbg_prefetch < 0) {
            dbg_prefetch = 0;
        }
        if (dbg_prefetch2 < 0) {
            dbg_prefetch2 = 0;
        }
        if (dbg_prefetch3 < 0) {
            dbg_prefetch3 = 0;
        }
    }
    return InjectPrefetch(env).mutate(s);
}

}
}
<|endoftext|>


<commit_before>/*******************************************************
 * Copyright (c) 2014, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/

#include <af/version.h>
#include <af/cuda.h>
#include <platform.hpp>
#include <defines.hpp>
#include <version.hpp>
#include <driver.h>
#include <vector>
#include <string>
#include <algorithm>
#include <sstream>
#include <stdexcept>
#include <cstdio>
#include <cstring>
#include <err_cuda.hpp>

using namespace std;

namespace cuda
{

///////////////////////////////////////////////////////////////////////////
// HELPERS
///////////////////////////////////////////////////////////////////////////
// pulled from CUTIL from CUDA SDK
static inline int compute2cores(int major, int minor)
{
    struct {
        int compute; // 0xMm (hex), M = major version, m = minor version
        int cores;
    } gpus[] = {
        { 0x10,  8 },
        { 0x11,  8 },
        { 0x12,  8 },
        { 0x13,  8 },
        { 0x20, 32 },
        { 0x21, 48 },
        { 0x30, 192 },
        { 0x32, 192 },
        { 0x35, 192 },
        { 0x37, 192 },
        { 0x50, 128 },
        { 0x52, 128 },
        { 0x53, 128 },
        {   -1, -1  },
    };

    for (int i = 0; gpus[i].compute != -1; ++i) {
        if (gpus[i].compute == (major << 4) + minor)
            return gpus[i].cores;
    }
    return 0;
}

// compare two cards based on (in order):
//   1. flops (theoretical)
//   2. total memory

#define COMPARE(a,b,f) do {                     \
        return ((a)->f >= (b)->f);              \
    } while (0);

static inline bool card_compare_compute(const cudaDevice_t &l, const cudaDevice_t &r)
{
    const cudaDevice_t *lc = &l;
    const cudaDevice_t *rc = &r;

    COMPARE(lc, rc, prop.major);
    COMPARE(lc, rc, prop.minor);
    COMPARE(lc, rc, flops);
    COMPARE(lc, rc, prop.totalGlobalMem);
    COMPARE(lc, rc, nativeId);
    return 0;
}

static inline bool card_compare_flops(const cudaDevice_t &l, const cudaDevice_t &r)
{
    const cudaDevice_t *lc = &l;
    const cudaDevice_t *rc = &r;

    COMPARE(lc, rc, flops);
    COMPARE(lc, rc, prop.totalGlobalMem);
    COMPARE(lc, rc, prop.major);
    COMPARE(lc, rc, prop.minor);
    COMPARE(lc, rc, nativeId);
    return 0;
}

static inline bool card_compare_mem(const cudaDevice_t &l, const cudaDevice_t &r)
{
    const cudaDevice_t *lc = &l;
    const cudaDevice_t *rc = &r;

    COMPARE(lc, rc, prop.totalGlobalMem);
    COMPARE(lc, rc, flops);
    COMPARE(lc, rc, prop.major);
    COMPARE(lc, rc, prop.minor);
    COMPARE(lc, rc, nativeId);
    return 0;
}

static inline bool card_compare_num(const cudaDevice_t &l, const cudaDevice_t &r)
{
    const cudaDevice_t *lc = &l;
    const cudaDevice_t *rc = &r;

    COMPARE(lc, rc, nativeId);
    return 0;
}

static const std::string get_system(void)
{
    std::string arch = (sizeof(void *) == 4) ? "32-bit " : "64-bit ";

    return arch +
#if defined(OS_LNX)
    "Linux";
#elif defined(OS_WIN)
    "Windows";
#elif defined(OS_MAC)
    "Mac OSX";
#endif
}

template <typename T>
static inline string toString(T val)
{
    stringstream s;
    s << val;
    return s.str();
}

///////////////////////////////////////////////////////////////////////////
// Wrapper Functions
///////////////////////////////////////////////////////////////////////////
int getBackend()
{
    return AF_BACKEND_CUDA;
}

string getInfo()
{
    ostringstream info;
    info << "ArrayFire v" << AF_VERSION
         << " (CUDA, " << get_system() << ", build " << AF_REVISION << ")" << std::endl;
    info << getPlatformInfo();
    for (int i = 0; i < getDeviceCount(); ++i) {
        info << getDeviceInfo(i);
    }
    return info.str();
}

string getDeviceInfo(int device)
{
    cudaDeviceProp dev = getDeviceProp(device);

    size_t mem_gpu_total = dev.totalGlobalMem;
    //double cc = double(dev.major) + double(dev.minor) / 10;

    bool show_braces = getActiveDeviceId() == device;

    string id = (show_braces ? string("[") : "-") + toString(device) +
                (show_braces ? string("]") : "-");
    string name(dev.name);
    string memory = toString((mem_gpu_total / (1024 * 1024))
                          + !!(mem_gpu_total % (1024 * 1024)))
                    + string(" MB");
    string compute = string("CUDA Compute ") + toString(dev.major) + string(".") + toString(dev.minor);

    string info = id + string(" ")  +
                name + string(", ") +
              memory + string(", ") +
             compute + string("\n");
    return info;
}

string getPlatformInfo()
{
    string driverVersion = getDriverVersion();
    std::string cudaRuntime = getCUDARuntimeVersion();
    string platform = "Platform: CUDA Toolkit " + cudaRuntime;
    if (!driverVersion.empty()) {
        platform.append(", Driver: ");
        platform.append(driverVersion);
    }
    platform.append("\n");
    return platform;
}

bool isDoubleSupported(int device)
{
    return true;
}

void devprop(char* d_name, char* d_platform, char *d_toolkit, char* d_compute)
{
    if (getDeviceCount() <= 0) {
        printf("No CUDA-capable devices detected.\n");
        return;
    }

    cudaDeviceProp dev = getDeviceProp(getActiveDeviceId());

    // Name
    snprintf(d_name, 32, "%s", dev.name);

    //Platform
    std::string cudaRuntime = getCUDARuntimeVersion();
    snprintf(d_platform, 10, "CUDA");
    snprintf(d_toolkit, 64, "v%s", cudaRuntime.c_str());

    // Compute Version
    snprintf(d_compute, 10, "%d.%d", dev.major, dev.minor);

    // Sanitize input
    for (int i = 0; i < 31; i++) {
        if (d_name[i] == ' ') {
            if (d_name[i + 1] == 0 || d_name[i + 1] == ' ') d_name[i] = 0;
            else d_name[i] = '_';
        }
    }
}

string getDriverVersion()
{
    char driverVersion[1024] = {" ",};
    int x = nvDriverVersion(driverVersion, sizeof(driverVersion));
    if (x != 1) {
        // Windows, OSX, Tegra Need a new way to fetch driver
        #if !defined(OS_WIN) && !defined(OS_MAC) && !defined(__arm__)
        throw runtime_error("Invalid driver");
        #endif
        int driver = 0;
        CUDA_CHECK(cudaDriverGetVersion(&driver));
        return string("CUDA Driver Version: ") + toString(driver);
    } else {
        return string(driverVersion);
    }
}

string getCUDARuntimeVersion()
{
    int runtime = 0;
    CUDA_CHECK(cudaRuntimeGetVersion(&runtime));
    if(runtime / 100.f > 0)
        return toString((runtime / 1000) + (runtime % 1000)/ 100.);
    else
        return toString(runtime / 1000) + string(".0");

}

int getDeviceCount()
{
    return DeviceManager::getInstance().nDevices;
}

int getActiveDeviceId()
{
    return DeviceManager::getInstance().activeDev;
}

int getDeviceNativeId(int device)
{
    if(device < (int)DeviceManager::getInstance().cuDevices.size())
        return DeviceManager::getInstance().cuDevices[device].nativeId;
    return -1;
}

int getDeviceIdFromNativeId(int nativeId)
{
    DeviceManager& mngr = DeviceManager::getInstance();

    int devId = 0;
    for(devId = 0; devId < mngr.nDevices; ++devId) {
        if (nativeId == mngr.cuDevices[devId].nativeId)
            break;
    }
    return devId;
}

cudaStream_t getStream(int device)
{
    return DeviceManager::getInstance().streams[device];
}

int setDevice(int device)
{
    return DeviceManager::getInstance().setActiveDevice(device);
}

cudaDeviceProp getDeviceProp(int device)
{
    if(device < (int)DeviceManager::getInstance().cuDevices.size())
        return DeviceManager::getInstance().cuDevices[device].prop;
    return DeviceManager::getInstance().cuDevices[0].prop;
}

///////////////////////////////////////////////////////////////////////////
// DeviceManager Class Functions
///////////////////////////////////////////////////////////////////////////
DeviceManager& DeviceManager::getInstance()
{
    static DeviceManager my_instance;
    return my_instance;
}

DeviceManager::DeviceManager()
    : cuDevices(0), activeDev(0), nDevices(0)
{
    CUDA_CHECK(cudaGetDeviceCount(&nDevices));
    if (nDevices == 0)
        throw runtime_error("No CUDA-Capable devices found");

    cuDevices.reserve(nDevices);

    for(int i = 0; i < nDevices; i++) {
        cudaDevice_t dev;
        cudaGetDeviceProperties(&dev.prop, i);
        dev.flops = dev.prop.multiProcessorCount * compute2cores(dev.prop.major, dev.prop.minor) * dev.prop.clockRate;
        dev.nativeId = i;
        cuDevices.push_back(dev);
    }

    sortDevices();

    for(int i = 0; i < nDevices; i++) {
        setActiveDevice(i, cuDevices[i].nativeId);
        CUDA_CHECK(cudaStreamCreate(&streams[i]));
    }

    const char* deviceENV = getenv("AF_CUDA_DEFAULT_DEVICE");
    if(!deviceENV) {
        setActiveDevice(0, cuDevices[0].nativeId);
    } else {
        stringstream s(deviceENV);
        int def_device = -1;
        s >> def_device;
        if(def_device < 0 || def_device >= nDevices) {
            printf("WARNING: AF_CUDA_DEFAULT_DEVICE is out of range\n");
            printf("Setting default device as 0\n");
            setActiveDevice(0, cuDevices[0].nativeId);
        } else {
            setActiveDevice(def_device, cuDevices[def_device].nativeId);
        }
    }
}

void DeviceManager::sortDevices(sort_mode mode)
{
    switch(mode) {
        case memory :
            sort(cuDevices.begin(), cuDevices.end(), card_compare_mem);
            break;
        case flops :
            sort(cuDevices.begin(), cuDevices.end(), card_compare_flops);
            break;
        case compute :
            sort(cuDevices.begin(), cuDevices.end(), card_compare_compute);
            break;
        case none : default :
            sort(cuDevices.begin(), cuDevices.end(), card_compare_num);
            break;
    }
}

int DeviceManager::setActiveDevice(int device, int nId)
{
    if(device > (int)cuDevices.size()) {
        return -1;
    } else {
        int old = activeDev;
        if(nId == -1) nId = getDeviceNativeId(device);
        CUDA_CHECK(cudaSetDevice(nId));
        activeDev = device;
        return old;
    }
}

void sync(int device)
{
    int currDevice = getActiveDeviceId();
    setDevice(device);
    CUDA_CHECK(cudaStreamSynchronize(getStream(getActiveDeviceId())));
    setDevice(currDevice);
}

}

af_err afcu_get_stream(cudaStream_t* stream, int id)
{
    *stream = cuda::getStream(id);
    return AF_SUCCESS;
}

af_err afcu_get_native_id(int* nativeid, int id)
{
    *nativeid = cuda::getDeviceNativeId(id);
    return AF_SUCCESS;
}

af_err afcu_set_native_id(int nativeid)
{
    cuda::setDevice(cuda::getDeviceIdFromNativeId(nativeid));
    return AF_SUCCESS;
}
<commit_msg>Creating streams for devices only when device is active<commit_after>/*******************************************************
 * Copyright (c) 2014, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/

#include <af/version.h>
#include <af/cuda.h>
#include <platform.hpp>
#include <defines.hpp>
#include <version.hpp>
#include <driver.h>
#include <vector>
#include <string>
#include <algorithm>
#include <sstream>
#include <stdexcept>
#include <cstdio>
#include <cstring>
#include <err_cuda.hpp>

using namespace std;

namespace cuda
{

///////////////////////////////////////////////////////////////////////////
// HELPERS
///////////////////////////////////////////////////////////////////////////
// pulled from CUTIL from CUDA SDK
static inline int compute2cores(int major, int minor)
{
    struct {
        int compute; // 0xMm (hex), M = major version, m = minor version
        int cores;
    } gpus[] = {
        { 0x10,  8 },
        { 0x11,  8 },
        { 0x12,  8 },
        { 0x13,  8 },
        { 0x20, 32 },
        { 0x21, 48 },
        { 0x30, 192 },
        { 0x32, 192 },
        { 0x35, 192 },
        { 0x37, 192 },
        { 0x50, 128 },
        { 0x52, 128 },
        { 0x53, 128 },
        {   -1, -1  },
    };

    for (int i = 0; gpus[i].compute != -1; ++i) {
        if (gpus[i].compute == (major << 4) + minor)
            return gpus[i].cores;
    }
    return 0;
}

// compare two cards based on (in order):
//   1. flops (theoretical)
//   2. total memory

#define COMPARE(a,b,f) do {                     \
        return ((a)->f >= (b)->f);              \
    } while (0);

static inline bool card_compare_compute(const cudaDevice_t &l, const cudaDevice_t &r)
{
    const cudaDevice_t *lc = &l;
    const cudaDevice_t *rc = &r;

    COMPARE(lc, rc, prop.major);
    COMPARE(lc, rc, prop.minor);
    COMPARE(lc, rc, flops);
    COMPARE(lc, rc, prop.totalGlobalMem);
    COMPARE(lc, rc, nativeId);
    return 0;
}

static inline bool card_compare_flops(const cudaDevice_t &l, const cudaDevice_t &r)
{
    const cudaDevice_t *lc = &l;
    const cudaDevice_t *rc = &r;

    COMPARE(lc, rc, flops);
    COMPARE(lc, rc, prop.totalGlobalMem);
    COMPARE(lc, rc, prop.major);
    COMPARE(lc, rc, prop.minor);
    COMPARE(lc, rc, nativeId);
    return 0;
}

static inline bool card_compare_mem(const cudaDevice_t &l, const cudaDevice_t &r)
{
    const cudaDevice_t *lc = &l;
    const cudaDevice_t *rc = &r;

    COMPARE(lc, rc, prop.totalGlobalMem);
    COMPARE(lc, rc, flops);
    COMPARE(lc, rc, prop.major);
    COMPARE(lc, rc, prop.minor);
    COMPARE(lc, rc, nativeId);
    return 0;
}

static inline bool card_compare_num(const cudaDevice_t &l, const cudaDevice_t &r)
{
    const cudaDevice_t *lc = &l;
    const cudaDevice_t *rc = &r;

    COMPARE(lc, rc, nativeId);
    return 0;
}

static const std::string get_system(void)
{
    std::string arch = (sizeof(void *) == 4) ? "32-bit " : "64-bit ";

    return arch +
#if defined(OS_LNX)
    "Linux";
#elif defined(OS_WIN)
    "Windows";
#elif defined(OS_MAC)
    "Mac OSX";
#endif
}

template <typename T>
static inline string toString(T val)
{
    stringstream s;
    s << val;
    return s.str();
}

///////////////////////////////////////////////////////////////////////////
// Wrapper Functions
///////////////////////////////////////////////////////////////////////////
int getBackend()
{
    return AF_BACKEND_CUDA;
}

string getInfo()
{
    ostringstream info;
    info << "ArrayFire v" << AF_VERSION
         << " (CUDA, " << get_system() << ", build " << AF_REVISION << ")" << std::endl;
    info << getPlatformInfo();
    for (int i = 0; i < getDeviceCount(); ++i) {
        info << getDeviceInfo(i);
    }
    return info.str();
}

string getDeviceInfo(int device)
{
    cudaDeviceProp dev = getDeviceProp(device);

    size_t mem_gpu_total = dev.totalGlobalMem;
    //double cc = double(dev.major) + double(dev.minor) / 10;

    bool show_braces = getActiveDeviceId() == device;

    string id = (show_braces ? string("[") : "-") + toString(device) +
                (show_braces ? string("]") : "-");
    string name(dev.name);
    string memory = toString((mem_gpu_total / (1024 * 1024))
                          + !!(mem_gpu_total % (1024 * 1024)))
                    + string(" MB");
    string compute = string("CUDA Compute ") + toString(dev.major) + string(".") + toString(dev.minor);

    string info = id + string(" ")  +
                name + string(", ") +
              memory + string(", ") +
             compute + string("\n");
    return info;
}

string getPlatformInfo()
{
    string driverVersion = getDriverVersion();
    std::string cudaRuntime = getCUDARuntimeVersion();
    string platform = "Platform: CUDA Toolkit " + cudaRuntime;
    if (!driverVersion.empty()) {
        platform.append(", Driver: ");
        platform.append(driverVersion);
    }
    platform.append("\n");
    return platform;
}

bool isDoubleSupported(int device)
{
    return true;
}

void devprop(char* d_name, char* d_platform, char *d_toolkit, char* d_compute)
{
    if (getDeviceCount() <= 0) {
        printf("No CUDA-capable devices detected.\n");
        return;
    }

    cudaDeviceProp dev = getDeviceProp(getActiveDeviceId());

    // Name
    snprintf(d_name, 32, "%s", dev.name);

    //Platform
    std::string cudaRuntime = getCUDARuntimeVersion();
    snprintf(d_platform, 10, "CUDA");
    snprintf(d_toolkit, 64, "v%s", cudaRuntime.c_str());

    // Compute Version
    snprintf(d_compute, 10, "%d.%d", dev.major, dev.minor);

    // Sanitize input
    for (int i = 0; i < 31; i++) {
        if (d_name[i] == ' ') {
            if (d_name[i + 1] == 0 || d_name[i + 1] == ' ') d_name[i] = 0;
            else d_name[i] = '_';
        }
    }
}

string getDriverVersion()
{
    char driverVersion[1024] = {" ",};
    int x = nvDriverVersion(driverVersion, sizeof(driverVersion));
    if (x != 1) {
        // Windows, OSX, Tegra Need a new way to fetch driver
        #if !defined(OS_WIN) && !defined(OS_MAC) && !defined(__arm__)
        throw runtime_error("Invalid driver");
        #endif
        int driver = 0;
        CUDA_CHECK(cudaDriverGetVersion(&driver));
        return string("CUDA Driver Version: ") + toString(driver);
    } else {
        return string(driverVersion);
    }
}

string getCUDARuntimeVersion()
{
    int runtime = 0;
    CUDA_CHECK(cudaRuntimeGetVersion(&runtime));
    if(runtime / 100.f > 0)
        return toString((runtime / 1000) + (runtime % 1000)/ 100.);
    else
        return toString(runtime / 1000) + string(".0");

}

int getDeviceCount()
{
    return DeviceManager::getInstance().nDevices;
}

int getActiveDeviceId()
{
    return DeviceManager::getInstance().activeDev;
}

int getDeviceNativeId(int device)
{
    if(device < (int)DeviceManager::getInstance().cuDevices.size())
        return DeviceManager::getInstance().cuDevices[device].nativeId;
    return -1;
}

int getDeviceIdFromNativeId(int nativeId)
{
    DeviceManager& mngr = DeviceManager::getInstance();

    int devId = 0;
    for(devId = 0; devId < mngr.nDevices; ++devId) {
        if (nativeId == mngr.cuDevices[devId].nativeId)
            break;
    }
    return devId;
}

cudaStream_t getStream(int device)
{
    return DeviceManager::getInstance().streams[device];
}

int setDevice(int device)
{
    return DeviceManager::getInstance().setActiveDevice(device);
}

cudaDeviceProp getDeviceProp(int device)
{
    if(device < (int)DeviceManager::getInstance().cuDevices.size())
        return DeviceManager::getInstance().cuDevices[device].prop;
    return DeviceManager::getInstance().cuDevices[0].prop;
}

///////////////////////////////////////////////////////////////////////////
// DeviceManager Class Functions
///////////////////////////////////////////////////////////////////////////
DeviceManager& DeviceManager::getInstance()
{
    static DeviceManager my_instance;
    return my_instance;
}

DeviceManager::DeviceManager()
    : cuDevices(0), activeDev(0), nDevices(0)
{
    CUDA_CHECK(cudaGetDeviceCount(&nDevices));
    if (nDevices == 0)
        throw runtime_error("No CUDA-Capable devices found");

    cuDevices.reserve(nDevices);

    for(int i = 0; i < nDevices; i++) {
        cudaDevice_t dev;
        cudaGetDeviceProperties(&dev.prop, i);
        dev.flops = dev.prop.multiProcessorCount * compute2cores(dev.prop.major, dev.prop.minor) * dev.prop.clockRate;
        dev.nativeId = i;
        cuDevices.push_back(dev);
    }

    sortDevices();

    // Initialize all streams to 0.
    // Streams will be created in setActiveDevice()
    for(int i = 0; i < (int)MAX_DEVICES; i++)
        streams[i] = (cudaStream_t)0;

    const char* deviceENV = getenv("AF_CUDA_DEFAULT_DEVICE");
    if(!deviceENV) {
        setActiveDevice(0, cuDevices[0].nativeId);
    } else {
        stringstream s(deviceENV);
        int def_device = -1;
        s >> def_device;
        if(def_device < 0 || def_device >= nDevices) {
            printf("WARNING: AF_CUDA_DEFAULT_DEVICE is out of range\n");
            printf("Setting default device as 0\n");
            setActiveDevice(0, cuDevices[0].nativeId);
        } else {
            setActiveDevice(def_device, cuDevices[def_device].nativeId);
        }
    }
}

void DeviceManager::sortDevices(sort_mode mode)
{
    switch(mode) {
        case memory :
            sort(cuDevices.begin(), cuDevices.end(), card_compare_mem);
            break;
        case flops :
            sort(cuDevices.begin(), cuDevices.end(), card_compare_flops);
            break;
        case compute :
            sort(cuDevices.begin(), cuDevices.end(), card_compare_compute);
            break;
        case none : default :
            sort(cuDevices.begin(), cuDevices.end(), card_compare_num);
            break;
    }
}

int DeviceManager::setActiveDevice(int device, int nId)
{
    if(device > (int)cuDevices.size()) {
        return -1;
    } else {
        int old = activeDev;
        if(nId == -1) nId = getDeviceNativeId(device);
        CUDA_CHECK(cudaSetDevice(nId));
        activeDev = device;

        if(!streams[device]) {
            CUDA_CHECK(cudaStreamCreate(&streams[device]));
        }

        return old;
    }
}

void sync(int device)
{
    int currDevice = getActiveDeviceId();
    setDevice(device);
    CUDA_CHECK(cudaStreamSynchronize(getStream(getActiveDeviceId())));
    setDevice(currDevice);
}

}

af_err afcu_get_stream(cudaStream_t* stream, int id)
{
    *stream = cuda::getStream(id);
    return AF_SUCCESS;
}

af_err afcu_get_native_id(int* nativeid, int id)
{
    *nativeid = cuda::getDeviceNativeId(id);
    return AF_SUCCESS;
}

af_err afcu_set_native_id(int nativeid)
{
    cuda::setDevice(cuda::getDeviceIdFromNativeId(nativeid));
    return AF_SUCCESS;
}
<|endoftext|>


<commit_before><commit_msg>[Android] Fix an error in extracting WebCL kernel parameter list.<commit_after><|endoftext|>


<commit_before>#include "examples/ace/ttcp/common.h"

#include "muduo/base/Logging.h"
#include <boost/asio.hpp>
#include <stdio.h>

using boost::asio::ip::tcp;

void transmit(const Options& opt)
{
  try
  {
  }
  catch (std::exception& e)
  {
    LOG_ERROR << e.what();
  }
}

class TtcpServerConnection : public std::enable_shared_from_this<TtcpServerConnection>,
                             muduo::noncopyable
{
 public:
#if BOOST_VERSION < 107000L
  TtcpServerConnection(boost::asio::io_service& io_service)
    : socket_(io_service), count_(0), payload_(NULL), ack_(0)
#else
  TtcpServerConnection(const boost::asio::executor& executor)
    : socket_(executor), count_(0), payload_(NULL), ack_(0)
#endif
  {
    sessionMessage_.number = 0;
    sessionMessage_.length = 0;
  }

  ~TtcpServerConnection()
  {
    ::free(payload_);
  }

  tcp::socket& socket() { return socket_; }

  void start()
  {
    std::ostringstream oss;
    oss << socket_.remote_endpoint();
    LOG_INFO << "Got connection from " << oss.str();
    doReadSession();
  }

 private:
  void doReadSession()
  {
    auto self(shared_from_this());
    boost::asio::async_read(
        socket_, boost::asio::buffer(&sessionMessage_, sizeof(sessionMessage_)),
        [this, self](const boost::system::error_code& error, size_t len)
        {
          if (!error && len == sizeof sessionMessage_)
          {
            sessionMessage_.number = ntohl(sessionMessage_.number);
            sessionMessage_.length = ntohl(sessionMessage_.length);
            printf("receive number = %d\nreceive length = %d\n",
                   sessionMessage_.number, sessionMessage_.length);
            const int total_len = static_cast<int>(sizeof(int32_t) + sessionMessage_.length);
            payload_ = static_cast<PayloadMessage*>(::malloc(total_len));
            doReadLength();
          }
          else
          {
            LOG_ERROR << "read session message: " << error.message();
          }
        });
  }

  void doReadLength()
  {
    auto self(shared_from_this());
    payload_->length = 0;
    boost::asio::async_read(
        socket_, boost::asio::buffer(&payload_->length, sizeof payload_->length),
        [this, self](const boost::system::error_code& error, size_t len)
        {
          if (!error && len == sizeof payload_->length)
          {
            payload_->length = ntohl(payload_->length);
            doReadPayload();
          }
          else
          {
            LOG_ERROR << "read length: " << error.message();
          }
        });
  }

  void doReadPayload()
  {
    assert(payload_->length == sessionMessage_.length);
    auto self(shared_from_this());
    boost::asio::async_read(
        socket_, boost::asio::buffer(&payload_->data, payload_->length),
        [this, self](const boost::system::error_code& error, size_t len)
        {
          if (!error && len == static_cast<size_t>(payload_->length))
          {
            doWriteAck();
          }
          else
          {
            LOG_ERROR << "read payload data: " << error.message();
          }
        });
  }

  void doWriteAck()
  {
    auto self(shared_from_this());
    ack_ = htonl(payload_->length);
    boost::asio::async_write(
        socket_, boost::asio::buffer(&ack_, sizeof ack_),
        [this, self](const boost::system::error_code& error, size_t len)
        {
          if (!error && len == sizeof ack_)
          {
            if (++count_ < sessionMessage_.number)
            {
              doReadLength();
            }
            else
            {
              LOG_INFO << "Done";
            }
          }
          else
          {
            LOG_ERROR << "write ack: " << error.message();
          }
        });
  }

  tcp::socket socket_;
  int count_;
  struct SessionMessage sessionMessage_;
  struct PayloadMessage* payload_;
  int32_t ack_;
};
typedef std::shared_ptr<TtcpServerConnection> TtcpServerConnectionPtr;

void doAccept(tcp::acceptor& acceptor)
{
#if BOOST_VERSION < 107000L
  // no need to pre-create new_connection if we use asio 1.12 or boost 1.66+
  TtcpServerConnectionPtr new_connection(new TtcpServerConnection(acceptor.get_io_service()));
#else
  TtcpServerConnectionPtr new_connection(new TtcpServerConnection(acceptor.get_executor()));
#endif
  acceptor.async_accept(
      new_connection->socket(),
      [&acceptor, new_connection](boost::system::error_code error)  // move new_connection in C++14
      {
        if (!error)
        {
          new_connection->start();
        }
        doAccept(acceptor);
      });
}

void receive(const Options& opt)
{
  try
  {
    boost::asio::io_service io_service;
    tcp::acceptor acceptor(io_service, tcp::endpoint(tcp::v4(), opt.port));
    doAccept(acceptor);
    io_service.run();
  }
  catch (std::exception& e)
  {
    LOG_ERROR << e.what();
  }
}

<commit_msg>fix for asio in boost 1.74 where executor_type becomes any_io_executor.<commit_after>#include "examples/ace/ttcp/common.h"

#include "muduo/base/Logging.h"
#include <boost/asio.hpp>
#include <stdio.h>

using boost::asio::ip::tcp;

void transmit(const Options& opt)
{
  try
  {
  }
  catch (std::exception& e)
  {
    LOG_ERROR << e.what();
  }
}

class TtcpServerConnection : public std::enable_shared_from_this<TtcpServerConnection>,
                             muduo::noncopyable
{
 public:
#if BOOST_VERSION < 107000L
  TtcpServerConnection(boost::asio::io_service& io_service)
    : socket_(io_service), count_(0), payload_(NULL), ack_(0)
#else
  TtcpServerConnection(const tcp::socket::executor_type& executor)
    : socket_(executor), count_(0), payload_(NULL), ack_(0)
#endif
  {
    sessionMessage_.number = 0;
    sessionMessage_.length = 0;
  }

  ~TtcpServerConnection()
  {
    ::free(payload_);
  }

  tcp::socket& socket() { return socket_; }

  void start()
  {
    std::ostringstream oss;
    oss << socket_.remote_endpoint();
    LOG_INFO << "Got connection from " << oss.str();
    doReadSession();
  }

 private:
  void doReadSession()
  {
    auto self(shared_from_this());
    boost::asio::async_read(
        socket_, boost::asio::buffer(&sessionMessage_, sizeof(sessionMessage_)),
        [this, self](const boost::system::error_code& error, size_t len)
        {
          if (!error && len == sizeof sessionMessage_)
          {
            sessionMessage_.number = ntohl(sessionMessage_.number);
            sessionMessage_.length = ntohl(sessionMessage_.length);
            printf("receive number = %d\nreceive length = %d\n",
                   sessionMessage_.number, sessionMessage_.length);
            const int total_len = static_cast<int>(sizeof(int32_t) + sessionMessage_.length);
            payload_ = static_cast<PayloadMessage*>(::malloc(total_len));
            doReadLength();
          }
          else
          {
            LOG_ERROR << "read session message: " << error.message();
          }
        });
  }

  void doReadLength()
  {
    auto self(shared_from_this());
    payload_->length = 0;
    boost::asio::async_read(
        socket_, boost::asio::buffer(&payload_->length, sizeof payload_->length),
        [this, self](const boost::system::error_code& error, size_t len)
        {
          if (!error && len == sizeof payload_->length)
          {
            payload_->length = ntohl(payload_->length);
            doReadPayload();
          }
          else
          {
            LOG_ERROR << "read length: " << error.message();
          }
        });
  }

  void doReadPayload()
  {
    assert(payload_->length == sessionMessage_.length);
    auto self(shared_from_this());
    boost::asio::async_read(
        socket_, boost::asio::buffer(&payload_->data, payload_->length),
        [this, self](const boost::system::error_code& error, size_t len)
        {
          if (!error && len == static_cast<size_t>(payload_->length))
          {
            doWriteAck();
          }
          else
          {
            LOG_ERROR << "read payload data: " << error.message();
          }
        });
  }

  void doWriteAck()
  {
    auto self(shared_from_this());
    ack_ = htonl(payload_->length);
    boost::asio::async_write(
        socket_, boost::asio::buffer(&ack_, sizeof ack_),
        [this, self](const boost::system::error_code& error, size_t len)
        {
          if (!error && len == sizeof ack_)
          {
            if (++count_ < sessionMessage_.number)
            {
              doReadLength();
            }
            else
            {
              LOG_INFO << "Done";
            }
          }
          else
          {
            LOG_ERROR << "write ack: " << error.message();
          }
        });
  }

  tcp::socket socket_;
  int count_;
  struct SessionMessage sessionMessage_;
  struct PayloadMessage* payload_;
  int32_t ack_;
};
typedef std::shared_ptr<TtcpServerConnection> TtcpServerConnectionPtr;

void doAccept(tcp::acceptor& acceptor)
{
#if BOOST_VERSION < 107000L
  // no need to pre-create new_connection if we use asio 1.12 or boost 1.66+
  TtcpServerConnectionPtr new_connection(new TtcpServerConnection(acceptor.get_io_service()));
#else
  TtcpServerConnectionPtr new_connection(new TtcpServerConnection(acceptor.get_executor()));
#endif
  acceptor.async_accept(
      new_connection->socket(),
      [&acceptor, new_connection](boost::system::error_code error)  // move new_connection in C++14
      {
        if (!error)
        {
          new_connection->start();
        }
        doAccept(acceptor);
      });
}

void receive(const Options& opt)
{
  try
  {
    boost::asio::io_service io_service;
    tcp::acceptor acceptor(io_service, tcp::endpoint(tcp::v4(), opt.port));
    doAccept(acceptor);
    io_service.run();
  }
  catch (std::exception& e)
  {
    LOG_ERROR << e.what();
  }
}

<|endoftext|>


<commit_before>/*
 * Copyright 2014 Milian Wolff <mail@milianw.de>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Library General Public License as
 * published by the Free Software Foundation; either version 2 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public
 * License along with this program; if not, write to the
 * Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/**
 * @file heaptrack_interpret.cpp
 *
 * @brief Interpret raw heaptrack data and add Dwarf based debug information.
 */

#include <iostream>
#include <sstream>
#include <unordered_map>
#include <vector>
#include <tuple>
#include <algorithm>

#include <cxxabi.h>
#include <boost/algorithm/string/predicate.hpp>

#include "libbacktrace/backtrace.h"
#include "linereader.h"

using namespace std;

namespace {

string demangle(const char* function)
{
    if (!function) {
        return {};
    } else if (function[0] != '_' || function[1] != 'Z') {
        return {function};
    }

    string ret;
    int status = 0;
    char* demangled = abi::__cxa_demangle(function, 0, 0, &status);
    if (demangled) {
        ret = demangled;
        free(demangled);
    }
    return ret;
}

struct AddressInformation
{
    string function;
    string file;
    int line = 0;
};

struct Module
{
    Module(uintptr_t addressStart, uintptr_t addressEnd, backtrace_state* backtraceState, size_t moduleIndex)
        : addressStart(addressStart)
        , addressEnd(addressEnd)
        , moduleIndex(moduleIndex)
        , backtraceState(backtraceState)
    {
    }

    AddressInformation resolveAddress(uintptr_t address) const
    {
        AddressInformation info;
        if (!backtraceState) {
            return info;
        }

        backtrace_pcinfo(backtraceState, address,
                         [] (void *data, uintptr_t /*addr*/, const char *file, int line, const char *function) -> int {
                            auto info = reinterpret_cast<AddressInformation*>(data);
                            info->function = demangle(function);
                            info->file = file ? file : "";
                            info->line = line;
                            return 0;
                         }, [] (void */*data*/, const char */*msg*/, int /*errnum*/) {}, &info);

        if (info.function.empty()) {
            backtrace_syminfo(backtraceState, address,
                              [] (void *data, uintptr_t /*pc*/, const char *symname, uintptr_t /*symval*/, uintptr_t /*symsize*/) {
                                if (symname) {
                                    reinterpret_cast<AddressInformation*>(data)->function = demangle(symname);
                                }
                              }, [] (void */*data*/, const char *msg, int errnum) {
                                cerr << "Module backtrace error (code " << errnum << "): " << msg << endl;
                              }, &info);
        }

        return info;
    }

    bool operator<(const Module& module) const
    {
        return make_tuple(addressStart, addressEnd, moduleIndex)
             < make_tuple(module.addressStart, module.addressEnd, module.moduleIndex);
    }

    bool operator!=(const Module& module) const
    {
        return make_tuple(addressStart, addressEnd, moduleIndex)
            != make_tuple(module.addressStart, module.addressEnd, module.moduleIndex);
    }

    uintptr_t addressStart;
    uintptr_t addressEnd;
    size_t moduleIndex;
    backtrace_state* backtraceState;
};

struct Allocation
{
    // backtrace entry point
    size_t ipIndex;
    // number of allocations
    size_t allocations;
    // amount of bytes leaked
    size_t leaked;
};

/**
 * Information for a single call to an allocation function
 */
struct AllocationInfo
{
    size_t ipIndex;
    size_t size;
};

struct ResolvedIP
{
    size_t moduleIndex = 0;
    size_t fileIndex = 0;
    size_t functionIndex = 0;
    int line = 0;
};

struct AccumulatedTraceData
{
    AccumulatedTraceData()
    {
        m_modules.reserve(256);
        m_backtraceStates.reserve(64);
        m_internedData.reserve(4096);
        m_encounteredIps.reserve(32768);
    }

    ~AccumulatedTraceData()
    {
        fprintf(stdout, "# strings: %lu\n# ips: %lu\n",
                m_internedData.size(), m_encounteredIps.size());
    }

    ResolvedIP resolve(const uintptr_t ip)
    {
        if (m_modulesDirty) {
            // sort by addresses, required for binary search below
            sort(m_modules.begin(), m_modules.end());

            for (size_t i = 0; i < m_modules.size(); ++i) {
                const auto& m1 = m_modules[i];
                for (size_t j = i + 1; j < m_modules.size(); ++j) {
                    if (i == j) {
                        continue;
                    }
                    const auto& m2 = m_modules[j];
                    if ((m1.addressStart <= m2.addressStart && m1.addressEnd > m2.addressStart) ||
                        (m1.addressStart < m2.addressEnd && m1.addressEnd >= m2.addressEnd))
                    {
                        cerr << "OVERLAPPING MODULES: " << hex
                             << m1.moduleIndex << " (" << m1.addressStart << " to " << m1.addressEnd << ") and "
                             << m1.moduleIndex << " (" << m2.addressStart << " to " << m2.addressEnd << ")\n"
                             << dec;
                    } else if (m2.addressStart >= m1.addressEnd) {
                        break;
                    }
                }
            }

            m_modulesDirty = false;
        }

        ResolvedIP data;
        // find module for this instruction pointer
        auto module = lower_bound(m_modules.begin(), m_modules.end(), ip,
                                    [] (const Module& module, const uintptr_t ip) -> bool {
                                        return module.addressEnd < ip;
                                    });
        if (module != m_modules.end() && module->addressStart <= ip && module->addressEnd >= ip) {
            data.moduleIndex = module->moduleIndex;
            const auto info = module->resolveAddress(ip);
            data.fileIndex = intern(info.file);
            data.functionIndex = intern(info.function);
            data.line = info.line;
        }
        return data;
    }

    size_t intern(const string& str)
    {
        if (str.empty()) {
            return 0;
        }

        auto it = m_internedData.find(str);
        if (it != m_internedData.end()) {
            return it->second;
        }
        const size_t id = m_internedData.size() + 1;
        m_internedData.insert(it, make_pair(str, id));
        fputs("s ", stdout);
        fputs(str.c_str(), stdout);
        fputc('\n', stdout);
        return id;
    }

    void addModule(backtrace_state* backtraceState, const size_t moduleIndex,
                   const uintptr_t addressStart, const uintptr_t addressEnd)
    {
        m_modules.emplace_back(addressStart, addressEnd, backtraceState, moduleIndex);
        m_modulesDirty = true;
    }

    void clearModules()
    {
        // TODO: optimize this, reuse modules that are still valid
        m_modules.clear();
        m_modulesDirty = true;
    }

    size_t addIp(const uintptr_t instructionPointer)
    {
        if (!instructionPointer) {
            return 0;
        }

        auto it = m_encounteredIps.find(instructionPointer);
        if (it != m_encounteredIps.end()) {
            return it->second;
        }

        const size_t ipId = m_encounteredIps.size() + 1;
        m_encounteredIps.insert(it, make_pair(instructionPointer, ipId));

        const auto ip = resolve(instructionPointer);
        fprintf(stdout, "i %lx %lx", instructionPointer, ip.moduleIndex);
        if (ip.functionIndex || ip.fileIndex) {
            fprintf(stdout, " %lx", ip.functionIndex);
            if (ip.fileIndex) {
                fprintf(stdout, " %lx %d", ip.fileIndex, ip.line);
            }
        }
        fputc('\n', stdout);
        return ipId;
    }

    /**
     * Prevent the same file from being initialized multiple times.
     * This drastically cuts the memory consumption down
     */
    backtrace_state* findBacktraceState(const string& fileName, uintptr_t addressStart, bool isExe)
    {
        if (boost::algorithm::starts_with(fileName, "linux-vdso.so")) {
            // prevent warning, since this will always fail
            return nullptr;
        }

        auto it = m_backtraceStates.find(fileName);
        if (it != m_backtraceStates.end()) {
            return it->second;
        }

        struct CallbackData {
            const string* fileName;
        };
        CallbackData data = {&fileName};

        auto errorHandler = [] (void *rawData, const char *msg, int errnum) {
            auto data = reinterpret_cast<const CallbackData*>(rawData);
            cerr << "Failed to create backtrace state for module " << *data->fileName << ": "
                 << msg << " (error code " << errnum << ")" << endl;
        };

        auto state = backtrace_create_state(fileName.c_str(), /* we are single threaded, so: not thread safe */ false,
                                            errorHandler, &data);

        if (state) {
            // when we could initialize the backtrace state, we initialize it with the first address
            // we get since that is the lowest one
            if (!addressStart && isExe) {
                // the address start should not be zero, that would lead to issues since
                // libbacktrace then tries to find the start address and fails
                // we just set it to 1 for the exe, which seems to work
                addressStart = 1;
            }
            backtrace_fileline_initialize(state, addressStart + isExe, isExe, errorHandler, &data);
        }

        m_backtraceStates.insert(it, make_pair(fileName, state));

        return state;
    }

private:
    vector<Module> m_modules;
    unordered_map<string, backtrace_state*> m_backtraceStates;
    bool m_modulesDirty = false;

    unordered_map<string, size_t> m_internedData;
    unordered_map<uintptr_t, size_t> m_encounteredIps;
};

}

int main(int /*argc*/, char** /*argv*/)
{
    // optimize: we only have a single thread
    ios_base::sync_with_stdio(false);

    AccumulatedTraceData data;

    LineReader reader;

    string exe;

    while (reader.getLine(cin)) {
        if (reader.mode() == 'x') {
            reader >> exe;
        } else if (reader.mode() == 'm') {
            string fileName;
            reader >> fileName;
            if (fileName == "-") {
                data.clearModules();
            } else {
                const bool isExe = (fileName == "x");
                if (isExe) {
                    fileName = exe;
                }
                const auto moduleIndex = data.intern(fileName);
                uintptr_t addressStart = 0;
                if (!(reader >> addressStart)) {
                    cerr << "failed to parse line: " << reader.line() << endl;
                    return 1;
                }
                auto state = data.findBacktraceState(fileName, addressStart, isExe);
                uintptr_t vAddr = 0;
                uintptr_t memSize = 0;
                while ((reader >> vAddr) && (reader >> memSize)) {
                    data.addModule(state, moduleIndex,
                                   addressStart + vAddr,
                                   addressStart + vAddr + memSize);
                }
            }
        } else if (reader.mode() == 't') {
            uintptr_t instructionPointer = 0;
            size_t parentIndex = 0;
            if (!(reader >> instructionPointer) || !(reader >> parentIndex)) {
                cerr << "failed to parse line: " << reader.line() << endl;
                return 1;
            }
            // ensure ip is encountered
            const auto ipId = data.addIp(instructionPointer);
            // trace point, map current output index to parent index
            fprintf(stdout, "t %lx %lx\n", ipId, parentIndex);
        } else {
            fputs(reader.line().c_str(), stdout);
            fputc('\n', stdout);
        }
    }

    return 0;
}
<commit_msg>Unset automatic locking of stdout and stdin.<commit_after>/*
 * Copyright 2014 Milian Wolff <mail@milianw.de>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Library General Public License as
 * published by the Free Software Foundation; either version 2 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public
 * License along with this program; if not, write to the
 * Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/**
 * @file heaptrack_interpret.cpp
 *
 * @brief Interpret raw heaptrack data and add Dwarf based debug information.
 */

#include <iostream>
#include <unordered_map>
#include <vector>
#include <tuple>
#include <algorithm>
#include <stdio_ext.h>

#include <cxxabi.h>
#include <boost/algorithm/string/predicate.hpp>

#include "libbacktrace/backtrace.h"
#include "linereader.h"

using namespace std;

namespace {

string demangle(const char* function)
{
    if (!function) {
        return {};
    } else if (function[0] != '_' || function[1] != 'Z') {
        return {function};
    }

    string ret;
    int status = 0;
    char* demangled = abi::__cxa_demangle(function, 0, 0, &status);
    if (demangled) {
        ret = demangled;
        free(demangled);
    }
    return ret;
}

struct AddressInformation
{
    string function;
    string file;
    int line = 0;
};

struct Module
{
    Module(uintptr_t addressStart, uintptr_t addressEnd, backtrace_state* backtraceState, size_t moduleIndex)
        : addressStart(addressStart)
        , addressEnd(addressEnd)
        , moduleIndex(moduleIndex)
        , backtraceState(backtraceState)
    {
    }

    AddressInformation resolveAddress(uintptr_t address) const
    {
        AddressInformation info;
        if (!backtraceState) {
            return info;
        }

        backtrace_pcinfo(backtraceState, address,
                         [] (void *data, uintptr_t /*addr*/, const char *file, int line, const char *function) -> int {
                            auto info = reinterpret_cast<AddressInformation*>(data);
                            info->function = demangle(function);
                            info->file = file ? file : "";
                            info->line = line;
                            return 0;
                         }, [] (void */*data*/, const char */*msg*/, int /*errnum*/) {}, &info);

        if (info.function.empty()) {
            backtrace_syminfo(backtraceState, address,
                              [] (void *data, uintptr_t /*pc*/, const char *symname, uintptr_t /*symval*/, uintptr_t /*symsize*/) {
                                if (symname) {
                                    reinterpret_cast<AddressInformation*>(data)->function = demangle(symname);
                                }
                              }, [] (void */*data*/, const char *msg, int errnum) {
                                cerr << "Module backtrace error (code " << errnum << "): " << msg << endl;
                              }, &info);
        }

        return info;
    }

    bool operator<(const Module& module) const
    {
        return make_tuple(addressStart, addressEnd, moduleIndex)
             < make_tuple(module.addressStart, module.addressEnd, module.moduleIndex);
    }

    bool operator!=(const Module& module) const
    {
        return make_tuple(addressStart, addressEnd, moduleIndex)
            != make_tuple(module.addressStart, module.addressEnd, module.moduleIndex);
    }

    uintptr_t addressStart;
    uintptr_t addressEnd;
    size_t moduleIndex;
    backtrace_state* backtraceState;
};

struct Allocation
{
    // backtrace entry point
    size_t ipIndex;
    // number of allocations
    size_t allocations;
    // amount of bytes leaked
    size_t leaked;
};

/**
 * Information for a single call to an allocation function
 */
struct AllocationInfo
{
    size_t ipIndex;
    size_t size;
};

struct ResolvedIP
{
    size_t moduleIndex = 0;
    size_t fileIndex = 0;
    size_t functionIndex = 0;
    int line = 0;
};

struct AccumulatedTraceData
{
    AccumulatedTraceData()
    {
        m_modules.reserve(256);
        m_backtraceStates.reserve(64);
        m_internedData.reserve(4096);
        m_encounteredIps.reserve(32768);
    }

    ~AccumulatedTraceData()
    {
        fprintf(stdout, "# strings: %lu\n# ips: %lu\n",
                m_internedData.size(), m_encounteredIps.size());
    }

    ResolvedIP resolve(const uintptr_t ip)
    {
        if (m_modulesDirty) {
            // sort by addresses, required for binary search below
            sort(m_modules.begin(), m_modules.end());

            for (size_t i = 0; i < m_modules.size(); ++i) {
                const auto& m1 = m_modules[i];
                for (size_t j = i + 1; j < m_modules.size(); ++j) {
                    if (i == j) {
                        continue;
                    }
                    const auto& m2 = m_modules[j];
                    if ((m1.addressStart <= m2.addressStart && m1.addressEnd > m2.addressStart) ||
                        (m1.addressStart < m2.addressEnd && m1.addressEnd >= m2.addressEnd))
                    {
                        cerr << "OVERLAPPING MODULES: " << hex
                             << m1.moduleIndex << " (" << m1.addressStart << " to " << m1.addressEnd << ") and "
                             << m1.moduleIndex << " (" << m2.addressStart << " to " << m2.addressEnd << ")\n"
                             << dec;
                    } else if (m2.addressStart >= m1.addressEnd) {
                        break;
                    }
                }
            }

            m_modulesDirty = false;
        }

        ResolvedIP data;
        // find module for this instruction pointer
        auto module = lower_bound(m_modules.begin(), m_modules.end(), ip,
                                    [] (const Module& module, const uintptr_t ip) -> bool {
                                        return module.addressEnd < ip;
                                    });
        if (module != m_modules.end() && module->addressStart <= ip && module->addressEnd >= ip) {
            data.moduleIndex = module->moduleIndex;
            const auto info = module->resolveAddress(ip);
            data.fileIndex = intern(info.file);
            data.functionIndex = intern(info.function);
            data.line = info.line;
        }
        return data;
    }

    size_t intern(const string& str)
    {
        if (str.empty()) {
            return 0;
        }

        auto it = m_internedData.find(str);
        if (it != m_internedData.end()) {
            return it->second;
        }
        const size_t id = m_internedData.size() + 1;
        m_internedData.insert(it, make_pair(str, id));
        fputs("s ", stdout);
        fputs(str.c_str(), stdout);
        fputc('\n', stdout);
        return id;
    }

    void addModule(backtrace_state* backtraceState, const size_t moduleIndex,
                   const uintptr_t addressStart, const uintptr_t addressEnd)
    {
        m_modules.emplace_back(addressStart, addressEnd, backtraceState, moduleIndex);
        m_modulesDirty = true;
    }

    void clearModules()
    {
        // TODO: optimize this, reuse modules that are still valid
        m_modules.clear();
        m_modulesDirty = true;
    }

    size_t addIp(const uintptr_t instructionPointer)
    {
        if (!instructionPointer) {
            return 0;
        }

        auto it = m_encounteredIps.find(instructionPointer);
        if (it != m_encounteredIps.end()) {
            return it->second;
        }

        const size_t ipId = m_encounteredIps.size() + 1;
        m_encounteredIps.insert(it, make_pair(instructionPointer, ipId));

        const auto ip = resolve(instructionPointer);
        fprintf(stdout, "i %lx %lx", instructionPointer, ip.moduleIndex);
        if (ip.functionIndex || ip.fileIndex) {
            fprintf(stdout, " %lx", ip.functionIndex);
            if (ip.fileIndex) {
                fprintf(stdout, " %lx %d", ip.fileIndex, ip.line);
            }
        }
        fputc('\n', stdout);
        return ipId;
    }

    /**
     * Prevent the same file from being initialized multiple times.
     * This drastically cuts the memory consumption down
     */
    backtrace_state* findBacktraceState(const string& fileName, uintptr_t addressStart, bool isExe)
    {
        if (boost::algorithm::starts_with(fileName, "linux-vdso.so")) {
            // prevent warning, since this will always fail
            return nullptr;
        }

        auto it = m_backtraceStates.find(fileName);
        if (it != m_backtraceStates.end()) {
            return it->second;
        }

        struct CallbackData {
            const string* fileName;
        };
        CallbackData data = {&fileName};

        auto errorHandler = [] (void *rawData, const char *msg, int errnum) {
            auto data = reinterpret_cast<const CallbackData*>(rawData);
            cerr << "Failed to create backtrace state for module " << *data->fileName << ": "
                 << msg << " (error code " << errnum << ")" << endl;
        };

        auto state = backtrace_create_state(fileName.c_str(), /* we are single threaded, so: not thread safe */ false,
                                            errorHandler, &data);

        if (state) {
            // when we could initialize the backtrace state, we initialize it with the first address
            // we get since that is the lowest one
            if (!addressStart && isExe) {
                // the address start should not be zero, that would lead to issues since
                // libbacktrace then tries to find the start address and fails
                // we just set it to 1 for the exe, which seems to work
                addressStart = 1;
            }
            backtrace_fileline_initialize(state, addressStart + isExe, isExe, errorHandler, &data);
        }

        m_backtraceStates.insert(it, make_pair(fileName, state));

        return state;
    }

private:
    vector<Module> m_modules;
    unordered_map<string, backtrace_state*> m_backtraceStates;
    bool m_modulesDirty = false;

    unordered_map<string, size_t> m_internedData;
    unordered_map<uintptr_t, size_t> m_encounteredIps;
};

}

int main(int /*argc*/, char** /*argv*/)
{
    // optimize: we only have a single thread
    ios_base::sync_with_stdio(false);

    AccumulatedTraceData data;

    LineReader reader;

    string exe;

    __fsetlocking(stdout, FSETLOCKING_BYCALLER);
    __fsetlocking(stdin, FSETLOCKING_BYCALLER);

    while (reader.getLine(cin)) {
        if (reader.mode() == 'x') {
            reader >> exe;
        } else if (reader.mode() == 'm') {
            string fileName;
            reader >> fileName;
            if (fileName == "-") {
                data.clearModules();
            } else {
                const bool isExe = (fileName == "x");
                if (isExe) {
                    fileName = exe;
                }
                const auto moduleIndex = data.intern(fileName);
                uintptr_t addressStart = 0;
                if (!(reader >> addressStart)) {
                    cerr << "failed to parse line: " << reader.line() << endl;
                    return 1;
                }
                auto state = data.findBacktraceState(fileName, addressStart, isExe);
                uintptr_t vAddr = 0;
                uintptr_t memSize = 0;
                while ((reader >> vAddr) && (reader >> memSize)) {
                    data.addModule(state, moduleIndex,
                                   addressStart + vAddr,
                                   addressStart + vAddr + memSize);
                }
            }
        } else if (reader.mode() == 't') {
            uintptr_t instructionPointer = 0;
            size_t parentIndex = 0;
            if (!(reader >> instructionPointer) || !(reader >> parentIndex)) {
                cerr << "failed to parse line: " << reader.line() << endl;
                return 1;
            }
            // ensure ip is encountered
            const auto ipId = data.addIp(instructionPointer);
            // trace point, map current output index to parent index
            fprintf(stdout, "t %lx %lx\n", ipId, parentIndex);
        } else {
            fputs(reader.line().c_str(), stdout);
            fputc('\n', stdout);
        }
    }

    return 0;
}
<|endoftext|>