FPU: Refine tests for FP registers protection

1. Directly call assembly functions with 'fp_func_jump_template' to
   make sure that FP registers are correctly protected.
2. Remove test case which is used to clear FP registers in service. FP
   registers in service are cleared in psa_close if they are correctly
   restored. It does not make sense to make it a single test case.
3. Extract common functions into fpu_tests_common.c.

Signed-off-by: Xinyu Zhang <xinyu.zhang@arm.com>
Change-Id: I87bf1f27c3e7755d206f28edbc43b88bf08f0b64
diff --git a/test/secure_fw/common_test_services/tfm_secure_client_service/tfm_secure_client_service.yaml b/test/secure_fw/common_test_services/tfm_secure_client_service/tfm_secure_client_service.yaml
index eed3dc3..0dacb05 100644
--- a/test/secure_fw/common_test_services/tfm_secure_client_service/tfm_secure_client_service.yaml
+++ b/test/secure_fw/common_test_services/tfm_secure_client_service/tfm_secure_client_service.yaml
@@ -54,8 +54,7 @@
     "TFM_FIRMWARE_UPDATE_SERVICE",
     "IPC_SERVICE_TEST_BASIC",
     "IPC_SERVICE_TEST_STATELESS_ROT",
-    "TFM_FPU_SERVICE_CLEAR_FP_REGISTER",
-    "TFM_FPU_SERVICE_CHECK_FP_REGISTER",
+    "TFM_FPU_SERVICE_CHECK_FP_CALLEE_REGISTER",
     "IPC_CLIENT_TEST_PSA_ACCESS_APP_MEM",
     "IPC_CLIENT_TEST_PSA_ACCESS_APP_READ_ONLY_MEM",
     "IPC_CLIENT_TEST_APP_ACCESS_PSA_MEM",
diff --git a/test/secure_fw/suites/fpu/fpu_tests_common.c b/test/secure_fw/suites/fpu/fpu_tests_common.c
index 1e22322..3c1821d 100644
--- a/test/secure_fw/suites/fpu/fpu_tests_common.c
+++ b/test/secure_fw/suites/fpu/fpu_tests_common.c
@@ -5,329 +5,103 @@
  *
  */
 
-#include <stdbool.h>
-#include "fpu_tests_common.h"
-#include "psa_manifest/sid.h"
 #include <string.h>
+#include "fpu_tests_common.h"
 
 
-/**
- * Change FP registers.
- */
-__attribute__((naked)) static void change_fp_in_client(void)
+__attribute__((naked))
+void fp_func_jump_template(uintptr_t *func_table, uintptr_t *func_return,
+                           uint32_t func_num)
 {
     __asm volatile(
-        "mov       r0, #0x000000C0         \n"
-        "vmov      s0, r0                  \n"
-        "mov       r0, #0x000000C1         \n"
-        "vmov      s1, r0                  \n"
-        "mov       r0, #0x000000C2         \n"
-        "vmov      s2, r0                  \n"
-        "mov       r0, #0x000000C3         \n"
-        "vmov      s3, r0                  \n"
-        "mov       r0, #0x000000C4         \n"
-        "vmov      s4, r0                  \n"
-        "mov       r0, #0x000000C5         \n"
-        "vmov      s5, r0                  \n"
-        "mov       r0, #0x000000C6         \n"
-        "vmov      s6, r0                  \n"
-        "mov       r0, #0x000000C7         \n"
-        "vmov      s7, r0                  \n"
-        "mov       r0, #0x000000C8         \n"
-        "vmov      s8, r0                  \n"
-        "mov       r0, #0x000000C9         \n"
-        "vmov      s9, r0                  \n"
-        "mov       r0, #0x000000CA         \n"
-        "vmov      s10, r0                 \n"
-        "mov       r0, #0x000000CB         \n"
-        "vmov      s11, r0                 \n"
-        "mov       r0, #0x000000CC         \n"
-        "vmov      s12, r0                 \n"
-        "mov       r0, #0x000000CD         \n"
-        "vmov      s13, r0                 \n"
-        "mov       r0, #0x000000CE         \n"
-        "vmov      s14, r0                 \n"
-        "mov       r0, #0x000000CF         \n"
-        "vmov      s15, r0                 \n"
-        "mov       r0, #0x000000D0         \n"
-        "vmov      s16, r0                 \n"
-        "mov       r0, #0x000000D1         \n"
-        "vmov      s17, r0                 \n"
-        "mov       r0, #0x000000D2         \n"
-        "vmov      s18, r0                 \n"
-        "mov       r0, #0x000000D3         \n"
-        "vmov      s19, r0                 \n"
-        "mov       r0, #0x000000D4         \n"
-        "vmov      s20, r0                 \n"
-        "mov       r0, #0x000000D5         \n"
-        "vmov      s21, r0                 \n"
-        "mov       r0, #0x000000D6         \n"
-        "vmov      s22, r0                 \n"
-        "mov       r0, #0x000000D7         \n"
-        "vmov      s23, r0                 \n"
-        "mov       r0, #0x000000D8         \n"
-        "vmov      s24, r0                 \n"
-        "mov       r0, #0x000000D9         \n"
-        "vmov      s25, r0                 \n"
-        "mov       r0, #0x000000DA         \n"
-        "vmov      s26, r0                 \n"
-        "mov       r0, #0x000000DB         \n"
-        "vmov      s27, r0                 \n"
-        "mov       r0, #0x000000DC         \n"
-        "vmov      s28, r0                 \n"
-        "mov       r0, #0x000000DD         \n"
-        "vmov      s29, r0                 \n"
-        "mov       r0, #0x000000DE         \n"
-        "vmov      s30, r0                 \n"
-        "mov       r0, #0x000000DF         \n"
-        "vmov      s31, r0                 \n"
-
-        "bx        lr                      \n"
+        "cmp       r2, #0           \n"
+        "beq       jump_exit        \n"
+        "push      {r4-r6, lr}      \n"
+        "mov       r4, r0           \n"
+        "mov       r5, r1           \n"
+        "mov       r6, r2           \n"
+    "jump_loop:                     \n"
+        /* Load func address to r2.
+         * Move r4 to the address of func param. */
+        "ldr       r2, [r4], #4     \n"
+        /* Load func param to r0.
+         * Move r4 to the address of next func.*/
+        "ldr       r0, [r4], #4     \n"
+        /* Jump to target func. */
+        "blx       r2               \n"
+        /* Store return value of finished func.
+         * Move r5 to store return value of next func. */
+        "str       r0, [r5], #4     \n"
+        "subs      r6, r6, #1       \n"
+        "bne       jump_loop        \n"
+        "pop       {r4-r6, pc}      \n"
+    "jump_exit:                     \n"
+        "bx        lr               \n"
     );
 }
 
