<commit_before>#include "Tracing.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "runtime/HalideRuntime.h"

namespace Halide {
namespace Internal {

int tracing_level() {
    char *trace = getenv("HL_TRACE");
    return trace ? atoi(trace) : 0;
}

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

class InjectTracing : public IRMutator {
public:
    const map<string, Function> &env;
    Function output;
    int global_level;
    InjectTracing(const map<string, Function> &e,
                  Function o) : env(e),
                                output(o),
                                global_level(tracing_level()) {}

private:
    using IRMutator::visit;

    void visit(const Call *op) {

        // Calls inside of an address_of don't count, but we want to
        // visit the args of the inner call.
        if (op->call_type == Call::Intrinsic && op->name == Call::address_of) {
            internal_assert(op->args.size() == 1);
            const Call *c = op->args[0].as<Call>();
            const Load *l = op->args[0].as<Load>();

            internal_assert(c || l);

            std::vector<Expr> args;
            if (c) {
                args = c->args;
            } else {
                args.push_back(l->index);
            }

            bool unchanged = true;
            vector<Expr> new_args(args.size());
            for (size_t i = 0; i < args.size(); i++) {
                new_args[i] = mutate(args[i]);
                unchanged = unchanged && (new_args[i].same_as(args[i]));
            }

            if (unchanged) {
                expr = op;
                return;
            } else {
                Expr inner;
                if (c) {
                    inner = Call::make(c->type, c->name, new_args, c->call_type,
                                       c->func, c->value_index, c->image, c->param);
                } else {
                    Expr inner = Load::make(l->type, l->name, new_args[0], l->image, l->param);
                }
                expr = Call::make(Handle(), Call::address_of, vec(inner), Call::Intrinsic);
                return;
            }
        }



        IRMutator::visit(op);
        op = expr.as<Call>();
        internal_assert(op);

        if (op->call_type != Call::Halide) {
            return;
        }

        Function f = op->func;
        bool inlined = !f.same_as(output) && f.schedule().compute_level().is_inline();

        if (f.is_tracing_loads() || (global_level > 2 && !inlined)) {

            // Wrap the load in a call to trace_load
            vector<Expr> args;
            args.push_back(f.name());
            args.push_back(halide_trace_load);
            args.push_back(Variable::make(Int(32), op->name + ".trace_id"));
            args.push_back(op->value_index);
            args.push_back(op);
            args.insert(args.end(), op->args.begin(), op->args.end());

            expr = Call::make(op->type, Call::trace_expr, args, Call::Intrinsic);
        }

    }

    void visit(const Provide *op) {
        IRMutator::visit(op);
        op = stmt.as<Provide>();
        internal_assert(op);

        map<string, Function>::const_iterator iter = env.find(op->name);
        if (iter == env.end()) return;
        Function f = iter->second;
        bool inlined = !f.same_as(output) && f.schedule().compute_level().is_inline();

        if (f.is_tracing_stores() || (global_level > 1 && !inlined)) {
            // Wrap each expr in a tracing call

            const vector<Expr> &values = op->values;
            vector<Expr> traces(op->values.size());
<commit_msg>tracing also reports loads from input images<commit_after>#include "Tracing.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "runtime/HalideRuntime.h"

namespace Halide {
namespace Internal {

int tracing_level() {
    char *trace = getenv("HL_TRACE");
    return trace ? atoi(trace) : 0;
}

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

class InjectTracing : public IRMutator {
public:
    const map<string, Function> &env;
    Function output;
    int global_level;
    InjectTracing(const map<string, Function> &e,
                  Function o) : env(e),
                                output(o),
                                global_level(tracing_level()) {}

private:
    using IRMutator::visit;

    void visit(const Call *op) {

        // Calls inside of an address_of don't count, but we want to
        // visit the args of the inner call.
        if (op->call_type == Call::Intrinsic && op->name == Call::address_of) {
            internal_assert(op->args.size() == 1);
            const Call *c = op->args[0].as<Call>();
            const Load *l = op->args[0].as<Load>();

            internal_assert(c || l);

            std::vector<Expr> args;
            if (c) {
                args = c->args;
            } else {
                args.push_back(l->index);
            }

            bool unchanged = true;
            vector<Expr> new_args(args.size());
            for (size_t i = 0; i < args.size(); i++) {
                new_args[i] = mutate(args[i]);
                unchanged = unchanged && (new_args[i].same_as(args[i]));
            }

            if (unchanged) {
                expr = op;
                return;
            } else {
                Expr inner;
                if (c) {
                    inner = Call::make(c->type, c->name, new_args, c->call_type,
                                       c->func, c->value_index, c->image, c->param);
                } else {
                    Expr inner = Load::make(l->type, l->name, new_args[0], l->image, l->param);
                }
                expr = Call::make(Handle(), Call::address_of, vec(inner), Call::Intrinsic);
                return;
            }
        }



        IRMutator::visit(op);
        op = expr.as<Call>();
        internal_assert(op);

        bool trace_it = false;
        Expr trace_parent;
        if (op->call_type == Call::Halide) {
            Function f = op->func;
            bool inlined = !f.same_as(output) && f.schedule().compute_level().is_inline();
            trace_it = f.is_tracing_loads() || (global_level > 2 && !inlined);
            trace_parent = Variable::make(Int(32), op->name + ".trace_id");
        } else if (op->call_type == Call::Image) {
            trace_it = global_level > 2;
            trace_parent = -1;
        }

        if (trace_it) {

            // Wrap the load in a call to trace_load
            vector<Expr> args;
            args.push_back(op->name);
            args.push_back(halide_trace_load);
            args.push_back(trace_parent);
            args.push_back(op->value_index);
            args.push_back(op);
            args.insert(args.end(), op->args.begin(), op->args.end());

            expr = Call::make(op->type, Call::trace_expr, args, Call::Intrinsic);
        }

    }

    void visit(const Provide *op) {
        IRMutator::visit(op);
        op = stmt.as<Provide>();
        internal_assert(op);

        map<string, Function>::const_iterator iter = env.find(op->name);
        if (iter == env.end()) return;
        Function f = iter->second;
        bool inlined = !f.same_as(output) && f.schedule().compute_level().is_inline();

        if (f.is_tracing_stores() || (global_level > 1 && !inlined)) {
            // Wrap each expr in a tracing call

            const vector<Expr> &values = op->values;
            vector<Expr> traces(op->values.size());

{'fraction_non_alphanumeric': 0.14423463238248294, 'fraction_numerical': 0.0029462970403107006, 'mean_word_length': 1.7455882352941177, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 92, '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': '3898221', 'n_tokens_mistral': 2316, 'n_tokens_neox': 2209, 'n_words': 493}