.arch   armv8-a+crypto
 
        /*
-        * void ce_aes_ccm_auth_data(u8 mac[], u8 const in[], u32 abytes,
-        *                           u32 *macp, u8 const rk[], u32 rounds);
+        * u32 ce_aes_ccm_auth_data(u8 mac[], u8 const in[], u32 abytes,
+        *                          u32 macp, u8 const rk[], u32 rounds);
         */
 SYM_FUNC_START(ce_aes_ccm_auth_data)
-       ldr     w8, [x3]                        /* leftover from prev round? */
        ld1     {v0.16b}, [x0]                  /* load mac */
-       cbz     w8, 1f
-       sub     w8, w8, #16
+       cbz     w3, 1f
+       sub     w3, w3, #16
        eor     v1.16b, v1.16b, v1.16b
 0:     ldrb    w7, [x1], #1                    /* get 1 byte of input */
        subs    w2, w2, #1
-       add     w8, w8, #1
+       add     w3, w3, #1
        ins     v1.b[0], w7
        ext     v1.16b, v1.16b, v1.16b, #1      /* rotate in the input bytes */
        beq     8f                              /* out of input? */
-       cbnz    w8, 0b
+       cbnz    w3, 0b
        eor     v0.16b, v0.16b, v1.16b
 1:     ld1     {v3.4s}, [x4]                   /* load first round key */
        prfm    pldl1strm, [x1]
        beq     10f
        adds    w2, w2, #16
        beq     10f
-       mov     w8, w2
+       mov     w3, w2
 7:     ldrb    w7, [x1], #1
        umov    w6, v0.b[0]
        eor     w6, w6, w7
        beq     10f
        ext     v0.16b, v0.16b, v0.16b, #1      /* rotate out the mac bytes */
        b       7b
-8:     cbz     w8, 91f
-       mov     w7, w8
-       add     w8, w8, #16
+8:     cbz     w3, 91f
+       mov     w7, w3
+       add     w3, w3, #16
 9:     ext     v1.16b, v1.16b, v1.16b, #1
        adds    w7, w7, #1
        bne     9b
 91:    eor     v0.16b, v0.16b, v1.16b
        st1     {v0.16b}, [x0]
-10:    str     w8, [x3]
+10:    mov     w0, w3
        ret
 SYM_FUNC_END(ce_aes_ccm_auth_data)
 
 
        return 6 + ctx->key_length / 4;
 }
 
-asmlinkage void ce_aes_ccm_auth_data(u8 mac[], u8 const in[], u32 abytes,
-                                    u32 *macp, u32 const rk[], u32 rounds);
+asmlinkage u32 ce_aes_ccm_auth_data(u8 mac[], u8 const in[], u32 abytes,
+                                   u32 macp, u32 const rk[], u32 rounds);
 
 asmlinkage void ce_aes_ccm_encrypt(u8 out[], u8 const in[], u32 cbytes,
                                   u32 const rk[], u32 rounds, u8 mac[],
        return 0;
 }
 
-static void ccm_update_mac(struct crypto_aes_ctx *key, u8 mac[], u8 const in[],
-                          u32 abytes, u32 *macp)
-{
-       ce_aes_ccm_auth_data(mac, in, abytes, macp, key->key_enc,
-                            num_rounds(key));
-}
-
 static void ccm_calculate_auth_mac(struct aead_request *req, u8 mac[])
 {
        struct crypto_aead *aead = crypto_aead_reqtfm(req);
                ltag.len = 6;
        }
 
-       ccm_update_mac(ctx, mac, (u8 *)<ag, ltag.len, &macp);
+       macp = ce_aes_ccm_auth_data(mac, (u8 *)<ag, ltag.len, macp,
+                                   ctx->key_enc, num_rounds(ctx));
        scatterwalk_start(&walk, req->src);
 
        do {
                }
                n = min_t(u32, n, SZ_4K); /* yield NEON at least every 4k */
                p = scatterwalk_map(&walk);
-               ccm_update_mac(ctx, mac, p, n, &macp);
+
+               macp = ce_aes_ccm_auth_data(mac, p, n, macp, ctx->key_enc,
+                                           num_rounds(ctx));
 
                if (len / SZ_4K > (len - n) / SZ_4K) {
                        kernel_neon_end();
                        kernel_neon_begin();
                }
-
                len -= n;
 
                scatterwalk_unmap(p);