cinn(test): add rmsnorm subgraph symbolic test (#61317)

PaddlePaddle · Jan 30, 2024 · 3d0601e · 3d0601e
1 parent ea2a78a
commit 3d0601e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/test/ir/pir/cinn/CMakeLists.txt b/test/ir/pir/cinn/CMakeLists.txt
@@ -35,15 +35,37 @@ if(WITH_GPU)
   set_tests_properties(test_subgraph_checker PROPERTIES LABELS "RUN_TYPE=CINN")
 
   add_test(
-    NAME test_rms_norm
+    NAME test_rms_norm_seq_len_symbolic
     COMMAND
       ${CMAKE_COMMAND} -E env
       PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/:$ENV{PYTHONPATH}
       FLAGS_cinn_convert_static_dim_to_dynamic_dim=2048:S0
       FLAGS_cinn_bucket_compile=True FLAGS_enable_pir_api=1 ${PYTHON_EXECUTABLE}
       ${CMAKE_CURRENT_SOURCE_DIR}/test_rms_norm.py
     WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
-  set_tests_properties(test_rms_norm PROPERTIES LABELS "RUN_TYPE=CINN")
+  set_tests_properties(test_rms_norm_seq_len_symbolic
+                       PROPERTIES LABELS "RUN_TYPE=CINN")
+  add_test(
+    NAME test_rms_norm_bs_symbolic
+    COMMAND
+      ${CMAKE_COMMAND} -E env
+      PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/:$ENV{PYTHONPATH}
+      FLAGS_cinn_convert_static_dim_to_dynamic_dim=7:S1
+      FLAGS_cinn_bucket_compile=True FLAGS_enable_pir_api=1 ${PYTHON_EXECUTABLE}
+      ${CMAKE_CURRENT_SOURCE_DIR}/test_rms_norm.py
+    WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
+  set_tests_properties(test_rms_norm_bs_symbolic PROPERTIES LABELS
+                                                            "RUN_TYPE=CINN")
+  add_test(
+    NAME test_rms_norm_symbolic
+    COMMAND
+      ${CMAKE_COMMAND} -E env
+      PYTHONPATH=${CMAKE_BINARY_DIR}:${CMAKE_BINARY_DIR}/python/:$ENV{PYTHONPATH}
+      FLAGS_cinn_convert_static_dim_to_dynamic_dim=2048:S0,7:S1
+      FLAGS_cinn_bucket_compile=True FLAGS_enable_pir_api=1 ${PYTHON_EXECUTABLE}
+      ${CMAKE_CURRENT_SOURCE_DIR}/test_rms_norm.py
+    WORKING_DIRECTORY ${CMAKE_BINARY_DIR})
+  set_tests_properties(test_rms_norm_symbolic PROPERTIES LABELS "RUN_TYPE=CINN")
   add_test(
     NAME test_rope_symbolic
     COMMAND

diff --git a/test/ir/pir/cinn/test_rms_norm.py b/test/ir/pir/cinn/test_rms_norm.py
@@ -41,7 +41,7 @@ def forward(self, hidden_states):
 
 class TestLlamaRMSNorm(TestCinnSubGraphBase):
     def prepare_data(self):
-        self.shape = [1, 2048, 768]
+        self.shape = [7, 2048, 768]
         self.hidden_states = paddle.randn(self.shape, dtype="float32")
         self.hidden_states.stop_gradient = False