arm64: filmgrain16: Add NEON implementation of gen_grain for 16 bpc

Relative speedup over C code:
                             Cortex A53    A72    A73   Apple M1
gen_grain_uv_ar0_16bpc_420_neon:   2.90   4.13   5.43   5.80
gen_grain_uv_ar0_16bpc_422_neon:   3.23   4.51   5.52   5.83
gen_grain_uv_ar0_16bpc_444_neon:   4.01   4.97   6.08   5.87
gen_grain_uv_ar1_16bpc_420_neon:   2.94   2.80   3.56   3.48
gen_grain_uv_ar1_16bpc_422_neon:   3.14   3.07   3.68   3.47
gen_grain_uv_ar1_16bpc_444_neon:   3.54   3.51   3.93   2.61
gen_grain_uv_ar2_16bpc_420_neon:   3.92   3.69   4.40   3.98
gen_grain_uv_ar2_16bpc_422_neon:   4.13   3.96   4.42   3.92
gen_grain_uv_ar2_16bpc_444_neon:   4.69   4.33   4.84   3.25
gen_grain_uv_ar3_16bpc_420_neon:   5.05   5.39   5.42   4.74
gen_grain_uv_ar3_16bpc_422_neon:   5.25   5.68   5.57   4.67
gen_grain_uv_ar3_16bpc_444_neon:   6.02   6.33   6.35   4.38
gen_grain_y_ar0_16bpc_neon:        4.67   5.23   5.22  10.11
gen_grain_y_ar1_16bpc_neon:        3.32   3.03   3.28   2.24
gen_grain_y_ar2_16bpc_neon:        4.59   3.95   4.64   3.52
gen_grain_y_ar3_16bpc_neon:        5.89   5.93   6.36   4.79

Absolute numbers:
                                 Cortex A53       A72       A73    Apple M1
gen_grain_uv_ar0_16bpc_420_neon:    19797.2    9725.0    9234.0    29.7
gen_grain_uv_ar0_16bpc_422_neon:    34899.4   16875.3   17021.6    57.7
gen_grain_uv_ar0_16bpc_444_neon:    53776.6   28470.1   28773.1   107.8
gen_grain_uv_ar1_16bpc_420_neon:    37998.2   24631.2   24754.0    84.2
gen_grain_uv_ar1_16bpc_422_neon:    70817.5   44642.5   46323.1   166.3
gen_grain_uv_ar1_16bpc_444_neon:   123333.0   77316.4   83523.1   427.5
gen_grain_uv_ar2_16bpc_420_neon:    49115.8   33053.7   33249.9    93.6
gen_grain_uv_ar2_16bpc_422_neon:    92965.3   59663.8   64741.9   187.9
gen_grain_uv_ar2_16bpc_444_neon:   160899.7  108845.6  115422.4   441.8
gen_grain_uv_ar3_16bpc_420_neon:    65786.6   41924.3   45562.1   108.1
gen_grain_uv_ar3_16bpc_422_neon:   126232.3   78691.6   87351.5   217.6
gen_grain_uv_ar3_16bpc_444_neon:   218702.6  140197.8  151294.8   454.3
gen_grain_y_ar0_16bpc_neon:         35867.9   17653.6   20770.7   108.0
gen_grain_y_ar1_16bpc_neon:        118781.8   74777.1   81338.6   426.0
gen_grain_y_ar2_16bpc_neon:        155919.9  102145.8  109698.1   438.5
gen_grain_y_ar3_16bpc_neon:        213348.1  133054.8  144726.0   447.9

Corresponding numbers for 8bpc:
                                 Cortex A53       A72       A73    Apple M1
gen_grain_uv_ar0_8bpc_420_neon:     15086.1    8384.7    8556.6    29.4
gen_grain_uv_ar0_8bpc_422_neon:     26800.6   14354.4   15526.5    56.6
gen_grain_uv_ar0_8bpc_444_neon:     43749.6   22408.6   24627.9   108.3
gen_grain_uv_ar1_8bpc_420_neon:     33706.3   21892.6   22835.9    87.1
gen_grain_uv_ar1_8bpc_422_neon:     63897.0   41820.1   43468.9   171.8
gen_grain_uv_ar1_8bpc_444_neon:    117345.1   76372.5   79938.3   370.0
gen_grain_uv_ar2_8bpc_420_neon:     42808.8   28493.8   29932.8    92.2
gen_grain_uv_ar2_8bpc_422_neon:     82282.5   53969.4   58191.1   181.8
gen_grain_uv_ar2_8bpc_444_neon:    147641.4   98136.4  103157.6   430.2
gen_grain_uv_ar3_8bpc_420_neon:     56784.3   36342.0   40812.3   102.2
gen_grain_uv_ar3_8bpc_422_neon:    110249.7   70215.6   79716.0   200.5
gen_grain_uv_ar3_8bpc_444_neon:    196461.7  125802.8  141781.5   440.1
gen_grain_y_ar0_8bpc_neon:          36451.7   17794.4   19839.3   109.5
gen_grain_y_ar1_8bpc_neon:         113155.6   71811.9   77296.8   370.2
gen_grain_y_ar2_8bpc_neon:         142812.3   95042.4  100434.4   431.8
gen_grain_y_ar3_8bpc_neon:         191608.6  121199.5  136946.4   437.2

2 files changed, 1112 insertions, 1 deletions

diff --git a/src/arm/64/film_grain16.S b/src/arm/64/film_grain16.S
index 18525c5..7bf703c 100644
--- a/src/arm/64/film_grain16.S
+++ b/src/arm/64/film_grain16.S
@@ -30,6 +30,1117 @@
 #include "src/arm/asm-offsets.h"
 
 #define GRAIN_WIDTH 82
