<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);<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);
{'fraction_non_alphanumeric': 0.14175769612711023, 'fraction_numerical': 0.005627275736511089, 'mean_word_length': 1.7108568864961866, 'pattern_counts': {'":': 1, '<': 225, '<?xml version=': 0, '>': 75, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4692275', 'n_tokens_mistral': 3802, 'n_tokens_neox': 3599, 'n_words': 886}