From 4812c707119f30d7c0e6b5d531091163d5f9ca33 Mon Sep 17 00:00:00 2001
From: hong <43953930+phlrain@users.noreply.github.com>
Date: Tue, 8 Aug 2023 10:59:22 +0800
Subject: [PATCH] [NewIR]enable dy2st new ir test in ci (#56034)

* support new ir dy2st

* revert code

* skip test with stride

* chang import file
---
 test/dygraph_to_static/dygraph_to_static_util.py |  2 ++
 test/dygraph_to_static/test_bert.py              | 13 ++++++++++++-
 test/dygraph_to_static/test_mobile_net.py        |  8 ++++++++
 test/dygraph_to_static/test_resnet.py            | 14 ++++++++++++++
 test/dygraph_to_static/test_resnet_v2.py         | 14 ++++++++++++++
 5 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/test/dygraph_to_static/dygraph_to_static_util.py b/test/dygraph_to_static/dygraph_to_static_util.py
index 7c19f865a03..f97857df568 100644
--- a/test/dygraph_to_static/dygraph_to_static_util.py
+++ b/test/dygraph_to_static/dygraph_to_static_util.py
@@ -121,6 +121,8 @@ def test_with_new_ir(func):
     @wraps(func)
     def impl(*args, **kwargs):
         ir_outs = None
+        if os.environ.get('FLAGS_use_stride_kernel', False):
+            return
         with static.scope_guard(static.Scope()):
             with static.program_guard(static.Program()):
                 try:
diff --git a/test/dygraph_to_static/test_bert.py b/test/dygraph_to_static/test_bert.py
index 4dda6c984f2..3b70956cb77 100644
--- a/test/dygraph_to_static/test_bert.py
+++ b/test/dygraph_to_static/test_bert.py
@@ -20,7 +20,7 @@ import unittest
 import numpy as np
 from bert_dygraph_model import PretrainModelLayer
 from bert_utils import get_bert_config, get_feed_data_reader
-from dygraph_to_static_util import ast_only_test
+from dygraph_to_static_util import ast_only_test, test_with_new_ir
 from predictor_utils import PredictorTools
 
 import paddle
@@ -263,6 +263,17 @@ class TestBert(unittest.TestCase):
         out = output()
         return out
 
+    @test_with_new_ir
+    def test_train_new_ir(self):
+        static_loss, static_ppl = self.train_static(
+            self.bert_config, self.data_reader
+        )
+        dygraph_loss, dygraph_ppl = self.train_dygraph(
+            self.bert_config, self.data_reader
+        )
+        np.testing.assert_allclose(static_loss, dygraph_loss, rtol=1e-05)
+        np.testing.assert_allclose(static_ppl, dygraph_ppl, rtol=1e-05)
+
     @ast_only_test
     def test_train(self):
         static_loss, static_ppl = self.train_static(
diff --git a/test/dygraph_to_static/test_mobile_net.py b/test/dygraph_to_static/test_mobile_net.py
index 220779f1477..c2fd027ea88 100644
--- a/test/dygraph_to_static/test_mobile_net.py
+++ b/test/dygraph_to_static/test_mobile_net.py
@@ -19,6 +19,7 @@ import time
 import unittest
 
 import numpy as np
+from dygraph_to_static_util import test_with_new_ir
 from predictor_utils import PredictorTools
 
 import paddle
@@ -731,6 +732,13 @@ class TestMobileNet(unittest.TestCase):
             ),
         )
 
+    @test_with_new_ir
+    def test_mobile_net_new_ir(self):
+        # MobileNet-V1
+        self.assert_same_loss("MobileNetV1")
+        # MobileNet-V2
+        self.assert_same_loss("MobileNetV2")
+
     def test_mobile_net(self):
         # MobileNet-V1
         self.assert_same_loss("MobileNetV1")
diff --git a/test/dygraph_to_static/test_resnet.py b/test/dygraph_to_static/test_resnet.py
index 5ed4adb205c..bcdc5008d25 100644
--- a/test/dygraph_to_static/test_resnet.py
+++ b/test/dygraph_to_static/test_resnet.py
@@ -19,6 +19,7 @@ import time
 import unittest
 
 import numpy as np
+from dygraph_to_static_util import test_with_new_ir
 from predictor_utils import PredictorTools
 
 import paddle
@@ -426,6 +427,19 @@ class TestResnet(unittest.TestCase):
             ),
         )
 
+    @test_with_new_ir
+    def test_resnet_new_ir(self):
+        static_loss = self.train(to_static=True)
+        dygraph_loss = self.train(to_static=False)
+        np.testing.assert_allclose(
+            static_loss,
+            dygraph_loss,
+            rtol=1e-05,
+            err_msg='static_loss: {} \n dygraph_loss: {}'.format(
+                static_loss, dygraph_loss
+            ),
+        )
+
     def test_resnet(self):
         static_loss = self.train(to_static=True)
         dygraph_loss = self.train(to_static=False)
diff --git a/test/dygraph_to_static/test_resnet_v2.py b/test/dygraph_to_static/test_resnet_v2.py
index ef7eff71bf1..29e3552b898 100644
--- a/test/dygraph_to_static/test_resnet_v2.py
+++ b/test/dygraph_to_static/test_resnet_v2.py
@@ -19,6 +19,7 @@ import time
 import unittest
 
 import numpy as np
+from dygraph_to_static_util import test_with_new_ir
 from predictor_utils import PredictorTools
 
 import paddle
@@ -431,6 +432,19 @@ class TestResnet(unittest.TestCase):
             ),
         )
 
+    @test_with_new_ir
+    def test_resnet_new_ir(self):
+        static_loss = self.train(to_static=True)
+        dygraph_loss = self.train(to_static=False)
+        np.testing.assert_allclose(
+            static_loss,
+            dygraph_loss,
+            rtol=1e-05,
+            err_msg='static_loss: {} \n dygraph_loss: {}'.format(
+                static_loss, dygraph_loss
+            ),
+        )
+
     def test_resnet(self):
         static_loss = self.train(to_static=True)
         dygraph_loss = self.train(to_static=False)
-- 
GitLab