[PIR] Open some ut for cf (#60966)

zhangbo9674 · web-flow · commit 41c280eaa71b · 2024-01-22T13:01:00.000+08:00
* fix

* fix

* fix

* fix

* fix
diff --git a/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc b/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
@@ -87,11 +87,8 @@ void ValueExecutionInfo::Add(::pir::Value value, const std::string& var_name) {
           "The size of variable_list and var_name_2_id map should be equal"));
 }
 
-void ValueExecutionInfo::Rename(pir::Value value,
-                                const std::string& new_name,
+void ValueExecutionInfo::Rename(const std::string& new_name,
                                 const std::string& orig_name) {
-  value_2_var_name_[value] = new_name;
-
   for (auto kv : value_2_var_name_) {
     if (kv.second == orig_name) {
       value_2_var_name_[kv.first] = new_name;
@@ -148,6 +145,11 @@ void ValueExecutionInfo::AddValue2VarName(::pir::Value value,
   value_2_var_name_.emplace(value, var_name);
 }
 
+void ValueExecutionInfo::UpdateValue2VarName(::pir::Value value,
+                                             const std::string& var_name) {
+  value_2_var_name_[value] = var_name;
+}
+
 const std::unordered_map<const Variable*, std::string>&
 ValueExecutionInfo::GetVar2VarName() const {
   return var_2_var_name_;
@@ -542,7 +544,7 @@ void HandleForSpecialOp(pir::Operation* op,
               << param_name;
     }
 
-    value_exe_info->Rename(value, param_name, orig_name);
+    value_exe_info->Rename(param_name, orig_name);
   } else if (op->isa<pir::ShadowOutputOp>()) {
     VLOG(6) << "Handle for builtin.shadow_ouptut";
     auto var_name = op->attributes()
@@ -558,14 +560,18 @@ void HandleForSpecialOp(pir::Operation* op,
       return;
     }
 
-    if (value_exe_info->GetScope()->FindVar(var_name) != nullptr) {
-      const_cast<Scope*>(value_exe_info->GetScope())->EraseVars({var_name});
-      VLOG(1) << "var " << var_name << " has been removed from scope";
+    if (value_exe_info->HasVar(var_name)) {
+      value_exe_info->UpdateValue2VarName(value, var_name);
+    } else {
+      if (value_exe_info->GetScope()->FindVar(var_name) != nullptr) {
+        const_cast<Scope*>(value_exe_info->GetScope())->EraseVars({var_name});
+        VLOG(1) << "var " << var_name << " has been removed from scope";
+      }
+      const_cast<Scope*>(value_exe_info->GetScope())
+          ->Rename(orig_name, var_name);
+      VLOG(8) << "var " << orig_name << " has been renamed to " << var_name;
+      value_exe_info->Rename(var_name, orig_name);
     }
-    const_cast<Scope*>(value_exe_info->GetScope())->Rename(orig_name, var_name);
-    VLOG(8) << "var " << orig_name << " has been renamed to " << var_name;
-
-    value_exe_info->Rename(value, var_name, orig_name);
   } else if (op->isa<pir::ParameterOp>()) {
     VLOG(6) << "Handle for builtin.parameter:";
     auto param_name = op->attributes()
diff --git a/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h b/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h
@@ -58,9 +58,7 @@ class ValueExecutionInfo {
 
   void Add(::pir::Value value, const std::string& var_name);
 
-  void Rename(pir::Value value,
-              const std::string& new_name,
-              const std::string& orig_name);
+  void Rename(const std::string& new_name, const std::string& orig_name);
 
   int GetIdByName(const std::string& name) const;
 
@@ -72,6 +70,8 @@ class ValueExecutionInfo {
 
   void AddValue2VarName(::pir::Value value, const std::string& var_name);
 
+  void UpdateValue2VarName(::pir::Value value, const std::string& var_name);
+
   const std::unordered_map<const paddle::framework::Variable*, std::string>&
   GetVar2VarName() const;
 
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -187,6 +187,7 @@ static bool NeedFallBackCpu(const pir::Operation* op,
 }
 
 static bool NeedFallBackFromGPUDNN2GPU(pir::Operation* op,
+                                       const std::string& kernel_name,
                                        const phi::KernelKey kernel_key) {
   // NOTE(phlrain): keep the same kernel select strategy with
   // GetExepectKernelKey
@@ -224,6 +225,21 @@ static bool NeedFallBackFromGPUDNN2GPU(pir::Operation* op,
 #endif
     return !use_cudnn;
   }
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  if (kernel_key.backend() == phi::Backend::GPUDNN) {
+    auto iter = phi::KernelFactory::Instance().kernels().find(kernel_name);
+    if (iter != phi::KernelFactory::Instance().kernels().end()) {
+      auto kernel_iter = iter->second.find({phi::Backend::GPUDNN,
+                                            phi::DataLayout::ALL_LAYOUT,
+                                            kernel_key.dtype()});
+      if (kernel_iter == iter->second.end()) {
+        return true;
+      }
+    }
+  }
+#endif
+
   return false;
 }
 
@@ -1069,7 +1085,7 @@ phi::KernelKey GetKernelKey(
     VLOG(8) << "kernel backend must be on CPU when need fallback";
   }
 
-  if (NeedFallBackFromGPUDNN2GPU(op, res)) {
+  if (NeedFallBackFromGPUDNN2GPU(op, kernel_fn_str, res)) {
     res.set_backend(phi::Backend::GPU);
     VLOG(8) << "kernel backend must be on GPU when need fallback from GPUDNN "
                "to GPU";
diff --git a/test/dygraph_to_static/test_break_continue.py b/test/dygraph_to_static/test_break_continue.py
@@ -17,9 +17,6 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
-    IrMode,
-    ToStaticMode,
-    disable_test_case,
     enable_to_static_guard,
     test_ast_only,
     test_legacy_and_pt,
@@ -300,7 +297,7 @@ def init_dygraph_func(self):
         self.dygraph_func = test_continue_in_while
 
     # TODO(dev): Remove this after fix PT Rename issue
-    @disable_test_case((ToStaticMode.AST, IrMode.PT))
+    @test_legacy_and_pt
     def test_transformed_static_result(self):
         self.init_dygraph_func()
         dygraph_res = self.run_dygraph_mode()
@@ -318,7 +315,7 @@ def init_dygraph_func(self):
         self.dygraph_func = test_break_in_while
 
     # TODO(dev): Remove this after fix PT Rename issue
-    @disable_test_case((ToStaticMode.AST, IrMode.PT))
+    @test_legacy_and_pt
     def test_transformed_static_result(self):
         self.init_dygraph_func()
         dygraph_res = self.run_dygraph_mode()
@@ -360,7 +357,7 @@ def init_dygraph_func(self):
         self.dygraph_func = test_optim_break_in_while
 
     # TODO(dev): Remove this after fix PT Rename issue
-    @disable_test_case((ToStaticMode.AST, IrMode.PT))
+    @test_legacy_and_pt
     def test_transformed_static_result(self):
         self.init_dygraph_func()
         dygraph_res = self.run_dygraph_mode()
diff --git a/test/dygraph_to_static/test_jit_setitem.py b/test/dygraph_to_static/test_jit_setitem.py
@@ -18,9 +18,7 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
-    IrMode,
-    ToStaticMode,
-    disable_test_case,
+    test_legacy_and_pt,
     test_legacy_and_pt_and_pir,
 )
 
@@ -206,8 +204,7 @@ def run_dygraph(self, func):
         return (y,)
 
     # TODO: Open PIR test when while_loop dy2st fixed
-    @disable_test_case((ToStaticMode.AST, IrMode.PIR))
-    @test_legacy_and_pt_and_pir
+    @test_legacy_and_pt
     def test_case(self):
         func = self.init_func()
         dy_res = self.run_dygraph(func)
