Func repeat_edge(const Func &source,
const std::vector<std::pair<Expr, Expr>> &bounds) {
std::vector<Var> args(source.args());
user_assert(args.size() >= bounds.size()) <<
"repeat_edge called with more bounds (" << bounds.size() <<
") than dimensions (" << args.size() << ") Func " <<
source.name() << "has.\n";
std::vector<Expr> actuals;
for (size_t i = 0; i < bounds.size(); i++) {
Var arg_var = args[i];
Expr min = bounds[i].first;
Expr extent = bounds[i].second;
if (min.defined() && extent.defined()) {
actuals.push_back(clamp(likely(arg_var), min, min + extent - 1));
} else if (!min.defined() && !extent.defined()) {
actuals.push_back(arg_var);
} else {
user_error << "Partially undefined bounds for dimension " << arg_var
<< " of Func " << source.name() << "\n";
}
}
// If there were fewer bounds than dimensions, regard the ones at the end as unbounded.
actuals.insert(actuals.end(), args.begin() + actuals.size(), args.end());
Func bounded("repeat_edge");
bounded(args) = source(actuals);
return bounded;
}
Internal::ReductionDomain build_domain(string name0, Expr min0, Expr extent0,
string name1, Expr min1, Expr extent1,
string name2, Expr min2, Expr extent2,
string name3, Expr min3, Expr extent3) {
vector<Internal::ReductionVariable> d;
if (min0.defined()) {
Internal::ReductionVariable v = {name0, min0, extent0};
d.push_back(v);
}
if (min1.defined()) {
Internal::ReductionVariable v = {name1, min1, extent1};
d.push_back(v);
}
if (min2.defined()) {
Internal::ReductionVariable v = {name2, min2, extent2};
d.push_back(v);
}
if (min3.defined()) {
Internal::ReductionVariable v = {name3, min3, extent3};
d.push_back(v);
}
Internal::ReductionDomain dom(d);
return dom;
}
Stmt Store::make(std::string name, Expr value, Expr index) {
internal_assert(value.defined()) << "Store of undefined\n";
internal_assert(index.defined()) << "Store of undefined\n";
Store *node = new Store;
node->name = name;
node->value = value;
node->index = index;
return node;
}
Expr GE::make(Expr a, Expr b) {
internal_assert(a.defined()) << "GE of undefined\n";
internal_assert(b.defined()) << "GE of undefined\n";
internal_assert(a.type() == b.type()) << "GE of mismatched types\n";
GE *node = new GE;
node->type = Bool(a.type().width);
node->a = a;
node->b = b;
return node;
}
// Order a pair of Exprs, treating undefined Exprs as infinity
void sort2(Expr &a, Expr &b) {
if (!a.defined()) {
std::swap(a, b);
} else if (!b.defined()) {
return;
} else {
Expr tmp = min(a, b);
b = max(a, b);
a = tmp;
}
}
Expr Mod::make(Expr a, Expr b) {
internal_assert(a.defined()) << "Mod of undefined\n";
internal_assert(b.defined()) << "Mod of undefined\n";
internal_assert(a.type() == b.type()) << "Mod of mismatched types\n";
Mod *node = new Mod;
node->type = a.type();
node->a = a;
node->b = b;
return node;
}
Expr Sub::make(Expr a, Expr b) {
internal_assert(a.defined()) << "Sub of undefined\n";
internal_assert(b.defined()) << "Sub of undefined\n";
internal_assert(a.type() == b.type()) << "Sub of mismatched types\n";
Sub *node = new Sub;
node->type = a.type();
node->a = a;
node->b = b;
return node;
}
Expr Let::make(std::string name, Expr value, Expr body) {
internal_assert(value.defined()) << "Let of undefined\n";
internal_assert(body.defined()) << "Let of undefined\n";
Let *node = new Let;
node->type = body.type();
node->name = name;
node->value = value;
node->body = body;
return node;
}
Expr Div::make(Expr a, Expr b) {
internal_assert(a.defined()) << "Div of undefined\n";
internal_assert(b.defined()) << "Div of undefined\n";
internal_assert(a.type() == b.type()) << "Div of mismatched types\n";
Div *node = new Div;
node->type = a.type();
node->a = a;
node->b = b;
return node;
}
Expr And::make(Expr a, Expr b) {
internal_assert(a.defined()) << "And of undefined\n";
internal_assert(b.defined()) << "And of undefined\n";
internal_assert(a.type().is_bool()) << "lhs of And is not a bool\n";
internal_assert(b.type().is_bool()) << "rhs of And is not a bool\n";
And *node = new And;
node->type = Bool(a.type().width);
node->a = a;
node->b = b;
return node;
}
Expr Or::make(Expr a, Expr b) {
internal_assert(a.defined()) << "Or of undefined\n";
internal_assert(b.defined()) << "Or of undefined\n";
internal_assert(a.type().is_bool()) << "lhs of Or is not a bool\n";
internal_assert(b.type().is_bool()) << "rhs of Or is not a bool\n";
Or *node = new Or;
node->type = Bool(a.type().width);
node->a = a;
node->b = b;
return node;
}
Expr box_size(const Box &b) {
Expr size = make_one(Int(64));
for (size_t i = 0; i < b.size(); i++) {
Expr extent = get_extent(b[i]);
if (extent.defined() && size.defined()) {
size *= extent;
} else if (is_zero(extent)) {
return make_zero(Int(64));
} else {
return Expr();
}
}
return simplify(size);
}
Stmt For::make(std::string name, Expr min, Expr extent, ForType for_type, Stmt body) {
internal_assert(min.defined()) << "For of undefined\n";
internal_assert(extent.defined()) << "For of undefined\n";
internal_assert(min.type().is_scalar()) << "For with vector min\n";
internal_assert(extent.type().is_scalar()) << "For with vector extent\n";
internal_assert(body.defined()) << "For of undefined\n";
For *node = new For;
node->name = name;
node->min = min;
node->extent = extent;
node->for_type = for_type;
node->body = body;
return node;
}
Expr Ramp::make(Expr base, Expr stride, int width) {
internal_assert(base.defined()) << "Ramp of undefined\n";
internal_assert(stride.defined()) << "Ramp of undefined\n";
internal_assert(base.type().is_scalar()) << "Ramp with vector base\n";
internal_assert(stride.type().is_scalar()) << "Ramp with vector stride\n";
internal_assert(width > 1) << "Ramp of width <= 1\n";
internal_assert(stride.type() == base.type()) << "Ramp of mismatched types\n";
Ramp *node = new Ramp;
node->type = base.type().vector_of(width);
node->base = base;
node->stride = stride;
node->width = width;
return node;
}
void visit(const FieldWrite *op) {
Expr elemOrSet = op->elementOrSet;
std::string fieldName = op->fieldName;
// If the same field is read and written in the same statement and the
// values are combined/reduced (e.g. multiplied) then we must introduce a
// temporary to avoid read/write interference.
Expr fieldRead = GetFieldRead(elemOrSet, fieldName).check(op->value);
if (!fieldRead.defined()) {
stmt = op;
return;
}
bool valsCombined = IsFieldReduced(fieldRead).check(op->value);
if (!valsCombined) {
stmt = op;
return;
}
Type fieldType = getFieldType(elemOrSet, fieldName);
Var tmp(names.getName(), fieldType);
Stmt tmpAssignment = AssignStmt::make(tmp, op->value);
Stmt writeTmpToField = FieldWrite::make(elemOrSet, fieldName, tmp);
stmt = Block::make(tmpAssignment, writeTmpToField);
}
void Function::define(const vector<string> &args, Expr value) {
assert(value.defined() && "Undefined expression in right-hand-side of function definition\n");
// Make sure all the vars in the value are either args or are
// attached to some parameter
CheckVars check;
check.pure_args = args;
value.accept(&check);
assert(!check.reduction_domain.defined() && "Reduction domain referenced in pure function definition");
if (!contents.defined()) {
contents = new FunctionContents;
contents.ptr->name = unique_name('f');
}
assert(!contents.ptr->value.defined() && "Function is already defined");
contents.ptr->value = value;
contents.ptr->args = args;
for (size_t i = 0; i < args.size(); i++) {
Schedule::Dim d = {args[i], For::Serial};
contents.ptr->schedule.dims.push_back(d);
}
}
Stmt Evaluate::make(Expr v) {
internal_assert(v.defined()) << "Evaluate of undefined\n";
Evaluate *node = new Evaluate;
node->value = v;
return node;
}
Expr Select::make(Expr condition, Expr true_value, Expr false_value) {
internal_assert(condition.defined()) << "Select of undefined\n";
internal_assert(true_value.defined()) << "Select of undefined\n";
internal_assert(false_value.defined()) << "Select of undefined\n";
internal_assert(condition.type().is_bool()) << "First argument to Select is not a bool\n";
internal_assert(false_value.type() == true_value.type()) << "Select of mismatched types\n";
internal_assert(condition.type().is_scalar() ||
condition.type().width == true_value.type().width)
<< "In Select, vector width of condition must either be 1, or equal to vector width of arguments\n";
Select *node = new Select;
node->type = true_value.type();
node->condition = condition;
node->true_value = true_value;
node->false_value = false_value;
return node;
}
Expr Cast::make(Type t, Expr v) {
internal_assert(v.defined()) << "Cast of undefined\n";
Cast *node = new Cast;
node->type = t;
node->value = v;
return node;
}
void visit(const Variable *v) {
Expr r = find_replacement(v->name);
if (r.defined()) {
expr = r;
} else {
expr = v;
}
}
Expr IRBuilder::binaryElwiseExpr(Expr l, BinaryOperator op, Expr r) {
const TensorType *ltype = l.type().toTensor();
const TensorType *rtype = r.type().toTensor();
Expr tensor = (ltype->order() > 0) ? l : r;
std::vector<IndexVar> indexVars;
const TensorType *tensorType = tensor.type().toTensor();
vector<IndexDomain> dimensions = tensorType->getDimensions();
for (unsigned int i=0; i < tensorType->order(); ++i) {
IndexDomain domain = dimensions[i];
indexVars.push_back(factory.createIndexVar(domain));
}
Expr a, b;
if (ltype->order() == 0 || rtype->order() == 0) {
std::vector<IndexVar> scalarIndexVars;
std::vector<IndexVar> *lIndexVars;
std::vector<IndexVar> *rIndexVars;
if (ltype->order() == 0) {
lIndexVars = &scalarIndexVars;
rIndexVars = &indexVars;
}
else {
lIndexVars = &indexVars;
rIndexVars = &scalarIndexVars;
}
a = IndexedTensor::make(l, *lIndexVars);
b = IndexedTensor::make(r, *rIndexVars);
}
else {
iassert(l.type() == r.type());
a = IndexedTensor::make(l, indexVars);
b = IndexedTensor::make(r, indexVars);
}
iassert(a.defined() && b.defined());
Expr val;
switch (op) {
case Add:
val = Add::make(a, b);
break;
case Sub:
val = Sub::make(a, b);
break;
case Mul:
val = Mul::make(a, b);
break;
case Div:
val = Div::make(a, b);
break;
}
iassert(val.defined());
const bool isColumnVector = tensor.type().toTensor()->isColumnVector;
return IndexExpr::make(indexVars, val, isColumnVector);
}
Func mirror_interior(const Func &source,
const std::vector<std::pair<Expr, Expr>> &bounds) {
std::vector<Var> args(source.args());
user_assert(args.size() >= bounds.size()) <<
"mirror_interior called with more bounds (" << bounds.size() <<
") than dimensions (" << args.size() << ") Func " <<
source.name() << "has.\n";
std::vector<Expr> actuals;
for (size_t i = 0; i < bounds.size(); i++) {
Var arg_var = args[i];
Expr min = bounds[i].first;
Expr extent = bounds[i].second;
if (min.defined() && extent.defined()) {
Expr limit = extent - 1;
Expr coord = arg_var - min; // Enforce zero origin.
coord = coord % (2 * limit); // Range is 0 to 2w-1
coord = coord - limit; // Range is -w, w
coord = abs(coord); // Range is 0, w
coord = limit - coord; // Range is 0, w
coord = coord + min; // Restore correct min
// The boundary condition probably doesn't apply
coord = select(arg_var < min || arg_var >= min + extent, coord,
clamp(likely(arg_var), min, min + extent - 1));
actuals.push_back(coord);
} else if (!min.defined() && !extent.defined()) {
actuals.push_back(arg_var);
} else {
user_error << "Partially undefined bounds for dimension " << arg_var
<< " of Func " << source.name() << "\n";
}
}
// If there were fewer bounds than dimensions, regard the ones at the end as unbounded.
actuals.insert(actuals.end(), args.begin() + actuals.size(), args.end());
Func bounded("mirror_interior");
bounded(args) = source(actuals);
return bounded;
}
Stmt AssertStmt::make(Expr condition, Expr message) {
internal_assert(condition.defined()) << "AssertStmt of undefined\n";
internal_assert(message.type() == Int(32)) << "AssertStmt message must be an int:" << message << "\n";
AssertStmt *node = new AssertStmt;
node->condition = condition;
node->message = message;
return node;
}
Expr IRMutator::mutate(const Expr &e) {
if (e.defined()) {
e.accept(this);
} else {
expr = Expr();
}
stmt = Stmt();
return std::move(expr);
}
Expr IRMutator::mutate(Expr e) {
if (e.defined()) {
e.accept(this);
} else {
expr = Expr();
}
stmt = Stmt();
return expr;
}
Expr Not::make(Expr a) {
internal_assert(a.defined()) << "Not of undefined\n";
internal_assert(a.type().is_bool()) << "argument of Not is not a bool\n";
Not *node = new Not;
node->type = Bool(a.type().width);
node->a = a;
return node;
}
Stmt AssertStmt::make(Expr condition, std::string message, const std::vector<Expr> &args) {
internal_assert(condition.defined()) << "AssertStmt of undefined\n";
AssertStmt *node = new AssertStmt;
node->condition = condition;
node->message = message;
node->args = args;
return node;
}
Expr Cast::make(Type t, Expr v) {
internal_assert(v.defined()) << "Cast of undefined\n";
internal_assert(t.width == v.type().width) << "Cast may not change vector widths\n";
Cast *node = new Cast;
node->type = t;
node->value = v;
return node;
}
Stmt IfThenElse::make(Expr condition, Stmt then_case, Stmt else_case) {
internal_assert(condition.defined() && then_case.defined()) << "IfThenElse of undefined\n";
// else_case may be null.
IfThenElse *node = new IfThenElse;
node->condition = condition;
node->then_case = then_case;
node->else_case = else_case;
return node;
}
Stmt visit(const Allocate *op) override {
Expr total_extent = make_const(Int(64), 1);
for (const Expr &e : op->extents) {
total_extent *= e;
}
Expr bound = find_constant_bound(total_extent, Direction::Upper, scope);
user_assert(bound.defined() ||
op->memory_type != MemoryType::Register)
<< "Allocation " << op->name << " has a dynamic size. "
<< "Only fixed-size allocations can be stored in registers. "
<< "Try storing on the heap or stack instead.";
user_assert(!in_thread_loop || bound.defined())
<< "Allocation " << op->name << " has a dynamic size. "
<< "Only fixed-size allocations are supported on the gpu. "
<< "Try storing into shared memory instead.";
const int64_t *size_ptr = bound.defined() ? as_const_int(bound) : nullptr;
int64_t size = size_ptr ? *size_ptr : 0;
if (size_ptr && size == 0 && !op->new_expr.defined()) {
// This allocation is dead
return Allocate::make(op->name, op->type, op->memory_type, {0}, const_false(),
mutate(op->body), op->new_expr, op->free_function);
}
// 128 bytes is a typical minimum allocation size in
// halide_malloc. For now we are very conservative, and only
// round sizes up to a constant if they're smaller than that.
int malloc_overhead = 128 / op->type.bytes();
if (size_ptr &&
(in_thread_loop ||
(op->memory_type == MemoryType::Stack && can_allocation_fit_on_stack(size)) ||
op->memory_type == MemoryType::Register ||
(op->memory_type == MemoryType::Auto && size <= malloc_overhead))) {
user_assert(size >= 0 && size < (int64_t)1 << 31)
<< "Allocation " << op->name << " has a size greater than 2^31: " << bound << "\n";
return Allocate::make(op->name, op->type, op->memory_type, {(int32_t)size}, op->condition,
mutate(op->body), op->new_expr, op->free_function);
} else {
return IRMutator::visit(op);
}
}