Stmt Pipeline::make(std::string name, Stmt produce, Stmt update, Stmt consume) {
internal_assert(produce.defined()) << "Pipeline of undefined\n";
// update is allowed to be null
internal_assert(consume.defined()) << "Pipeline of undefined\n";
Pipeline *node = new Pipeline;
node->name = name;
node->produce = produce;
node->update = update;
node->consume = consume;
return node;
}
Stmt IRRewriter::rewrite(Stmt s) {
if (s.defined()) {
s.accept(this);
Stmt spilledStmts = getSpilledStmts();
if (spilledStmts.defined()) {
stmt = Block::make(spilledStmts, stmt);
}
s = stmt;
}
else {
s = Stmt();
}
expr = Expr();
stmt = Stmt();
func = Func();
return s;
}
void visit(const IfThenElse *op) {
Expr start = rewrite(op->condition);
Stmt thenBody = rewrite(op->thenBody);
Stmt elseBody = rewrite(op->elseBody);
stmt = !elseBody.defined() ? IfThenElse::make(op->condition, thenBody) :
IfThenElse::make(op->condition, thenBody, elseBody);
}
Stmt Block::make(Stmt first, Stmt rest) {
internal_assert(first.defined()) << "Block of undefined\n";
// rest is allowed to be null
Block *node = new Block;
node->first = first;
node->rest = rest;
return node;
}
Stmt IRMutator::mutate(Stmt s) {
if (s.defined()) {
s.accept(this);
} else {
stmt = Stmt();
}
expr = Expr();
return stmt;
}
Stmt IRMutator::mutate(const Stmt &s) {
if (s.defined()) {
s.accept(this);
} else {
stmt = Stmt();
}
expr = Expr();
return std::move(stmt);
}
void IRRewriter::visit(const While *op) {
Expr condition = rewrite(op->condition);
Stmt spilledCond = getSpilledStmts();
Stmt body = rewrite(op->body);
if (condition == op->condition && body == op->body) {
stmt = op;
}
else {
if (spilledCond.defined()) {
body = Block::make(body, spilledCond);
}
stmt = While::make(condition, body);
if (spilledCond.defined()) {
stmt = Block::make(spilledCond, stmt);
}
}
}
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 LetStmt::make(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 = value;
node->body = body;
return node;
}
void IRRewriter::visit(const IfThenElse *op) {
Expr condition = rewrite(op->condition);
Stmt spilledCond = getSpilledStmts();
Stmt thenBody = rewrite(op->thenBody);
Stmt elseBody = rewrite(op->elseBody);
if (condition == op->condition && thenBody == op->thenBody &&
elseBody == op->elseBody) {
stmt = op;
}
else {
if (elseBody.defined()) {
stmt = IfThenElse::make(condition, thenBody, elseBody);
} else {
stmt = IfThenElse::make(condition, thenBody);
}
if (spilledCond.defined()) {
stmt = Block::make(spilledCond, stmt);
}
}
}
void IRRewriter::visit(const Func *f) {
Stmt body = rewrite(f->getBody());
if (body == f->getBody()) {
func = *f;
}
else {
if (!body.defined()) {
body = Pass::make();
}
func = Func(*f, body);
}
}
Func::Func(const std::string& name, const std::vector<Var>& arguments,
const std::vector<Var>& results, const Stmt& body,
const Environment& environment, Kind kind)
: IntrusivePtr(new FuncContent) {
ptr->kind = kind;
ptr->name = name;
ptr->arguments = arguments;
ptr->results = results;
ptr->env = environment;
ptr->body = new Stmt();
if (body.defined()) {
*ptr->body = body;
}
}
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;
}
Stmt lowerTranspose(Var target, const IndexExpr* iexpr,
Environment* env, Storage* storage) {
simit_iassert(isa<IndexedTensor>(iexpr->value));
simit_iassert(isa<VarExpr>(to<IndexedTensor>(iexpr->value)->tensor));
Var source = to<VarExpr>(to<IndexedTensor>(iexpr->value)->tensor)->var;
auto sourceIndex = storage->getStorage(source).getTensorIndex();
auto targetIndex = storage->getStorage(target).getTensorIndex();
auto sourceType = source.getType().toTensor();
auto iRange = sourceType->getOuterDimensions()[0];
Var i("i", Int);
Var ij("ij", Int);
Var j("j", Int);
Var locVar(INTERNAL_PREFIX("locVar"), Int);
Stmt locStmt = CallStmt::make({locVar}, intrinsics::loc(),
{Load::make(sourceIndex.getColidxArray(),ij), i,
targetIndex.getRowptrArray(),
targetIndex.getColidxArray()});
Stmt body;
auto blockType = *sourceType->getBlockType().toTensor();
if (blockType.order() == 0) { // Not blocked
body = Store::make(target, locVar, Load::make(source, ij));
}
else { // Blocked
simit_iassert(blockType.order() == 2);
Var ii("ii", Int);
Var jj("jj", Int);
auto d1 = blockType.getOuterDimensions()[0];
auto d2 = blockType.getOuterDimensions()[1];
Expr l1 = Length::make(d1);
Expr l2 = Length::make(d2);
body = Store::make(target, locVar*l1*l2 + ii*l2 + jj,
Load::make(source, ij*l1*l2 + ii*l2 + jj));
body = For::make(jj, ForDomain(d2), body);
body = For::make(ii, ForDomain(d1), body);
}
simit_iassert(body.defined());
Expr start = Load::make(sourceIndex.getRowptrArray(), i);
Expr stop = Load::make(sourceIndex.getRowptrArray(), i+1);
Stmt innerLoop = ForRange::make(ij, start, stop, Block::make(locStmt, body));
return For::make(i, iRange, innerLoop);
}
void IRRewriter::visit(const ForRange *op) {
Expr start = rewrite(op->start);
Expr end = rewrite(op->end);
Stmt spilledBounds = getSpilledStmts();
Stmt body = rewrite(op->body);
if (body == op->body && start == op->start && end == op->end) {
stmt = op;
}
else {
stmt = ForRange::make(op->var, start, end, body);
if (spilledBounds.defined()) {
stmt = Block::make(spilledBounds, stmt);
}
}
}
Stmt Allocate::make(std::string name, Type type, const std::vector<Expr> &extents,
Expr condition, Stmt body) {
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";
Allocate *node = new Allocate;
node->name = name;
node->type = type;
node->extents = extents;
node->condition = condition;
node->body = body;
return node;
}
Stmt Realize::make(const std::string &name, const std::vector<Type> &types, 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";
Realize *node = new Realize;
node->name = name;
node->types = types;
node->bounds = bounds;
node->condition = condition;
node->body = body;
return node;
}
void visit(const Block *op) {
/* First we dig into the block traversing down the 'first'
* stmt until we find one that is not a block. We push all of
* the rest stmt's into the 'rest' stmt of the top-level
* block, and then fix up the 'rest' stmt recursively at the
* end. The result of this mutation is an equivalent Block
* node that does not contain any Block nodes in a 'first' stmt.
*/
Stmt first = op->first;
Stmt rest = op->rest;
while(const Block *first_block = first.as<Block>()) {
first = first_block->first;
if (first_block->rest.defined()) {
rest = rest.defined()? Block::make(first_block->rest, rest): first_block->rest;
}
}
if (first.same_as(op->first)) {
rest = mutate(rest);
stmt = rest.same_as(op->rest)? op: Block::make(first, rest);
} else {
stmt = Block::make(first, mutate(rest));
}
}
void IRRewriter::visit(const Block *op) {
Stmt first = rewrite(op->first);
Stmt rest = rewrite(op->rest);
if (first == op->first && rest == op->rest) {
stmt = op;
}
else {
if (first.defined() && rest.defined()) {
stmt = Block::make(first, rest);
}
else if (first.defined() && !rest.defined()) {
stmt = first;
}
else if (!first.defined() && rest.defined()) {
stmt = rest;
}
else {
stmt = Stmt();
}
}
}
Stmt IRMutator2::mutate(const Stmt &s) {
return s.defined() ? ((const BaseStmtNode *)s.get())->mutate_stmt(this) : Stmt();
}
