<commit_before>#include "StoreForwarding.h"
#include "IRMutator.h"
#include "Substitute.h"
#include "IROperator.h"
#include "Simplify.h"
#include "IREquality.h"

namespace Halide {
namespace Internal {

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

class FindLoadsAndStores : public IRVisitor {
    void visit(const Load *op) {
        debug(0) << "Got a load: " << op->name << "\n";
        loads.push_back(op);
        IRVisitor::visit(op);
    }

    void visit(const Store *op) {
        debug(0) << "Got a load: " << op->name << "\n";
        stores.push_back(op);
        IRVisitor::visit(op);
    }

    void visit(const IfThenElse *op) {
        // Conservatively don't enter a conditional
        op->condition.accept(this);
    }

    void visit(const For *op) {
        // We don't want to lift loads and stores that depend on inner
        // loop variables. For now we just don't descend into inner
        // loops (TODO: fix this).
    }

    // Consider lets

public:
    vector<const Load *> loads;
    vector<const Store *> stores;


};

class UseScratch : public IRMutator {
    const Load *load;
    const Store *store;
    const string &name;

    void visit(const Load *op) {
        if (equal(Expr(op), Expr(load))) {
            expr = Load::make(op->type, name, make_zero(Int(32, op->type.width)),
                              Buffer(), Parameter());
        } else {
            IRMutator::visit(op);
        }
    }    

    void visit(const Store *op) {
        if (op->name == store->name &&
            equal(op->index, store->index)) {
            Expr value = mutate(op->value);

            string var_name = unique_name('t');
            Expr var = Variable::make(value.type(), var_name);
            Stmt store_to_orig = Store::make(op->name, var, op->index);
            Stmt store_to_scratch = Store::make(name, var, make_zero(Int(32, value.type().width)));
            stmt = Block::make(store_to_orig, store_to_scratch);
            stmt = LetStmt::make(var_name, value, stmt);

            // TODO: What if the load is in the store index itself?
        } else {
            IRMutator::visit(op);
        }
    }
    
public:
    UseScratch(const Load *l, const Store *s, const string &n) : load(l), store(s), name(n) {}
};

class StoreForwarding : public IRMutator {
    void visit(const For *op) {
        Stmt body = mutate(op->body);

        // Find all the unconditional loads and stores
        FindLoadsAndStores f;<commit_msg>Add another todo<commit_after>#include "StoreForwarding.h"
#include "IRMutator.h"
#include "Substitute.h"
#include "IROperator.h"
#include "Simplify.h"
#include "IREquality.h"

namespace Halide {
namespace Internal {

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

class FindLoadsAndStores : public IRVisitor {
    void visit(const Load *op) {
        debug(0) << "Got a load: " << op->name << "\n";
        loads.push_back(op);
        IRVisitor::visit(op);
    }

    void visit(const Store *op) {
        debug(0) << "Got a load: " << op->name << "\n";
        stores.push_back(op);
        IRVisitor::visit(op);
    }

    void visit(const IfThenElse *op) {
        // Conservatively don't enter a conditional
        op->condition.accept(this);
    }

    void visit(const For *op) {
        // We don't want to lift loads and stores that depend on inner
        // loop variables. For now we just don't descend into inner
        // loops (TODO: fix this).
    }

    // Consider lets

public:
    vector<const Load *> loads;
    vector<const Store *> stores;


};

class UseScratch : public IRMutator {
    const Load *load;
    const Store *store;
    const string &name;

    void visit(const Load *op) {
        if (equal(Expr(op), Expr(load))) {
            expr = Load::make(op->type, name, make_zero(Int(32, op->type.width)),
                              Buffer(), Parameter());
        } else {
            IRMutator::visit(op);
        }
    }    

    void visit(const Store *op) {
        if (op->name == store->name &&
            equal(op->index, store->index)) {
            Expr value = mutate(op->value);

            string var_name = unique_name('t');
            Expr var = Variable::make(value.type(), var_name);
            Stmt store_to_orig = Store::make(op->name, var, op->index);
            Stmt store_to_scratch = Store::make(name, var, make_zero(Int(32, value.type().width)));
            stmt = Block::make(store_to_orig, store_to_scratch);
            stmt = LetStmt::make(var_name, value, stmt);

            // TODO: What if the load is in the store index itself?
        } else {
            IRMutator::visit(op);
        }
    }
    
public:
    UseScratch(const Load *l, const Store *s, const string &n) : load(l), store(s), name(n) {}
};

class StoreForwarding : public IRMutator {
    void visit(const For *op) {
        Stmt body = mutate(op->body);

        // TODO: skip parallel loops

        // Find all the unconditional loads and stores
        FindLoadsAndStores f;
{'fraction_non_alphanumeric': 0.13556939144406507, 'fraction_numerical': 0.002410122514561157, 'mean_word_length': 2.265573770491803, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5684833', 'n_tokens_mistral': 1598, 'n_tokens_neox': 1533, 'n_words': 423}