diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp --- a/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp +++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelper.cpp @@ -3928,33 +3928,30 @@ // 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; } diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir b/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir --- a/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir +++ b/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir @@ -29,6 +29,40 @@ $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