From ab1e9984dd87dcf0e83a9b7482db07b58ae10856 Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <kparzysz@quicinc.com>
Date: Mon, 27 Sep 2021 19:38:26 -0500
Subject: [PATCH] [Codegen] Swap out analyzer when outlining (#9117)

Problem: the `analyzer_` in `CodeGenLLVM` and derived classes
can generate invalid code for outlined functions.

Consider code like this:

  let x = y in
    // attr compute_scope
    blah = x

Then it gets outlined in codegen_cpu (for example):

  let x = y in
    call foo(x)

  foo(x) {
    blah = x
  }

Now, if `analyzer_->Simplify` was run on the body of `foo`, it
would produce:

  foo(x) {
    blah = y
  }

Because the `analyzer_` knows that `x` is same as `y` (because
of the `Let` statemement), but doesn't know that `y` is no longer
available in the outlined function `foo`.

See
https://discuss.tvm.apache.org/t/compute-scope-issue-with-analyzer-invalid-simplification/11111
---
 src/target/llvm/codegen_cpu.cc | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index e67dee3c37c4..40e6245e8beb 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -473,14 +473,17 @@ void CodeGenCPU::CreateComputeScope(const AttrStmtNode* op) {
     }
 #endif
   }
+  auto new_analyzer = std::make_unique<arith::Analyzer>();
   std::swap(function_, fcompute);
-  std::swap(new_vmap, var_map_);
+  std::swap(analyzer_, new_analyzer);
+  std::swap(var_map_, new_vmap);
   BasicBlock* compute_entry = BasicBlock::Create(*ctx_, "entry", function_);
   builder_->SetInsertPoint(compute_entry);
   this->VisitStmt(op->body);
   builder_->CreateRet(ConstInt32(0));
   // swap the var map back, now we are back on track.
-  std::swap(new_vmap, var_map_);
+  std::swap(var_map_, new_vmap);
+  std::swap(analyzer_, new_analyzer);
   std::swap(function_, fcompute);
   builder_->SetInsertPoint(compute_call_end);
 }
@@ -551,13 +554,16 @@ void CodeGenCPU::CreateParallelLaunch(const Stmt& body, int num_task) {
   new_vmap[par_env.num_task.get()] =
       builder_->CreateLoad(builder_->CreateInBoundsGEP(penv, {ConstInt32(0), ConstInt32(1)}));
   par_env.penv = penv;
+  auto new_analyzer = std::make_unique<arith::Analyzer>();
   std::swap(function_, f);
   std::swap(parallel_env_, par_env);
+  std::swap(analyzer_, new_analyzer);
   std::swap(var_map_, new_vmap);
   this->VisitStmt(body);
   builder_->CreateRet(ConstInt32(0));
   // swap the var map back, now we are back on track.
   std::swap(var_map_, new_vmap);
+  std::swap(analyzer_, new_analyzer);
   std::swap(parallel_env_, par_env);
   std::swap(function_, f);
   ICHECK_NE(par_env.parallel_loop_count, 0) << "Cannot find parallel loop within parallel launch";
@@ -604,12 +610,15 @@ void CodeGenCPU::CreateStaticInit(const std::string& init_fname, const Stmt& bod
   std::unordered_map<const VarNode*, llvm::Value*> new_vmap;
   UnpackClosureData(cdata, vfields, &new_vmap);
   ICHECK(parallel_env_.penv == nullptr);
+  auto new_analyzer = std::make_unique<arith::Analyzer>();
   std::swap(function_, f);
+  std::swap(analyzer_, new_analyzer);
   std::swap(var_map_, new_vmap);
   this->VisitStmt(body);
   builder_->CreateRet(ConstInt32(0));
   // swap the var map back, now we are back on track.
   std::swap(var_map_, new_vmap);
+  std::swap(analyzer_, new_analyzer);
   std::swap(function_, f);
   builder_->SetInsertPoint(init_end);
 }