[GISel] Handle G_TRUNC in `matchExtractVecEltBuildVec`

[Pierre-vh]

Summary

Spotted some cases in D134354 where this was an issue.

[arsenm]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

3930–3931

I don't see why this legality check (or the one below for build_vector_trunc) are here. The equivalent DAG combine does not check for the legality, and you're not creating a new one in the apply

3937

getVRegDef should be adequate here

[Pierre-vh]

getDefIgnoringCopies -> getVRegDef

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

3930–3931

@aemerson I see you implemented this one, is it okay to remove the legalizer checks?

[aemerson]

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

3930–3931

Sure, I don’t remember why it’s there either.

[Pierre-vh]

Remove legalizer checks, simplify code

llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp

 }
 
 bool CombinerHelper::matchExtractVecEltBuildVec(MachineInstr &MI,
                                                 Register &Reg) {
   assert(MI.getOpcode() == TargetOpcode::G_EXTRACT_VECTOR_ELT);
   // If we have a constant index, look for a G_BUILD_VECTOR source
   // and find the source register that the index maps to.
   Register SrcVec = MI.getOperand(1).getReg();
   LLT SrcTy = MRI.getType(SrcVec);
-  if (!isLegalOrBeforeLegalizer(
-          {TargetOpcode::G_BUILD_VECTOR, {SrcTy, SrcTy.getElementType()}}))
-    return false;
 
   auto Cst = getIConstantVRegValWithLookThrough(MI.getOperand(2).getReg(), MRI);
   if (!Cst || Cst->Value.getZExtValue() >= SrcTy.getNumElements())
     return false;
 
   unsigned VecIdx = Cst->Value.getZExtValue();
-  MachineInstr *BuildVecMI =
-      getOpcodeDef(TargetOpcode::G_BUILD_VECTOR, SrcVec, MRI);
-  if (!BuildVecMI) {
-    BuildVecMI = getOpcodeDef(TargetOpcode::G_BUILD_VECTOR_TRUNC, SrcVec, MRI);
-    if (!BuildVecMI)
-      return false;
-    LLT ScalarTy = MRI.getType(BuildVecMI->getOperand(1).getReg());
-    if (!isLegalOrBeforeLegalizer(
-            {TargetOpcode::G_BUILD_VECTOR_TRUNC, {SrcTy, ScalarTy}}))
-      return false;
+
+  // Check if we have a build_vector or build_vector_trunc with an optional
+  // trunc in front.
+  MachineInstr *SrcVecMI = MRI.getVRegDef(SrcVec);
+  if (SrcVecMI->getOpcode() == TargetOpcode::G_TRUNC) {
+    SrcVecMI = MRI.getVRegDef(SrcVecMI->getOperand(1).getReg());
   }
 
+  if (SrcVecMI->getOpcode() != TargetOpcode::G_BUILD_VECTOR &&
+      SrcVecMI->getOpcode() != TargetOpcode::G_BUILD_VECTOR_TRUNC)
+    return false;
+
   EVT Ty(getMVTForLLT(SrcTy));
   if (!MRI.hasOneNonDBGUse(SrcVec) &&
       !getTargetLowering().aggressivelyPreferBuildVectorSources(Ty))
     return false;
 
-  Reg = BuildVecMI->getOperand(VecIdx + 1).getReg();
+  Reg = SrcVecMI->getOperand(VecIdx + 1).getReg();
   return true;
 }
 
 void CombinerHelper::applyExtractVecEltBuildVec(MachineInstr &MI,
                                                 Register &Reg) {
   // Check the type of the register, since it may have come from a
   // G_BUILD_VECTOR_TRUNC.
   LLT ScalarTy = MRI.getType(Reg);

llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir

bb.0:
     %zero:_(s32) = G_CONSTANT i32 0
     %bv:_(<2 x s64>) = G_BUILD_VECTOR %arg1(s64), %arg2(s64)
     %extract:_(s64) = G_EXTRACT_VECTOR_ELT %bv(<2 x s64>), %zero(s32)
     $x0 = COPY %extract(s64)
     RET_ReallyLR implicit $x0
 
 ...
 ---
+name:            extract_from_trunc_build_vector
+alignment:       4
+tracksRegLiveness: true
+liveins:
+  - { reg: '$x0' }
+  - { reg: '$x1' }
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+
+    ; CHECK-LABEL: name: extract_from_trunc_build_vector
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %arg1:_(s64) = COPY $x0
+    ; CHECK-NEXT: %extract:_(s32) = G_TRUNC %arg1(s64)
+    ; CHECK-NEXT: %zext:_(s64) = G_ZEXT %extract(s32)
+    ; CHECK-NEXT: $x0 = COPY %zext(s64)
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %arg1:_(s64) = COPY $x0
+    %arg2:_(s64) = COPY $x1
+    %zero:_(s32) = G_CONSTANT i32 0
+    %bv:_(<2 x s64>) = G_BUILD_VECTOR %arg1(s64), %arg2(s64)
+    %truncbv:_(<2 x s32>) = G_TRUNC %bv
+    %extract:_(s32) = G_EXTRACT_VECTOR_ELT %truncbv(<2 x s32>), %zero(s32)
+    %zext:_(s64) = G_ZEXT %extract
+    $x0 = COPY %zext(s64)
+    RET_ReallyLR implicit $x0
+
+...
+---
 name:            extract_from_build_vector_idx1
 alignment:       4
 tracksRegLiveness: true
 liveins:
   - { reg: '$x0' }
   - { reg: '$x1' }
 frameInfo:
   maxAlignment:    1