IR.cpp
#include "IR.h"
#include "IRMutator.h"
#include "IRPrinter.h"
#include "IRVisitor.h"
#include <numeric>
#include <utility>
namespace Halide {
namespace Internal {
Expr Cast::make(Type t, Expr v) {
internal_assert(v.defined()) << "Cast of undefined\n";
internal_assert(t.lanes() == v.type().lanes()) << "Cast may not change vector widths\n";
Cast *node = new Cast;
node->type = t;
node->value = std::move(v);
return node;
}
Expr Add::make(Expr a, Expr b) {
internal_assert(a.defined()) << "Add of undefined\n";
internal_assert(b.defined()) << "Add of undefined\n";
internal_assert(a.type() == b.type()) << "Add of mismatched types\n";
Add *node = new Add;
node->type = a.type();
node->a = std::move(a);
node->b = std::move(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 = std::move(a);
node->b = std::move(b);
return node;
}
Expr Mul::make(Expr a, Expr b) {
internal_assert(a.defined()) << "Mul of undefined\n";
internal_assert(b.defined()) << "Mul of undefined\n";
internal_assert(a.type() == b.type()) << "Mul of mismatched types\n";
Mul *node = new Mul;
node->type = a.type();
node->a = std::move(a);
node->b = std::move(b);
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 = std::move(a);
node->b = std::move(b);
return node;
}
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 = std::move(a);
node->b = std::move(b);
return node;
}
Expr Min::make(Expr a, Expr b) {
internal_assert(a.defined()) << "Min of undefined\n";
internal_assert(b.defined()) << "Min of undefined\n";
internal_assert(a.type() == b.type()) << "Min of mismatched types\n";
Min *node = new Min;
node->type = a.type();
node->a = std::move(a);
node->b = std::move(b);
return node;
}
Expr Max::make(Expr a, Expr b) {
internal_assert(a.defined()) << "Max of undefined\n";
internal_assert(b.defined()) << "Max of undefined\n";
internal_assert(a.type() == b.type()) << "Max of mismatched types\n";
Max *node = new Max;
node->type = a.type();
node->a = std::move(a);
node->b = std::move(b);
return node;
}
Expr EQ::make(Expr a, Expr b) {
internal_assert(a.defined()) << "EQ of undefined\n";
internal_assert(b.defined()) << "EQ of undefined\n";
internal_assert(a.type() == b.type()) << "EQ of mismatched types\n";
EQ *node = new EQ;
node->type = Bool(a.type().lanes());
node->a = std::move(a);
node->b = std::move(b);
return node;
}
Expr NE::make(Expr a, Expr b) {
internal_assert(a.defined()) << "NE of undefined\n";
internal_assert(b.defined()) << "NE of undefined\n";
internal_assert(a.type() == b.type()) << "NE of mismatched types\n";
NE *node = new NE;
node->type = Bool(a.type().lanes());
node->a = std::move(a);
node->b = std::move(b);
return node;
}
Expr LT::make(Expr a, Expr b) {
internal_assert(a.defined()) << "LT of undefined\n";
internal_assert(b.defined()) << "LT of undefined\n";
internal_assert(a.type() == b.type()) << "LT of mismatched types\n";
LT *node = new LT;
node->type = Bool(a.type().lanes());
node->a = std::move(a);
node->b = std::move(b);
return node;
}
Expr LE::make(Expr a, Expr b) {
internal_assert(a.defined()) << "LE of undefined\n";
internal_assert(b.defined()) << "LE of undefined\n";
internal_assert(a.type() == b.type()) << "LE of mismatched types\n";
LE *node = new LE;
node->type = Bool(a.type().lanes());
node->a = std::move(a);
node->b = std::move(b);
return node;
}
Expr GT::make(Expr a, Expr b) {
internal_assert(a.defined()) << "GT of undefined\n";
internal_assert(b.defined()) << "GT of undefined\n";
internal_assert(a.type() == b.type()) << "GT of mismatched types\n";
GT *node = new GT;
node->type = Bool(a.type().lanes());
node->a = std::move(a);
node->b = std::move(b);
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().lanes());
node->a = std::move(a);
node->b = std::move(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";
internal_assert(a.type() == b.type()) << "And of mismatched types\n";
And *node = new And;
node->type = Bool(a.type().lanes());
node->a = std::move(a);
node->b = std::move(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";
internal_assert(a.type() == b.type()) << "Or of mismatched types\n";
Or *node = new Or;
node->type = Bool(a.type().lanes());
node->a = std::move(a);
node->b = std::move(b);
return node;
}
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().lanes());
node->a = std::move(a);
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: " << condition.type() << "\n";
internal_assert(false_value.type() == true_value.type()) << "Select of mismatched types\n";
internal_assert(condition.type().is_scalar() ||
condition.type().lanes() == true_value.type().lanes())
<< "In Select, vector lanes of condition must either be 1, or equal to vector lanes of arguments\n";
Select *node = new Select;
node->type = true_value.type();
node->condition = std::move(condition);
node->true_value = std::move(true_value);
node->false_value = std::move(false_value);
return node;
}
Expr Load::make(Type type, const std::string &name, Expr index, Buffer<> image, Parameter param, Expr predicate, ModulusRemainder alignment) {
internal_assert(predicate.defined()) << "Load with undefined predicate\n";
internal_assert(index.defined()) << "Load of undefined\n";
internal_assert(type.lanes() == index.type().lanes()) << "Vector lanes of Load must match vector lanes of index\n";
internal_assert(type.lanes() == predicate.type().lanes())
<< "Vector lanes of Load must match vector lanes of predicate\n";
Load *node = new Load;
node->type = type;
node->name = name;
node->predicate = std::move(predicate);
node->index = std::move(index);
node->image = std::move(image);
node->param = std::move(param);
node->alignment = alignment;
return node;
}
Expr Ramp::make(Expr base, Expr stride, int lanes) {
internal_assert(base.defined()) << "Ramp of undefined\n";
internal_assert(stride.defined()) << "Ramp of undefined\n";
internal_assert(lanes > 1) << "Ramp of lanes <= 1\n";
internal_assert(stride.type() == base.type()) << "Ramp of mismatched types\n";
Ramp *node = new Ramp;
node->type = base.type().with_lanes(lanes * base.type().lanes());
node->base = std::move(base);
node->stride = std::move(stride);
node->lanes = lanes;
return node;
}
Expr Broadcast::make(Expr value, int lanes) {
internal_assert(value.defined()) << "Broadcast of undefined\n";
internal_assert(lanes != 1) << "Broadcast of lanes 1\n";
Broadcast *node = new Broadcast;
node->type = value.type().with_lanes(lanes * value.type().lanes());
node->value = std::move(value);
node->lanes = lanes;
return node;
}
Expr Let::make(const 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 = std::move(value);
node->body = std::move(body);
return node;
}
Stmt LetStmt::make(const std::string &name, Expr value, Stmt body) {
internal_assert(value.defined()) << "Let of undefined\n";
internal_assert(body.defined()) << "Let of undefined\n";
LetStmt *node = new LetStmt;
node->name = name;
node->value = std::move(value);
node->body = std::move(body);
return node;
}
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 = std::move(condition);
node->message = std::move(message);
return node;
}
Stmt ProducerConsumer::make(const std::string &name, bool is_producer, Stmt body) {
internal_assert(body.defined()) << "ProducerConsumer of undefined\n";
ProducerConsumer *node = new ProducerConsumer;
node->name = name;
node->is_producer = is_producer;
node->body = std::move(body);
return node;
}
Stmt ProducerConsumer::make_produce(const std::string &name, Stmt body) {
return ProducerConsumer::make(name, true, std::move(body));
}
Stmt ProducerConsumer::make_consume(const std::string &name, Stmt body) {
return ProducerConsumer::make(name, false, std::move(body));
}
Stmt For::make(const std::string &name, Expr min, Expr extent, ForType for_type, DeviceAPI device_api, Stmt body) {
internal_assert(min.defined()) << "For of undefined\n";
internal_assert(extent.defined()) << "For of undefined\n";
internal_assert(min.type() == Int(32)) << "For with non-integer min\n";
internal_assert(extent.type() == Int(32)) << "For with non-integer extent\n";
internal_assert(body.defined()) << "For of undefined\n";
For *node = new For;
node->name = name;
node->min = std::move(min);
node->extent = std::move(extent);
node->for_type = for_type;
node->device_api = device_api;
node->body = std::move(body);
return node;
}
Stmt Acquire::make(Expr semaphore, Expr count, Stmt body) {
internal_assert(semaphore.defined()) << "Acquire with undefined semaphore\n";
internal_assert(body.defined()) << "Acquire with undefined body\n";
Acquire *node = new Acquire;
node->semaphore = std::move(semaphore);
node->count = std::move(count);
node->body = std::move(body);
return node;
}
Stmt Store::make(const std::string &name, Expr value, Expr index, Parameter param, Expr predicate, ModulusRemainder alignment) {
internal_assert(predicate.defined()) << "Store with undefined predicate\n";
internal_assert(value.defined()) << "Store of undefined\n";
internal_assert(index.defined()) << "Store of undefined\n";
internal_assert(value.type().lanes() == index.type().lanes()) << "Vector lanes of Store must match vector lanes of index\n";
internal_assert(value.type().lanes() == predicate.type().lanes())
<< "Vector lanes of Store must match vector lanes of predicate\n";
Store *node = new Store;
node->name = name;
node->predicate = std::move(predicate);
node->value = std::move(value);
node->index = std::move(index);
node->param = std::move(param);
node->alignment = alignment;
return node;
}
Stmt Provide::make(const std::string &name, const std::vector<Expr> &values, const std::vector<Expr> &args, const Expr &predicate) {
internal_assert(predicate.defined()) << "Provide with undefined predicate\n";
internal_assert(!values.empty()) << "Provide of no values\n";
for (size_t i = 0; i < values.size(); i++) {
internal_assert(values[i].defined()) << "Provide of undefined value\n";
}
for (size_t i = 0; i < args.size(); i++) {
internal_assert(args[i].defined()) << "Provide to undefined location\n";
}
Provide *node = new Provide;
node->name = name;
node->values = values;
node->args = args;
node->predicate = predicate;
return node;
}
Stmt Allocate::make(const std::string &name, Type type, MemoryType memory_type,
const std::vector<Expr> &extents,
Expr condition, Stmt body,
Expr new_expr, const std::string &free_function) {
for (size_t i = 0; i < extents.size(); i++) {
internal_assert(extents[i].defined()) << "Allocate of undefined extent\n";
internal_assert(extents[i].type().is_scalar() == 1) << "Allocate of vector extent\n";
}
internal_assert(body.defined()) << "Allocate of undefined\n";
internal_assert(condition.defined()) << "Allocate with undefined condition\n";
internal_assert(condition.type().is_bool()) << "Allocate condition is not boolean\n";
Allocate *node = new Allocate;
node->name = name;
node->type = type;
node->memory_type = memory_type;
node->extents = extents;
node->new_expr = std::move(new_expr);
node->free_function = free_function;
node->condition = std::move(condition);
node->body = std::move(body);
return node;
}
int32_t Allocate::constant_allocation_size(const std::vector<Expr> &extents, const std::string &name) {
int64_t result = 1;
for (size_t i = 0; i < extents.size(); i++) {
if (const IntImm *int_size = extents[i].as<IntImm>()) {
// Check if the individual dimension is > 2^31 - 1. Not
// currently necessary because it's an int32_t, which is
// always smaller than 2^31 - 1. If we ever upgrade the
// type of IntImm but not the maximum allocation size, we
// should re-enable this.
/*
if ((int64_t)int_size->value > (((int64_t)(1)<<31) - 1)) {
user_error
<< "Dimension " << i << " for allocation " << name << " has size " <<
int_size->value << " which is greater than 2^31 - 1.";
}
*/
result *= int_size->value;
if (result > (static_cast<int64_t>(1) << 31) - 1) {
user_error
<< "Total size for allocation " << name
<< " is constant but exceeds 2^31 - 1.\n";
}
} else {
return 0;
}
}
return static_cast<int32_t>(result);
}
int32_t Allocate::constant_allocation_size() const {
return Allocate::constant_allocation_size(extents, name);
}
Stmt Free::make(const std::string &name) {
Free *node = new Free;
node->name = name;
return node;
}
Stmt Realize::make(const std::string &name, const std::vector<Type> &types, MemoryType memory_type, const Region &bounds, Expr condition, Stmt body) {
for (size_t i = 0; i < bounds.size(); i++) {
internal_assert(bounds[i].min.defined()) << "Realize of undefined\n";
internal_assert(bounds[i].extent.defined()) << "Realize of undefined\n";
internal_assert(bounds[i].min.type().is_scalar()) << "Realize of vector size\n";
internal_assert(bounds[i].extent.type().is_scalar()) << "Realize of vector size\n";
}
internal_assert(body.defined()) << "Realize of undefined\n";
internal_assert(!types.empty()) << "Realize has empty type\n";
internal_assert(condition.defined()) << "Realize with undefined condition\n";
internal_assert(condition.type().is_bool()) << "Realize condition is not boolean\n";
Realize *node = new Realize;
node->name = name;
node->types = types;
node->memory_type = memory_type;
node->bounds = bounds;
node->condition = std::move(condition);
node->body = std::move(body);
return node;
}
Stmt Prefetch::make(const std::string &name, const std::vector<Type> &types,
const Region &bounds,
const PrefetchDirective &prefetch,
Expr condition, Stmt body) {
for (size_t i = 0; i < bounds.size(); i++) {
internal_assert(bounds[i].min.defined()) << "Prefetch of undefined\n";
internal_assert(bounds[i].extent.defined()) << "Prefetch of undefined\n";
internal_assert(bounds[i].min.type().is_scalar()) << "Prefetch of vector size\n";
internal_assert(bounds[i].extent.type().is_scalar()) << "Prefetch of vector size\n";
}
internal_assert(!types.empty()) << "Prefetch has empty type\n";
internal_assert(body.defined()) << "Prefetch of undefined\n";
internal_assert(condition.defined()) << "Prefetch with undefined condition\n";
internal_assert(condition.type().is_bool()) << "Prefetch condition is not boolean\n";
Prefetch *node = new Prefetch;
node->name = name;
node->types = types;
node->bounds = bounds;
node->prefetch = prefetch;
node->condition = std::move(condition);
node->body = std::move(body);
return node;
}
Stmt Block::make(Stmt first, Stmt rest) {
internal_assert(first.defined()) << "Block of undefined\n";
internal_assert(rest.defined()) << "Block of undefined\n";
Block *node = new Block;
if (const Block *b = first.as<Block>()) {
// Use a canonical block nesting order
node->first = b->first;
node->rest = Block::make(b->rest, std::move(rest));
} else {
node->first = std::move(first);
node->rest = std::move(rest);
}
return node;
}
Stmt Block::make(const std::vector<Stmt> &stmts) {
if (stmts.empty()) {
return Stmt();
}
Stmt result = stmts.back();
for (size_t i = stmts.size() - 1; i > 0; i--) {
result = Block::make(stmts[i - 1], result);
}
return result;
}
Stmt Fork::make(Stmt first, Stmt rest) {
internal_assert(first.defined()) << "Fork of undefined\n";
internal_assert(rest.defined()) << "Fork of undefined\n";
Fork *node = new Fork;
if (const Fork *b = first.as<Fork>()) {
// Use a canonical fork nesting order
node->first = b->first;
node->rest = Fork::make(b->rest, std::move(rest));
} else {
node->first = std::move(first);
node->rest = std::move(rest);
}
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 = std::move(condition);
node->then_case = std::move(then_case);
node->else_case = std::move(else_case);
return node;
}
Stmt Evaluate::make(Expr v) {
internal_assert(v.defined()) << "Evaluate of undefined\n";
Evaluate *node = new Evaluate;
node->value = std::move(v);
return node;
}
Expr Call::make(const Function &func, const std::vector<Expr> &args, int idx) {
internal_assert(idx >= 0 &&
idx < func.outputs())
<< "Value index out of range in call to halide function\n";
internal_assert(func.has_pure_definition() || func.has_extern_definition())
<< "Call to undefined halide function\n";
return make(func.output_types()[(size_t)idx], func.name(), args, Halide,
func.get_contents(), idx, Buffer<>(), Parameter());
}
namespace {
const char *const intrinsic_op_names[] = {
"abs",
"absd",
"add_image_checks_marker",
"alloca",
"bitwise_and",
"bitwise_not",
"bitwise_or",
"bitwise_xor",
"bool_to_mask",
"bundle",
"call_cached_indirect_function",
"cast_mask",
"count_leading_zeros",
"count_trailing_zeros",
"declare_box_touched",
"debug_to_file",
"div_round_to_zero",
"dynamic_shuffle",
"extract_mask_element",
"gpu_thread_barrier",
"halving_add",
"halving_sub",
"hvx_gather",
"hvx_scatter",
"hvx_scatter_acc",
"hvx_scatter_release",
"if_then_else",
"if_then_else_mask",
"image_load",
"image_store",
"lerp",
"likely",
"likely_if_innermost",
"make_struct",
"memoize_expr",
"mod_round_to_zero",
"mul_shift_right",
"mux",
"popcount",
"prefetch",
"promise_clamped",
"random",
"register_destructor",
"reinterpret",
"require",
"require_mask",
"return_second",
"rewrite_buffer",
"rounding_halving_add",
"rounding_halving_sub",
"rounding_mul_shift_right",
"rounding_shift_left",
"rounding_shift_right",
"saturating_add",
"saturating_sub",
"scatter_gather",
"select_mask",
"shift_left",
"shift_right",
"signed_integer_overflow",
"size_of_halide_buffer_t",
"sorted_avg",
"strict_float",
"stringify",
"undef",
"unreachable",
"unsafe_promise_clamped",
"widening_add",
"widening_mul",
"widening_shift_left",
"widening_shift_right",
"widening_sub",
};
static_assert(sizeof(intrinsic_op_names) / sizeof(intrinsic_op_names[0]) == Call::IntrinsicOpCount,
"intrinsic_op_names needs attention");
} // namespace
const char *Call::get_intrinsic_name(IntrinsicOp op) {
return intrinsic_op_names[op];
}
Expr Call::make(Type type, Call::IntrinsicOp op, const std::vector<Expr> &args, CallType call_type,
FunctionPtr func, int value_index,
const Buffer<> &image, Parameter param) {
internal_assert(call_type == Call::Intrinsic || call_type == Call::PureIntrinsic);
return Call::make(type, intrinsic_op_names[op], args, call_type, std::move(func), value_index, image, std::move(param));
}
Expr Call::make(Type type, const std::string &name, const std::vector<Expr> &args, CallType call_type,
FunctionPtr func, int value_index,
Buffer<> image, Parameter param) {
if (name == intrinsic_op_names[Call::prefetch] && call_type == Call::Intrinsic) {
internal_assert(args.size() % 2 == 0)
<< "Number of args to a prefetch call should be even: {base, offset, extent0, stride0, extent1, stride1, ...}\n";
}
for (size_t i = 0; i < args.size(); i++) {
internal_assert(args[i].defined()) << "Call of " << name << " with argument " << i << " undefined.\n";
}
if (call_type == Halide) {
for (size_t i = 0; i < args.size(); i++) {
internal_assert(args[i].type() == Int(32))
<< "Args to call to halide function must be type Int(32)\n";
}
} else if (call_type == Image) {
internal_assert((param.defined() || image.defined()))
<< "Call node to undefined image\n";
for (size_t i = 0; i < args.size(); i++) {
internal_assert(args[i].type() == Int(32))
<< "Args to load from image must be type Int(32)\n";
}
}
Call *node = new Call;
node->type = type;
node->name = name;
node->args = args;
node->call_type = call_type;
node->func = std::move(func);
node->value_index = value_index;
node->image = std::move(image);
node->param = std::move(param);
return node;
}
Expr Variable::make(Type type, const std::string &name, Buffer<> image, Parameter param, ReductionDomain reduction_domain) {
Variable *node = new Variable;
node->type = type;
node->name = name;
node->image = std::move(image);
node->param = std::move(param);
node->reduction_domain = std::move(reduction_domain);
return node;
}
Expr Shuffle::make(const std::vector<Expr> &vectors,
const std::vector<int> &indices) {
internal_assert(!vectors.empty()) << "Shuffle of zero vectors.\n";
internal_assert(!indices.empty()) << "Shufle with zero indices.\n";
Type element_ty = vectors.front().type().element_of();
int input_lanes = 0;
for (const Expr &i : vectors) {
internal_assert(i.type().element_of() == element_ty) << "Shuffle of vectors of mismatched types.\n";
input_lanes += i.type().lanes();
}
for (int i : indices) {
internal_assert(0 <= i && i < input_lanes) << "Shuffle vector index out of range: " << i << "\n";
}
Shuffle *node = new Shuffle;
node->type = element_ty.with_lanes((int)indices.size());
node->vectors = vectors;
node->indices = indices;
return node;
}
Expr Shuffle::make_interleave(const std::vector<Expr> &vectors) {
internal_assert(!vectors.empty()) << "Interleave of zero vectors.\n";
if (vectors.size() == 1) {
return vectors.front();
}
int lanes = vectors.front().type().lanes();
for (const Expr &i : vectors) {
internal_assert(i.type().lanes() == lanes)
<< "Interleave of vectors with different sizes.\n";
}
std::vector<int> indices;
for (int i = 0; i < lanes; i++) {
for (int j = 0; j < (int)vectors.size(); j++) {
indices.push_back(j * lanes + i);
}
}
return make(vectors, indices);
}
Expr Shuffle::make_concat(const std::vector<Expr> &vectors) {
internal_assert(!vectors.empty()) << "Concat of zero vectors.\n";
if (vectors.size() == 1) {
return vectors.front();
}
std::vector<int> indices;
int lane = 0;
for (int i = 0; i < (int)vectors.size(); i++) {
for (int j = 0; j < vectors[i].type().lanes(); j++) {
indices.push_back(lane++);
}
}
return make(vectors, indices);
}
Expr Shuffle::make_broadcast(Expr vector, int factor) {
std::vector<int> indices(factor * vector.type().lanes());
for (int ix = 0; ix < factor; ix++) {
std::iota(indices.begin() + ix * vector.type().lanes(),
indices.begin() + (ix + 1) * vector.type().lanes(), 0);
}
return make({std::move(vector)}, indices);
}
Expr Shuffle::make_slice(Expr vector, int begin, int stride, int size) {
if (begin == 0 && size == vector.type().lanes() && stride == 1) {
return vector;
}
std::vector<int> indices;
for (int i = 0; i < size; i++) {
indices.push_back(begin + i * stride);
}
return make({std::move(vector)}, indices);
}
Expr Shuffle::make_extract_element(Expr vector, int i) {
return make_slice(std::move(vector), i, 1, 1);
}
bool Shuffle::is_broadcast() const {
int lanes = indices.size();
int factor = broadcast_factor();
if (factor == 0 || factor >= lanes) {
return false;
}
int broadcasted_lanes = lanes / factor;
if (broadcasted_lanes < 2 || broadcasted_lanes >= lanes || lanes % broadcasted_lanes != 0) {
return false;
}
for (int i = 0; i < lanes; i++) {
if (indices[i % broadcasted_lanes] != indices[i]) {
return false;
}
}
return true;
}
int Shuffle::broadcast_factor() const {
int lanes = indices.size();
int broadcasted_lanes = 0;
for (; broadcasted_lanes < lanes; broadcasted_lanes++) {
if (indices[broadcasted_lanes] != broadcasted_lanes) {
break;
}
}
if (broadcasted_lanes > 0) {
return lanes / broadcasted_lanes;
} else {
return 0;
}
}
bool Shuffle::is_interleave() const {
int lanes = vectors.front().type().lanes();
// Don't consider concat of scalars as an interleave.
if (lanes == 1) {
return false;
}
for (const Expr &i : vectors) {
if (i.type().lanes() != lanes) {
return false;
}
}
// Require that we are a complete interleaving.
if (lanes * vectors.size() != indices.size()) {
return false;
}
for (int i = 0; i < (int)vectors.size(); i++) {
for (int j = 0; j < lanes; j++) {
if (indices[j * (int)vectors.size() + i] != i * lanes + j) {
return false;
}
}
}
return true;
}
Stmt Atomic::make(const std::string &producer_name,
const std::string &mutex_name,
Stmt body) {
Atomic *node = new Atomic;
node->producer_name = producer_name;
node->mutex_name = mutex_name;
internal_assert(body.defined()) << "Atomic must have a body statement.\n";
node->body = std::move(body);
return node;
}
Expr VectorReduce::make(VectorReduce::Operator op,
Expr vec,
int lanes) {
if (vec.type().is_bool()) {
internal_assert(op == VectorReduce::And || op == VectorReduce::Or)
<< "The only legal operators for VectorReduce on a Bool"
<< "vector are VectorReduce::And and VectorReduce::Or\n";
}
internal_assert(!vec.type().is_handle()) << "VectorReduce of handle type";
// Check the output lanes is a factor of the input lanes. They can
// also both be zero if we're constructing a wildcard expression.
internal_assert((lanes == 0 && vec.type().lanes() == 0) ||
(lanes != 0 && (vec.type().lanes() % lanes == 0)))
<< "Vector reduce output lanes must be a divisor of the number of lanes in the argument "
<< lanes << " " << vec.type().lanes() << "\n";
VectorReduce *node = new VectorReduce;
node->type = vec.type().with_lanes(lanes);
node->op = op;
node->value = std::move(vec);
return node;
}
namespace {
// Helper function to determine if a sequence of indices is a
// contiguous ramp.
bool is_ramp(const std::vector<int> &indices, int stride = 1) {
for (size_t i = 0; i + 1 < indices.size(); i++) {
if (indices[i + 1] != indices[i] + stride) {
return false;
}
}
return true;
}
} // namespace
bool Shuffle::is_concat() const {
size_t input_lanes = 0;
for (const Expr &i : vectors) {
input_lanes += i.type().lanes();
}
// A concat is a ramp where the output has the same number of
// lanes as the input.
return indices.size() == input_lanes && is_ramp(indices);
}
bool Shuffle::is_slice() const {
size_t input_lanes = 0;
for (const Expr &i : vectors) {
input_lanes += i.type().lanes();
}
// A slice is a ramp where the output does not contain all of the
// lanes of the input.
return indices.size() < input_lanes && is_ramp(indices, slice_stride());
}
bool Shuffle::is_extract_element() const {
return indices.size() == 1;
}
template<>
void ExprNode<IntImm>::accept(IRVisitor *v) const {
v->visit((const IntImm *)this);
}
template<>
void ExprNode<UIntImm>::accept(IRVisitor *v) const {
v->visit((const UIntImm *)this);
}
template<>
void ExprNode<FloatImm>::accept(IRVisitor *v) const {
v->visit((const FloatImm *)this);
}
template<>
void ExprNode<StringImm>::accept(IRVisitor *v) const {
v->visit((const StringImm *)this);
}
template<>
void ExprNode<Cast>::accept(IRVisitor *v) const {
v->visit((const Cast *)this);
}
template<>
void ExprNode<Variable>::accept(IRVisitor *v) const {
v->visit((const Variable *)this);
}
template<>
void ExprNode<Add>::accept(IRVisitor *v) const {
v->visit((const Add *)this);
}
template<>
void ExprNode<Sub>::accept(IRVisitor *v) const {
v->visit((const Sub *)this);
}
template<>
void ExprNode<Mul>::accept(IRVisitor *v) const {
v->visit((const Mul *)this);
}
template<>
void ExprNode<Div>::accept(IRVisitor *v) const {
v->visit((const Div *)this);
}
template<>
void ExprNode<Mod>::accept(IRVisitor *v) const {
v->visit((const Mod *)this);
}
template<>
void ExprNode<Min>::accept(IRVisitor *v) const {
v->visit((const Min *)this);
}
template<>
void ExprNode<Max>::accept(IRVisitor *v) const {
v->visit((const Max *)this);
}
template<>
void ExprNode<EQ>::accept(IRVisitor *v) const {
v->visit((const EQ *)this);
}
template<>
void ExprNode<NE>::accept(IRVisitor *v) const {
v->visit((const NE *)this);
}
template<>
void ExprNode<LT>::accept(IRVisitor *v) const {
v->visit((const LT *)this);
}
template<>
void ExprNode<LE>::accept(IRVisitor *v) const {
v->visit((const LE *)this);
}
template<>
void ExprNode<GT>::accept(IRVisitor *v) const {
v->visit((const GT *)this);
}
template<>
void ExprNode<GE>::accept(IRVisitor *v) const {
v->visit((const GE *)this);
}
template<>
void ExprNode<And>::accept(IRVisitor *v) const {
v->visit((const And *)this);
}
template<>
void ExprNode<Or>::accept(IRVisitor *v) const {
v->visit((const Or *)this);
}
template<>
void ExprNode<Not>::accept(IRVisitor *v) const {
v->visit((const Not *)this);
}
template<>
void ExprNode<Select>::accept(IRVisitor *v) const {
v->visit((const Select *)this);
}
template<>
void ExprNode<Load>::accept(IRVisitor *v) const {
v->visit((const Load *)this);
}
template<>
void ExprNode<Ramp>::accept(IRVisitor *v) const {
v->visit((const Ramp *)this);
}
template<>
void ExprNode<Broadcast>::accept(IRVisitor *v) const {
v->visit((const Broadcast *)this);
}
template<>
void ExprNode<Call>::accept(IRVisitor *v) const {
v->visit((const Call *)this);
}
template<>
void ExprNode<Shuffle>::accept(IRVisitor *v) const {
v->visit((const Shuffle *)this);
}
template<>
void ExprNode<VectorReduce>::accept(IRVisitor *v) const {
v->visit((const VectorReduce *)this);
}
template<>
void ExprNode<Let>::accept(IRVisitor *v) const {
v->visit((const Let *)this);
}
template<>
void StmtNode<LetStmt>::accept(IRVisitor *v) const {
v->visit((const LetStmt *)this);
}
template<>
void StmtNode<AssertStmt>::accept(IRVisitor *v) const {
v->visit((const AssertStmt *)this);
}
template<>
void StmtNode<ProducerConsumer>::accept(IRVisitor *v) const {
v->visit((const ProducerConsumer *)this);
}
template<>
void StmtNode<For>::accept(IRVisitor *v) const {
v->visit((const For *)this);
}
template<>
void StmtNode<Store>::accept(IRVisitor *v) const {
v->visit((const Store *)this);
}
template<>
void StmtNode<Provide>::accept(IRVisitor *v) const {
v->visit((const Provide *)this);
}
template<>
void StmtNode<Allocate>::accept(IRVisitor *v) const {
v->visit((const Allocate *)this);
}
template<>
void StmtNode<Free>::accept(IRVisitor *v) const {
v->visit((const Free *)this);
}
template<>
void StmtNode<Realize>::accept(IRVisitor *v) const {
v->visit((const Realize *)this);
}
template<>
void StmtNode<Block>::accept(IRVisitor *v) const {
v->visit((const Block *)this);
}
template<>
void StmtNode<IfThenElse>::accept(IRVisitor *v) const {
v->visit((const IfThenElse *)this);
}
template<>
void StmtNode<Evaluate>::accept(IRVisitor *v) const {
v->visit((const Evaluate *)this);
}
template<>
void StmtNode<Prefetch>::accept(IRVisitor *v) const {
v->visit((const Prefetch *)this);
}
template<>
void StmtNode<Acquire>::accept(IRVisitor *v) const {
v->visit((const Acquire *)this);
}
template<>
void StmtNode<Fork>::accept(IRVisitor *v) const {
v->visit((const Fork *)this);
}
template<>
void StmtNode<Atomic>::accept(IRVisitor *v) const {
v->visit((const Atomic *)this);
}
template<>
Expr ExprNode<IntImm>::mutate_expr(IRMutator *v) const {
return v->visit((const IntImm *)this);
}
template<>
Expr ExprNode<UIntImm>::mutate_expr(IRMutator *v) const {
return v->visit((const UIntImm *)this);
}
template<>
Expr ExprNode<FloatImm>::mutate_expr(IRMutator *v) const {
return v->visit((const FloatImm *)this);
}
template<>
Expr ExprNode<StringImm>::mutate_expr(IRMutator *v) const {
return v->visit((const StringImm *)this);
}
template<>
Expr ExprNode<Cast>::mutate_expr(IRMutator *v) const {
return v->visit((const Cast *)this);
}
template<>
Expr ExprNode<Variable>::mutate_expr(IRMutator *v) const {
return v->visit((const Variable *)this);
}
template<>
Expr ExprNode<Add>::mutate_expr(IRMutator *v) const {
return v->visit((const Add *)this);
}
template<>
Expr ExprNode<Sub>::mutate_expr(IRMutator *v) const {
return v->visit((const Sub *)this);
}
template<>
Expr ExprNode<Mul>::mutate_expr(IRMutator *v) const {
return v->visit((const Mul *)this);
}
template<>
Expr ExprNode<Div>::mutate_expr(IRMutator *v) const {
return v->visit((const Div *)this);
}
template<>
Expr ExprNode<Mod>::mutate_expr(IRMutator *v) const {
return v->visit((const Mod *)this);
}
template<>
Expr ExprNode<Min>::mutate_expr(IRMutator *v) const {
return v->visit((const Min *)this);
}
template<>
Expr ExprNode<Max>::mutate_expr(IRMutator *v) const {
return v->visit((const Max *)this);
}
template<>
Expr ExprNode<EQ>::mutate_expr(IRMutator *v) const {
return v->visit((const EQ *)this);
}
template<>
Expr ExprNode<NE>::mutate_expr(IRMutator *v) const {
return v->visit((const NE *)this);
}
template<>
Expr ExprNode<LT>::mutate_expr(IRMutator *v) const {
return v->visit((const LT *)this);
}
template<>
Expr ExprNode<LE>::mutate_expr(IRMutator *v) const {
return v->visit((const LE *)this);
}
template<>
Expr ExprNode<GT>::mutate_expr(IRMutator *v) const {
return v->visit((const GT *)this);
}
template<>
Expr ExprNode<GE>::mutate_expr(IRMutator *v) const {
return v->visit((const GE *)this);
}
template<>
Expr ExprNode<And>::mutate_expr(IRMutator *v) const {
return v->visit((const And *)this);
}
template<>
Expr ExprNode<Or>::mutate_expr(IRMutator *v) const {
return v->visit((const Or *)this);
}
template<>
Expr ExprNode<Not>::mutate_expr(IRMutator *v) const {
return v->visit((const Not *)this);
}
template<>
Expr ExprNode<Select>::mutate_expr(IRMutator *v) const {
return v->visit((const Select *)this);
}
template<>
Expr ExprNode<Load>::mutate_expr(IRMutator *v) const {
return v->visit((const Load *)this);
}
template<>
Expr ExprNode<Ramp>::mutate_expr(IRMutator *v) const {
return v->visit((const Ramp *)this);
}
template<>
Expr ExprNode<Broadcast>::mutate_expr(IRMutator *v) const {
return v->visit((const Broadcast *)this);
}
template<>
Expr ExprNode<Call>::mutate_expr(IRMutator *v) const {
return v->visit((const Call *)this);
}
template<>
Expr ExprNode<Shuffle>::mutate_expr(IRMutator *v) const {
return v->visit((const Shuffle *)this);
}
template<>
Expr ExprNode<VectorReduce>::mutate_expr(IRMutator *v) const {
return v->visit((const VectorReduce *)this);
}
template<>
Expr ExprNode<Let>::mutate_expr(IRMutator *v) const {
return v->visit((const Let *)this);
}
template<>
Stmt StmtNode<LetStmt>::mutate_stmt(IRMutator *v) const {
return v->visit((const LetStmt *)this);
}
template<>
Stmt StmtNode<AssertStmt>::mutate_stmt(IRMutator *v) const {
return v->visit((const AssertStmt *)this);
}
template<>
Stmt StmtNode<ProducerConsumer>::mutate_stmt(IRMutator *v) const {
return v->visit((const ProducerConsumer *)this);
}
template<>
Stmt StmtNode<For>::mutate_stmt(IRMutator *v) const {
return v->visit((const For *)this);
}
template<>
Stmt StmtNode<Store>::mutate_stmt(IRMutator *v) const {
return v->visit((const Store *)this);
}
template<>
Stmt StmtNode<Provide>::mutate_stmt(IRMutator *v) const {
return v->visit((const Provide *)this);
}
template<>
Stmt StmtNode<Allocate>::mutate_stmt(IRMutator *v) const {
return v->visit((const Allocate *)this);
}
template<>
Stmt StmtNode<Free>::mutate_stmt(IRMutator *v) const {
return v->visit((const Free *)this);
}
template<>
Stmt StmtNode<Realize>::mutate_stmt(IRMutator *v) const {
return v->visit((const Realize *)this);
}
template<>
Stmt StmtNode<Block>::mutate_stmt(IRMutator *v) const {
return v->visit((const Block *)this);
}
template<>
Stmt StmtNode<IfThenElse>::mutate_stmt(IRMutator *v) const {
return v->visit((const IfThenElse *)this);
}
template<>
Stmt StmtNode<Evaluate>::mutate_stmt(IRMutator *v) const {
return v->visit((const Evaluate *)this);
}
template<>
Stmt StmtNode<Prefetch>::mutate_stmt(IRMutator *v) const {
return v->visit((const Prefetch *)this);
}
template<>
Stmt StmtNode<Acquire>::mutate_stmt(IRMutator *v) const {
return v->visit((const Acquire *)this);
}
template<>
Stmt StmtNode<Fork>::mutate_stmt(IRMutator *v) const {
return v->visit((const Fork *)this);
}
template<>
Stmt StmtNode<Atomic>::mutate_stmt(IRMutator *v) const {
return v->visit((const Atomic *)this);
}
Call::ConstString Call::buffer_get_dimensions = "_halide_buffer_get_dimensions";
Call::ConstString Call::buffer_get_min = "_halide_buffer_get_min";
Call::ConstString Call::buffer_get_extent = "_halide_buffer_get_extent";
Call::ConstString Call::buffer_get_stride = "_halide_buffer_get_stride";
Call::ConstString Call::buffer_get_max = "_halide_buffer_get_max";
Call::ConstString Call::buffer_get_host = "_halide_buffer_get_host";
Call::ConstString Call::buffer_get_device = "_halide_buffer_get_device";
Call::ConstString Call::buffer_get_device_interface = "_halide_buffer_get_device_interface";
Call::ConstString Call::buffer_get_shape = "_halide_buffer_get_shape";
Call::ConstString Call::buffer_get_host_dirty = "_halide_buffer_get_host_dirty";
Call::ConstString Call::buffer_get_device_dirty = "_halide_buffer_get_device_dirty";
Call::ConstString Call::buffer_get_type = "_halide_buffer_get_type";
Call::ConstString Call::buffer_set_host_dirty = "_halide_buffer_set_host_dirty";
Call::ConstString Call::buffer_set_device_dirty = "_halide_buffer_set_device_dirty";
Call::ConstString Call::buffer_is_bounds_query = "_halide_buffer_is_bounds_query";
Call::ConstString Call::buffer_init = "_halide_buffer_init";
Call::ConstString Call::buffer_init_from_buffer = "_halide_buffer_init_from_buffer";
Call::ConstString Call::buffer_crop = "_halide_buffer_crop";
Call::ConstString Call::buffer_set_bounds = "_halide_buffer_set_bounds";
Call::ConstString Call::trace = "halide_trace_helper";
} // namespace Internal
} // namespace Halide