+#define GRAIN_HEIGHT 73
+
+#define SUB_GRAIN_WIDTH 44
+#define SUB_GRAIN_HEIGHT 38
+
+.macro increment_seed steps, shift=1
+        lsr             w11, w2,  #3
+        lsr             w12, w2,  #12
+        lsr             w13, w2,  #1
+        eor             w11, w2,  w11                     // (r >> 0) ^ (r >> 3)
+        eor             w12, w12, w13                     // (r >> 12) ^ (r >> 1)
+        eor             w11, w11, w12                     // (r >> 0) ^ (r >> 3) ^ (r >> 12) ^ (r >> 1)
+.if \shift
+        lsr             w2,  w2,  #\steps
+.endif
+        and             w11, w11, #((1 << \steps) - 1)    // bit
+.if \shift
+        orr             w2,  w2,  w11, lsl #(16 - \steps) // *state
+.else
+        orr             w2,  w2,  w11, lsl #16            // *state
+.endif
+.endm
+
+.macro read_rand dest, bits, age
+        ubfx            \dest,  x2,   #16 - \bits - \age, #\bits
+.endm
+
+.macro read_shift_rand dest, bits
+        ubfx            \dest,  x2,   #17 - \bits, #\bits
+        lsr             w2,  w2,  #1
+.endm
+
+// special calling convention:
+// w2 holds seed
+// x3 holds dav1d_gaussian_sequence
+// clobbers x11-x15
+// returns in v0.8h
+function get_gaussian_neon
+        increment_seed  4
+        read_rand       x14, 11,  3
+        read_rand       x15, 11,  2
+        add             x14, x3,  x14, lsl #1
+        add             x15, x3,  x15, lsl #1
+        ld1             {v0.h}[0], [x14]
+        read_rand       x14, 11,  1
+        ld1             {v0.h}[1], [x15]
+        add             x14, x3,  x14, lsl #1
+        read_rand       x15, 11,  0
+        increment_seed  4
+        add             x15, x3,  x15, lsl #1
+        ld1             {v0.h}[2], [x14]
+        read_rand       x14, 11,  3
+        ld1             {v0.h}[3], [x15]
+        add             x14, x3,  x14, lsl #1
+        read_rand       x15, 11,  2
+        ld1             {v0.h}[4], [x14]
+        add             x15, x3,  x15, lsl #1
+        read_rand       x14, 11,  1
+        ld1             {v0.h}[5], [x15]
+        read_rand       x15, 11,  0
+        add             x14, x3,  x14, lsl #1
+        add             x15, x3,  x15, lsl #1
+        ld1             {v0.h}[6], [x14]
+        ld1             {v0.h}[7], [x15]
+        ret
+endfunc
+
+.macro store_grain_row r0, r1, r2, r3, r4, r5
+        st1             {\r0\().16b,\r1\().16b}, [x0], #32
+        st1             {\r2\().16b,\r3\().16b}, [x0], #32
+        st1             {\r4\().16b},  [x0], #16
+        st1             {\r5\().h}[0], [x0], #2
+.endm
+
+.macro store_grain_row_44 r0, r1, r2
+        st1             {\r0\().16b,\r1\().16b}, [x0], #32
+        st1             {\r2\().16b},  [x0]
+        add             x0,  x0,  #GRAIN_WIDTH-32
+.endm
+
+function get_grain_2_neon
+        increment_seed  2
+        read_rand       x14, 11,  1
+        read_rand       x15, 11,  0
+        add             x14, x3,  x14, lsl #1
+        add             x15, x3,  x15, lsl #1
+        ld1             {v0.h}[0], [x14]
+        ld1             {v0.h}[1], [x15]
+        srshl           v0.4h,   v0.4h,   v31.4h
+        ret
+endfunc
+
+.macro get_grain_2 dst
+        bl              get_grain_2_neon
+.ifnc \dst, v0
+        mov             \dst\().8b, v0.8b
+.endif
+.endm
+
+function get_grain_4_neon
+        increment_seed  4
+        read_rand       x14, 11,  3
+        read_rand       x15, 11,  2
+        add             x14, x3,  x14, lsl #1
+        add             x15, x3,  x15, lsl #1
+        ld1             {v0.h}[0], [x14]
+        read_rand       x14, 11,  1
+        ld1             {v0.h}[1], [x15]
+        add             x14, x3,  x14, lsl #1
+        read_rand       x15, 11,  0
+        add             x15, x3,  x15, lsl #1
+        ld1             {v0.h}[2], [x14]
+        ld1             {v0.h}[3], [x15]
+        srshl           v0.4h,   v0.4h,   v31.4h
+        ret
+endfunc
+
+.macro get_grain_4 dst
+        bl              get_grain_4_neon
+.ifnc \dst, v0
+        mov             \dst\().8b, v0.8b
+.endif
+.endm
+
+// w15 holds the number of entries to produce
+// w14, w16 and w17 hold the previous output entries
+// v0 holds the vector of produced entries
+// v1 holds the input vector of sums from above
+.macro output_lag n
+function output_lag\n\()_neon
+1:
+        read_shift_rand x13, 11
+        mov             w11, v1.s[0]
+        ldrsh           w12, [x3, x13, lsl #1]
+        ext             v0.16b,  v0.16b,  v0.16b,  #2
+.if \n == 1
+        madd            w11, w14, w4,  w11        // sum (above) + *coeff * prev output
+.elseif \n == 2
+        madd            w11, w16, w4,  w11        // sum (above) + *coeff * prev output 1
+        madd            w11, w14, w17, w11        // += *coeff * prev output 2
+        mov             w16, w14
+.else
+        madd            w11, w17, w4,  w11        // sum (above) + *coeff * prev output 1
+        madd            w11, w16, w20, w11        // sum (above) + *coeff * prev output 2
+        madd            w11, w14, w21, w11        // += *coeff * prev output 3
+        mov             w17, w16
+        mov             w16, w14
+.endif
+        add             w14, w11, w8              // 1 << (ar_coeff_shift - 1)
+        asr             w14, w14, w7              // >> ar_coeff_shift
+        add             w12, w12, w10
+        asr             w12, w12, w9              // >> (4 - bitdepth_min_8 + grain_scale_shift)
+        add             w14, w14, w12
+        cmp             w14, w5
+        csel            w14, w14, w5,  le
+        cmp             w14, w6
+        csel            w14, w14, w6,  ge
+        subs            w15, w15, #1
+        ext             v1.16b,  v1.16b,  v1.16b,  #4
+        ins             v0.h[7], w14
+        b.gt            1b
+        ret
+endfunc
+.endm
+
+output_lag 1
+output_lag 2
+output_lag 3
+
+
+function sum_lag1_above_neon
+        sub             x12, x0,  #1*GRAIN_WIDTH*2 - 16
+        ld1             {v18.8h}, [x12] // load top right
+
+        ext             v0.16b,  v16.16b, v17.16b, #14 // top left, top mid
+        ext             v1.16b,  v17.16b, v18.16b, #2  // top mid, top right
+
+        smull           v4.4s,   v17.4h,  v28.4h
+        smlal           v4.4s,   v0.4h,   v27.4h
+        smlal           v4.4s,   v1.4h,   v29.4h
+        smull2          v5.4s,   v17.8h,  v28.8h
+        smlal2          v5.4s,   v0.8h,   v27.8h
+        smlal2          v5.4s,   v1.8h,   v29.8h
+
+        mov             v16.16b, v17.16b
+        mov             v17.16b, v18.16b
+
+        ret
+endfunc
+
+.macro sum_lag_n_body lag, type, uv_layout, edge, elems, uv_coeff
+        bl              sum_\lag\()_above_neon
+.ifc \type, uv_420
+        add             x12, x19, #GRAIN_WIDTH*2
+        ld1             {v22.8h, v23.8h}, [x19], #32
+        ld1             {v24.8h, v25.8h}, [x12]
+        addp            v22.8h,  v22.8h,  v23.8h
+        addp            v23.8h,  v24.8h,  v25.8h
+        add             v22.8h,  v22.8h,  v23.8h
+        srshr           v0.8h,   v22.8h,  #2
+.endif
+.ifc \type, uv_422
+        ld1             {v22.8h, v23.8h}, [x19], #32
+        addp            v22.8h,  v22.8h,  v23.8h
+        srshr           v0.8h,   v22.8h,  #1
+.endif
+.ifc \type, uv_444
+        ld1             {v0.8h}, [x19], #16
+.endif
+.if \uv_layout
+.ifnb \uv_coeff
+        dup             v1.8b,   \uv_coeff
+        sxtl            v1.8h,   v1.8b
+        smlal           v4.4s,   v0.4h,   v1.4h
+        smlal2          v5.4s,   v0.8h,   v1.8h
+.else
+        smlal           v4.4s,   v0.4h,   v30.4h
+        smlal2          v5.4s,   v0.8h,   v30.8h
+.endif
+.endif
+.if \uv_layout && \elems == 8
+        b               sum_\lag\()_y_\edge\()_start
+.elseif \uv_layout == 444 && \elems == 7
+        b               sum_\lag\()_y_\edge\()_start
+.elseif \uv_layout == 422 && \elems == 1
+        b               sum_\lag\()_uv_420_\edge\()_start
+.else
+sum_\lag\()_\type\()_\edge\()_start:
+.if \elems > 4
+.ifc \edge, left
+        increment_seed  4
+        read_rand       x12, 11,  3
+        read_rand       x13, 11,  2
+        read_rand       x14, 11,  1
+        add             x12, x3,  x12, lsl #1
+        add             x13, x3,  x13, lsl #1
+        add             x14, x3,  x14, lsl #1
+        ld1             {v0.h}[5], [x12]
+        ld1             {v0.h}[6], [x13]
+        ld1             {v0.h}[7], [x14]
+        lsl             x2,  x2,  #1             // shift back the state as if we'd done increment_seed with shift=0
+        srshl           v0.8h,   v0.8h,   v31.8h
+        ext             v4.16b,  v4.16b,  v4.16b,  #12
+.ifc \lag, lag3
+        smov            w17, v0.h[5]
+.endif
+.ifnc \lag, lag1
+        smov            w16, v0.h[6]
+.endif
+        smov            w14, v0.h[7]
+
+        mov             v1.16b,  v4.16b
+        mov             w15, #1
+        bl              output_\lag\()_neon
+.else
+        increment_seed  4, shift=0
+        mov             v1.16b,  v4.16b
+        mov             w15, #4
+        bl              output_\lag\()_neon
+.endif
+
+        increment_seed  4, shift=0
+        mov             v1.16b,  v5.16b
+.ifc \edge, right
+        mov             w15, #3
+        bl              output_\lag\()_neon
+        read_shift_rand x15, 11
+        add             x15, x3,  x15, lsl #1
+        ld1             {v1.h}[0], [x15]
+        srshl           v1.4h,   v1.4h,   v31.4h
+        ext             v0.16b,  v0.16b,  v1.16b,  #2
+.else
+        mov             w15, #4
+        bl              output_\lag\()_neon
+.endif
+.else
+        // elems == 1
+        increment_seed  4, shift=0
+        mov             v1.16b,  v4.16b
+        mov             w15, #1
+        bl              output_\lag\()_neon
+        lsr             w2,  w2,  #3
+
+        read_rand       x12, 11,  2
+        read_rand       x13, 11,  1
+        read_rand       x14, 11,  0
+        add             x12, x3,  x12, lsl #1
+        add             x13, x3,  x13, lsl #1
+        add             x14, x3,  x14, lsl #1
+        ld1             {v1.h}[0], [x12]
+        ld1             {v1.h}[1], [x13]
+        ld1             {v1.h}[2], [x14]
+        srshl           v1.4h,   v1.4h,   v31.4h
+        ext             v0.16b,  v0.16b,  v1.16b,  #14
+.endif
+        st1             {v0.8h}, [x0], #16
+        ldr             x30,     [sp], #16
+        ret
+.endif
+.endm
+
+.macro sum_lag1_func type, uv_layout, edge, elems=8
+function sum_\type\()_lag1_\edge\()_neon
+        str             x30, [sp, #-16]!
+.ifc \edge, left
+        sub             x12, x0,  #1*GRAIN_WIDTH*2
+        ld1             {v17.8h}, [x12] // load the previous block right above
+.endif
+        sum_lag_n_body  lag1, \type, \uv_layout, \edge, \elems
+endfunc
+.endm
+
+sum_lag1_func y,      0,   left
+sum_lag1_func y,      0,   mid
+sum_lag1_func y,      0,   right, 7
+sum_lag1_func uv_444, 444, left
+sum_lag1_func uv_444, 444, mid
+sum_lag1_func uv_444, 444, right, 7
+sum_lag1_func uv_422, 422, left
+sum_lag1_func uv_422, 422, mid
+sum_lag1_func uv_422, 422, right, 1
+sum_lag1_func uv_420, 420, left
+sum_lag1_func uv_420, 420, mid
+sum_lag1_func uv_420, 420, right, 1
+
+
+function sum_lag2_above_neon
+        sub             x12, x0,  #2*GRAIN_WIDTH*2 - 16
+        sub             x13, x0,  #1*GRAIN_WIDTH*2 - 16
+        ld1             {v18.8h}, [x12] // load top right
+        ld1             {v21.8h}, [x13]
+
+        dup             v26.8b,  v30.b[0]
+        ext             v22.16b, v16.16b, v17.16b, #12 // top left, top mid
+        dup             v27.8b,  v30.b[1]
+        ext             v23.16b, v16.16b, v17.16b, #14
+        sxtl            v26.8h,  v26.8b
+        dup             v28.8b,  v30.b[3]
+        ext             v0.16b,  v17.16b, v18.16b, #2  // top mid, top right
+        sxtl            v27.8h,  v27.8b
+        dup             v29.8b,  v30.b[4]
+        ext             v1.16b,  v17.16b, v18.16b, #4
+        sxtl            v28.8h,  v28.8b
+        sxtl            v29.8h,  v29.8b
+
+        smull           v4.4s,   v22.4h,  v26.4h
+        smlal           v4.4s,   v23.4h,  v27.4h
+        smlal           v4.4s,   v0.4h,   v28.4h
+        smlal           v4.4s,   v1.4h,   v29.4h
+        smull2          v5.4s,   v22.8h,  v26.8h
+        smlal2          v5.4s,   v23.8h,  v27.8h
+        smlal2          v5.4s,   v0.8h,   v28.8h
+        smlal2          v5.4s,   v1.8h,   v29.8h
+
+        dup             v26.16b, v30.b[5]
+        ext             v22.16b, v19.16b, v20.16b, #12 // top left, top mid
+        dup             v27.16b, v30.b[6]
+        ext             v23.16b, v19.16b, v20.16b, #14
+        sxtl            v26.8h,  v26.8b
+        dup             v28.16b, v30.b[8]
+        ext             v0.16b,  v20.16b, v21.16b, #2  // top mid, top right
+        sxtl            v27.8h,  v27.8b
+        dup             v29.16b, v30.b[9]
+        ext             v1.16b,  v20.16b, v21.16b, #4
+        sxtl            v28.8h,  v28.8b
+        sxtl            v29.8h,  v29.8b
+
+        smlal           v4.4s,   v22.4h,  v26.4h
+        smlal           v4.4s,   v23.4h,  v27.4h
+        smlal           v4.4s,   v0.4h,   v28.4h
+        smlal           v4.4s,   v1.4h,   v29.4h
+        smlal2          v5.4s,   v22.8h,  v26.8h
+        smlal2          v5.4s,   v23.8h,  v27.8h
+        smlal2          v5.4s,   v0.8h,   v28.8h
+        smlal2          v5.4s,   v1.8h,   v29.8h
+
+        dup             v26.16b, v30.b[2]
+        dup             v27.16b, v30.b[7]
+        sxtl            v26.8h,  v26.8b
+        sxtl            v27.8h,  v27.8b
+
+        smlal           v4.4s,   v17.4h,  v26.4h
+        smlal           v4.4s,   v20.4h,  v27.4h
+        smlal2          v5.4s,   v17.8h,  v26.8h
+        smlal2          v5.4s,   v20.8h,  v27.8h
+        mov             v16.16b, v17.16b
+        mov             v17.16b, v18.16b
+
+        mov             v19.16b, v20.16b
+        mov             v20.16b, v21.16b
+        ret
+endfunc
+
+.macro sum_lag2_func type, uv_layout, edge, elems=16
+function sum_\type\()_lag2_\edge\()_neon
+        str             x30, [sp, #-16]!
+.ifc \edge, left
+        sub             x12, x0,  #2*GRAIN_WIDTH*2
+        sub             x13, x0,  #1*GRAIN_WIDTH*2
+        ld1             {v17.8h}, [x12] // load the previous block right above
+        ld1             {v20.8h}, [x13]
+.endif
+        sum_lag_n_body  lag2, \type, \uv_layout, \edge, \elems, v30.b[12]
+endfunc
+.endm
+
+sum_lag2_func y,      0,   left
+sum_lag2_func y,      0,   mid
+sum_lag2_func y,      0,   right, 7
+sum_lag2_func uv_444, 444, left
+sum_lag2_func uv_444, 444, mid
+sum_lag2_func uv_444, 444, right, 7
+sum_lag2_func uv_422, 422, left
+sum_lag2_func uv_422, 422, mid
+sum_lag2_func uv_422, 422, right, 1
+sum_lag2_func uv_420, 420, left
+sum_lag2_func uv_420, 420, mid
+sum_lag2_func uv_420, 420, right, 1
+
+
+function sum_lag3_above_neon
+        sub             x11, x0,  #3*GRAIN_WIDTH*2 - 16
+        sub             x12, x0,  #2*GRAIN_WIDTH*2 - 16
+        sub             x13, x0,  #1*GRAIN_WIDTH*2 - 16
+        ld1             {v15.8h}, [x11] // load top right
+        ld1             {v18.8h}, [x12]
+        ld1             {v21.8h}, [x13]
+
+        dup             v22.8b,  v29.b[0]
+        ext             v8.16b,  v13.16b, v14.16b, #10 // top left, top mid
+        dup             v23.8b,  v29.b[1]
+        ext             v9.16b,  v13.16b, v14.16b, #12
+        sxtl            v22.8h,  v22.8b
+        dup             v24.8b,  v29.b[2]
+        sxtl            v23.8h,  v23.8b
+        dup             v25.8b,  v29.b[3]
+        ext             v10.16b, v13.16b, v14.16b, #14
+        sxtl            v24.8h,  v24.8b
+        dup             v26.8b,  v29.b[4]
+        ext             v11.16b, v14.16b, v15.16b, #2  // top mid, top right
+        sxtl            v25.8h,  v25.8b
+        dup             v27.8b,  v29.b[5]
+        ext             v12.16b, v14.16b, v15.16b, #4
+        sxtl            v26.8h,  v26.8b
+        dup             v28.8b,  v29.b[6]
+        ext             v13.16b, v14.16b, v15.16b, #6
+        sxtl            v27.8h,  v27.8b
+        sxtl            v28.8h,  v28.8b
+
+        smull           v4.4s,   v8.4h,   v22.4h
+        smlal           v4.4s,   v9.4h,   v23.4h
+        smlal           v4.4s,   v10.4h,  v24.4h
+        smlal           v4.4s,   v11.4h,  v26.4h
+        smlal           v4.4s,   v12.4h,  v27.4h
+        smlal           v4.4s,   v13.4h,  v28.4h
+        smlal           v4.4s,   v14.4h,  v25.4h
+        smull2          v5.4s,   v8.8h,   v22.8h
+        smlal2          v5.4s,   v9.8h,   v23.8h
+        smlal2          v5.4s,   v10.8h,  v24.8h
+        smlal2          v5.4s,   v11.8h,  v26.8h
+        smlal2          v5.4s,   v12.8h,  v27.8h
+        smlal2          v5.4s,   v13.8h,  v28.8h
+        smlal2          v5.4s,   v14.8h,  v25.8h
+
+        dup             v22.8b,  v29.b[7]
+        ext             v8.16b,  v16.16b, v17.16b, #10 // top left, top mid
+        dup             v23.8b,  v29.b[8]
+        ext             v9.16b,  v16.16b, v17.16b, #12
+        sxtl            v22.8h,  v22.8b
+        dup             v24.8b,  v29.b[9]
+        sxtl            v23.8h,  v23.8b
+        dup             v25.8b,  v29.b[10]
+        ext             v10.16b, v16.16b, v17.16b, #14
+        sxtl            v24.8h,  v24.8b
+        dup             v26.8b,  v29.b[11]
+        ext             v11.16b, v17.16b, v18.16b, #2  // top mid, top right
+        sxtl            v25.8h,  v25.8b
+        dup             v27.8b,  v29.b[12]
+        ext             v12.16b, v17.16b, v18.16b, #4
+        sxtl            v26.8h,  v26.8b
+        dup             v28.8b,  v29.b[13]
+        ext             v13.16b, v17.16b, v18.16b, #6
+        sxtl            v27.8h,  v27.8b
+        sxtl            v28.8h,  v28.8b
+
+        smlal           v4.4s,   v8.4h,   v22.4h
+        smlal           v4.4s,   v9.4h,   v23.4h
+        smlal           v4.4s,   v10.4h,  v24.4h
+        smlal           v4.4s,   v11.4h,  v26.4h
+        smlal           v4.4s,   v12.4h,  v27.4h
+        smlal           v4.4s,   v13.4h,  v28.4h
+        smlal           v4.4s,   v17.4h,  v25.4h
+        smlal2          v5.4s,   v8.8h,   v22.8h
+        smlal2          v5.4s,   v9.8h,   v23.8h
+        smlal2          v5.4s,   v10.8h,  v24.8h
+        smlal2          v5.4s,   v11.8h,  v26.8h
+        smlal2          v5.4s,   v12.8h,  v27.8h
+        smlal2          v5.4s,   v13.8h,  v28.8h
+        smlal2          v5.4s,   v17.8h,  v25.8h
+
+        dup             v22.8b,  v29.b[14]
+        ext             v8.16b,  v19.16b, v20.16b, #10 // top left, top mid
+        dup             v23.8b,  v29.b[15]
+        ext             v9.16b,  v19.16b, v20.16b, #12
+        sxtl            v22.8h,  v22.8b
+        dup             v24.8b,  v30.b[0]
+        sxtl            v23.8h,  v23.8b
+        dup             v25.8b,  v30.b[1]
+        ext             v10.16b, v19.16b, v20.16b, #14
+        sxtl            v24.8h,  v24.8b
+        dup             v26.8b,  v30.b[2]
+        ext             v11.16b, v20.16b, v21.16b, #2  // top mid, top right
+        sxtl            v25.8h,  v25.8b
+        dup             v27.8b,  v30.b[3]
+        ext             v12.16b, v20.16b, v21.16b, #4
+        sxtl            v26.8h,  v26.8b
+        dup             v28.8b,  v30.b[4]
+        ext             v13.16b, v20.16b, v21.16b, #6
+        sxtl            v27.8h,  v27.8b
+        sxtl            v28.8h,  v28.8b
+
+        smlal           v4.4s,   v8.4h,   v22.4h
+        smlal           v4.4s,   v9.4h,   v23.4h
+        smlal           v4.4s,   v10.4h,  v24.4h
+        smlal           v4.4s,   v11.4h,  v26.4h
+        smlal           v4.4s,   v12.4h,  v27.4h
+        smlal           v4.4s,   v13.4h,  v28.4h
+        smlal           v4.4s,   v20.4h,  v25.4h
+        mov             v16.16b, v17.16b
+        mov             v17.16b, v18.16b
+        smlal2          v5.4s,   v8.8h,   v22.8h
+        smlal2          v5.4s,   v9.8h,   v23.8h
+        smlal2          v5.4s,   v10.8h,  v24.8h
+        smlal2          v5.4s,   v11.8h,  v26.8h
+        smlal2          v5.4s,   v12.8h,  v27.8h
+        smlal2          v5.4s,   v13.8h,  v28.8h
+        smlal2          v5.4s,   v20.8h,  v25.8h
+
+        mov             v13.16b, v14.16b
+        mov             v14.16b, v15.16b
+
+        mov             v19.16b, v20.16b
+        mov             v20.16b, v21.16b
+        ret
+endfunc
+
+.macro sum_lag3_func type, uv_layout, edge, elems=16
+function sum_\type\()_lag3_\edge\()_neon
+        str             x30, [sp, #-16]!
+.ifc \edge, left
+        sub             x11, x0,  #3*GRAIN_WIDTH*2
+        sub             x12, x0,  #2*GRAIN_WIDTH*2
+        sub             x13, x0,  #1*GRAIN_WIDTH*2
+        ld1             {v14.8h}, [x11] // load the previous block right above
+        ld1             {v17.8h}, [x12]
+        ld1             {v20.8h}, [x13]
+.endif
+        sum_lag_n_body  lag3, \type, \uv_layout, \edge, \elems, v30.b[8]
+endfunc
+.endm
+
+sum_lag3_func y,      0,   left
+sum_lag3_func y,      0,   mid
+sum_lag3_func y,      0,   right, 7
+sum_lag3_func uv_444, 444, left
+sum_lag3_func uv_444, 444, mid
+sum_lag3_func uv_444, 444, right, 7
+sum_lag3_func uv_422, 422, left
+sum_lag3_func uv_422, 422, mid
+sum_lag3_func uv_422, 422, right, 1
+sum_lag3_func uv_420, 420, left
+sum_lag3_func uv_420, 420, mid
+sum_lag3_func uv_420, 420, right, 1
+
+function generate_grain_rows_neon
+        str             x30, [sp, #-16]!
+1:
+        mov             w16, #80
+2:
+        bl              get_gaussian_neon
+        srshl           v0.8h,   v0.8h,   v31.8h
+        subs            w16, w16, #8
+        st1             {v0.8h}, [x0], #16
+        b.gt            2b
+        get_grain_2     v0
+        subs            w1,  w1,  #1
+        st1             {v0.s}[0], [x0], #4
+        b.gt            1b
+        ldr             x30, [sp], #16
+        ret
+endfunc
+
+function generate_grain_rows_44_neon
+        str             x30, [sp, #-16]!
+1:
+        mov             w16, #40
+2:
+        bl              get_gaussian_neon
+        srshl           v0.8h,   v0.8h,   v31.8h
+        subs            w16, w16, #8
+        st1             {v0.8h}, [x0], #16
+        b.gt            2b
+        get_grain_4     v0
+        subs            w1,  w1,  #1
+        st1             {v0.4h}, [x0]
+        add             x0,  x0,  #GRAIN_WIDTH*2-80
+        b.gt            1b
+        ldr             x30, [sp], #16
+        ret
+endfunc
+
+function gen_grain_uv_444_lag0_neon
+        str             x30, [sp, #-16]!
+        ld1             {v4.8h}, [x19], #16
+gen_grain_uv_lag0_8_start:
+        bl              get_gaussian_neon
+        srshl           v0.8h,   v0.8h,   v31.8h
+gen_grain_uv_lag0_8_add:
+        and             v4.16b,  v4.16b,  v1.16b
+        smull           v2.4s,   v4.4h,   v27.4h
+        smull2          v3.4s,   v4.8h,   v27.8h
+        srshl           v2.4s,   v2.4s,   v28.4s
+        srshl           v3.4s,   v3.4s,   v28.4s
+        sqxtn           v2.4h,   v2.4s
+        sqxtn2          v2.8h,   v3.4s
+        sqadd           v2.8h,   v2.8h,   v0.8h
+        smin            v2.8h,   v2.8h,   v25.8h
+        smax            v2.8h,   v2.8h,   v26.8h
+        st1             {v2.8h}, [x0], #16
+        ldr             x30, [sp], #16
+        ret
+endfunc
+
+function gen_grain_uv_420_lag0_8_neon
+        add             x12, x19, #GRAIN_WIDTH*2
+        str             x30, [sp, #-16]!
+        ld1             {v16.8h, v17.8h}, [x19], #32
+        ld1             {v18.8h, v19.8h}, [x12]
+        addp            v16.8h,  v16.8h,  v17.8h
+        addp            v17.8h,  v18.8h,  v19.8h
+        add             v16.8h,  v16.8h,  v17.8h
+        srshr           v4.8h,   v16.8h,  #2
+        b               gen_grain_uv_lag0_8_start
+endfunc
+
+function gen_grain_uv_422_lag0_8_neon
+        str             x30, [sp, #-16]!
+        ld1             {v16.8h, v17.8h}, [x19], #32
+        addp            v16.8h,  v16.8h,  v17.8h
+        srshr           v4.8h,   v16.8h,  #1
+        b               gen_grain_uv_lag0_8_start
+endfunc
+
+function gen_grain_uv_420_lag0_4_neon
+        add             x12, x19, #GRAIN_WIDTH*2
+        str             x30, [sp, #-16]!
+        ld1             {v16.4h, v17.4h}, [x19]
+        ld1             {v18.4h, v19.4h}, [x12]
+        add             x19,  x19,  #32
+        addp            v16.4h,  v16.4h,  v17.4h
+        addp            v17.4h,  v18.4h,  v19.4h
+        add             v16.4h,  v16.4h,  v17.4h
+        srshr           v4.4h,   v16.4h,  #2
+        get_grain_4     v0
+        b               gen_grain_uv_lag0_8_add
+endfunc
+
+function gen_grain_uv_422_lag0_4_neon
+        str             x30, [sp, #-16]!
+        ld1             {v16.4h, v17.4h}, [x19]
+        add             x19,  x19,  #32
+        addp            v16.4h,  v16.4h,  v17.4h
+        srshr           v4.4h,   v16.4h,  #1
+        get_grain_4     v0
+        b               gen_grain_uv_lag0_8_add
+endfunc
+
+.macro gen_grain_82 type
+function generate_grain_\type\()_16bpc_neon, export=1
+        stp             x30, x19, [sp, #-96]!
+
+.ifc \type, uv_444
+        mov             w13, w3
+        mov             w14, #28
+        add             x19, x1,  #3*GRAIN_WIDTH*2
+        mov             x1,  x2
+        mul             w13, w13, w14
+        clz             w15, w4
+.else
+        clz             w15, w2
+.endif
+        movrel          x3,  X(gaussian_sequence)
+        sub             w15, w15, #24 // -bitdepth_min_8
+        ldr             w2,  [x1, #FGD_SEED]
+        ldr             w9,  [x1, #FGD_GRAIN_SCALE_SHIFT]
+.ifc \type, y
+        add             x4,  x1,  #FGD_AR_COEFFS_Y
+.else
+        add             x4,  x1,  #FGD_AR_COEFFS_UV
+.endif
+        add             w9,  w9,  w15 // grain_scale_shift - bitdepth_min_8
+        adr             x16, L(gen_grain_\type\()_tbl)
+        ldr             w17, [x1, #FGD_AR_COEFF_LAG]
+        add             w9,  w9,  #4
+        ldrh            w17, [x16, w17, uxtw #1]
+        dup             v31.8h,  w9    // 4 - bitdepth_min_8 + data->grain_scale_shift
+        sub             x16, x16, w17, uxtw
+        neg             v31.8h,  v31.8h
+
+.ifc \type, uv_444
+        cmp             w13, #0
+        mov             w11, #0x49d8
+        mov             w14, #0xb524
+        add             x4,  x4,  w13, uxtw // Add offset to ar_coeffs_uv[1]
+        csel            w11, w11, w14, ne
+.endif
+
+        ldr             w7,  [x1, #FGD_AR_COEFF_SHIFT]
+        neg             w15, w15            // bitdepth_min_8
+        mov             w8,  #1
+        mov             w10, #1
+        lsl             w8,  w8,  w7        // 1 << ar_coeff_shift
+        lsl             w10, w10, w9        // 1 << (4 + data->grain_scale_shift)
+        lsr             w8,  w8,  #1        // 1 << (ar_coeff_shift - 1)
+        lsr             w10, w10, #1        // 1 << (4 + data->grain_scale_shift - 1)
+        mov             w5,  #128
+        lsl             w5,  w5,  w15       //   128 << bitdepth_min_8
+        neg             w6,  w5             // -(128 << bitpdeth_min_8)
+        sub             w5,  w5,  #1        //  (128 << bitdepth_min_8) - 1
+
+.ifc \type, uv_444
+        eor             w2,  w2,  w11
+.endif
+
+        br              x16
+
+L(generate_grain_\type\()_lag0):
+.ifc \type, y
+        mov             w1,  #GRAIN_HEIGHT
+        bl              generate_grain_rows_neon
+.else
+        dup             v28.4s,  w7
+        ld1r            {v27.8b}, [x4]      // ar_coeffs_uv[0]
+        movi            v0.16b,  #0
+        movi            v1.16b,  #255
+        dup             v25.8h,  w5
+        dup             v26.8h,  w6
+        ext             v29.16b, v0.16b,  v1.16b,  #10
+        ext             v30.16b, v1.16b,  v0.16b,  #2
+        neg             v28.4s,  v28.4s
+        sxtl            v27.8h,  v27.8b
+
+        mov             w1,  #3
+        bl              generate_grain_rows_neon
+        mov             w1,  #GRAIN_HEIGHT-3
+1:
+        mov             v1.16b,  v29.16b
+        bl              gen_grain_uv_444_lag0_neon // 8
+        movi            v1.16b,  #255
+        bl              gen_grain_uv_444_lag0_neon // 16
+        bl              gen_grain_uv_444_lag0_neon // 24
+        bl              gen_grain_uv_444_lag0_neon // 32
+        bl              gen_grain_uv_444_lag0_neon // 40
+        bl              gen_grain_uv_444_lag0_neon // 48
+        bl              gen_grain_uv_444_lag0_neon // 56
+        bl              gen_grain_uv_444_lag0_neon // 64
+        bl              gen_grain_uv_444_lag0_neon // 72
+        mov             v1.16b,  v30.16b
+        bl              gen_grain_uv_444_lag0_neon // 80
+        get_grain_2     v16
+        subs            w1,  w1,  #1
+        add             x19, x19, #4
+        st1             {v16.s}[0], [x0], #4
+        b.gt            1b
+.endif
+        ldp             x30, x19, [sp], #96
+        ret
+
+L(generate_grain_\type\()_lag1):
+        ld1r            {v27.8b}, [x4], #1  // ar_coeffs_y[0]
+        ld1r            {v28.8b}, [x4], #1  // ar_coeffs_y[1]
+        ld1r            {v29.8b}, [x4]      // ar_coeffs_y[2]
+.ifc \type, y
+        ldrsb           w4,  [x4, #1]       // ar_coeffs_y[3]
+.else
+        add             x4,  x4,  #2
+.endif
+
+        mov             w1,  #3
+.ifc \type, uv_444
+        ld1r            {v30.8b}, [x4]      // ar_coeffs_uv[4]
+        ldursb          w4,  [x4, #-1]      // ar_coeffs_uv[3]
+.endif
+        bl              generate_grain_rows_neon
+        sxtl            v27.8h,  v27.8b
+        sxtl            v28.8h,  v28.8b
+        sxtl            v29.8h,  v29.8b
+.ifc \type, uv_444
+        sxtl            v30.8h,  v30.8b
+.endif
+
+        mov             w1,  #GRAIN_HEIGHT - 3
+1:
+        bl              sum_\type\()_lag1_left_neon  // 8
+        bl              sum_\type\()_lag1_mid_neon   // 16
+        bl              sum_\type\()_lag1_mid_neon   // 24
+        bl              sum_\type\()_lag1_mid_neon   // 32
+        bl              sum_\type\()_lag1_mid_neon   // 40
+        bl              sum_\type\()_lag1_mid_neon   // 48
+        bl              sum_\type\()_lag1_mid_neon   // 56
+        bl              sum_\type\()_lag1_mid_neon   // 64
+        bl              sum_\type\()_lag1_mid_neon   // 72
+        bl              sum_\type\()_lag1_right_neon // 80
+        get_grain_2     v16
+        subs            w1,  w1,  #1
+.ifc \type, uv_444
+        add             x19, x19, #4
+.endif
+        st1             {v16.s}[0], [x0], #4
+        b.gt            1b
+
+        ldp             x30, x19, [sp], #96
+        ret
+
+L(generate_grain_\type\()_lag2):
+        ld1             {v30.16b}, [x4]     // ar_coeffs_y[0-11], ar_coeffs_uv[0-12]
+
+        smov            w4,  v30.b[10]
+        smov            w17, v30.b[11]
+
+        mov             w1,  #3
+        bl              generate_grain_rows_neon
+
+        mov             w1,  #GRAIN_HEIGHT - 3
+1:
+        bl              sum_\type\()_lag2_left_neon  // 8
+        bl              sum_\type\()_lag2_mid_neon   // 16
+        bl              sum_\type\()_lag2_mid_neon   // 24
+        bl              sum_\type\()_lag2_mid_neon   // 32
+        bl              sum_\type\()_lag2_mid_neon   // 40
+        bl              sum_\type\()_lag2_mid_neon   // 48
+        bl              sum_\type\()_lag2_mid_neon   // 56
+        bl              sum_\type\()_lag2_mid_neon   // 64
+        bl              sum_\type\()_lag2_mid_neon   // 72
+        bl              sum_\type\()_lag2_right_neon // 80
+        get_grain_2     v16
+        subs            w1,  w1,  #1
+.ifc \type, uv_444
+        add             x19, x19, #4
+.endif
+        st1             {v16.s}[0], [x0], #4
+        b.gt            1b
+
+        ldp             x30, x19, [sp], #96
+        ret
+
+L(generate_grain_\type\()_lag3):
+        ldr             q29,      [x4]      // ar_coeffs_y[0-15]
+        ldr             q30,      [x4, #16] // ar_coeffs_y[16-23], ar_coeffs_uv[16-24]
+        stp             d8,  d9,  [sp, #16]
+        stp             d10, d11, [sp, #32]
+        stp             d12, d13, [sp, #48]
+        stp             d14, d15, [sp, #64]
+        stp             x20, x21, [sp, #80]
+
+        smov            w4,  v30.b[5]
+        smov            w20, v30.b[6]
+        smov            w21, v30.b[7]
+
+        mov             w1,  #3
+        bl              generate_grain_rows_neon
+
+        mov             w1,  #GRAIN_HEIGHT - 3
+1:
+        bl              sum_\type\()_lag3_left_neon  // 8
+        bl              sum_\type\()_lag3_mid_neon   // 16
+        bl              sum_\type\()_lag3_mid_neon   // 24
+        bl              sum_\type\()_lag3_mid_neon   // 32
+        bl              sum_\type\()_lag3_mid_neon   // 40
+        bl              sum_\type\()_lag3_mid_neon   // 48
+        bl              sum_\type\()_lag3_mid_neon   // 56
+        bl              sum_\type\()_lag3_mid_neon   // 64
+        bl              sum_\type\()_lag3_mid_neon   // 72
+        bl              sum_\type\()_lag3_right_neon // 80
+        get_grain_2     v16
+        subs            w1,  w1,  #1
+.ifc \type, uv_444
+        add             x19, x19, #4
+.endif
+        st1             {v16.s}[0], [x0], #4
+        b.gt            1b
+
+        ldp             x20, x21, [sp, #80]
+        ldp             d14, d15, [sp, #64]
+        ldp             d12, d13, [sp, #48]
+        ldp             d10, d11, [sp, #32]
+        ldp             d8,  d9,  [sp, #16]
+        ldp             x30, x19, [sp], #96
+        ret
+
+L(gen_grain_\type\()_tbl):
+        .hword L(gen_grain_\type\()_tbl) - L(generate_grain_\type\()_lag0)
+        .hword L(gen_grain_\type\()_tbl) - L(generate_grain_\type\()_lag1)
+        .hword L(gen_grain_\type\()_tbl) - L(generate_grain_\type\()_lag2)
+        .hword L(gen_grain_\type\()_tbl) - L(generate_grain_\type\()_lag3)
+endfunc
+.endm
+
+gen_grain_82 y
+gen_grain_82 uv_444
+
+.macro set_height dst, type
+.ifc \type, uv_420
+        mov             \dst,  #SUB_GRAIN_HEIGHT-3
+.else
+        mov             \dst,  #GRAIN_HEIGHT-3
+.endif
+.endm
+
+.macro increment_y_ptr reg, type
+.ifc \type, uv_420
+        add             \reg, \reg, #2*GRAIN_WIDTH*2-(6*32)
+.else
+        sub             \reg, \reg, #6*32-GRAIN_WIDTH*2
+.endif
+.endm
+
+.macro gen_grain_44 type
+function generate_grain_\type\()_16bpc_neon, export=1
+        stp             x30, x19, [sp, #-96]!
+
+        mov             w13, w3
+        mov             w14, #28
+        add             x19, x1,  #(3*GRAIN_WIDTH-3)*2
+        mov             x1,  x2
+        mul             w13, w13, w14
+        clz             w15, w4
+
+        movrel          x3,  X(gaussian_sequence)
+        sub             w15, w15, #24 // -bitdepth_min_8
+        ldr             w2,  [x1, #FGD_SEED]
+        ldr             w9,  [x1, #FGD_GRAIN_SCALE_SHIFT]
+        add             x4,  x1,  #FGD_AR_COEFFS_UV
+        add             w9,  w9,  w15 // grain_scale_shift - bitdepth_min_8
+        adr             x16, L(gen_grain_\type\()_tbl)
+        ldr             w17, [x1, #FGD_AR_COEFF_LAG]
+        add             w9,  w9,  #4
+        ldrh            w17, [x16, w17, uxtw #1]
+        dup             v31.8h,  w9    // 4 - bitdepth_min_8 + data->grain_scale_shift
+        sub             x16, x16, w17, uxtw
+        neg             v31.8h,  v31.8h
+
+        cmp             w13, #0
+        mov             w11, #0x49d8
+        mov             w14, #0xb524
+        add             x4,  x4,  w13, uxtw // Add offset to ar_coeffs_uv[1]
+        csel            w11, w11, w14, ne
+
+        ldr             w7,  [x1, #FGD_AR_COEFF_SHIFT]
+        neg             w15, w15            // bitdepth_min_8
+        mov             w8,  #1
+        mov             w10, #1
+        lsl             w8,  w8,  w7        // 1 << ar_coeff_shift
+        lsl             w10, w10, w9        // 1 << (4 + data->grain_scale_shift)
+        lsr             w8,  w8,  #1        // 1 << (ar_coeff_shift - 1)
+        lsr             w10, w10, #1        // 1 << (4 + data->grain_scale_shift - 1)
+        mov             w5,  #128
+        lsl             w5,  w5,  w15       //   128 << bitdepth_min_8
+        neg             w6,  w5             // -(128 << bitpdeth_min_8)
+        sub             w5,  w5,  #1        //  (128 << bitdepth_min_8) - 1
+
+        eor             w2,  w2,  w11
+
+        br              x16
+
+L(generate_grain_\type\()_lag0):
+        dup             v28.4s,  w7
+        ld1r            {v27.8b}, [x4]      // ar_coeffs_uv[0]
+        movi            v0.16b,  #0
+        movi            v1.16b,  #255
+        dup             v25.8h,  w5
+        dup             v26.8h,  w6
+        ext             v29.16b, v0.16b,  v1.16b,  #10
+        ext             v30.16b, v1.16b,  v0.16b,  #14
+        neg             v28.4s,  v28.4s
+        sxtl            v27.8h,  v27.8b
+
+        mov             w1,  #3
+        bl              generate_grain_rows_44_neon
+        set_height      w1,  \type
+1:
+        mov             v1.16b,  v29.16b
+        bl              gen_grain_\type\()_lag0_8_neon // 8
+        movi            v1.16b,  #255
+        bl              gen_grain_\type\()_lag0_8_neon // 16
+        bl              gen_grain_\type\()_lag0_8_neon // 24
+        bl              gen_grain_\type\()_lag0_8_neon // 32
+        bl              gen_grain_\type\()_lag0_8_neon // 40
+        mov             v1.16b,  v30.16b
+        bl              gen_grain_\type\()_lag0_4_neon // 44
+        subs            w1,  w1,  #1
+        increment_y_ptr x19, \type
+        add             x0,  x0,  #GRAIN_WIDTH*2-6*16
+        b.gt            1b
+
+        ldp             x30, x19, [sp], #96
+        ret
+
+L(generate_grain_\type\()_lag1):
+        ld1r            {v27.8b}, [x4], #1  // ar_coeffs_uv[0]
+        ld1r            {v28.8b}, [x4], #1  // ar_coeffs_uv[1]
+        ld1r            {v29.8b}, [x4]      // ar_coeffs_uv[2]
+        add             x4,  x4,  #2
+
+        mov             w1,  #3
+        ld1r            {v30.8b}, [x4]      // ar_coeffs_u4[4]
+        ldursb          w4,  [x4, #-1]      // ar_coeffs_uv[3]
+        bl              generate_grain_rows_44_neon
+
+        sxtl            v27.8h,  v27.8b
+        sxtl            v28.8h,  v28.8b
+        sxtl            v29.8h,  v29.8b
+        sxtl            v30.8h,  v30.8b
+        set_height      w1,  \type
+1:
+        bl              sum_\type\()_lag1_left_neon  // 8
+        bl              sum_\type\()_lag1_mid_neon   // 16
+        bl              sum_\type\()_lag1_mid_neon   // 24
+        bl              sum_\type\()_lag1_mid_neon   // 32
+        bl              sum_\type\()_lag1_mid_neon   // 40
+        bl              sum_\type\()_lag1_right_neon // 44
+        subs            w1,  w1,  #1
+        increment_y_ptr x19, \type
+        add             x0,  x0,  #GRAIN_WIDTH*2-6*16
+        b.gt            1b
+
+        ldp             x30, x19, [sp], #96
+        ret
+
+L(generate_grain_\type\()_lag2):
+        ld1             {v30.16b}, [x4]     // ar_coeffs_uv[0-12]
+
+        smov            w4,  v30.b[10]
+        smov            w17, v30.b[11]
+
+        mov             w1,  #3
+        bl              generate_grain_rows_44_neon
+
+        set_height      w1,  \type
+1:
+        bl              sum_\type\()_lag2_left_neon  // 8
+        bl              sum_\type\()_lag2_mid_neon   // 16
+        bl              sum_\type\()_lag2_mid_neon   // 24
+        bl              sum_\type\()_lag2_mid_neon   // 32
+        bl              sum_\type\()_lag2_mid_neon   // 40
+        bl              sum_\type\()_lag2_right_neon // 44
+        subs            w1,  w1,  #1
+        increment_y_ptr x19, \type
+        add             x0,  x0,  #GRAIN_WIDTH*2-6*16
+        b.gt            1b
+
+        ldp             x30, x19, [sp], #96
+        ret
+
+L(generate_grain_\type\()_lag3):
+        ldr             q29,      [x4]      // ar_coeffs_uv[0-15]
+        ldr             q30,      [x4, #16] // ar_coeffs_uv[16-24]
+        stp             d8,  d9,  [sp, #16]
+        stp             d10, d11, [sp, #32]
+        stp             d12, d13, [sp, #48]
+        stp             d14, d15, [sp, #64]
+        stp             x20, x21, [sp, #80]
+
+        smov            w4,  v30.b[5]
+        smov            w20, v30.b[6]
+        smov            w21, v30.b[7]
+
+        mov             w1,  #3
+        bl              generate_grain_rows_44_neon
+
+        set_height      w1,  \type
+1:
+        bl              sum_\type\()_lag3_left_neon  // 8
+        bl              sum_\type\()_lag3_mid_neon   // 16
+        bl              sum_\type\()_lag3_mid_neon   // 24
+        bl              sum_\type\()_lag3_mid_neon   // 32
+        bl              sum_\type\()_lag3_mid_neon   // 40
+        bl              sum_\type\()_lag3_right_neon // 44
+        subs            w1,  w1,  #1
+        increment_y_ptr x19, \type
+        add             x0,  x0,  #GRAIN_WIDTH*2-6*16
+        b.gt            1b
+
+        ldp             x20, x21, [sp, #80]
+        ldp             d14, d15, [sp, #64]
+        ldp             d12, d13, [sp, #48]
+        ldp             d10, d11, [sp, #32]
+        ldp             d8,  d9,  [sp, #16]
+        ldp             x30, x19, [sp], #96
+        ret
+
+L(gen_grain_\type\()_tbl):
+        .hword L(gen_grain_\type\()_tbl) - L(generate_grain_\type\()_lag0)
+        .hword L(gen_grain_\type\()_tbl) - L(generate_grain_\type\()_lag1)
+        .hword L(gen_grain_\type\()_tbl) - L(generate_grain_\type\()_lag2)
+        .hword L(gen_grain_\type\()_tbl) - L(generate_grain_\type\()_lag3)
+endfunc
+.endm
+
+gen_grain_44 uv_420
+gen_grain_44 uv_422
 
 .macro gather_interleaved dst1, dst2, src1, src2, off
         umov            w14, \src1[0]
diff --git a/src/arm/film_grain_init_tmpl.c b/src/arm/film_grain_init_tmpl.c
index 5e81e7c..9f7d93b 100644
--- a/src/arm/film_grain_init_tmpl.c
+++ b/src/arm/film_grain_init_tmpl.c
@@ -209,7 +209,7 @@ COLD void bitfn(dav1d_film_grain_dsp_init_arm)(Dav1dFilmGrainDSPContext *const c
 
     if (!(flags & DAV1D_ARM_CPU_FLAG_NEON)) return;
 
-#if ARCH_AARCH64 && BITDEPTH == 8
+#if ARCH_AARCH64
     c->generate_grain_y = BF(dav1d_generate_grain_y, neon);
     c->generate_grain_uv[DAV1D_PIXEL_LAYOUT_I420 - 1] = BF(dav1d_generate_grain_uv_420, neon);
     c->generate_grain_uv[DAV1D_PIXEL_LAYOUT_I422 - 1] = BF(dav1d_generate_grain_uv_422, neon);