RemoveDeadAllocations.cpp
#include "RemoveDeadAllocations.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "Scope.h"
namespace Halide {
namespace Internal {
namespace {
class RemoveDeadAllocations : public IRMutator {
using IRMutator::visit;
Scope<int> allocs;
Expr visit(const Call *op) override {
if (op->is_extern()) {
for (const auto &arg : op->args) {
const Variable *var = arg.as<Variable>();
if (var && ends_with(var->name, ".buffer")) {
std::string func = var->name.substr(0, var->name.find_first_of('.'));
if (allocs.contains(func)) {
allocs.pop(func);
}
}
}
}
return IRMutator::visit(op);
}
Expr visit(const Load *op) override {
if (allocs.contains(op->name)) {
allocs.pop(op->name);
}
return IRMutator::visit(op);
}
Stmt visit(const Store *op) override {
if (allocs.contains(op->name)) {
allocs.pop(op->name);
}
return IRMutator::visit(op);
}
Expr visit(const Variable *op) override {
if (allocs.contains(op->name)) {
allocs.pop(op->name);
}
return op;
}
Stmt visit(const Allocate *op) override {
allocs.push(op->name, 1);
Stmt body = mutate(op->body);
if (allocs.contains(op->name) && op->free_function.empty()) {
allocs.pop(op->name);
return body;
} else if (body.same_as(op->body)) {
return op;
} else {
return Allocate::make(op->name, op->type, op->memory_type, op->extents,
op->condition, body, op->new_expr, op->free_function);
}
}
Stmt visit(const Free *op) override {
if (allocs.contains(op->name)) {
// We have reached a Free Stmt without ever using this buffer, do nothing.
return Evaluate::make(0);
} else {
return op;
}
}
Stmt visit(const Atomic *op) override {
if (allocs.contains(op->mutex_name)) {
allocs.pop(op->mutex_name);
}
return IRMutator::visit(op);
}
};
} // namespace
Stmt remove_dead_allocations(const Stmt &s) {
return RemoveDeadAllocations().mutate(s);
}
} // namespace Internal
} // namespace Halide