C++ Stmt::defined方法代码示例

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 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);
    }
  }
}

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 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 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));
        }
    }

Stmt IRMutator2::mutate(const Stmt &s) {
    return s.defined() ? ((const BaseStmtNode *)s.get())->mutate_stmt(this) : Stmt();
}