[Z80] Fix legalization of non-power-of-2 and large merged types

Sightem · Sightem · commit d7d42ee50632 · 2025-12-27T17:36:12.000-08:00
fix two legalization issues:

-  non-power-of-2 types (like s9) produced by -Oz optimizations caused
   crashes. the optimizer's closed form summation transformation
   (n*(n-1)/2) widens 8 bit inputs to 9 bit to prevent overflow,
   but the legalizer had no rule to handle these types.

-  G_MERGE_VALUES with s96 destination (4×s24) failed because only
   s32/s48/s64 were in LegalLargeScalars. this was seemingly a latent bug that
   became exposed once the s9 widening fix allowed legalization to
   proceed further.

The s9 crash also affected v15, not just v17
diff --git a/llvm/lib/Target/Z80/GISel/Z80LegalizerInfo.cpp b/llvm/lib/Target/Z80/GISel/Z80LegalizerInfo.cpp
@@ -106,6 +106,17 @@ Z80LegalizerInfo::Z80LegalizerInfo(const Z80Subtarget &STI,
       .legalForCartesianProduct(NotMin, NotMax)
       // allow creating s32/s48/s64 from legal scalar sources (needed for libcall args)
       .legalForCartesianProduct(LegalLargeScalars, LegalScalars)
+      // allow merging legal scalars into larger types (handles s72, s96, s128, etc.)
+      .legalIf([=](const LegalityQuery &Q) {
+        LLT DstTy = Q.Types[0];
+        LLT SrcTy = Q.Types[1];
+        if (!SrcTy.isScalar() || !DstTy.isScalar())
+          return false;
+        if (DstTy.getSizeInBits() % SrcTy.getSizeInBits() != 0)
+          return false;
+        unsigned SrcSize = SrcTy.getSizeInBits();
+        return SrcSize == 8 || SrcSize == 16 || (Is24Bit && SrcSize == 24);
+      })
       .customIf([=](const LegalityQuery &Q) {
         // custom legalize G_MERGE_VALUES when source type is s1, expand into shifts+ORs
         return Q.Types[1] == s1;
@@ -117,6 +128,17 @@ Z80LegalizerInfo::Z80LegalizerInfo(const Z80Subtarget &STI,
       .legalForCartesianProduct(NotMax, NotMin)
       // allow splitting wide scalars into legal scalar parts (s32->s16, s48->s24, s64->s16)
       .legalForCartesianProduct(LegalScalars, LegalLargeScalars)
+      // allow unmerging larger types into legal scalar parts (handles s72, s96, s128, etc.)
+      .legalIf([=](const LegalityQuery &Q) {
+        LLT DstTy = Q.Types[0];
+        LLT SrcTy = Q.Types[1];
+        if (!SrcTy.isScalar() || !DstTy.isScalar())
+          return false;
+        if (SrcTy.getSizeInBits() % DstTy.getSizeInBits() != 0)
+          return false;
+        unsigned DstSize = DstTy.getSizeInBits();
+        return DstSize == 8 || DstSize == 16 || (Is24Bit && DstSize == 24);
+      })
       .clampScalar(1, *NotMin.begin(), *std::prev(NotMin.end()))
       .clampScalar(0, *NotMax.begin(), *std::prev(NotMax.end()));
 
@@ -126,6 +148,8 @@ Z80LegalizerInfo::Z80LegalizerInfo(const Z80Subtarget &STI,
 
   getActionDefinitionsBuilder({G_ZEXT, G_ANYEXT})
       .legalForCartesianProduct(LegalScalars, NotMaxWithOne)
+      .widenScalarToNextPow2(0, 8)
+      .widenScalarToNextPow2(1, 1)
       .clampScalar(0, *LegalScalars.begin(), *std::prev(LegalScalars.end()))
       .clampScalar(1, *NotMaxWithOne.begin(), *std::prev(NotMaxWithOne.end()));
 
@@ -174,11 +198,14 @@ Z80LegalizerInfo::Z80LegalizerInfo(const Z80Subtarget &STI,
 
   getActionDefinitionsBuilder(G_TRUNC)
       .legalForCartesianProduct(NotMaxWithOne, LegalScalars)
+      .widenScalarToNextPow2(0, 1)
+      .widenScalarToNextPow2(1, 8)
       .clampScalar(1, *LegalScalars.begin(), *std::prev(LegalScalars.end()))
       .clampScalar(0, *NotMaxWithOne.begin(), *std::prev(NotMaxWithOne.end()));
 
   getActionDefinitionsBuilder({G_FREEZE, G_PHI, G_CONSTANT})
       .legalFor(LegalTypes)
+      .widenScalarToNextPow2(0, 8)
       .clampScalar(0, s8, sMax);
 
   getActionDefinitionsBuilder(G_FCONSTANT)
@@ -219,6 +246,7 @@ Z80LegalizerInfo::Z80LegalizerInfo(const Z80Subtarget &STI,
   getActionDefinitionsBuilder(G_MUL)
       .legalIf(all(predZ180Ops, typeIs(0, s8)))
       .libcallFor(LegalLibcallScalars)
+      .widenScalarToNextPow2(0, 8)
       .minScalar(0, s8)
       .minScalar(0, s16)
       .minScalarIf(pred24Bit, 0, s24)
@@ -239,6 +267,7 @@ Z80LegalizerInfo::Z80LegalizerInfo(const Z80Subtarget &STI,
   getActionDefinitionsBuilder({G_SHL, G_LSHR, G_ASHR})
       .customForCartesianProduct(LegalLibcallScalars, {s8})
       .clampScalar(1, s8, s8)
+      .widenScalarToNextPow2(0, 8)
       .minScalar(0, s8)
       .minScalar(0, s16)
       .minScalarIf(pred24Bit, 0, s24)
diff --git a/llvm/lib/Target/Z80/Z80MachineLateOptimization.cpp b/llvm/lib/Target/Z80/Z80MachineLateOptimization.cpp
@@ -142,6 +142,11 @@ class Z80MachineLateOptimization : public MachineFunctionPass {
   // after SBC A,A: A is a member of {0, 0xFF} and Z = (A == 0), so OR A,A is redundant
   bool ZFlagReflectsAZero = false;
 
+  // track redundant A<->reg copies, after "ld X, a", X mirrors A
+  // if A hasnt changed when we see "ld a, X", we can eliminate the copy
+  // AMirroredInReg = X means the value currently in A is also in X
+  MCRegister AMirroredInReg = Z80::NoRegister;
+
   template <typename... Args> void assign(MCRegister Reg, Args &&...args) {
     RegVals[Reg] = RegVal(std::forward<Args>(args)..., Reg, *TRI);
   }
@@ -452,6 +457,7 @@ bool Z80MachineLateOptimization::runOnMachineFunction(MachineFunction &MF) {
     LiveUnits.addLiveIns(MBB);
     clobberAll();
     ZFlagReflectsAZero = false;
+    AMirroredInReg = Z80::NoRegister;
     for (MachineBasicBlock::iterator I = MBB.begin(), E = MBB.end(); I != E;) {
       MachineInstrBuilder MIB(MF, I);
       LiveUnits.stepForward(*I);
@@ -624,6 +630,45 @@ bool Z80MachineLateOptimization::runOnMachineFunction(MachineFunction &MF) {
           continue;
         }
         break;
+      // redundant copy elimination: track COPY from A and eliminate COPY back if A unchanged
+      // this runs before pseudo expansion, so we see COPY, not LD8gg/xx/yy
+      case TargetOpcode::COPY: {
+        MCRegister CopyDst = MIB->getOperand(0).getReg();
+        MCRegister CopySrc = MIB->getOperand(1).getReg();
+        // only handle 8-bit A-related copies
+        if (!Z80::R8RegClass.contains(CopyDst) || !Z80::R8RegClass.contains(CopySrc))
+          break;
+        // if copying TO A from the register that mirrors A, the copy is redundant
+        if (CopyDst == Z80::A && CopySrc == AMirroredInReg) {
+          LLVM_DEBUG(dbgs() << "Erasing redundant copy to A from " 
+                            << TRI->getName(CopySrc) << " (A mirrors "
+                            << TRI->getName(AMirroredInReg) << "): ";
+                     MIB->dump());
+          MIB->eraseFromParent();
+          Changed = true;
+          continue;
+        }
+        // if copying FROM A to X, record that X now mirrors A
+        if (CopySrc == Z80::A) {
+          AMirroredInReg = CopyDst;
+        }
+        break;
+      }
+      // 8 bit ALU pseudos: ADD8_gisel, SUB8_gisel, etc. have an explicit 8 bit def
+      // and implicit def A (because Z80 ALU ops always write to A)
+      // after these instructions, A and the dest register hold the same value
+      case Z80::ADD8_gisel:
+      case Z80::SUB8_gisel:
+      case Z80::AND8_gisel:
+      case Z80::OR8_gisel:
+      case Z80::XOR8_gisel: {
+        MCRegister ExplicitDst = MIB->getOperand(0).getReg();
+        // after ALU pseudo, explicit dest and A have the same value if dest is not A itself, record that dest mirrors A
+        if (ExplicitDst != Z80::A) {
+          AMirroredInReg = ExplicitDst;
+        }
+        break;
+      }
       case Z80::Sub16ao:
       case Z80::Sub24ao:
       case Z80::Cmp16ao:
@@ -734,16 +779,46 @@ bool Z80MachineLateOptimization::runOnMachineFunction(MachineFunction &MF) {
         }
       }
       
-      // after SBC A,A, Z flag correctly reflects (A == 0)
-      // track this so we can eliminate redundant OR A,A
+      // SBC A,A sets Z=(A==0). INC/DEC A preserve this. Track to eliminate redundant OR A,A
       unsigned Opc = MIB->getOpcode();
       if (Opc == Z80::SBC8ar && MIB->getOperand(0).getReg() == Z80::A) {
         ZFlagReflectsAZero = true;
+      } else if (ZFlagReflectsAZero && 
+                 (Opc == Z80::INC8r || Opc == Z80::DEC8r) &&
+                 MIB->getOperand(0).getReg() == Z80::A) {
+        // INC/DEC A preserves Z=(A==0)
       } else if (ClobberedA || ClobberedF) {
-        // if A or F is modified by something other than SBC A,A, reset tracking
         ZFlagReflectsAZero = false;
       }
 
+      // invalidate A<->reg mirroring if A or mirrored reg is clobbered (except by tracked ops)
+      if (ClobberedA) {
+        bool isCopyFromA = Opc == TargetOpcode::COPY &&
+                           MIB->getOperand(1).getReg() == Z80::A;
+        bool isALUPseudo = (Opc == Z80::ADD8_gisel || Opc == Z80::SUB8_gisel ||
+                            Opc == Z80::AND8_gisel || Opc == Z80::OR8_gisel ||
+                            Opc == Z80::XOR8_gisel);
+        if (!isCopyFromA && !isALUPseudo)
+          AMirroredInReg = Z80::NoRegister;
+      }
+      if (AMirroredInReg != Z80::NoRegister) {
+        bool isALUPseudoOrCopy = (Opc == Z80::ADD8_gisel || Opc == Z80::SUB8_gisel ||
+                                  Opc == Z80::AND8_gisel || Opc == Z80::OR8_gisel ||
+                                  Opc == Z80::XOR8_gisel || Opc == TargetOpcode::COPY);
+        if (!isALUPseudoOrCopy) {
+          for (MachineOperand &MO : MIB->operands()) {
+            if (MO.isReg() && MO.isDef()) {
+              for (MCRegAliasIterator AI(MO.getReg(), TRI, true); AI.isValid(); ++AI) {
+                if (*AI == AMirroredInReg) {
+                  AMirroredInReg = Z80::NoRegister;
+                  break;
+                }
+              }
+            }
+          }
+        }
+      }
+
       // Apply KnownFlags after clobbering defs.
       if (KnownFlags)
         assign(Z80::F, KnownFlags.getKnownVal(KnownFlagsVal),
