[tvm] branch main updated: [AutoScheduler] Skip useless calls to RewriteLayout (#6993)

[lm...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

lmzheng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new d14e779  [AutoScheduler] Skip useless calls to RewriteLayout (#6993)
d14e779 is described below

commit d14e779fd6fd5e3501bdee3880e7f914520e2f48
Author: Lianmin Zheng <li...@gmail.com>
AuthorDate: Sun Nov 29 13:36:56 2020 -0800

    [AutoScheduler] Skip useless calls to RewriteLayout (#6993)
    
    * [AutoScheduler] Skip useless calls of RewriteLayout
    
    * fix lint
    
    * fix lint
---
 src/auto_scheduler/compute_dag.cc | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/auto_scheduler/compute_dag.cc b/src/auto_scheduler/compute_dag.cc
index e57fc8c..caaed6f 100755
--- a/src/auto_scheduler/compute_dag.cc
+++ b/src/auto_scheduler/compute_dag.cc
@@ -1121,10 +1121,25 @@ ComputeDAG ComputeDAG::RewriteLayout(Array<Step>* transform_steps,
   return new_dag;
 }
 
+// Return whether a DAG has placeholders that are marked as "layout free".
+bool HasLayoutFreeTensors(const ComputeDAG& dag) {
+  for (const auto& op : dag->ops) {
+    if (!op->IsInstance<te::ComputeOpNode>()) {
+      continue;
+    }
+    if (op->attrs.count(ComputeDAG::layout_free_placeholders_key)) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
 std::pair<te::Schedule, Array<te::Tensor>> ComputeDAG::ApplySteps(
     const Array<Step>& transform_steps, Array<te::Stage>* stages, StageToAxesMap* stage_to_axes,
     LayoutRewriteOption layout_rewrite) const {
-  if (layout_rewrite != LayoutRewriteOption::NoRewrite && !transform_steps.empty()) {
+  if (layout_rewrite != LayoutRewriteOption::NoRewrite && HasLayoutFreeTensors(*this) &&
+      !transform_steps.empty()) {
     Array<Step> steps = transform_steps;
     const auto& dag = RewriteLayout(&steps, layout_rewrite);
     return dag.ApplySteps(steps);