diff --git a/llvm/include/llvm/CodeGen/GlobalISel/Utils.h b/llvm/include/llvm/CodeGen/GlobalISel/Utils.h --- a/llvm/include/llvm/CodeGen/GlobalISel/Utils.h +++ b/llvm/include/llvm/CodeGen/GlobalISel/Utils.h @@ -263,6 +263,8 @@ Optional ConstantFoldBinOp(unsigned Opcode, const Register Op1, const Register Op2, const MachineRegisterInfo &MRI); +Optional ConstantFoldUnaryOp(unsigned Opcode, const Register Op1, + const MachineRegisterInfo &MRI); Optional ConstantFoldFPBinOp(unsigned Opcode, const Register Op1, const Register Op2, const MachineRegisterInfo &MRI); @@ -274,6 +276,14 @@ Optional ConstantFoldVectorBinop(unsigned Opcode, const Register Op1, const Register Op2, const MachineRegisterInfo &MRI, MachineIRBuilder &MIB); +/// Tries to constant fold a vector unary operation with source \p Op. +/// If successful, returns the G_BUILD_VECTOR representing the folded vector +/// constant. \p MIB should have an insertion point already set to create new +/// G_CONSTANT instructions as needed. +Optional +ConstantFoldVectorUnaryOp(unsigned Opcode, const Register Op, + const MachineRegisterInfo &MRI, + MachineIRBuilder &MIB); Optional ConstantFoldExtOp(unsigned Opcode, const Register Op1, uint64_t Imm, const MachineRegisterInfo &MRI); diff --git a/llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp b/llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp --- a/llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp +++ b/llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp @@ -226,18 +226,18 @@ case TargetOpcode::G_CTLZ: { assert(SrcOps.size() == 1 && "Expected one source"); assert(DstOps.size() == 1 && "Expected one dest"); - auto MaybeCsts = ConstantFoldCTLZ(SrcOps[0].getReg(), *getMRI()); - if (!MaybeCsts) + if (SrcOps[0].getLLTTy(*getMRI()).isVector()) { + // Try to constant fold vector constants. + auto VecCst = + ConstantFoldVectorUnaryOp(Opc, SrcOps[0].getReg(), *getMRI(), *this); + if (VecCst) + return MachineInstrBuilder(getMF(), *VecCst); break; - if (MaybeCsts->size() == 1) - return buildConstant(DstOps[0], (*MaybeCsts)[0]); - // This was a vector constant. Build a G_BUILD_VECTOR for them. - SmallVector ConstantRegs; - LLT VecTy = DstOps[0].getLLTTy(*getMRI()); - for (unsigned Cst : *MaybeCsts) - ConstantRegs.emplace_back( - buildConstant(VecTy.getScalarType(), Cst).getReg(0)); - return buildBuildVector(DstOps[0], ConstantRegs); + } + if (Optional Cst = + ConstantFoldUnaryOp(Opc, SrcOps[0].getReg(), *getMRI())) + return buildConstant(DstOps[0], *Cst); + break; } } bool CanCopy = checkCopyToDefsPossible(DstOps); diff --git a/llvm/lib/CodeGen/GlobalISel/Utils.cpp b/llvm/lib/CodeGen/GlobalISel/Utils.cpp --- a/llvm/lib/CodeGen/GlobalISel/Utils.cpp +++ b/llvm/lib/CodeGen/GlobalISel/Utils.cpp @@ -530,6 +530,46 @@ return None; } +Optional llvm::ConstantFoldUnaryOp(unsigned Opcode, const Register Op1, + + const MachineRegisterInfo &MRI) { + Optional MaybeCst = getIConstantVRegVal(Op1, MRI); + if (!MaybeCst) + return None; + + APInt &C = *MaybeCst; + switch (Opcode) { + default: + return None; + case TargetOpcode::G_CTLZ: + return APInt(C.getBitWidth(), C.countLeadingZeros()); + } +} + +Optional +llvm::ConstantFoldVectorUnaryOp(unsigned Opcode, const Register Op, + const MachineRegisterInfo &MRI, + MachineIRBuilder &MIB) { + auto *SrcVec = getOpcodeDef(Op, MRI); + if (!SrcVec) + return None; + + const LLT EltTy = MRI.getType(SrcVec->getSourceReg(0)); + + SmallVector FoldedElements; + for (unsigned Idx = 0, E = SrcVec->getNumSources(); Idx < E; ++Idx) { + auto MaybeCst = ConstantFoldUnaryOp(Opcode, SrcVec->getSourceReg(Idx), MRI); + if (!MaybeCst) + return None; + auto FoldedCstReg = MIB.buildConstant(EltTy, *MaybeCst).getReg(0); + FoldedElements.emplace_back(FoldedCstReg); + } + // Create the new vector constant. + auto CstVec = + MIB.buildBuildVector(MRI.getType(SrcVec->getReg(0)), FoldedElements); + return &*CstVec; +} + Optional llvm::ConstantFoldFPBinOp(unsigned Opcode, const Register Op1, const Register Op2, const MachineRegisterInfo &MRI) {