[llvm] [AArch64] Fix correctness of FORM_TRANSPOSED_REG_TUPLE expansions (PR #205528)

[Benjamin Maxwell via llvm-commits]

================
@@ -1282,26 +1282,114 @@ bool AArch64ExpandPseudoImpl::expandMultiVecPseudo(
   return true;
 }
 
+struct Copy {
+  Register Dst;
+  Register Src;
+  int SrcTupleIdx = -1;
+};
+
 bool AArch64ExpandPseudoImpl::expandFormTuplePseudo(
     MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
     MachineBasicBlock::iterator &NextMBBI, unsigned Size) {
   assert((Size == 2 || Size == 4) && "Invalid Tuple Size");
-  MachineInstr &MI = *MBBI;
-  Register ReturnTuple = MI.getOperand(0).getReg();
 
   const TargetRegisterInfo *TRI =
       MBB.getParent()->getSubtarget().getRegisterInfo();
+
+  MachineInstr &MI = *MBBI;
+  DebugLoc DL = MI.getDebugLoc();
+  Register Tuple = MI.getOperand(0).getReg();
+
+  // Collect the copies required to form the tuple. Count if any tuple members
+  // appear as operands to the FORM_TRANSPOSED_REG_TUPLE. If a destination
+  // also appears as a source operand, we can't write to that register until
+  // we've handled the other copies that use it.
+
+  Copy Copies[4] = {};
+  unsigned Users[4] = {};
+
+  for (unsigned I = 0; I < Size; ++I) {
+    Register Dst = TRI->getSubReg(Tuple, AArch64::zsub0 + I);
+    Register Src = MI.getOperand(I + 1).getReg();
+
+    Copies[I] = {Dst, Src};
+
+    if (Dst != Src && TRI->isSubRegister(Tuple, Src)) {
+      int SrcTupleIdx = TRI->getSubRegIndex(Tuple, Src) - AArch64::zsub0;
+      Copies[I].SrcTupleIdx = SrcTupleIdx;
+      ++Users[SrcTupleIdx];
----------------
MacDue wrote:

I don't think so. 

>i.e, where the source register is in the destination tuple and the destination sub-reg is in source tuple

This does not imply a swap is required (this exists in all cases that can be reordered). See `form_transposed_reg_tuple_x4_overlap`. There's three copies where the source/dest are both tuple members, but they can be scheduled. 

What this is doing with the `Users` counts is building a small graph of copies. The copies can be scheduled without swaps if that graph is a DAG, i.e. there's no cycles in the graph. This is what the first half the the lowering does.

If there's a cycle in the graph, e.g. a simple swap: `(A <- B, B <- A) `, then we have to emit that using swaps. A cycle is a permutation.

This gets a little mathematical, but a permutation e.g., `perm = [2, 3, 1, 0]` (where `result[i] = src[perm[i]])`,  is one or more cycles.  In this case, `0 -> 2 -> 1 -> 3` (to get this, we see perm[0] goes to 2, then perm[2] goes to 1, and so on until we've back at 0, which happens at perm[3]). 

To apply decompose a cycle into swaps, you go right to left. So in this case, the cycle is `0 -> 2 -> 1 -> 3`:

```
tup = src // (keep a copy of the original tuple for reference)
swap tup[0], tup[3] // tup[3] = src[0], tup[0] = src[3]
swap tup[0], tup[1] // tup[1] = src[3], tup[0] = src[1]
swap tup[0], tup[2] // tup[2] = src[1], tup[0] = src[2]
```

Note that at each swap, `tup[0]` contains the final value for the higher index. This works because it's a cycle. Once we've swapped `tup[0]` and `tup[3]`, `tup[0]` now contains the value for the next step in the cycle. 


https://github.com/llvm/llvm-project/pull/205528