WIP: dav1d ITX AVX2 Assembly Port

libavcodec/vvc/vvc_ctu.c

@@ -268,6 +268,8 @@ static TransformBlock* add_tb(TransformUnit *tu, VVCLocalContext *lc,
     tb->ts = 0;
     tb->coeffs = lc->coeffs;
     lc->coeffs += tb_width * tb_height;
+    tb->pixels = lc->pixels;
+    lc->pixels += tb_width * tb_height;
     return tb;
 }
 
@@ -2382,6 +2384,7 @@ int ff_vvc_coding_tree_unit(VVCLocalContext *lc,
     }
 
     lc->coeffs = fc->tab.coeffs + rs * ctb_size * VVC_MAX_SAMPLE_ARRAYS;
+    lc->pixels = fc->tab.pixels + rs * ctb_size * VVC_MAX_SAMPLE_ARRAYS;
     lc->cu     = NULL;
 
     ff_vvc_cabac_init(lc, ctu_idx, rx, ry);
@@ -2474,4 +2477,4 @@ void ff_vvc_ep_init_stat_coeff(EntryPoint *ep,
         ep->stat_coeff[i] =
             persistent_rice_adaptation_enabled_flag ? 2 * (av_log2(bit_depth - 10)) : 0;
     }
-}
+}

libavcodec/vvc/vvc_ctu.h

