Revision 69e4dffe991aafe28a7457fbe7cc3978a8074020 authored by Andrew Adams on 21 February 2021, 17:31:49 UTC, committed by Andrew Adams on 21 February 2021, 17:31:49 UTC
Inline.cpp
#include <set>
#include "CSE.h"
#include "Debug.h"
#include "ExternFuncArgument.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "IRPrinter.h"
#include "Inline.h"
#include "Qualify.h"
#include "Substitute.h"
namespace Halide {
namespace Internal {
using std::string;
using std::vector;
// Sanity check that this is a reasonable function to inline
void validate_schedule_inlined_function(Function f) {
const FuncSchedule &func_s = f.schedule();
const StageSchedule &stage_s = f.definition().schedule();
if (!func_s.store_level().is_inlined()) {
user_error << "Function " << f.name() << " is scheduled to be computed inline, "
<< "but is not scheduled to be stored inline. A storage schedule "
<< "is meaningless for functions computed inline.\n";
}
// Inlining is allowed only if there is no specialization.
user_assert(f.definition().specializations().empty())
<< "Function " << f.name() << " is scheduled inline, so it"
<< " must not have any specializations. Specialize on the"
<< " scheduled function instead.\n";
if (func_s.memoized()) {
user_error << "Cannot memoize function "
<< f.name() << " because the function is scheduled inline.\n";
}
for (size_t i = 0; i < stage_s.dims().size(); i++) {
Dim d = stage_s.dims()[i];
if (d.is_unordered_parallel()) {
user_error << "Cannot parallelize dimension "
<< d.var << " of function "
<< f.name() << " because the function is scheduled inline.\n";
} else if (d.for_type == ForType::Unrolled) {
user_error << "Cannot unroll dimension "
<< d.var << " of function "
<< f.name() << " because the function is scheduled inline.\n";
} else if (d.for_type == ForType::Vectorized) {
user_error << "Cannot vectorize dimension "
<< d.var << " of function "
<< f.name() << " because the function is scheduled inline.\n";
}
}
for (size_t i = 0; i < stage_s.splits().size(); i++) {
if (stage_s.splits()[i].is_rename()) {
user_warning << "It is meaningless to rename variable "
<< stage_s.splits()[i].old_var << " of function "
<< f.name() << " to " << stage_s.splits()[i].outer
<< " because " << f.name() << " is scheduled inline.\n";
} else if (stage_s.splits()[i].is_fuse()) {
user_warning << "It is meaningless to fuse variables "
<< stage_s.splits()[i].inner << " and " << stage_s.splits()[i].outer
<< " because " << f.name() << " is scheduled inline.\n";
} else {
user_warning << "It is meaningless to split variable "
<< stage_s.splits()[i].old_var << " of function "
<< f.name() << " into "
<< stage_s.splits()[i].outer << " * "
<< stage_s.splits()[i].factor << " + "
<< stage_s.splits()[i].inner << " because "
<< f.name() << " is scheduled inline.\n";
}
}
for (size_t i = 0; i < func_s.bounds().size(); i++) {
if (func_s.bounds()[i].min.defined()) {
user_warning << "It is meaningless to bound dimension "
<< func_s.bounds()[i].var << " of function "
<< f.name() << " to be within ["
<< func_s.bounds()[i].min << ", "
<< func_s.bounds()[i].extent << "] because the function is scheduled inline.\n";
} else if (func_s.bounds()[i].modulus.defined()) {
user_warning << "It is meaningless to align the bounds of dimension "
<< func_s.bounds()[i].var << " of function "
<< f.name() << " to have modulus/remainder ["
<< func_s.bounds()[i].modulus << ", "
<< func_s.bounds()[i].remainder << "] because the function is scheduled inline.\n";
}
}
}
class Inliner : public IRMutator {
using IRMutator::visit;
Function func;
Expr visit(const Call *op) override {
if (op->name == func.name()) {
// Mutate the args
vector<Expr> args(op->args.size());
for (size_t i = 0; i < args.size(); i++) {
args[i] = mutate(op->args[i]);
}
// Grab the body
Expr body = qualify(func.name() + ".", func.values()[op->value_index]);
const vector<string> func_args = func.args();
// Bind the args using Let nodes
internal_assert(args.size() == func_args.size());
for (size_t i = 0; i < args.size(); i++) {
if (is_const(args[i]) || args[i].as<Variable>()) {
body = substitute(func.name() + "." + func_args[i], args[i], body);
} else {
body = Let::make(func.name() + "." + func_args[i], args[i], body);
}
}
found++;
return body;
} else {
return IRMutator::visit(op);
}
}
Expr visit(const Variable *op) override {
if (op->name == func.name() + ".buffer") {
const Call *call = func.is_wrapper();
internal_assert(call);
// Do a whole-image inline. Substitute the .buffer symbol
// for the wrapped object's .buffer symbol.
string buf_name;
if (call->call_type == Call::Halide) {
buf_name = call->name;
if (Function(call->func).outputs() > 1) {
buf_name += "." + std::to_string(call->value_index);
}
buf_name += ".buffer";
return Variable::make(type_of<halide_buffer_t *>(), buf_name);
} else if (call->param.defined()) {
return Variable::make(type_of<halide_buffer_t *>(), call->name + ".buffer", call->param);
} else {
internal_assert(call->image.defined());
return Variable::make(type_of<halide_buffer_t *>(), call->name + ".buffer", call->image);
}
} else {
return op;
}
}
Stmt visit(const Provide *op) override {
ScopedValue<int> old_found(found, 0);
Stmt stmt = IRMutator::visit(op);
// TODO: making this > 1 should be desirable,
// but explodes compiletimes in some situations.
if (found > 0) {
stmt = common_subexpression_elimination(stmt);
}
return stmt;
}
public:
int found = 0;
Inliner(const Function &f)
: func(f) {
internal_assert(f.can_be_inlined()) << "Illegal to inline " << f.name() << "\n";
validate_schedule_inlined_function(f);
}
};
Stmt inline_function(Stmt s, const Function &f) {
Inliner i(f);
s = i.mutate(s);
return s;
}
Expr inline_function(Expr e, const Function &f) {
Inliner i(f);
e = i.mutate(e);
// TODO: making this > 1 should be desirable,
// but explodes compiletimes in some situations.
if (i.found > 0) {
e = common_subexpression_elimination(e);
}
return e;
}
// Inline all calls to 'f' inside 'caller'
void inline_function(Function caller, const Function &f) {
Inliner i(f);
caller.mutate(&i);
if (caller.has_extern_definition()) {
for (ExternFuncArgument &arg : caller.extern_arguments()) {
if (arg.is_func() && arg.func.same_as(f.get_contents())) {
const Call *call = f.is_wrapper();
internal_assert(call);
arg.func = call->func;
}
}
}
}
} // namespace Internal
} // namespace Halide
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...