[llvm] [AArch64][GlobalISel] Combine G_CONCAT_VECTORS with Illegal Form (PR #85047)

[David Green via llvm-commits]

================
@@ -359,6 +359,108 @@ class LegalizationArtifactCombiner {
     return false;
   }
 
+  // Try to combine illegal G_CONCAT_VECTORS instructions
+  // Combine if the G_CONCAT_VECTORS instruction is illegal and
+  // Source Registers are:
+  //  - Previously defined by a G_BITCAST instruction
+  //  - Defined by G_IMPLICIT_DEF with look through
+  // ===============
+  //
+  // %0(s32) = G_LOAD %ptr
+  // %1(<4 x s8>) = G_BITCAST %0(s32)
+  // %2(s8) = G_IMPLICIT_DEF
+  // %3(<4 x s8>) = G_BUILD_VECTOR %2(s8), %2(s8), %2(s8), %2(s8)
+  // %4(<16 x s8>) = G_CONCAT_VECTORS %1(<4 x s8>), %3(<4 x s8>), %3(<4 x s8>),
+  // %3(<4 x s8>)
+  //
+  // ====>
+  //
+  // %0(s32) = G_LOAD %ptr
+  // %1(s32) = G_IMPLICIT_DEF
+  // %2(<4 x s32>) = G_BUILD_VECTOR %0(s32), %1(s32), %1(s32), %1(s32)
+  // %3(<16 x s8>) = G_BITCAST %2(<4 x s32>)
+  bool tryCombineConcat(GConcatVectors &MI,
+                        SmallVectorImpl<MachineInstr *> &DeadInsts,
+                        SmallVectorImpl<Register> &UpdatedDefs,
+                        GISelObserverWrapper &WrapperObserver) {
+    Register SrcReg0 = MI.getSourceReg(0);
+    LLT SrcTy = MRI.getType(SrcReg0);
+    unsigned NumSrc = MI.getNumSources();
+
+    Register DstReg = MI.getReg(0);
+    LLT DstTy = MRI.getType(DstReg);
+
+    // Do not attempt combine if it is already legal
+    if (isInstLegal({TargetOpcode::G_CONCAT_VECTORS, {DstTy, SrcTy}}))
+      return false;
+
+    // CHECK if this instruction should be combined before building instructions
+    SmallVector<Register> SrcRegs;
+    unsigned NumUndef = 0;
+    for (unsigned i = 0; i < NumSrc; i++) {
+      Register CurrReg = MI.getSourceReg(i);
+      LLT CurrTy = MRI.getType(CurrReg);
+
+      // Only handle cases where all source types are the same type
+      if (CurrTy.getSizeInBits() != SrcTy.getSizeInBits())
+        return false;
+
+      // Register defined with G_IMPLIT_DEF
+      if (isImpDefVRegValWithLookThrough(CurrReg, MRI)) {
+        NumUndef++;
+        SrcRegs.push_back(CurrReg);
+      }
+
+      else if (MachineInstr *DefMI = getDefIgnoringCopies(CurrReg, MRI);
+               DefMI->getOpcode() == TargetOpcode::G_BITCAST) {
+        Register DefSrcReg = DefMI->getOperand(1).getReg();
+        LLT DefSrcTy = MRI.getType(DefSrcReg);
+
+        // Only handle cases where G_BITCAST source type is scalar
+        if (DefSrcTy.isScalar())
+          SrcRegs.push_back(DefSrcReg);
+        else
+          return false;
+      } else {
+        return false;
+      }
+    }
+
+    // Ignore if all sources are G_IMPLICIT_DEF
+    if (NumUndef == NumSrc)
+      return false;
+
+    // Combine the instruction, matches the condition
+    Builder.setInstrAndDebugLoc(MI);
+    SmallVector<Register> BuildRegs;
+    MachineInstr *UndefMI = nullptr;
+
+    for (Register CurrReg : SrcRegs) {
+      LLT CurrTy = MRI.getType(CurrReg);
+      if (CurrTy == LLT::scalar(SrcTy.getSizeInBits())) {
+        BuildRegs.push_back(CurrReg);
+      } else if (isImpDefVRegValWithLookThrough(CurrReg, MRI)) {
+        UndefMI = !UndefMI
+                      ? Builder.buildUndef(LLT::scalar(SrcTy.getSizeInBits()))
+                      : UndefMI;
----------------
davemgreen wrote:

I think I would just do
```
if (!UndefMI)
  UndefMI = ...
```

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