void visit(const LetStmt *op) {
is_impure = false;
Expr value = mutate(op->value);
Stmt body = op->body;
bool should_pop = false;
bool should_erase = false;
if (!is_impure) {
map<Expr, string, IRDeepCompare>::iterator iter = scope.find(value);
if (iter == scope.end()) {
scope[value] = op->name;
should_pop = true;
} else {
value = Variable::make(value.type(), iter->second);
rewrites[op->name] = iter->second;
should_erase = true;
}
}
body = mutate(op->body);
if (should_pop) {
scope.erase(value);
}
if (should_erase) {
rewrites.erase(op->name);
}
if (value.same_as(op->value) && body.same_as(op->body)) {
stmt = op;
} else {
stmt = LetStmt::make(op->name, value, body);
}
}
void visit(const LetStmt *let) {
Expr var = canonicalize(Variable::make(let->value.type(), let->name));
Expr new_value = mutate(let->value);
enter_scope();
add_replacement(var, new_value);
stmt = mutate(let->body);
leave_scope();
}
void visit(const LetStmt *op) {
Expr new_value = mutate(op->value);
hidden.push(op->name, 0);
Stmt new_body = mutate(op->body);
hidden.pop(op->name);
if (new_value.same_as(op->value) &&
new_body.same_as(op->body)) {
stmt = op;
} else {
stmt = LetStmt::make(op->name, new_value, new_body);
}
}
void visit(const For *op) {
Expr new_min = mutate(op->min);
Expr new_extent = mutate(op->extent);
hidden.push(op->name, 0);
Stmt new_body = mutate(op->body);
hidden.pop(op->name);
if (new_min.same_as(op->min) &&
new_extent.same_as(op->extent) &&
new_body.same_as(op->body)) {
stmt = op;
} else {
stmt = For::make(op->name, new_min, new_extent, op->for_type, new_body);
}
}
void visit(const For *op) {
bool old_in_gpu_loop = in_gpu_loop;
in_gpu_loop =
(op->for_type == ForType::GPUBlock ||
op->for_type == ForType::GPUThread);
if (old_in_gpu_loop && in_gpu_loop) {
// Don't lift lets to in-between gpu blocks/threads
IRMutator::visit(op);
} else if (op->device_api == DeviceAPI::GLSL ||
op->device_api == DeviceAPI::OpenGLCompute) {
// Don't lift anything out of OpenGL loops
IRMutator::visit(op);
} else {
// Lift invariants
LiftLoopInvariants lifter;
Stmt new_stmt = lifter.mutate(op);
// Recurse
const For *loop = new_stmt.as<For>();
internal_assert(loop);
new_stmt = For::make(loop->name, loop->min, loop->extent,
loop->for_type, loop->device_api, mutate(loop->body));
// Wrap lets for the lifted invariants
for (const auto &p : lifter.lifted) {
new_stmt = LetStmt::make(p.second, p.first, new_stmt);
}
stmt = new_stmt;
}
in_gpu_loop = old_in_gpu_loop;
}