// Load and save register values with pauses for ptrace
 //
-// x0 - SVE in use
-// x1 - SME in use
-// x2 - SME2 in use
-// x3 - FA64 supported
+// x0 - HAVE_ flags indicating which features are in use
 
 .globl load_and_save
 load_and_save:
        ldp     q30, q31, [x7, #16 * 30]
 
        // SME?
-       cbz     x1, check_sve_in
+       tbz     x0, #HAVE_SME_SHIFT, check_sve_in
 
        adrp    x7, svcr_in
        ldr     x7, [x7, :lo12:svcr_in]
        bne     1b
 
        // ZT?
-       cbz     x2, check_sm_in
+       tbz     x0, #HAVE_SME2_SHIFT, check_sm_in
        adrp    x6, zt_in
        add     x6, x6, :lo12:zt_in
        _ldr_zt 6
        // In streaming mode?
 check_sm_in:
        tbz     x7, #SVCR_SM_SHIFT, check_sve_in
-       mov     x4, x3          // Load FFR if we have FA64
+
+       // Load FFR if we have FA64
+       mov     x4, #0
+       tbz     x0, #HAVE_FA64_SHIFT, load_sve
+       mov     x4, #1
        b       load_sve
 
        // SVE?
 check_sve_in:
-       cbz     x0, wait_for_writes
+       tbz     x0, #HAVE_SVE_SHIFT, wait_for_writes
        mov     x4, #1
 
 load_sve:
        stp     q28, q29, [x7, #16 * 28]
        stp     q30, q31, [x7, #16 * 30]
 
-       // SME?
-       cbz     x1, check_sve_out
+       tbz     x0, #HAVE_SME_SHIFT, check_sve_out
 
        rdsvl   11, 1
        adrp    x6, sme_vl_out
        bne     1b
 
        // ZT?
-       cbz     x2, check_sm_out
+       tbz     x0, #HAVE_SME2_SHIFT, check_sm_out
        adrp    x6, zt_out
        add     x6, x6, :lo12:zt_out
        _str_zt 6
        // In streaming mode?
 check_sm_out:
        tbz     x7, #SVCR_SM_SHIFT, check_sve_out
-       mov     x4, x3                          // FFR?
+
+       // Do we have FA64 and FFR?
+       mov     x4, #0
+       tbz     x0, #HAVE_FA64_SHIFT, read_sve
+       mov     x4, #1
        b       read_sve
 
        // SVE?
 check_sve_out:
-       cbz     x0, wait_for_reads
+       tbz     x0, #HAVE_SVE_SHIFT, wait_for_reads
        mov     x4, #1
 
        rdvl    x7, #1
        brk #0
 
        // Ensure we don't leave ourselves in streaming mode
-       cbz     x1, out
+       tbz     x0, #HAVE_SME_SHIFT, out
        msr     S3_3_C4_C2_2, xzr
 
 out:
 
 uint64_t sme_vl_out;
 uint64_t svcr_in, svcr_expected, svcr_out;
 
-void load_and_save(int sve, int sme, int sme2, int fa64);
+void load_and_save(int flags);
 
 static bool got_alarm;
 
 
 static void run_child(struct test_config *config)
 {
-       int ret;
+       int ret, flags;
 
        /* Let the parent attach to us */
        ret = ptrace(PTRACE_TRACEME, 0, 0, 0);
        }
 
        /* Load values and wait for the parent */
-       load_and_save(sve_supported(), sme_supported(),
-                     sme2_supported(), fa64_supported());
+       flags = 0;
+       if (sve_supported())
+               flags |= HAVE_SVE;
+       if (sme_supported())
+               flags |= HAVE_SME;
+       if (sme2_supported())
+               flags |= HAVE_SME2;
+       if (fa64_supported())
+               flags |= HAVE_FA64;
+
+       load_and_save(flags);
 
        exit(0);
 }
 
 #define SVCR_SM (1 << SVCR_SM_SHIFT)
 #define SVCR_ZA (1 << SVCR_ZA_SHIFT)
 
+#define HAVE_SVE_SHIFT         0
+#define HAVE_SME_SHIFT         1
+#define HAVE_SME2_SHIFT                2
+#define HAVE_FA64_SHIFT                3
+
+#define HAVE_SVE       (1 << HAVE_SVE_SHIFT)
+#define HAVE_SME       (1 << HAVE_SME_SHIFT)
+#define HAVE_SME2      (1 << HAVE_SME2_SHIFT)
+#define HAVE_FA64      (1 << HAVE_FA64_SHIFT)
+
 #endif