void Function::define_reduction(const vector<Expr> &args, Expr value) {
assert(defined() && "Can't add a reduction definition without a regular definition first");
assert(!is_reduction() && "Function already has a reduction definition");
assert(value.defined() && "Undefined expression in right-hand-side of reduction");
// Check the dimensionality matches
assert(args.size() == contents.ptr->args.size() &&
"Dimensionality of reduction definition must match dimensionality of pure definition");
// The pure args are those naked vars in the args that are not in
// a reduction domain and are not parameters
vector<string> pure_args;
for (size_t i = 0; i < args.size(); i++) {
assert(args[i].defined() && "Undefined expression in left-hand-side of reduction");
if (const Variable *var = args[i].as<Variable>()) {
if (!var->param.defined() && !var->reduction_domain.defined()) {
assert(var->name == contents.ptr->args[i] && "Pure argument to update step must have the same name as pure argument to initialization step in the same dimension");
pure_args.push_back(var->name);
}
}
}
// Make sure all the vars in the args and the value are either
// pure args, in the reduction domain, or a parameter
CheckVars check;
check.pure_args = pure_args;
value.accept(&check);
for (size_t i = 0; i < args.size(); i++) {
args[i].accept(&check);
}
assert(check.reduction_domain.defined() && "No reduction domain referenced in reduction definition");
contents.ptr->reduction_args = args;
contents.ptr->reduction_value = value;
contents.ptr->reduction_domain = check.reduction_domain;
// First add the pure args in order
for (size_t i = 0; i < pure_args.size(); i++) {
Schedule::Dim d = {pure_args[i], For::Serial};
contents.ptr->reduction_schedule.dims.push_back(d);
}
// Then add the reduction domain outside of that
for (size_t i = 0; i < check.reduction_domain.domain().size(); i++) {
Schedule::Dim d = {check.reduction_domain.domain()[i].var, For::Serial};
contents.ptr->reduction_schedule.dims.push_back(d);
}
}
inline bool implemented_as_kernel(scheduler::op_element const & op){
return op.type_subfamily==scheduler::OPERATION_FUNCTION_TYPE_SUBFAMILY || is_reduction(op);
}
