ReduceNode in TIR

xqdan · October 30, 2020, 1:10am

Hi,

ReduceNode only exisits in ScheduleOps for the tvm 0.6, is it still true in latest tvm? if it is true, I’m confused two passes after ScheduleOps have ReduceNode logical. Can anyone explain a little bit? Thanks in advance.

github.com

apache/incubator-tvm/blob/1831c17998b29f3797f364410980809bfef554ca/src/tir/transforms/narrow_datatype.cc#L117


    IterVar iv = Downcast<IterVar>(op->node);
    ICHECK_NE(iv->thread_tag.length(), 0U);
    analyzer_.Bind(iv->var, Range::FromMinExtent(0, op->value));
    vextent_[iv->var.as<VarNode>()] = op->value.dtype();
    StmtExprVisitor::VisitStmt_(op);
  } else {
    StmtExprVisitor::VisitStmt_(op);
  }
}

void VisitExpr_(const ReduceNode* op) {
  // Setup the domain information before simplification.
  for (const IterVar& iv : op->axis) {
    analyzer_.Bind(iv->var, iv->dom);
    vextent_[iv->var.as<VarNode>()] = iv->dom->extent.dtype();
  }
  // Recursively call simplification when necessary.
  StmtExprVisitor::VisitExpr_(op);
}

void VisitExpr_(const VarNode* op) {

github.com

apache/incubator-tvm/blob/1831c17998b29f3797f364410980809bfef554ca/src/tir/transforms/split_host_device.cc#L137


      return Let(op->var, value, body);
    }
  }
}

PrimExpr VisitExpr_(const VarNode* op) final {
  this->HandleUse(GetRef<PrimExpr>(op));
  return StmtExprMutator::VisitExpr_(op);
}

PrimExpr VisitExpr_(const ReduceNode* op) final {
  for (const auto& iv : op->axis) {
    this->HandleDef(iv->var.get());
  }
  return StmtExprMutator::VisitExpr_(op);
}

PrimExpr VisitExpr_(const LoadNode* op) final {
  this->HandleUse(op->buffer_var);
  return StmtExprMutator::VisitExpr_(op);
}

@tqchen @yzhliu @spectrometerHBH