-/**
- * Check whether FP registers are restored correctly.
- * Return:
- *   True - FP registers are restored correctly
- *   False - FP registers are not restored correctly
- */
-static bool check_fp_restored_client(void)
-{
-    uint32_t fp_buffer[NR_FP_REG] = {0};
-    const uint32_t fp_expect[NR_FP_REG] =
-                            {0xC0, 0xC1, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7,
-                             0xC8, 0xC9, 0xCA, 0xCB, 0xCC, 0xCD, 0xCE, 0xCF,
-                             0xD0, 0xD1, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7,
-                             0xD8, 0xD9, 0xDA, 0xDB, 0xDC, 0xDD, 0xDE, 0xDF};
-
-    __asm volatile(
-        "vstm      %0, {S0-S31}            \n"
-        :
-        :"r"(fp_buffer)
-        :"memory"
-    );
-
-    if (!memcmp(fp_buffer, fp_expect, FP_BUF_SIZE)) {
-        return true;
-    }
-
-    return false;
-}
-
-/**
- * Clear FP registers.
- */
-__attribute__((naked)) static void fpu_client_fp_clear_test(void)
+__attribute__((naked))
+void dump_fp_caller(uint32_t *fp_caller_buffer)
 {
     __asm volatile(
-        "eor       r0, r0, r0              \n"
-        "vmov      s0, r0                  \n"
-        "vmov      s1, r0                  \n"
-        "vmov      s2, r0                  \n"
-        "vmov      s3, r0                  \n"
-        "vmov      s4, r0                  \n"
-        "vmov      s5, r0                  \n"
-        "vmov      s6, r0                  \n"
-        "vmov      s7, r0                  \n"
-        "vmov      s8, r0                  \n"
-        "vmov      s9, r0                  \n"
-        "vmov      s10, r0                 \n"
-        "vmov      s11, r0                 \n"
-        "vmov      s12, r0                 \n"
-        "vmov      s13, r0                 \n"
-        "vmov      s14, r0                 \n"
-        "vmov      s15, r0                 \n"
-        "vmov      s16, r0                 \n"
-        "vmov      s17, r0                 \n"
-        "vmov      s18, r0                 \n"
-        "vmov      s19, r0                 \n"
-        "vmov      s20, r0                 \n"
-        "vmov      s21, r0                 \n"
-        "vmov      s22, r0                 \n"
-        "vmov      s23, r0                 \n"
-        "vmov      s24, r0                 \n"
-        "vmov      s25, r0                 \n"
-        "vmov      s26, r0                 \n"
-        "vmov      s27, r0                 \n"
-        "vmov      s28, r0                 \n"
-        "vmov      s29, r0                 \n"
-        "vmov      s30, r0                 \n"
-        "vmov      s31, r0                 \n"
-
-        "bx        lr                      \n"
+        "vstm      r0, {s0-s15}     \n"
+        "bx        lr               \n"
     );
 }
 
-/**
- * Check invalidation of FP registers.
- * Return:
- *   True - FP registers are invalidated
- *   False - FP registers are not invalidated
- */
-static bool check_fp_invalidated(void)
+__attribute__((naked))
+void dump_fp_callee(uint32_t *fp_callee_buffer)
 {
-    uint32_t fp_buffer[NR_FP_REG] = {0};
-    const uint32_t fp_expect[NR_FP_REG] = {0};
-
     __asm volatile(
-        "vstm      %0, {S0-S31}            \n"
-        :
-        :"r"(fp_buffer)
-        :"memory"
+        "vstm      r0, {s16-s31}    \n"
+        "bx        lr               \n"
     );
-
-#if DOMAIN_NS == 1
-    if (!memcmp(fp_buffer, fp_expect, FP_BUF_SIZE))
-#else
-    if (!memcmp(fp_buffer, fp_expect, FP_BUF_SIZE))
-#endif
-    {
-        return true;
-    }
-
-    return false;
 }
 
-/**
- * Description: Clear FP registers to check basic FP register write/read
- * functionality.
- * Expectation: FP registers in FPU client partition should be cleared.
- */
+__attribute__((naked))
+void populate_caller_fp_regs(const uint32_t *expecting_caller_content)
+{
+    __asm volatile(
+        "vldm      r0, {s0-s15}     \n"
+        "bx        lr               \n"
+    );
+}
+
+__attribute__((naked))
+void populate_callee_fp_regs(const uint32_t *expecting_callee_content)
+{
+    __asm volatile(
+        "vldm      r0, {s16-s31}    \n"
+        "bx        lr               \n"
+    );
+}
+
 void tfm_fpu_test_clear_client_fp_data(struct test_result_t *ret)
 {
-    ret->val = TEST_FAILED;
+    uint32_t fp_caller_buffer[NR_FP_CALLER_REG] = {0};
+    uint32_t fp_callee_buffer[NR_FP_CALLEE_REG] = {0};
+    const uint32_t expecting_caller_content[NR_FP_CALLER_REG] = {0};
+    const uint32_t expecting_callee_content[NR_FP_CALLEE_REG] = {0};
 
-    fpu_client_fp_clear_test();
+    uintptr_t func_table[] = {
+        (uintptr_t)populate_caller_fp_regs, (uintptr_t)expecting_caller_content,
+        (uintptr_t)populate_callee_fp_regs, (uintptr_t)expecting_callee_content,
+        (uintptr_t)dump_fp_caller, (uintptr_t)fp_caller_buffer,
+        (uintptr_t)dump_fp_callee, (uintptr_t)fp_callee_buffer
+    };
+    uintptr_t func_return[ARRAY_SIZE(func_table)/2] = {0};
 
-    if (check_fp_invalidated()) {
-        ret->val = TEST_PASSED;
-    } else {
-        ret->val = TEST_FAILED;
-    }
-}
+    ret->val = TEST_PASSED;
 
