pan/lower_fb: Re-order components when dealing with raw formats

The output swizzle defined in the render-target descriptor is ignored
when the format is RAW. In that case, we have to swap the components
when lowering FB stores/loads if we want to get the right color.

Signed-off-by: Boris Brezillon <boris.brezillon@collabora.com>
Acked-by: Alyssa Rosenzweig <alyssa.rosenzweig@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/12793>

src/gallium/drivers/panfrost/pan_screen.h

@@ -90,7 +90,7 @@ struct panfrost_vtable {
         /* Shader compilation methods */
         const nir_shader_compiler_options *(*get_compiler_options)(void);
         void (*compile_shader)(nir_shader *s,
-                               const struct panfrost_compile_inputs *inputs,
+                               struct panfrost_compile_inputs *inputs,
                                struct util_dynarray *binary,
                                struct pan_shader_info *info);
 };

src/panfrost/lib/pan_shader.c

@@ -24,6 +24,7 @@
 
 #include "pan_device.h"
 #include "pan_shader.h"
+#include "pan_format.h"
 
 #if PAN_ARCH <= 5
 #include "panfrost/midgard/midgard_compile.h"
@@ -177,7 +178,7 @@ bifrost_blend_type_from_nir(nir_alu_type nir_type)
 
 void
 GENX(pan_shader_compile)(nir_shader *s,
-                         const struct panfrost_compile_inputs *inputs,
+                         struct panfrost_compile_inputs *inputs,
                          struct util_dynarray *binary,
                          struct pan_shader_info *info)
 {
@@ -186,6 +187,14 @@ GENX(pan_shader_compile)(nir_shader *s,
 #if PAN_ARCH >= 6
         bifrost_compile_shader_nir(s, inputs, binary, info);
 #else
+        for (unsigned i = 0; i < ARRAY_SIZE(inputs->rt_formats); i++) {
+                enum pipe_format fmt = inputs->rt_formats[i];
+                unsigned wb_fmt = panfrost_blendable_formats_v6[fmt].writeback;
+
+                if (wb_fmt <= MALI_MFBD_COLOR_FORMAT_RAW2048)
+                        inputs->raw_fmt_mask |= BITFIELD_BIT(i);
+        }
+
         midgard_compile_shader_nir(s, inputs, binary, info);
 #endif
 

src/panfrost/lib/pan_shader.h

@@ -39,7 +39,7 @@ GENX(pan_shader_get_compiler_options)(void);
 
 void
 GENX(pan_shader_compile)(nir_shader *nir,
-                         const struct panfrost_compile_inputs *inputs,
+                         struct panfrost_compile_inputs *inputs,
                          struct util_dynarray *binary,
                          struct pan_shader_info *info);
 

src/panfrost/midgard/midgard_compile.c

@@ -3088,7 +3088,8 @@ midgard_compile_shader_nir(nir_shader *nir,
 
         unsigned pan_quirks = panfrost_get_quirks(inputs->gpu_id, 0);
         NIR_PASS_V(nir, pan_lower_framebuffer,
-                   inputs->rt_formats, inputs->is_blend, pan_quirks);
+                   inputs->rt_formats, inputs->raw_fmt_mask,
+                   inputs->is_blend, pan_quirks);
 
         NIR_PASS_V(nir, nir_lower_io, nir_var_shader_in | nir_var_shader_out,
                         glsl_type_size, 0);

src/panfrost/util/pan_ir.h

@@ -131,6 +131,7 @@ struct panfrost_compile_inputs {
         bool no_ubo_to_push;
 
         enum pipe_format rt_formats[8];
+        uint8_t raw_fmt_mask;
         unsigned nr_cbufs;
 
         union {

src/panfrost/util/pan_lower_framebuffer.c

@@ -196,6 +196,36 @@ pan_unpack_pure_16(nir_builder *b, nir_ssa_def *pack, unsigned num_components)
         return nir_pad_vec4(b, nir_vec(b, unpacked, num_components));
 }
 
+static nir_ssa_def *
+pan_pack_reorder(nir_builder *b,
+                 const struct util_format_description *desc,
+                 nir_ssa_def *v)
+{
+        unsigned swizzle[4] = { 0, 1, 2, 3 };
+
+        for (unsigned i = 0; i < v->num_components; i++) {
+                if (desc->swizzle[i] <= PIPE_SWIZZLE_W)
+                        swizzle[i] = desc->swizzle[i];
+        }
+
+        return nir_swizzle(b, v, swizzle, v->num_components);
+}
+
+static nir_ssa_def *
+pan_unpack_reorder(nir_builder *b,
+                   const struct util_format_description *desc,
+                   nir_ssa_def *v)
+{
+        unsigned swizzle[4] = { 0, 1, 2, 3 };
+
+        for (unsigned i = 0; i < v->num_components; i++) {
+                if (desc->swizzle[i] <= PIPE_SWIZZLE_W)
+                        swizzle[desc->swizzle[i]] = i;
+        }
+
+        return nir_swizzle(b, v, swizzle, v->num_components);
+}
+
 static nir_ssa_def *
 pan_replicate_4(nir_builder *b, nir_ssa_def *v)
 {
@@ -475,10 +505,16 @@ pan_lower_fb_store(nir_shader *shader,
                 nir_builder *b,
                 nir_intrinsic_instr *intr,
                 const struct util_format_description *desc,
+                bool reorder_comps,
                 unsigned quirks)
 {
         /* For stores, add conversion before */
         nir_ssa_def *unpacked = nir_ssa_for_src(b, intr->src[1], 4);
+
+        /* Re-order the components */
+        if (reorder_comps)
+                unpacked = pan_pack_reorder(b, desc, unpacked);
+
         nir_ssa_def *packed = pan_pack(b, desc, unpacked);
 
         nir_store_raw_output_pan(b, packed);
@@ -495,6 +531,7 @@ pan_lower_fb_load(nir_shader *shader,
                 nir_builder *b,
                 nir_intrinsic_instr *intr,
                 const struct util_format_description *desc,
+                bool reorder_comps,
                 unsigned base, int sample, unsigned quirks)
 {
         nir_ssa_def *packed =
@@ -524,12 +561,16 @@ pan_lower_fb_load(nir_shader *shader,
         unpacked = nir_convert_to_bit_size(b, unpacked, src_type, bits);
         unpacked = nir_pad_vector(b, unpacked, nir_dest_num_components(intr->dest));
 
+        /* Reorder the components */
+        if (reorder_comps)
+                unpacked = pan_unpack_reorder(b, desc, unpacked);
+
         nir_ssa_def_rewrite_uses_after(&intr->dest.ssa, unpacked, &intr->instr);
 }
 
 bool
 pan_lower_framebuffer(nir_shader *shader, const enum pipe_format *rt_fmts,
-                      bool is_blend, unsigned quirks)
+                      uint8_t raw_fmt_mask, bool is_blend, unsigned quirks)
 {
         if (shader->info.stage != MESA_SHADER_FRAGMENT)
                return false;
@@ -579,16 +620,17 @@ pan_lower_framebuffer(nir_shader *shader, const enum pipe_format *rt_fmts,
                                  * MSAA blend shaders are not yet handled, so
                                  * for now always load sample 0. */
                                 int sample = is_blend ? 0 : -1;
+                                bool reorder_comps = raw_fmt_mask & BITFIELD_BIT(rt);
 
                                 nir_builder b;
                                 nir_builder_init(&b, func->impl);
 
                                 if (is_store) {
                                         b.cursor = nir_before_instr(instr);
-                                        pan_lower_fb_store(shader, &b, intr, desc, quirks);
+                                        pan_lower_fb_store(shader, &b, intr, desc, reorder_comps, quirks);
                                 } else {
                                         b.cursor = nir_after_instr(instr);
-                                        pan_lower_fb_load(shader, &b, intr, desc, base, sample, quirks);
+                                        pan_lower_fb_load(shader, &b, intr, desc, reorder_comps, base, sample, quirks);
                                 }
 
                                 nir_instr_remove(instr);

src/panfrost/util/pan_lower_framebuffer.h

@@ -41,7 +41,9 @@ enum pan_format_class {
 
 nir_alu_type pan_unpacked_type_for_format(const struct util_format_description *desc);
 
-bool pan_lower_framebuffer(nir_shader *shader, const enum pipe_format *rt_fmts,
+bool pan_lower_framebuffer(nir_shader *shader,
+                           const enum pipe_format *rt_fmts,
+                           uint8_t raw_fmt_mask,
                            bool is_blend, unsigned quirks);
 
 #endif