@@ -96,6 +96,7 @@ typedef struct TransformBlock {
     int bd_offset;
 
     int *coeffs;
+    int16_t *pixels;
 } TransformBlock;
 
 typedef enum VVCTreeType {
@@ -369,6 +370,7 @@ struct VVCLocalContext {
     VVCFrameContext *fc;
     EntryPoint *ep;
     int *coeffs;
+    int16_t *pixels;
 } ;
 
 typedef struct VVCAllowedSplit {

libavcodec/vvc/vvc_intra.c

@@ -176,11 +176,11 @@ static void add_residual_for_joint_coding_chroma(VVCLocalContext *lc,
     uint8_t *dst = &fc->frame->data[c_idx][(tb->y0 >> vs) * stride +
                                           ((tb->x0 >> hs) << fc->ps.sps->pixel_shift)];
     if (chroma_scale) {
-        fc->vvcdsp.itx.pred_residual_joint(tb->coeffs, tb->tb_width, tb->tb_height, c_sign, shift);
-        fc->vvcdsp.intra.lmcs_scale_chroma(lc, tb->coeffs, tb->coeffs, tb->tb_width, tb->tb_height, cu->x0, cu->y0);
-        fc->vvcdsp.itx.add_residual(dst, tb->coeffs, tb->tb_width, tb->tb_height, stride);
+        fc->vvcdsp.itx.pred_residual_joint(tb->pixels, tb->tb_width, tb->tb_height, c_sign, shift);
+        fc->vvcdsp.intra.lmcs_scale_chroma(lc, tb->pixels, tb->pixels, tb->tb_width, tb->tb_height, cu->x0, cu->y0);
+        fc->vvcdsp.itx.add_residual(dst, tb->pixels, tb->tb_width, tb->tb_height, stride);
     } else {
-        fc->vvcdsp.itx.add_residual_joint(dst, tb->coeffs, tb->tb_width, tb->tb_height, stride, c_sign, shift);
+        fc->vvcdsp.itx.add_residual_joint(dst, tb->pixels, tb->tb_width, tb->tb_height, stride, c_sign, shift);
     }
 }
 
@@ -272,32 +272,6 @@ static void predict_intra(VVCLocalContext *lc, const TransformUnit *tu, const in
     }
 }
 
-static void scale_clip(int *coeff, const int nzw, const int w, const int h,
-    const int shift, const int log2_transform_range)
-{
-    const int add = 1 << (shift - 1);
-    for (int y = 0; y < h; y++) {
-        int *p = coeff + y * w;
-        for (int x = 0; x < nzw; x++) {
-            *p = av_clip_intp2((*p + add) >> shift, log2_transform_range);
-            p++;
-        }
-        memset(p, 0, sizeof(*p) * (w - nzw));
-    }
-}
-
-static void scale(int *out, const int *in, const int w, const int h, const int shift)
-{
-    const int add = 1 << (shift - 1);
-    for (int y = 0; y < h; y++) {
-        for (int x = 0; x < w; x++) {
-            int *o = out + y * w + x;
-            const int *i = in + y * w + x;
-            *o = (*i + add) >> shift;
-        }
-    }
-}
-
 // part of 8.7.3 Scaling process for transform coefficients
 static void derive_qp(const VVCLocalContext *lc, const TransformUnit *tu, TransformBlock *tb)
 {
@@ -441,35 +415,6 @@ static void dequant(const VVCLocalContext *lc, const TransformUnit *tu, Transfor
     }
 }
 
-static void itx_2d(const VVCFrameContext *fc, TransformBlock *tb, const enum TxType trh, const enum TxType trv, int *temp)
-{
-    const VVCSPS *sps   = fc->ps.sps;
-    const int w         = tb->tb_width;
-    const int h         = tb->tb_height;
-    const int nzw       = tb->max_scan_x + 1;
-
-    for (int x = 0; x < nzw; x++)
-        fc->vvcdsp.itx.itx[trv][tb->log2_tb_height - 1](temp + x, w, tb->coeffs + x, w);
-    scale_clip(temp, nzw, w, h, 7, sps->log2_transform_range);
-
-    for (int y = 0; y < h; y++)
-        fc->vvcdsp.itx.itx[trh][tb->log2_tb_width - 1](tb->coeffs + y * w, 1, temp + y * w, 1);
-    scale(tb->coeffs, tb->coeffs, w, h, 5 + sps->log2_transform_range - sps->bit_depth);
-}
-
-static void itx_1d(const VVCFrameContext *fc, TransformBlock *tb, const enum TxType trh, const enum TxType trv, int  *temp)
-{
-    const VVCSPS *sps   = fc->ps.sps;
-    const int w         = tb->tb_width;
-    const int h         = tb->tb_height;
-
-    if (w > 1)
-        fc->vvcdsp.itx.itx[trh][tb->log2_tb_width - 1](temp, 1, tb->coeffs, 1);
-    else
-        fc->vvcdsp.itx.itx[trv][tb->log2_tb_height - 1](temp, 1, tb->coeffs, 1);
-    scale(tb->coeffs, temp, w, h, 6 + sps->log2_transform_range - sps->bit_depth);
-}
-
 static void transform_bdpcm(TransformBlock *tb, const VVCLocalContext *lc, const CodingUnit *cu)
 {
     const VVCSPS *sps        = lc->fc->ps.sps;
@@ -490,7 +435,7 @@ static void itransform(VVCLocalContext *lc, TransformUnit *tu, const int tu_idx,
     const VVCSH *sh             = &lc->sc->sh;
     const CodingUnit *cu        = lc->cu;
     const int ps                = fc->ps.sps->pixel_shift;
-    DECLARE_ALIGNED(32, int, temp)[MAX_TB_SIZE * MAX_TB_SIZE];
+    DECLARE_ALIGNED(32, int16_t, temp)[MAX_TB_SIZE * MAX_TB_SIZE];
 
     for (int i = 0; i < tu->nb_tbs; i++) {
         TransformBlock *tb  = &tu->tbs[i];
@@ -511,19 +456,26 @@ static void itransform(VVCLocalContext *lc, TransformUnit *tu, const int tu_idx,
             dequant(lc, tu, tb);
             if (!tb->ts) {
                 enum TxType trh, trv;
+                int nzw;
 
                 if (cu->apply_lfnst_flag[c_idx])
                     ilfnst_transform(lc, tb);
                 derive_transform_type(fc, lc, tb, &trh, &trv);
-                if (w > 1 && h > 1)
-                    itx_2d(fc, tb, trh, trv, temp);
-                else
-                    itx_1d(fc, tb, trh, trv, temp);
+
+                nzw = tb->max_scan_x + 1;
+                fc->vvcdsp.itx.itx[trh][trv][tb->log2_tb_width][tb->log2_tb_height](
+                        tb->pixels, tb->coeffs, nzw, sps->log2_transform_range);
+            } else {
+                for (int x = 0; x < w; ++x) {
+                    for (int y = 0; y < h; ++y) {
+                        tb->pixels[x * h + y] = tb->coeffs[y * w + x];
+                    }
+                }
             }
 
             if (chroma_scale)
-                fc->vvcdsp.intra.lmcs_scale_chroma(lc, temp, tb->coeffs, w, h, cu->x0, cu->y0);
-            fc->vvcdsp.itx.add_residual(dst, chroma_scale ? temp : tb->coeffs, w, h, stride);
+                fc->vvcdsp.intra.lmcs_scale_chroma(lc, temp, tb->pixels, w, h, cu->x0, cu->y0);
+            fc->vvcdsp.itx.add_residual(dst, chroma_scale ? temp : tb->pixels, w, h, stride);
 
             if (tu->joint_cbcr_residual_flag && tb->c_idx)
                 add_residual_for_joint_coding_chroma(lc, tu, tb, chroma_scale);

libavcodec/vvc/vvc_intra_template.c

@@ -429,7 +429,7 @@ static int FUNC(lmcs_derive_chroma_scale)(VVCLocalContext *lc, const int x0, con
 }
 
 // 8.7.5.3 Picture reconstruction with luma dependent chroma residual scaling process for chroma samples
-static void FUNC(lmcs_scale_chroma)(VVCLocalContext *lc, int *dst, const int *coeff,
+static void FUNC(lmcs_scale_chroma)(VVCLocalContext *lc, int16_t *dst, const int16_t *coeff,
     const int width, const int height, const int x0_cu, const int y0_cu)
 {
     const int chroma_scale = FUNC(lmcs_derive_chroma_scale)(lc, x0_cu, y0_cu);

libavcodec/vvc/vvcdec.c

@@ -98,6 +98,9 @@ static int ctb_arrays_init(VVCFrameContext *fc, const int ctu_count, const int c
         fc->tab.coeffs = av_malloc(ctu_count * sizeof(*fc->tab.coeffs) * ctu_size * VVC_MAX_SAMPLE_ARRAYS);
         if (!fc->tab.coeffs)
             return AVERROR(ENOMEM);
+        fc->tab.pixels = av_malloc(ctu_count * sizeof(*fc->tab.pixels) * ctu_size * VVC_MAX_SAMPLE_ARRAYS);
+        if (!fc->tab.pixels)
+            return AVERROR(ENOMEM);
         fc->rpl_tab_pool = av_buffer_pool_init(ctu_count * sizeof(RefPicListTab), av_buffer_allocz);
         if (!fc->rpl_tab_pool)
             return AVERROR(ENOMEM);

libavcodec/vvc/vvcdec.h

@@ -252,6 +252,7 @@ struct VVCFrameContext {
         uint8_t *alf_pixel_buffer_v[VVC_MAX_SAMPLE_ARRAYS][2];
 
         int     *coeffs;
+        int16_t *pixels;
         CTU     *ctus;
 
         //used in arrays_init only

libavcodec/vvc/vvcdsp.c

@@ -254,23 +254,31 @@ static int vvc_sad(const int16_t *src0, const int16_t *src1, int dx, int dy,
     return sad;
 }
 
-#define itx_fn(type, s)                                                                         \
-static void itx_##type##_##s(int *out, ptrdiff_t out_step, const int *in, ptrdiff_t in_step)    \
-{                                                                                               \
-  ff_vvc_inv_##type##_##s(out, out_step, in, in_step);                                          \
+static void scale_clip(int *coeff, const int nzw, const int w, const int h,
+    const int shift, const int log2_transform_range)
+{
+    const int add = 1 << (shift - 1);
+    for (int y = 0; y < h; y++) {
+        int *p = coeff + y * w;
+        for (int x = 0; x < nzw; x++) {
+            *p = av_clip_intp2((*p + add) >> shift, log2_transform_range);
+            p++;
+        }
+        memset(p, 0, sizeof(*p) * (w - nzw));
+    }
 }
 
-#define itx_fn_common(type) \
-    itx_fn(type, 4);        \
-    itx_fn(type, 8);        \
-    itx_fn(type, 16);       \
-    itx_fn(type, 32);       \
-
-itx_fn_common(dct2);
-itx_fn_common(dst7);
-itx_fn_common(dct8);
-itx_fn(dct2, 2);
-itx_fn(dct2, 64);
+static void scale(int16_t *out, const int *in, const int w, const int h, const int shift)
+{
+    const int add = 1 << (shift - 1);
+    for (int y = 0; y < h; y++) {
+        for (int x = 0; x < w; x++) {
+            int16_t *o = out + y * w + x;
+            const int *i = in + y * w + x;
+            *o = (*i + add) >> shift;
+        }
+    }
+}
 
 typedef struct IntraEdgeParams {
     uint8_t* top;

libavcodec/vvc/vvcdsp.h

@@ -35,6 +35,7 @@ enum TxType {
 };
 
 enum TxSize {
+    TX_SIZE_1 = 0,
     TX_SIZE_2,
     TX_SIZE_4,
     TX_SIZE_8,
@@ -101,7 +102,7 @@ struct VVCLocalContext;
 
 typedef struct VVCIntraDSPContext {
     void (*intra_cclm_pred)(const struct VVCLocalContext *lc, int x0, int y0, int w, int h);
-    void (*lmcs_scale_chroma)(struct VVCLocalContext *lc, int *dst, const int *coeff, int w, int h, int x0_cu, int y0_cu);
+    void (*lmcs_scale_chroma)(struct VVCLocalContext *lc, int16_t *dst, const int16_t *coeff, int w, int h, int x0_cu, int y0_cu);
     void (*intra_pred)(const struct VVCLocalContext *lc, int x0, int y0, int w, int h, int c_idx);
     void (*pred_planar)(uint8_t *src, const uint8_t *top, const uint8_t *left, int w, int h, ptrdiff_t stride);
     void (*pred_mip)(uint8_t *src, const uint8_t *top, const uint8_t *left, int w, int h, ptrdiff_t stride,
@@ -116,11 +117,11 @@ typedef struct VVCIntraDSPContext {
 } VVCIntraDSPContext;
 
 typedef struct VVCItxDSPContext {
-    void (*add_residual)(uint8_t *dst, const int *res, int width, int height, ptrdiff_t stride);
-    void (*add_residual_joint)(uint8_t *dst, const int *res, int width, int height, ptrdiff_t stride, int c_sign, int shift);
-    void (*pred_residual_joint)(int *buf, int width, int height, int c_sign, int shift);
+    void (*add_residual)(uint8_t *dst, const int16_t *res, int width, int height, ptrdiff_t stride);
+    void (*add_residual_joint)(uint8_t *dst, const int16_t *res, int width, int height, ptrdiff_t stride, int c_sign, int shift);
+    void (*pred_residual_joint)(int16_t *buf, int width, int height, int c_sign, int shift);
 
-    void (*itx[N_TX_TYPE][N_TX_SIZE])(int *out, ptrdiff_t out_step, const int *in, ptrdiff_t in_step);
+    void (*itx[N_TX_TYPE][N_TX_TYPE][N_TX_SIZE][N_TX_SIZE])(int16_t *dst, const int *coeff, int nzw, int log2_transform_range);
     void (*transform_bdpcm)(int *coeffs, int width, int height, int vertical, int log2_transform_range);
 } VVCItxDSPContext;