!31938 fix endless infer of the high-order differential function.

Merge pull request !31938 from lanzhineng/func_closure

mindspore/ccsrc/pipeline/jit/static_analysis/evaluator.cc

@@ -319,7 +319,7 @@ AbstractBasePtrList FuncGraphEvaluator::NormalizeArgs(const AbstractBasePtrList
     AbstractBasePtrList broaded_list;
     BroadenArgs(args_spec_list, &broaded_list);
     MS_LOG(DEBUG) << func_graph_->ToString() << ", original: " << mindspore::ToString(args_spec_list)
-                  << ", broaded: " << mindspore::ToString(broaded_list);
+                  << ", broadened: " << mindspore::ToString(broaded_list);
     return broaded_list;
   }
   return args_spec_list;
@@ -349,6 +349,10 @@ FuncGraphPtr FuncGraphEvaluator::GetFuncGraph(AnalysisEnginePtr engine, const Ab
     MS_EXCEPTION_IF_NULL(fg);
     FuncGraphPtr generated_graph = fg->GenerateGraph(args_spec_list);
     func_graph_cache_[args_spec_list] = generated_graph;
+    MS_LOG(DEBUG) << "Generate special instance of function graph: " << ToString()
+                  << ", special function: " << generated_graph->ToString()
+                  << ", args: " << ArgsToString(args_spec_list);
+
     MS_EXCEPTION_IF_NULL(engine);
     engine->func_graph_manager()->AddFuncGraph(generated_graph);
     res = generated_graph;

mindspore/ccsrc/pipeline/jit/static_analysis/static_analysis.cc

@@ -661,26 +661,23 @@ void AnalysisEngine::SetUndeterminedFlag(const EvaluatorPtr &evaluator, const Fu
   MS_EXCEPTION_IF_NULL(evaluator);
   static std::mutex fg_lock;
   std::lock_guard<std::mutex> infer_lock(fg_lock);
+  if (possible_parent_fg != nullptr) {
+    possible_parent_fg->set_flag(kFuncGraphFlagUndetermined, true);
+    MS_LOG(DEBUG) << "Set graph undetermined: " << possible_parent_fg->ToString();
+  }
   auto fg_eval = evaluator->cast<FuncGraphEvaluatorPtr>();
   if (fg_eval == nullptr) {
     return;
   }
-
   auto fg = fg_eval->func_graph();
   MS_EXCEPTION_IF_NULL(fg);
-  auto undetermined_fgs = fg->recursive();
-  if (undetermined_fgs) {
-    auto fg_parent = fg->parent();
-    if (fg_parent != nullptr) {
-      fg_parent->set_flag(kFuncGraphFlagUndetermined, true);
-      MS_LOG(DEBUG) << "Set graph undetermined: " << fg_parent->ToString() << " for fg: " << fg->ToString();
-      return;
-    } else if (possible_parent_fg != nullptr) {
-      possible_parent_fg->set_flag(kFuncGraphFlagUndetermined, true);
-      MS_LOG(DEBUG) << "Set graph undetermined: " << possible_parent_fg->ToString() << " for fg: " << fg->ToString();
-    } else {
-      MS_LOG(EXCEPTION) << "cannot find parent for fg: " << fg->ToString();
-    }
+  auto fg_parent = fg->parent();
+  if (fg_parent != nullptr) {
+    fg_parent->set_flag(kFuncGraphFlagUndetermined, true);
+    MS_LOG(DEBUG) << "Set graph undetermined: " << fg_parent->ToString() << " for fg: " << fg->ToString();
+    return;
+  } else {
+    MS_LOG(DEBUG) << "cannot find parent for fg: " << fg->ToString();
   }
 }
 
@@ -960,7 +957,6 @@ EvalResultPtr AnalysisEngine::ExecuteMultipleEvaluatorsMultiThread(const std::ve
     return std::make_shared<EvalResult>(eval_result, nullptr);
   }
   auto possible_parent_fg = out_conf->node()->func_graph();
-
   // Eval result of the main.
   AsyncAbstractPtr async_result_main = std::make_shared<AsyncAbstract>();
   // Eval result of the branches

tests/st/control/test_high_order_control.py

@@ -0,0 +1,84 @@
+# Copyright 2021-2022 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+""" test high order control flow """
+import pytest
+from mindspore.nn import Cell
+from mindspore.common import Tensor, dtype
+import mindspore.ops.functional as F
+
+
+@pytest.mark.level0
+@pytest.mark.platform_x86_gpu_training
+@pytest.mark.platform_arm_ascend_training
+@pytest.mark.platform_x86_ascend_training
+@pytest.mark.env_onecard
+def test_high_control_while():
+    """
+    Feature: High-order differential function.
+    Description: Infer of the high-order differential function.
+    Expectation: Null.
+    """
+
+    class Net(Cell):
+        def construct(self, x):
+            while x < 10:
+                x = (x * 2)
+            return x
+
+    net = Net()
+    x = Tensor(1, dtype.float32)
+    grad_net = F.grad(net)
+    order_grad_net = F.grad(grad_net)
+    order_grad = order_grad_net(x)
+    assert order_grad == 0.0
+
+
+@pytest.mark.level0
+@pytest.mark.platform_x86_gpu_training
+@pytest.mark.platform_arm_ascend_training
+@pytest.mark.platform_x86_ascend_training
+@pytest.mark.env_onecard
+def test_high_control_for_while():
+    """
+    Feature: High-order differential function.
+    Description: Infer of the complex high-order differential function.
+    Expectation: Null.
+    """
+
+    class Net(Cell):
+        def construct(self, x):
+            for _ in [2]:
+                for _ in [2]:
+                    while x > 1:
+                        x = (x / 3)
+                        x = (x / 2)
+                for _ in [2]:
+                    x = (x / 1)
+                    x = (x + 1)
+            for _ in [3]:
+                for _ in [4]:
+                    x = (x / 1)
+                    x = (x + 3)
+                for _ in [5]:
+                    x = (x / 3)
+                    x = (x / 2)
+            return x
+
+    net = Net()
+    x = Tensor(4, dtype.float32)
+    grad_net = F.grad(net)
+    grad_grad_net = F.grad(grad_net)
+    result = grad_grad_net(x)
+    assert result == 0.0