-/**
- * Description: Test FP context protection after psa calls. Change FP registers
- * in FPU client/service partition separately, then check FP registers after
- * psa calls.
- * Expectation: FP registers in FPU client/service partition should be saved
- * and restored correctly.
- */
-void tfm_fpu_test_fp_protection_psa_call(struct test_result_t *ret)
-{
-    psa_handle_t handle;
-    psa_status_t status;
+    fp_func_jump_template(func_table, func_return, ARRAY_SIZE(func_table)/2);
 
-    ret->val = TEST_FAILED;
-
-    handle = psa_connect(TFM_FPU_SERVICE_CHECK_FP_REGISTER_SID,
-                         TFM_FPU_SERVICE_CHECK_FP_REGISTER_VERSION);
-    if (!PSA_HANDLE_IS_VALID(handle)) {
-        ret->val = TEST_FAILED;
-        return;
+    if (memcmp(fp_caller_buffer, expecting_caller_content,
+               FP_CALLER_BUF_SIZE)) {
+        TEST_FAIL("FP caller registers are not correctly cleared!");
     }
 
-    /* Change FP registers */
-    change_fp_in_client();
-
-    status = psa_call(handle, PSA_IPC_CALL, NULL, 0, NULL, 0);
-    if (status == PSA_SUCCESS) {
-        /* FP registers should be restored */
-        if (check_fp_restored_client()) {
-            ret->val = TEST_PASSED;
-        } else {
-            ret->val = TEST_FAILED;
-        }
-    }
-
-    psa_close(handle);
-}
-
-/**
- * Description: Clear FP registers in FPU service partition for next test.
- * Expectation: FP registers in FPU service partition should be cleared.
- */
-void tfm_fpu_test_clear_service_fp_data(struct test_result_t *ret)
-{
-    psa_handle_t handle;
-    psa_status_t status;
-
-    ret->val = TEST_FAILED;
-
-    handle = psa_connect(TFM_FPU_SERVICE_CLEAR_FP_REGISTER_SID,
-                         TFM_FPU_SERVICE_CLEAR_FP_REGISTER_VERSION);
-    if (!PSA_HANDLE_IS_VALID(handle)) {
-        ret->val = TEST_FAILED;
-        return;
-    }
-
-    status = psa_call(handle, PSA_IPC_CALL, NULL, 0, NULL, 0);
-    if (status == PSA_SUCCESS) {
-        ret->val = TEST_PASSED;
-    } else {
-        ret->val = TEST_FAILED;
-    }
-
-    psa_close(handle);
-}
-
-/**
- * Description: Test reliability of FP context protection after psa calls by
- * loops. Change FP registers in FPU client/service partition separately, then
- * check FP registers after psa calls.
- * Expectation: FP registers in FPU client/service partition should be saved
- * and restored correctly.
- */
-void tfm_fpu_test_fp_protection_psa_call_loop(struct test_result_t *ret)
-{
-    psa_handle_t handle;
-    psa_status_t status;
-    uint32_t itr;
-
-    ret->val = TEST_FAILED;
-
-    for (itr = 0; itr < LOOP_ITERATIONS; itr++) {
-        TEST_LOG("  > Iteration %d of %d\r", itr + 1, LOOP_ITERATIONS);
-
-        handle = psa_connect(TFM_FPU_SERVICE_CHECK_FP_REGISTER_SID,
-                             TFM_FPU_SERVICE_CHECK_FP_REGISTER_VERSION);
-        if (!PSA_HANDLE_IS_VALID(handle)) {
-            ret->val = TEST_FAILED;
-            return;
-        }
-
-        /* Change FP registers */
-        change_fp_in_client();
-
-        status = psa_call(handle, PSA_IPC_CALL, NULL, 0, NULL, 0);
-        if (status == PSA_SUCCESS) {
-            /* FP registers should be restored */
-            if (check_fp_restored_client()) {
-                ret->val = TEST_PASSED;
-            } else {
-                ret->val = TEST_FAILED;
-            }
-        } else {
-            ret->val = TEST_FAILED;
-        }
-
-        psa_close(handle);
-
-        handle = psa_connect(TFM_FPU_SERVICE_CLEAR_FP_REGISTER_SID,
-                             TFM_FPU_SERVICE_CLEAR_FP_REGISTER_VERSION);
-        if (!PSA_HANDLE_IS_VALID(handle)) {
-            ret->val = TEST_FAILED;
-            return;
-        }
-
-        status = psa_call(handle, PSA_IPC_CALL, NULL, 0, NULL, 0);
-        if (status == PSA_SUCCESS) {
-            ret->val = TEST_PASSED;
-        } else {
-            ret->val = TEST_FAILED;
-        }
-
-        psa_close(handle);
+    if (memcmp(fp_callee_buffer, expecting_callee_content,
+               FP_CALLEE_BUF_SIZE)) {
+        TEST_FAIL("FP callee registers are not correctly cleared!");
     }
 }
diff --git a/test/secure_fw/suites/fpu/fpu_tests_common.h b/test/secure_fw/suites/fpu/fpu_tests_common.h
index 794ae90..105678d 100644
--- a/test/secure_fw/suites/fpu/fpu_tests_common.h
+++ b/test/secure_fw/suites/fpu/fpu_tests_common.h
@@ -19,6 +19,9 @@
 #include "psa_manifest/sid.h"
 #endif
 
+#ifndef ARRAY_SIZE
+#define ARRAY_SIZE(arr) (sizeof(arr)/sizeof(arr[0]))
+#endif
 
 #define LOOP_ITERATIONS             (300U)
 #define LOOPS_S_INT_TEST            (50U)
@@ -37,26 +40,44 @@
 #define S_TIMER_TRIGGERED                    (0x19)
 #define S_TIMER_NOT_TRIGGERED                (0x91)
 
-/**
- * Test FP context protection after psa calls.
- */
-void tfm_fpu_test_fp_protection_psa_call(struct test_result_t *ret);
+#define PSA_CALL_PARAM_LEN                   (4U)
 
-/**
- * Clear FP registers in FPU service partition.
- */
-void tfm_fpu_test_clear_service_fp_data(struct test_result_t *ret);
-
-/**
- * Test reliability of FP context protection after psa calls by loops.
- */
-void tfm_fpu_test_fp_protection_psa_call_loop(struct test_result_t *ret);
-
-/**
- * Clear FP registers in FPU client partition.
+/*
+ * Description: Clear FP registers to check basic FP register write/read
+ * functionality. This is a shared test case for both S and NS sides.
+ * Expectation: FP registers in FPU client partition should be cleared.
  */
 void tfm_fpu_test_clear_client_fp_data(struct test_result_t *ret);
 
+/*
+ * Introduction: Run functions in func_table sequentially with assembly
+ * instructions.
+ * Input:
+ *  - func_table:  Table containing function addresses and input parameters.
+ *  - func_return: Array containing function return values.
+ *  - func_num:    Number of functions to be run.
+ * Description: The format of func_table should be a function address followed
+ * by its input parameter. If the function needs more than 1 input parameters,
+ * input parameters need to be put into a struct. Then put the pointer of the
+ * struct into func_table.
+ * The length of func_return should be same as func_num. The return value of
+ * each funtion will be stored in func_return after the function is finished.
+ */
+void fp_func_jump_template(uintptr_t *func_table, uintptr_t *func_return,
+                           uint32_t func_num);
+
+/* Dump FP caller registers to fp_caller_buffer. */
+void dump_fp_caller(uint32_t *fp_caller_buffer);
+
+/* Dump FP callee registers to fp_callee_buffer. */
+void dump_fp_callee(uint32_t *fp_callee_buffer);
+
+/* Change FP caller registers. */
+void populate_caller_fp_regs(const uint32_t *expecting_caller_content);
+
+/* Change FP callee registers. */
+void populate_callee_fp_regs(const uint32_t *expecting_callee_content);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/test/secure_fw/suites/fpu/non_secure/fpu_ns_interface_testsuite.c b/test/secure_fw/suites/fpu/non_secure/fpu_ns_interface_testsuite.c
index 42b6e56..0f9afca 100644
--- a/test/secure_fw/suites/fpu/non_secure/fpu_ns_interface_testsuite.c
+++ b/test/secure_fw/suites/fpu/non_secure/fpu_ns_interface_testsuite.c
@@ -5,19 +5,22 @@
  *
  */
 
-#include <stdio.h>
-#include <stdbool.h>
 #include <string.h>
 #include "fpu_ns_tests.h"
 #include "../fpu_tests_common.h"
 #include "os_wrapper/delay.h"
+#include "tfm_psa_call_pack.h"
 
 static void tfm_fpu_test_fp_protection_secure_interrupt(
                                                 struct test_result_t *ret);
 static void tfm_fpu_test_fp_protection_non_secure_interrupt(
                                                 struct test_result_t *ret);
-static bool check_fp_caller_restored(void);
-static bool check_fp_callee_restored(void);
+/* Test FP context protection after psa calls. */
+static void tfm_fpu_test_fp_protection_ns_psa_call(struct test_result_t *ret);
+
+/* Test reliability of FP context protection after psa calls by loops. */
+static void tfm_fpu_test_fp_protection_ns_psa_call_loop(
+                                                struct test_result_t *ret);
 
 static struct test_t fpu_ns_tests[] = {
     {
@@ -25,8 +28,8 @@
         "Clear FP registers in FPU client partition"
     },
     {
-        &tfm_fpu_test_fp_protection_psa_call, "TFM_NS_FPU_TEST_1002",
-        "Test FP context protection after psa calls"
+        &tfm_fpu_test_fp_protection_ns_psa_call_loop, "TFM_NS_FPU_TEST_1002",
+        "Test reliability of FP context protection after psa calls"
     },
     {
         &tfm_fpu_test_fp_protection_secure_interrupt, "TFM_NS_FPU_TEST_1003",
@@ -50,141 +53,6 @@
 }
 
 /**
- * Change FP registers.
- */
-__attribute__((naked)) void change_fp_in_ns_thread(void)
-{
-    __asm volatile(
-        "mov       r0, #0xC0000000         \n"
-        "vmov      s0, r0                  \n"
-        "mov       r0, #0xC1000000         \n"
-        "vmov      s1, r0                  \n"
-        "mov       r0, #0xC2000000         \n"
-        "vmov      s2, r0                  \n"
-        "mov       r0, #0xC3000000         \n"
-        "vmov      s3, r0                  \n"
-        "mov       r0, #0xC4000000         \n"
-        "vmov      s4, r0                  \n"
-        "mov       r0, #0xC5000000         \n"
-        "vmov      s5, r0                  \n"
-        "mov       r0, #0xC6000000         \n"
-        "vmov      s6, r0                  \n"
-        "mov       r0, #0xC7000000         \n"
-        "vmov      s7, r0                  \n"
-        "mov       r0, #0xC8000000         \n"
-        "vmov      s8, r0                  \n"
-        "mov       r0, #0xC9000000         \n"
-        "vmov      s9, r0                  \n"
-        "mov       r0, #0xCA000000         \n"
-        "vmov      s10, r0                 \n"
-        "mov       r0, #0xCB000000         \n"
-        "vmov      s11, r0                 \n"
-        "mov       r0, #0xCC000000         \n"
-        "vmov      s12, r0                 \n"
-        "mov       r0, #0xCD000000         \n"
-        "vmov      s13, r0                 \n"
-        "mov       r0, #0xCE000000         \n"
-        "vmov      s14, r0                 \n"
-        "mov       r0, #0xCF000000         \n"
-        "vmov      s15, r0                 \n"
-        "mov       r0, #0xD0000000         \n"
-        "vmov      s16, r0                 \n"
-        "mov       r0, #0xD1000000         \n"
-        "vmov      s17, r0                 \n"
-        "mov       r0, #0xD2000000         \n"
-        "vmov      s18, r0                 \n"
-        "mov       r0, #0xD3000000         \n"
-        "vmov      s19, r0                 \n"
-        "mov       r0, #0xD4000000         \n"
-        "vmov      s20, r0                 \n"
-        "mov       r0, #0xD5000000         \n"
-        "vmov      s21, r0                 \n"
-        "mov       r0, #0xD6000000         \n"
-        "vmov      s22, r0                 \n"
-        "mov       r0, #0xD7000000         \n"
-        "vmov      s23, r0                 \n"
-        "mov       r0, #0xD8000000         \n"
-        "vmov      s24, r0                 \n"
-        "mov       r0, #0xD9000000         \n"
-        "vmov      s25, r0                 \n"
-        "mov       r0, #0xDA000000         \n"
-        "vmov      s26, r0                 \n"
-        "mov       r0, #0xDB000000         \n"
-        "vmov      s27, r0                 \n"
-        "mov       r0, #0xDC000000         \n"
-        "vmov      s28, r0                 \n"
-        "mov       r0, #0xDD000000         \n"
-        "vmov      s29, r0                 \n"
-        "mov       r0, #0xDE000000         \n"
-        "vmov      s30, r0                 \n"
-        "mov       r0, #0xDF000000         \n"
-        "vmov      s31, r0                 \n"
-
-        "bx        lr                      \n"
-    );
-}
-
-/**
- * Check whether FP caller registers are restored correctly.
- * Return:
- *   True - FP caller registers are restored correctly
- *   False - FP caller registers are not restored correctly
- */
-static bool check_fp_caller_restored(void)
-{
-    static uint32_t fp_buffer[NR_FP_CALLER_REG] = {0};
-    const uint32_t fp_expect[NR_FP_CALLER_REG] = {
-        0xC0000000, 0xC1000000, 0xC2000000, 0xC3000000,
-        0xC4000000, 0xC5000000, 0xC6000000, 0xC7000000,
-        0xC8000000, 0xC9000000, 0xCA000000, 0xCB000000,
-        0xCC000000, 0xCD000000, 0xCE000000, 0xCF000000
-    };
-
-    __asm volatile(
-        "vstm      %0, {S0-S15}            \n"
-        :
-        :"r"(fp_buffer)
-        :"memory"
-    );
-
-    if (!memcmp(fp_buffer, fp_expect, FP_CALLER_BUF_SIZE)) {
-        return true;
-    }
-
-    return false;
-}
-
-/**
- * Check whether FP callee registers are restored correctly.
- * Return:
- *   True - FP callee registers are restored correctly
- *   False - FP callee registers are not restored correctly
- */
-static bool check_fp_callee_restored(void)
-{
-    static uint32_t fp_buffer[NR_FP_CALLEE_REG] = {0};
-    const uint32_t fp_expect[NR_FP_CALLEE_REG] = {
-        0xD0000000, 0xD1000000, 0xD2000000, 0xD3000000,
-        0xD4000000, 0xD5000000, 0xD6000000, 0xD7000000,
-        0xD8000000, 0xD9000000, 0xDA000000, 0xDB000000,
-        0xDC000000, 0xDD000000, 0xDE000000, 0xDF000000
-    };
-
-    __asm volatile(
-        "vstm      %0, {S16-S31}            \n"
-        :
-        :"r"(fp_buffer)
-        :"memory"
-    );
-
-    if (!memcmp(fp_buffer, fp_expect, FP_CALLEE_BUF_SIZE)) {
-        return true;
-    }
-
-    return false;
-}
-
-/**
  * \brief Test FP context protection in S interrupt. Change FP registers in
  * non-secure thread first, then change them in S interrupt. After interrupt
  * return, check FP context protection in non-secure thread.
@@ -200,10 +68,26 @@
     struct psa_outvec outvecs[1] = {{outvec_data, sizeof(outvec_data[0])}};
     static uint8_t i;
 
+    uint32_t fp_caller_buffer[NR_FP_CALLER_REG] = {0};
+    uint32_t fp_callee_buffer[NR_FP_CALLEE_REG] = {0};
+    const uint32_t expecting_caller_content[NR_FP_CALLER_REG] = {
+        0xC0000000, 0xC1000000, 0xC2000000, 0xC3000000,
+        0xC4000000, 0xC5000000, 0xC6000000, 0xC7000000,
+        0xC8000000, 0xC9000000, 0xCA000000, 0xCB000000,
+        0xCC000000, 0xCD000000, 0xCE000000, 0xCF000000
+    };
+    const uint32_t expecting_callee_content[NR_FP_CALLEE_REG] = {
+        0xD0000000, 0xD1000000, 0xD2000000, 0xD3000000,
+        0xD4000000, 0xD5000000, 0xD6000000, 0xD7000000,
+        0xD8000000, 0xD9000000, 0xDA000000, 0xDB000000,
+        0xDC000000, 0xDD000000, 0xDE000000, 0xDF000000
+    };
+
     ret->val = TEST_FAILED;
 
     /* Change FP regs */
-    change_fp_in_ns_thread();
+    populate_caller_fp_regs(expecting_caller_content);
+    populate_callee_fp_regs(expecting_callee_content);
 
     /* Start the timer */
     handle = psa_connect(TFM_FPU_SERVICE_START_S_TIMER_SID,
@@ -248,14 +132,17 @@
 
     psa_close(handle);
 
-    /* FP caller registers should be restored correctly after S interrupt */
-    if (!check_fp_caller_restored()) {
-        return;
-    }
+    /* FP registers should be restored correctly after S interrupt */
+    dump_fp_caller(fp_caller_buffer);
+    dump_fp_callee(fp_callee_buffer);
 
-    /* FP callee registers should be restored correctly */
-    if (check_fp_callee_restored()) {
+    if ((!memcmp(fp_caller_buffer, expecting_caller_content,
+                 FP_CALLER_BUF_SIZE)) &&
+        (!memcmp(fp_callee_buffer, expecting_callee_content,
+                 FP_CALLEE_BUF_SIZE))) {
         ret->val = TEST_PASSED;
+    } else {
+        ret->val = TEST_FAILED;
     }
 }
 
@@ -288,3 +175,96 @@
 
     psa_close(handle);
 }
+
+/*
+ * Description: Directly jump to tfm_psa_call_veneer from NS side.
+ * Expectation: Return psa_status_t status code.
+ */
+__attribute__((naked)) static uint32_t fpu_ns_call(uintptr_t psa_call_param)
+{
+        __asm volatile(
+#if !defined(__ICCARM__)
+        ".syntax unified                   \n"
+#endif
+        "   push    {r12, lr}              \n"
+        "   mov     r12, r0                \n"
+        /* Load params of tfm_psa_call_veneer into r0~r4. */
+        "   ldr     r0, [r12], #4          \n"  /* psa handle */
+        "   ldr     r1, [r12], #4          \n"  /* ctrl_param */
+        "   ldr     r2, [r12], #4          \n"  /* in_vec */
+        "   ldr     r3, [r12]              \n"  /* out_vec */
+        "   blx     tfm_psa_call_veneer    \n"
+        "   pop     {r12, pc}              \n"
+    );
+}
+
+/*
+ * Description: Test FP context protection after psa calls. Change FP registers
+ * in FPU client/service partition separately, then check FP registers after
+ * psa calls.
+ * Expectation: FP registers in FPU client/service partition should be saved
+ * and restored correctly.
+ */
+static void tfm_fpu_test_fp_protection_ns_psa_call(struct test_result_t *ret)
+{
+    psa_handle_t handle;
+    uint32_t param;
+    uint32_t psa_call_param[PSA_CALL_PARAM_LEN] = {0};
+    uint32_t fp_callee_buffer[NR_FP_CALLEE_REG] = {0};
+    const uint32_t expecting_callee_content[NR_FP_CALLEE_REG] = {
+        0xD0, 0xD1, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7,
+        0xD8, 0xD9, 0xDA, 0xDB, 0xDC, 0xDD, 0xDE, 0xDF
+    };
+
+    uintptr_t func_table[] = {
+        (uintptr_t)populate_callee_fp_regs, (uintptr_t)expecting_callee_content,
+        (uintptr_t)fpu_ns_call, (uintptr_t)psa_call_param,
+        (uintptr_t)dump_fp_callee, (uintptr_t)fp_callee_buffer
+    };
+    uintptr_t func_return[ARRAY_SIZE(func_table)/2] = {0};
+
+    ret->val = TEST_PASSED;
+
+    handle = psa_connect(TFM_FPU_SERVICE_CHECK_FP_CALLEE_REGISTER_SID,
+                         TFM_FPU_SERVICE_CHECK_FP_CALLEE_REGISTER_VERSION);
+    if (!PSA_HANDLE_IS_VALID(handle)) {
+        TEST_FAIL("PSA Connect fail!");
+        return;
+    }
+
+    param = PARAM_PACK(PSA_IPC_CALL, 0, 0);
+    psa_call_param[0] = (uint32_t)handle;
+    psa_call_param[1] = param;
+
+    fp_func_jump_template(func_table, func_return, ARRAY_SIZE(func_table)/2);
+
+    if ((psa_status_t)func_return[1] != PSA_SUCCESS) {
+        TEST_FAIL("FP callee registers check fail in service!");
+    }
+
+    if (memcmp(fp_callee_buffer, expecting_callee_content,
+               FP_CALLEE_BUF_SIZE)) {
+        TEST_FAIL("FP callee registers are not correctly restored in client!");
+    }
+
+    psa_close(handle);
+}
+
+/*
+ * Description: Test reliability of FP context protection after psa calls by
+ * loops. Change FP callee registers in FPU client/service partition separately,
+ * then check FP callee registers after psa calls.
+ * Expectation: FP callee registers in FPU client/service partition should be
+ * saved and restored correctly.
+ */
+static void tfm_fpu_test_fp_protection_ns_psa_call_loop(
+                                                struct test_result_t *ret)
+{
+    uint32_t itr;
+
+    for (itr = 0; itr < LOOP_ITERATIONS; itr++) {
+        TEST_LOG("  > Iteration %d of %d\r", itr + 1, LOOP_ITERATIONS);
+
+        tfm_fpu_test_fp_protection_ns_psa_call(ret);
+    }
+}
diff --git a/test/secure_fw/suites/fpu/secure/fpu_s_interface_testsuite.c b/test/secure_fw/suites/fpu/secure/fpu_s_interface_testsuite.c
index 35a9190..748d450 100755
--- a/test/secure_fw/suites/fpu/secure/fpu_s_interface_testsuite.c
+++ b/test/secure_fw/suites/fpu/secure/fpu_s_interface_testsuite.c
@@ -7,6 +7,15 @@
 
 #include "fpu_s_tests.h"
 #include "../fpu_tests_common.h"
+#include "tfm_psa_call_pack.h"
+#include "tfm_secure_api.h"
+
+/* Test FP context protection after psa calls. */
+static void tfm_fpu_test_fp_protection_s_psa_call(struct test_result_t *ret);
+
+/* Test reliability of FP context protection after psa calls by loops. */
+static void tfm_fpu_test_fp_protection_s_psa_call_loop(
+                                                    struct test_result_t *ret);
 
 static struct test_t fpu_s_tests[] = {
     {
@@ -14,15 +23,7 @@
         "Clear FP registers in FPU client partition"
     },
     {
-        &tfm_fpu_test_fp_protection_psa_call, "TFM_S_FPU_TEST_1002",
-        "Test FP context protection after psa calls"
-    },
-    {
-        &tfm_fpu_test_clear_service_fp_data, "TFM_S_FPU_TEST_1003",
-        "Clear FP registers in FPU service partition for next test"
-    },
-    {
-        &tfm_fpu_test_fp_protection_psa_call_loop, "TFM_S_FPU_TEST_1004",
+        &tfm_fpu_test_fp_protection_s_psa_call_loop, "TFM_S_FPU_TEST_1002",
         "Test reliability of FP context protection after psa calls"
     }
 };
@@ -36,3 +37,96 @@
     set_testsuite("FPU secure interface test (TFM_S_FPU_TEST_1XXX)",
                   fpu_s_tests, list_size, p_test_suite);
 }
+
+/*
+ * Description: Directly jump to tfm_psa_call_pack from S side.
+ * Expectation: Return psa_status_t status code.
+ */
+__attribute__((naked)) static uint32_t fpu_s_call(uintptr_t psa_call_param)
+{
+        __asm volatile(
+#if !defined(__ICCARM__)
+        ".syntax unified                        \n"
+#endif
+        "   push    {r12, lr}                   \n"
+        "   mov     r12, r0                     \n"
+        /* Load params of tfm_psa_call_pack into r0~r4. */
+        "   ldr     r0, [r12], #4               \n"  /* psa handle */
+        "   ldr     r1, [r12], #4               \n"  /* ctrl_param */
+        "   ldr     r2, [r12], #4               \n"  /* in_vec */
+        "   ldr     r3, [r12]                   \n"  /* out_vec */
+        "   blx     "M2S(tfm_psa_call_pack)"    \n"
+        "   pop     {r12, pc}                   \n"
+    );
+}
+
+/*
+ * Description: Test FP context protection after psa calls. Change FP callee
+ * registers in FPU client/service partition separately, then check FP callee
+ * registers after psa calls.
+ * Expectation: FP callee registers in FPU client/service partition should be
+ * saved and restored correctly.
+ */
+static void tfm_fpu_test_fp_protection_s_psa_call(struct test_result_t *ret)
+{
+    psa_handle_t handle;
+    uint32_t param;
+    uint32_t psa_call_param[PSA_CALL_PARAM_LEN] = {0};
+    uint32_t fp_callee_buffer[NR_FP_CALLEE_REG] = {0};
+    const uint32_t expecting_callee_content[NR_FP_CALLEE_REG] = {
+        0xD0, 0xD1, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7,
+        0xD8, 0xD9, 0xDA, 0xDB, 0xDC, 0xDD, 0xDE, 0xDF
+    };
+
+    uintptr_t func_table[] = {
+        (uintptr_t)populate_callee_fp_regs, (uintptr_t)expecting_callee_content,
+        (uintptr_t)fpu_s_call, (uintptr_t)psa_call_param,
+        (uintptr_t)dump_fp_callee, (uintptr_t)fp_callee_buffer
+    };
+    uintptr_t func_return[ARRAY_SIZE(func_table)/2] = {0};
+
+    ret->val = TEST_PASSED;
+
+    handle = psa_connect(TFM_FPU_SERVICE_CHECK_FP_CALLEE_REGISTER_SID,
+                         TFM_FPU_SERVICE_CHECK_FP_CALLEE_REGISTER_VERSION);
+    if (!PSA_HANDLE_IS_VALID(handle)) {
+        TEST_FAIL("PSA Connect fail!");
+        return;
+    }
+
+    param = PARAM_PACK(PSA_IPC_CALL, 0, 0);
+    psa_call_param[0] = (uint32_t)handle;
+    psa_call_param[1] = param;
+
+    fp_func_jump_template(func_table, func_return, ARRAY_SIZE(func_table)/2);
+
+    if ((psa_status_t)func_return[1] != PSA_SUCCESS) {
+        TEST_FAIL("FP callee registers check fail in service!");
+    }
+
+    if (memcmp(fp_callee_buffer, expecting_callee_content,
+               FP_CALLEE_BUF_SIZE)) {
+        TEST_FAIL("FP callee registers are not correctly restored in client!");
+    }
+
+    psa_close(handle);
+}
+
+/*
+ * Description: Test reliability of FP context protection after psa calls by
+ * loops. Change FP callee registers in FPU client/service partition separately,
+ * then check FP callee registers after psa calls.
+ * Expectation: FP callee registers in FPU client/service partition should be
+ * saved and restored correctly.
+ */
+static void tfm_fpu_test_fp_protection_s_psa_call_loop(
+                                                    struct test_result_t *ret)
+{
+    uint32_t itr;
+
+    for (itr = 0; itr < LOOP_ITERATIONS; itr++) {
+        TEST_LOG("  > Iteration %d of %d\r", itr + 1, LOOP_ITERATIONS);
+
+        tfm_fpu_test_fp_protection_s_psa_call(ret);
+    }
+}
diff --git a/test/secure_fw/suites/fpu/service/tfm_fpu_service_test.c b/test/secure_fw/suites/fpu/service/tfm_fpu_service_test.c
index 60fa58d..2f17a36 100644
--- a/test/secure_fw/suites/fpu/service/tfm_fpu_service_test.c
+++ b/test/secure_fw/suites/fpu/service/tfm_fpu_service_test.c
@@ -5,8 +5,6 @@
  *
  */
 
-#include <stdbool.h>
-#include <string.h>
 #include "psa/client.h"
 #include "psa/service.h"
 #include "psa_manifest/tfm_fpu_service_test.h"
@@ -18,348 +16,41 @@
 #include "device_definition.h"
 #include "fpu_tests_common.h"
 
-/* Define the return status */
-#define FPU_SP_TEST_SUCCESS     (0)
-#define FPU_SP_TEST_FAILED      (-1)
-
-/**
- * Clear FP registers.
- */
-__attribute__((naked)) static void clear_fp_regs(void)
-{
-    __asm volatile(
-        "eor       r0, r0, r0              \n"
-        "vmov      s0, r0                  \n"
-        "vmov      s1, r0                  \n"
-        "vmov      s2, r0                  \n"
-        "vmov      s3, r0                  \n"
-        "vmov      s4, r0                  \n"
-        "vmov      s5, r0                  \n"
-        "vmov      s6, r0                  \n"
-        "vmov      s7, r0                  \n"
-        "vmov      s8, r0                  \n"
-        "vmov      s9, r0                  \n"
-        "vmov      s10, r0                 \n"
-        "vmov      s11, r0                 \n"
-        "vmov      s12, r0                 \n"
-        "vmov      s13, r0                 \n"
-        "vmov      s14, r0                 \n"
-        "vmov      s15, r0                 \n"
-        "vmov      s16, r0                 \n"
-        "vmov      s17, r0                 \n"
-        "vmov      s18, r0                 \n"
-        "vmov      s19, r0                 \n"
-        "vmov      s20, r0                 \n"
-        "vmov      s21, r0                 \n"
-        "vmov      s22, r0                 \n"
-        "vmov      s23, r0                 \n"
-        "vmov      s24, r0                 \n"
-        "vmov      s25, r0                 \n"
-        "vmov      s26, r0                 \n"
-        "vmov      s27, r0                 \n"
-        "vmov      s28, r0                 \n"
-        "vmov      s29, r0                 \n"
-        "vmov      s30, r0                 \n"
-        "vmov      s31, r0                 \n"
-        "bx        lr                      \n"
-    );
-}
-
-/**
- * Check whether FP registers are restored correctly.
- * Return:
- *   True - FP registers are restored correctly
- *   False - FP registers are not restored correctly
- */
-bool check_fp_restored_service(void)
-{
-    uint32_t fp_buffer[NR_FP_REG] = {0};
-    const uint32_t fp_expect[NR_FP_REG] = {
-        0xE0, 0xE1, 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7,
-        0xE8, 0xE9, 0xEA, 0xEB, 0xEC, 0xED, 0xEE, 0xEF,
-        0xF0, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7,
-        0xF8, 0xF9, 0xFA, 0xFB, 0xFC, 0xFD, 0xFE, 0xFF
-     };
-
-    /* Dump FP data from FP registers to buffer */
-    __asm volatile(
-        "vstm      %0, {S0-S31}            \n"
-        :
-        :"r"(fp_buffer)
-        :"memory"
-    );
-
-    if (!memcmp(fp_buffer, fp_expect, FP_BUF_SIZE)) {
-        return true;
-    }
-
-    return false;
-}
-
-/**
- * Change FP registers in FP service partition.
- */
-__attribute__((naked)) static void change_fp_in_service(void)
-{
-    __asm volatile(
-        "mov       r0, #0x000000E0         \n"
-        "vmov      s0, r0                  \n"
-        "mov       r0, #0x000000E1         \n"
-        "vmov      s1, r0                  \n"
-        "mov       r0, #0x000000E2         \n"
-        "vmov      s2, r0                  \n"
-        "mov       r0, #0x000000E3         \n"
-        "vmov      s3, r0                  \n"
-        "mov       r0, #0x000000E4         \n"
-        "vmov      s4, r0                  \n"
-        "mov       r0, #0x000000E5         \n"
-        "vmov      s5, r0                  \n"
-        "mov       r0, #0x000000E6         \n"
-        "vmov      s6, r0                  \n"
-        "mov       r0, #0x000000E7         \n"
-        "vmov      s7, r0                  \n"
-        "mov       r0, #0x000000E8         \n"
-        "vmov      s8, r0                  \n"
-        "mov       r0, #0x000000E9         \n"
-        "vmov      s9, r0                  \n"
-        "mov       r0, #0x000000EA         \n"
-        "vmov      s10, r0                 \n"
-        "mov       r0, #0x000000EB         \n"
-        "vmov      s11, r0                 \n"
-        "mov       r0, #0x000000EC         \n"
-        "vmov      s12, r0                 \n"
-        "mov       r0, #0x000000ED         \n"
-        "vmov      s13, r0                 \n"
-        "mov       r0, #0x000000EE         \n"
-        "vmov      s14, r0                 \n"
-        "mov       r0, #0x000000EF         \n"
-        "vmov      s15, r0                 \n"
-        "mov       r0, #0x000000F0         \n"
-        "vmov      s16, r0                 \n"
-        "mov       r0, #0x000000F1         \n"
-        "vmov      s17, r0                 \n"
-        "mov       r0, #0x000000F2         \n"
-        "vmov      s18, r0                 \n"
-        "mov       r0, #0x000000F3         \n"
-        "vmov      s19, r0                 \n"
-        "mov       r0, #0x000000F4         \n"
-        "vmov      s20, r0                 \n"
-        "mov       r0, #0x000000F5         \n"
-        "vmov      s21, r0                 \n"
-        "mov       r0, #0x000000F6         \n"
-        "vmov      s22, r0                 \n"
-        "mov       r0, #0x000000F7         \n"
-        "vmov      s23, r0                 \n"
-        "mov       r0, #0x000000F8         \n"
-        "vmov      s24, r0                 \n"
-        "mov       r0, #0x000000F9         \n"
-        "vmov      s25, r0                 \n"
-        "mov       r0, #0x000000FA         \n"
-        "vmov      s26, r0                 \n"
-        "mov       r0, #0x000000FB         \n"
-        "vmov      s27, r0                 \n"
-        "mov       r0, #0x000000FC         \n"
-        "vmov      s28, r0                 \n"
-        "mov       r0, #0x000000FD         \n"
-        "vmov      s29, r0                 \n"
-        "mov       r0, #0x000000FE         \n"
-        "vmov      s30, r0                 \n"
-        "mov       r0, #0x000000FF         \n"
-        "vmov      s31, r0                 \n"
-
-        "bx        lr                      \n"
-    );
-}
-
-/**
- * Check whether FP registers is invalidated.
- */
-bool is_fp_cleaned(void)
-{
-    uint32_t fp_buffer[NR_FP_REG] = {0};
-    uint32_t i;
-
-    /* Dump FP data from FP registers to buffer */
-    __asm volatile(
-        "vstm      %0, {S0-S31}            \n"
-        :
-        :"r"(fp_buffer)
-        :"memory"
-    );
-
-    for (i = 0; i < NR_FP_REG; i++) {
-        if (fp_buffer[i] != 0) {
-            return false;
-        }
-    }
-
-    return true;
-}
-
-/**
- * Change FP registers in secure thread.
- */
-__attribute__((naked)) void change_fp_in_s_thread(void)
-{
-    __asm volatile(
-        "push      {r4, lr}                    \n"
-
-        "mov       r0, #0xB0000000         \n"
-        "vmov      s0, r0                  \n"
-        "mov       r0, #0xB1000000         \n"
-        "vmov      s1, r0                  \n"
-        "mov       r0, #0xB2000000         \n"
-        "vmov      s2, r0                  \n"
-        "mov       r0, #0xB3000000         \n"
-        "vmov      s3, r0                  \n"
-        "mov       r0, #0xB4000000         \n"
-        "vmov      s4, r0                  \n"
-        "mov       r0, #0xB5000000         \n"
-        "vmov      s5, r0                  \n"
-        "mov       r0, #0xB6000000         \n"
-        "vmov      s6, r0                  \n"
-        "mov       r0, #0xB7000000         \n"
-        "vmov      s7, r0                  \n"
-        "mov       r0, #0xB8000000         \n"
-        "vmov      s8, r0                  \n"
-        "mov       r0, #0xB9000000         \n"
-        "vmov      s9, r0                  \n"
-        "mov       r0, #0xBA000000         \n"
-        "vmov      s10, r0                 \n"
-        "mov       r0, #0xBB000000         \n"
-        "vmov      s11, r0                 \n"
-        "mov       r0, #0xBC000000         \n"
-        "vmov      s12, r0                 \n"
-        "mov       r0, #0xBD000000         \n"
-        "vmov      s13, r0                 \n"
-        "mov       r0, #0xBE000000         \n"
-        "vmov      s14, r0                 \n"
-        "mov       r0, #0xBF000000         \n"
-        "vmov      s15, r0                 \n"
-        "mov       r0, #0xC0000000         \n"
-        "vmov      s16, r0                 \n"
-        "mov       r0, #0xC1000000         \n"
-        "vmov      s17, r0                 \n"
-        "mov       r0, #0xC2000000         \n"
-        "vmov      s18, r0                 \n"
-        "mov       r0, #0xC3000000         \n"
-        "vmov      s19, r0                 \n"
-        "mov       r0, #0xC4000000         \n"
-        "vmov      s20, r0                 \n"
-        "mov       r0, #0xC5000000         \n"
-        "vmov      s21, r0                 \n"
-        "mov       r0, #0xC6000000         \n"
-        "vmov      s22, r0                 \n"
-        "mov       r0, #0xC7000000         \n"
-        "vmov      s23, r0                 \n"
-        "mov       r0, #0xC8000000         \n"
-        "vmov      s24, r0                 \n"
-        "mov       r0, #0xC9000000         \n"
-        "vmov      s25, r0                 \n"
-        "mov       r0, #0xCA000000         \n"
-        "vmov      s26, r0                 \n"
-        "mov       r0, #0xCB000000         \n"
-        "vmov      s27, r0                 \n"
-        "mov       r0, #0xCC000000         \n"
-        "vmov      s28, r0                 \n"
-        "mov       r0, #0xCD000000         \n"
-        "vmov      s29, r0                 \n"
-        "mov       r0, #0xCE000000         \n"
-        "vmov      s30, r0                 \n"
-        "mov       r0, #0xCF000000         \n"
-        "vmov      s31, r0                 \n"
-
-        "pop       {r4, pc}                \n"
-    );
-}
-
-/**
- * Check whether FP registers are restored correctly.
- * Return:
- *   True - FP registers are restored correctly
- *   False - FP registers are not restored correctly
- */
-bool check_fp_restored_s(void)
-{
-    uint32_t fp_buffer[NR_FP_REG] = {0};
-    const uint32_t fp_expect[NR_FP_REG] = {
-        0xB0000000, 0xB1000000, 0xB2000000, 0xB3000000,
-        0xB4000000, 0xB5000000, 0xB6000000, 0xB7000000,
-        0xB8000000, 0xB9000000, 0xBA000000, 0xBB000000,
-        0xBC000000, 0xBD000000, 0xBE000000, 0xBF000000,
-        0xC0000000, 0xC1000000, 0xC2000000, 0xC3000000,
-        0xC4000000, 0xC5000000, 0xC6000000, 0xC7000000,
-        0xC8000000, 0xC9000000, 0xCA000000, 0xCB000000,
-        0xCC000000, 0xCD000000, 0xCE000000, 0xCF000000
-    };
-
-    /* Dump FP data from FP registers to buffer */
-    __asm volatile(
-        "vstm      %0, {S0-S31}            \n"
-        :
-        :"r"(fp_buffer)
-        :"memory"
-    );
-
-    if (!memcmp(fp_buffer, fp_expect, FP_BUF_SIZE)) {
-        return true;
-    }
-
-    return false;
-}
-
-/**
- * Service handler for clear FP register.
- */
-static void fpu_service_clear_fp_register(void)
-{
-    psa_msg_t msg;
-    psa_status_t r;
-
-    psa_get(TFM_FPU_SERVICE_CLEAR_FP_REGISTER_SIGNAL, &msg);
-    switch (msg.type) {
-    case PSA_IPC_CONNECT:
-        psa_reply(msg.handle, PSA_SUCCESS);
-        break;
-    case PSA_IPC_CALL:
-        if (check_fp_restored_service()) {
-            clear_fp_regs();
-            r = PSA_SUCCESS;
-        } else {
-            r = PSA_ERROR_INVALID_ARGUMENT;
-        }
-        psa_reply(msg.handle, r);
-        break;
-    case PSA_IPC_DISCONNECT:
-        psa_reply(msg.handle, PSA_SUCCESS);
-        break;
-    default:
-        psa_panic();
-        break;
-    }
-}
-
-/**
- * Service handler for checking FP register.
+/*
+ * Description: Service handler for checking FP register.
+ * Expectation: FP callee registers should be restored correctly in service.
  */
 static void fpu_service_check_fp_register(void)
 {
     psa_msg_t msg;
+    uint32_t fp_callee_buffer[NR_FP_CALLEE_REG] = {0};
+    const uint32_t fp_clear_callee_content[NR_FP_CALLEE_REG] = {0};
+    const uint32_t expecting_callee_content[NR_FP_CALLEE_REG] = {
+        0xF0, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7,
+        0xF8, 0xF9, 0xFA, 0xFB, 0xFC, 0xFD, 0xFE, 0xFF
+    };
 
-    psa_get(TFM_FPU_SERVICE_CHECK_FP_REGISTER_SIGNAL, &msg);
+    psa_get(TFM_FPU_SERVICE_CHECK_FP_CALLEE_REGISTER_SIGNAL, &msg);
     switch (msg.type) {
     case PSA_IPC_CONNECT:
         psa_reply(msg.handle, PSA_SUCCESS);
         break;
     case PSA_IPC_CALL:
-        if (is_fp_cleaned()) {
-            change_fp_in_service();
+        dump_fp_callee(fp_callee_buffer);
+        if (!memcmp(fp_callee_buffer, fp_clear_callee_content,
+                    FP_CALLEE_BUF_SIZE)) {
+            populate_callee_fp_regs(expecting_callee_content);
             psa_reply(msg.handle, PSA_SUCCESS);
         } else {
             psa_reply(msg.handle, PSA_ERROR_GENERIC_ERROR);
         }
         break;
     case PSA_IPC_DISCONNECT:
+        dump_fp_callee(fp_callee_buffer);
+        if ((!memcmp(fp_callee_buffer, expecting_callee_content,
+                     FP_CALLEE_BUF_SIZE))) {
+            populate_callee_fp_regs(fp_clear_callee_content);
+        }
         psa_reply(msg.handle, PSA_SUCCESS);
         break;
     default:
@@ -443,6 +134,21 @@
     psa_status_t r;
     static uint32_t i;
 
+    uint32_t fp_caller_buffer[NR_FP_CALLER_REG] = {0};
+    const uint32_t expecting_caller_content[NR_FP_CALLER_REG] = {
+        0xB0000000, 0xB1000000, 0xB2000000, 0xB3000000,
+        0xB4000000, 0xB5000000, 0xB6000000, 0xB7000000,
+        0xB8000000, 0xB9000000, 0xBA000000, 0xBB000000,
+        0xBC000000, 0xBD000000, 0xBE000000, 0xBF000000
+    };
+    uint32_t fp_callee_buffer[NR_FP_CALLEE_REG] = {0};
+    const uint32_t expecting_callee_content[NR_FP_CALLEE_REG] = {
+        0xC0000000, 0xC1000000, 0xC2000000, 0xC3000000,
+        0xC4000000, 0xC5000000, 0xC6000000, 0xC7000000,
+        0xC8000000, 0xC9000000, 0xCA000000, 0xCB000000,
+        0xCC000000, 0xCD000000, 0xCE000000, 0xCF000000
+    };
+
     r = psa_get(TFM_FPU_SERVICE_CHECK_NS_INTERRUPT_S_TEST_SIGNAL, &msg);
     switch (msg.type) {
     case PSA_IPC_CONNECT:
@@ -451,7 +157,8 @@
         break;
     case PSA_IPC_CALL:
         /* Change FP regs */
-        change_fp_in_s_thread();
+        populate_caller_fp_regs(expecting_caller_content);
+        populate_callee_fp_regs(expecting_callee_content);
         /* Start the timer */
         tfm_plat_test_non_secure_timer_start();
         LOG_DBGFMT("Wait for NS timer interrupt!\r\n");
@@ -475,7 +182,12 @@
             r = PSA_ERROR_GENERIC_ERROR;
         } else {
             /* FP register should be restored after NS interrupt. */
-            if (check_fp_restored_s()) {
+            dump_fp_caller(fp_caller_buffer);
+            dump_fp_callee(fp_callee_buffer);
+            if ((!memcmp(fp_caller_buffer, expecting_caller_content,
+                         FP_CALLER_BUF_SIZE)) &&
+                (!memcmp(fp_callee_buffer, expecting_callee_content,
+                         FP_CALLEE_BUF_SIZE))) {
                 r = PSA_SUCCESS;
             } else {
                 r = PSA_ERROR_GENERIC_ERROR;
@@ -490,19 +202,14 @@
     }
 }
 
-/**
- * FP service partition main thread.
- */
+/* FP service partition main thread. */
 void fpu_service_test_main(void *param)
 {
     uint32_t signals = 0;
 
-    clear_fp_regs();
     while (1) {
         signals = psa_wait(PSA_WAIT_ANY, PSA_BLOCK);
-        if (signals & TFM_FPU_SERVICE_CLEAR_FP_REGISTER_SIGNAL) {
-            fpu_service_clear_fp_register();
-        } else if (signals & TFM_FPU_SERVICE_CHECK_FP_REGISTER_SIGNAL) {
+        if (signals & TFM_FPU_SERVICE_CHECK_FP_CALLEE_REGISTER_SIGNAL) {
             fpu_service_check_fp_register();
         } else if (signals & TFM_FPU_SERVICE_START_S_TIMER_SIGNAL) {
             fpu_client_start_secure_timer();
diff --git a/test/secure_fw/suites/fpu/service/tfm_fpu_service_test.yaml b/test/secure_fw/suites/fpu/service/tfm_fpu_service_test.yaml
index e826b01..1adc9f8 100644
--- a/test/secure_fw/suites/fpu/service/tfm_fpu_service_test.yaml
+++ b/test/secure_fw/suites/fpu/service/tfm_fpu_service_test.yaml
@@ -17,15 +17,7 @@
   ],
   "services" : [
     {
-      "name": "TFM_FPU_SERVICE_CLEAR_FP_REGISTER",
-      "sid": "0x0000F090",
-      "connection_based": true,
-      "non_secure_clients": true,
-      "version": 1,
-      "version_policy": "STRICT"
-    },
-    {
-      "name": "TFM_FPU_SERVICE_CHECK_FP_REGISTER",
+      "name": "TFM_FPU_SERVICE_CHECK_FP_CALLEE_REGISTER",
       "sid": "0x0000F091",
       "connection_based": true,
       "non_secure_clients": true,