pan/midgard: Handle get/set_swizzle for load/store arguments

Load/store's  main "argument 0" already has its swizzle handled
correctly (for stores, that is). But the tinier arguments, the compact
ones with a component select but not a full swizzle, those are not yet
handled. Let's do something about that!

2 files changed, 83 insertions, 3 deletions

diff --git a/src/panfrost/midgard/helpers.h b/src/panfrost/midgard/helpers.h
index c44783e3d68..f1f502372d8 100644
--- a/src/panfrost/midgard/helpers.h
+++ b/src/panfrost/midgard/helpers.h
@@ -358,4 +358,41 @@ midgard_ldst_reg(unsigned reg, unsigned component)
         return packed;
 }
 
+/* Unpacks a load/store argument */
+
+static inline midgard_ldst_register_select
+midgard_ldst_select(uint8_t u)
+{
+        midgard_ldst_register_select sel;
+        memcpy(&sel, &u, sizeof(u));
+        return sel;
+}
+
+static inline uint8_t
+midgard_ldst_pack(midgard_ldst_register_select sel)
+{
+        uint8_t packed;
+        memcpy(&packed, &sel, sizeof(packed));
+        return packed;
+}
+
+/* Gets a swizzle like yyyy and returns y */
+
+static inline unsigned
+swizzle_to_component(unsigned swizzle)
+{
+        unsigned c = swizzle & 3;
+        assert(((swizzle >> 2) & 3) == c);
+        assert(((swizzle >> 4) & 3) == c);
+        assert(((swizzle >> 6) & 3) == c);
+        return  c;
+}
+
+
+static inline unsigned
+component_to_swizzle(unsigned c)
+{
+        return SWIZZLE(c, c, c, c);
+}
+
 #endif
diff --git a/src/panfrost/midgard/mir.c b/src/panfrost/midgard/mir.c
index 045948a74b5..0198430953b 100644
--- a/src/panfrost/midgard/mir.c
+++ b/src/panfrost/midgard/mir.c
@@ -44,8 +44,23 @@ mir_get_swizzle(midgard_instruction *ins, unsigned idx)
 
                 return s.swizzle;
         } else if (ins->type == TAG_LOAD_STORE_4) {
-                assert(idx == 0);
-                return ins->load_store.swizzle;
+                /* Main swizzle of a load is on the destination */
+                if (!OP_IS_STORE(ins->load_store.op))
+                        idx++;
+
+                switch (idx) {
+                case 0:
+                        return ins->load_store.swizzle;
+                case 1:
+                case 2: {
+                        uint8_t raw =
+                                (idx == 2) ? ins->load_store.arg_2 : ins->load_store.arg_1;
+
+                        return component_to_swizzle(midgard_ldst_select(raw).component);
+                }
+                default:
+                        unreachable("Unknown load/store source");
+                }
         } else if (ins->type == TAG_TEXTURE_4) {
                 switch (idx) {
                 case 0:
@@ -78,7 +93,35 @@ mir_set_swizzle(midgard_instruction *ins, unsigned idx, unsigned new)
                 else
                         ins->alu.src2 = pack;
         } else if (ins->type == TAG_LOAD_STORE_4) {
-                ins->load_store.swizzle = new;
+                /* Main swizzle of a load is on the destination */
+                if (!OP_IS_STORE(ins->load_store.op))
+                        idx++;
+
+                switch (idx) {
+                case 0:
+                        ins->load_store.swizzle = new;
+                        break;
+                case 1:
+                case 2: {
+                        uint8_t raw =
+                                (idx == 2) ? ins->load_store.arg_2 : ins->load_store.arg_1;
+
+                        midgard_ldst_register_select sel
+                                = midgard_ldst_select(raw);
+                        sel.component = swizzle_to_component(new);
+                        uint8_t packed = midgard_ldst_pack(sel);
+
+                        if (idx == 2)
+                                ins->load_store.arg_2 = packed;
+                        else
+                                ins->load_store.arg_1 = packed;
+
+                        break;
+                }
+                default:
+                        assert(new == 0);
+                        break;
+                }
         } else if (ins->type == TAG_TEXTURE_4) {
                 switch (idx) {
                 case 0: