GlobalISel: Fix CSE handling of buildConstant

[arsenm]

Summary

This fixes two problems with CSE done in buildConstant. First, this
would hit an assert when used with a vector result type. Solve this by
allowing CSE on the vector elements, but not on the result vector for
now.

Second, this was also performing the CSE based on the input
ConstantInt pointer. The underlying buildConstant could potentially
convert the constant depending on the result type, giving in a
different ConstantInt*. Stop allowing the APInt and ConstantInt forms
from automatically casting to the result type to avoid any similar
problems in the future.

[aditya_nandakumar]

LGTM. Please move the EXPECT_EQ changes to the other commit.

include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

738

Should be Src and not Src0 here.

unittests/CodeGen/GlobalISel/CSETest.cpp

33

Please make the EXPECT_EQ changes along with https://reviews.llvm.org/D57652

[arsenm]

r353077

include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

public:
   /// \pre The entire register \p Res (and no more) must be covered by the
   ///      input scalar registers.
   /// \pre The type of all \p Ops registers must be identical.
   ///
   /// \return a MachineInstrBuilder for the newly created instruction.
   MachineInstrBuilder buildBuildVector(const DstOp &Res,
                                        ArrayRef<unsigned> Ops);
 
+  /// Build and insert \p Res = G_BUILD_VECTOR with \p Src0 replicated to fill
+  /// the number of elements
+  MachineInstrBuilder buildSplatVector(const DstOp &Res,
+                                       const SrcOp &Src);
+
   /// Build and insert \p Res = G_BUILD_VECTOR_TRUNC \p Op0, ...
   ///
   /// G_BUILD_VECTOR_TRUNC creates a vector value from multiple scalar registers
   /// which have types larger than the destination vector element type, and
   /// truncates the values to fit.
   ///
   /// If the operands given are already the same size as the vector elt type,
   /// then this method will instead create a G_BUILD_VECTOR instruction.

lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp

MachineInstrBuilder CSEMIRBuilder::buildInstr(unsigned Opc,
   MachineInstrBuilder NewMIB =
       MachineIRBuilder::buildInstr(Opc, DstOps, SrcOps, Flag);
   return memoizeMI(NewMIB, InsertPos);
 }
 
 MachineInstrBuilder CSEMIRBuilder::buildConstant(const DstOp &Res,
                                                  const ConstantInt &Val) {
   constexpr unsigned Opc = TargetOpcode::G_CONSTANT;
+  LLT Ty = Res.getLLTTy(*getMRI()); // TODO: Handle vectors
   if (!canPerformCSEForOpc(Opc))
     return MachineIRBuilder::buildConstant(Res, Val);
+
+  // For vectors, CSE the element only for now.
+  if (Ty.isVector())
+    return buildSplatVector(Res, buildConstant(Ty.getElementType(), Val));
+
   FoldingSetNodeID ID;
   GISelInstProfileBuilder ProfBuilder(ID, *getMRI());
   void *InsertPos = nullptr;
   profileMBBOpcode(ProfBuilder, Opc);
   profileDstOp(Res, ProfBuilder);
   ProfBuilder.addNodeIDMachineOperand(MachineOperand::CreateCImm(&Val));
   MachineInstrBuilder MIB = getDominatingInstrForID(ID, InsertPos);
   if (MIB) {
     // Handle generating copies here.
     return generateCopiesIfRequired({Res}, MIB);
   }
+
   MachineInstrBuilder NewMIB = MachineIRBuilder::buildConstant(Res, Val);
   return memoizeMI(NewMIB, InsertPos);
 }
 
 MachineInstrBuilder CSEMIRBuilder::buildFConstant(const DstOp &Res,
                                                   const ConstantFP &Val) {
   constexpr unsigned Opc = TargetOpcode::G_FCONSTANT;
+  LLT Ty = Res.getLLTTy(*getMRI()); // TODO: Handle vectors
   if (!canPerformCSEForOpc(Opc))
     return MachineIRBuilder::buildFConstant(Res, Val);
+
+  // For vectors, CSE the element only for now.
+  if (Ty.isVector())
+    return buildSplatVector(Res, buildFConstant(Ty.getElementType(), Val));
+
   FoldingSetNodeID ID;
   GISelInstProfileBuilder ProfBuilder(ID, *getMRI());
   void *InsertPos = nullptr;
   profileMBBOpcode(ProfBuilder, Opc);
   profileDstOp(Res, ProfBuilder);
   ProfBuilder.addNodeIDMachineOperand(MachineOperand::CreateFPImm(&Val));
   MachineInstrBuilder MIB = getDominatingInstrForID(ID, InsertPos);
   if (MIB) {
     // Handle generating copies here.
     return generateCopiesIfRequired({Res}, MIB);
   }
   MachineInstrBuilder NewMIB = MachineIRBuilder::buildFConstant(Res, Val);
   return memoizeMI(NewMIB, InsertPos);
 }

lib/CodeGen/GlobalISel/MachineIRBuilder.cpp

MachineInstrBuilder MachineIRBuilder::buildCopy(const DstOp &Res,
                                                 const SrcOp &Op) {
   return buildInstr(TargetOpcode::COPY, Res, Op);
 }
 
 MachineInstrBuilder MachineIRBuilder::buildConstant(const DstOp &Res,
                                                     const ConstantInt &Val) {
   LLT Ty = Res.getLLTTy(*getMRI());
   LLT EltTy = Ty.getScalarType();
-
-  const ConstantInt *NewVal = &Val;
-  if (EltTy.getSizeInBits() != Val.getBitWidth()) {
-    NewVal = ConstantInt::get(
-      getMF().getFunction().getContext(),
-      Val.getValue().sextOrTrunc(EltTy.getSizeInBits()));
-  }
+  assert(EltTy.getScalarSizeInBits() == Val.getBitWidth() &&
+         "creating constant with the wrong size");
 
   if (Ty.isVector()) {
-    unsigned EltReg = getMRI()->createGenericVirtualRegister(EltTy);
-    buildInstr(TargetOpcode::G_CONSTANT)
-      .addDef(EltReg)
-      .addCImm(NewVal);
-
-    auto MIB = buildInstr(TargetOpcode::G_BUILD_VECTOR);
-    Res.addDefToMIB(*getMRI(), MIB);
-
-    for (unsigned I = 0, E = Ty.getNumElements(); I != E; ++I)
-      MIB.addUse(EltReg);
-    return MIB;
+    auto Const = buildInstr(TargetOpcode::G_CONSTANT)
+    .addDef(getMRI()->createGenericVirtualRegister(EltTy))
+    .addCImm(&Val);
+    return buildSplatVector(Res, Const);
   }
 
-  auto MIB = buildInstr(TargetOpcode::G_CONSTANT);
-  Res.addDefToMIB(*getMRI(), MIB);
-  MIB.addCImm(NewVal);
-  return MIB;
+  auto Const = buildInstr(TargetOpcode::G_CONSTANT);
+  Res.addDefToMIB(*getMRI(), Const);
+  Const.addCImm(&Val);
+  return Const;
 }
 
 MachineInstrBuilder MachineIRBuilder::buildConstant(const DstOp &Res,
                                                     int64_t Val) {
   auto IntN = IntegerType::get(getMF().getFunction().getContext(),
-                               Res.getLLTTy(*getMRI()).getSizeInBits());
+                               Res.getLLTTy(*getMRI()).getScalarSizeInBits());
   ConstantInt *CI = ConstantInt::get(IntN, Val, true);
   return buildConstant(Res, *CI);
 }
 
 MachineInstrBuilder MachineIRBuilder::buildConstant(const DstOp &Res,
                                                     const APInt &Val) {
   ConstantInt *CI = ConstantInt::get(getMF().getFunction().getContext(), Val);
   return buildConstant(Res, *CI);
 }
 
 MachineInstrBuilder MachineIRBuilder::buildFConstant(const DstOp &Res,
                                                      const ConstantFP &Val) {
   LLT Ty = Res.getLLTTy(*getMRI());
+  LLT EltTy = Ty.getScalarType();
+
+  assert(APFloat::getSizeInBits(Val.getValueAPF().getSemantics())
+         == EltTy.getSizeInBits() &&
+         "creating fconstant with the wrong size");
 
   assert(!Ty.isPointer() && "invalid operand type");
 
   if (Ty.isVector()) {
-    unsigned EltReg
-      = getMRI()->createGenericVirtualRegister(Ty.getElementType());
-    buildInstr(TargetOpcode::G_FCONSTANT)
-      .addDef(EltReg)
-      .addFPImm(&Val);
+    auto Const = buildInstr(TargetOpcode::G_FCONSTANT)
+    .addDef(getMRI()->createGenericVirtualRegister(EltTy))
 
-    auto MIB = buildInstr(TargetOpcode::G_BUILD_VECTOR);
-    Res.addDefToMIB(*getMRI(), MIB);
-
-    for (unsigned I = 0, E = Ty.getNumElements(); I != E; ++I)
-      MIB.addUse(EltReg);
-    return MIB;
+    return buildSplatVector(Res, Const);
   }
 
-  auto MIB = buildInstr(TargetOpcode::G_FCONSTANT);
-  Res.addDefToMIB(*getMRI(), MIB);
-  MIB.addFPImm(&Val);
-  return MIB;
+  auto Const = buildInstr(TargetOpcode::G_FCONSTANT);
+  Res.addDefToMIB(*getMRI(), Const);
+  Const.addFPImm(&Val);
+  return Const;
 }
 
 MachineInstrBuilder MachineIRBuilder::buildFConstant(const DstOp &Res,
                                                      double Val) {
   LLT DstTy = Res.getLLTTy(*getMRI());
   auto &Ctx = getMF().getFunction().getContext();
   auto *CFP =
-      ConstantFP::get(Ctx, getAPFloatFromSize(Val, DstTy.getSizeInBits()));
+      ConstantFP::get(Ctx, getAPFloatFromSize(Val, DstTy.getScalarSizeInBits()));
   return buildFConstant(Res, *CFP);
 }
 
 MachineInstrBuilder MachineIRBuilder::buildBrCond(unsigned Tst,
                                                   MachineBasicBlock &Dest) {
   assert(getMRI()->getType(Tst).isScalar() && "invalid operand type");
 
   return buildInstr(TargetOpcode::G_BRCOND).addUse(Tst).addMBB(&Dest);
MachineInstrBuilder MachineIRBuilder::buildBuildVector(const DstOp &Res,
                                                        ArrayRef<unsigned> Ops) {
   // Unfortunately to convert from ArrayRef<unsigned> to ArrayRef<SrcOp>,
   // we need some temporary storage for the DstOp objects. Here we use a
   // sufficiently large SmallVector to not go through the heap.
   SmallVector<SrcOp, 8> TmpVec(Ops.begin(), Ops.end());
   return buildInstr(TargetOpcode::G_BUILD_VECTOR, Res, TmpVec);
 }
 
+MachineInstrBuilder MachineIRBuilder::buildSplatVector(const DstOp &Res,
+                                                       const SrcOp &Src) {
+  SmallVector<SrcOp, 8> TmpVec(Res.getLLTTy(*getMRI()).getNumElements(), Src);
+  return buildInstr(TargetOpcode::G_BUILD_VECTOR, Res, TmpVec);
+}
+
 MachineInstrBuilder
 MachineIRBuilder::buildBuildVectorTrunc(const DstOp &Res,
                                         ArrayRef<unsigned> Ops) {
   // Unfortunately to convert from ArrayRef<unsigned> to ArrayRef<SrcOp>,
   // we need some temporary storage for the DstOp objects. Here we use a
   // sufficiently large SmallVector to not go through the heap.
   SmallVector<SrcOp, 8> TmpVec(Ops.begin(), Ops.end());
   return buildInstr(TargetOpcode::G_BUILD_VECTOR_TRUNC, Res, TmpVec);

unittests/CodeGen/GlobalISel/CSETest.cpp

TEST_F(GISelMITest, TestCSE) {
   auto MIBInput = B.buildInstr(TargetOpcode::G_TRUNC, {s16}, {Copies[0]});
   auto MIBInput1 = B.buildInstr(TargetOpcode::G_TRUNC, {s16}, {Copies[1]});
   auto MIBAdd = B.buildInstr(TargetOpcode::G_ADD, {s16}, {MIBInput, MIBInput});
   GISelCSEInfo CSEInfo;
   CSEInfo.setCSEConfig(make_unique<CSEConfig>());
   CSEInfo.analyze(*MF);
   B.setCSEInfo(&CSEInfo);
   CSEMIRBuilder CSEB(B.getState());
+
   CSEB.setInsertPt(*EntryMBB, EntryMBB->begin());
   unsigned AddReg = MRI->createGenericVirtualRegister(s16);
   auto MIBAddCopy =
       CSEB.buildInstr(TargetOpcode::G_ADD, {AddReg}, {MIBInput, MIBInput});
-  ASSERT_EQ(MIBAddCopy->getOpcode(), TargetOpcode::COPY);
+  EXPECT_EQ(MIBAddCopy->getOpcode(), TargetOpcode::COPY);
   auto MIBAdd2 =
       CSEB.buildInstr(TargetOpcode::G_ADD, {s16}, {MIBInput, MIBInput});
-  ASSERT_TRUE(&*MIBAdd == &*MIBAdd2);
+  EXPECT_TRUE(&*MIBAdd == &*MIBAdd2);
   auto MIBAdd4 =
       CSEB.buildInstr(TargetOpcode::G_ADD, {s16}, {MIBInput, MIBInput});
-  ASSERT_TRUE(&*MIBAdd == &*MIBAdd4);
+  EXPECT_TRUE(&*MIBAdd == &*MIBAdd4);
   auto MIBAdd5 =
       CSEB.buildInstr(TargetOpcode::G_ADD, {s16}, {MIBInput, MIBInput1});
-  ASSERT_TRUE(&*MIBAdd != &*MIBAdd5);
+  EXPECT_TRUE(&*MIBAdd != &*MIBAdd5);
 
   // Try building G_CONSTANTS.
   auto MIBCst = CSEB.buildConstant(s32, 0);
   auto MIBCst1 = CSEB.buildConstant(s32, 0);
-  ASSERT_TRUE(&*MIBCst == &*MIBCst1);
+  EXPECT_TRUE(&*MIBCst == &*MIBCst1);
   // Try the CFing of BinaryOps.
   auto MIBCF1 = CSEB.buildInstr(TargetOpcode::G_ADD, {s32}, {MIBCst, MIBCst});
-  ASSERT_TRUE(&*MIBCF1 == &*MIBCst);
+  EXPECT_TRUE(&*MIBCF1 == &*MIBCst);
+
 
   // Try out building FCONSTANTs.
   auto MIBFP0 = CSEB.buildFConstant(s32, 1.0);
   auto MIBFP0_1 = CSEB.buildFConstant(s32, 1.0);
-  ASSERT_TRUE(&*MIBFP0 == &*MIBFP0_1);
+  EXPECT_TRUE(&*MIBFP0 == &*MIBFP0_1);
   CSEInfo.print();
 
+  // Make sure buildConstant with a vector type doesn't crash, and the elements
+  // CSE.
+  auto Splat0 = CSEB.buildConstant(LLT::vector(2, s32), 0);
+  EXPECT_EQ(TargetOpcode::G_BUILD_VECTOR, Splat0->getOpcode());
+  EXPECT_EQ(Splat0->getOperand(1).getReg(), Splat0->getOperand(2).getReg());
+  EXPECT_EQ(&*MIBCst, MRI->getVRegDef(Splat0->getOperand(1).getReg()));
+
+  auto FSplat = CSEB.buildFConstant(LLT::vector(2, s32), 1.0);
+  EXPECT_EQ(TargetOpcode::G_BUILD_VECTOR, FSplat->getOpcode());
+  EXPECT_EQ(FSplat->getOperand(1).getReg(), FSplat->getOperand(2).getReg());
+  EXPECT_EQ(&*MIBFP0, MRI->getVRegDef(FSplat->getOperand(1).getReg()));
+
   // Check G_UNMERGE_VALUES
   auto MIBUnmerge = CSEB.buildUnmerge({s32, s32}, Copies[0]);
   auto MIBUnmerge2 = CSEB.buildUnmerge({s32, s32}, Copies[0]);
-  ASSERT_TRUE(&*MIBUnmerge == &*MIBUnmerge2);
+  EXPECT_TRUE(&*MIBUnmerge == &*MIBUnmerge2);
 }
 
 TEST_F(GISelMITest, TestCSEConstantConfig) {
   if (!TM)
     return;
 
   LLT s16{LLT::scalar(16)};
   auto MIBInput = B.buildInstr(TargetOpcode::G_TRUNC, {s16}, {Copies[0]});
   auto MIBAdd = B.buildInstr(TargetOpcode::G_ADD, {s16}, {MIBInput, MIBInput});
   auto MIBZero = B.buildConstant(s16, 0);
   GISelCSEInfo CSEInfo;
   CSEInfo.setCSEConfig(make_unique<CSEConfigConstantOnly>());
   CSEInfo.analyze(*MF);
   B.setCSEInfo(&CSEInfo);
   CSEMIRBuilder CSEB(B.getState());
   CSEB.setInsertPt(*EntryMBB, EntryMBB->begin());
   auto MIBAdd1 =
       CSEB.buildInstr(TargetOpcode::G_ADD, {s16}, {MIBInput, MIBInput});
   // We should CSE constants only. Adds should not be CSEd.
-  ASSERT_TRUE(MIBAdd1->getOpcode() != TargetOpcode::COPY);
-  ASSERT_TRUE(&*MIBAdd1 != &*MIBAdd);
+  EXPECT_TRUE(MIBAdd1->getOpcode() != TargetOpcode::COPY);
+  EXPECT_TRUE(&*MIBAdd1 != &*MIBAdd);
   // We should CSE constant.
   auto MIBZeroTmp = CSEB.buildConstant(s16, 0);
-  ASSERT_TRUE(&*MIBZero == &*MIBZeroTmp);
+  EXPECT_TRUE(&*MIBZero == &*MIBZeroTmp);
 }
 } // namespace

unittests/CodeGen/GlobalISel/MachineIRBuilderTest.cpp

 //===- MachineIRBuilderTest.cpp -------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "GISelMITest.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
 
 TEST_F(GISelMITest, TestBuildConstantFConstant) {
   if (!TM)
     return;
 
-  MachineIRBuilder B(*MF);
-  B.setInsertPt(*EntryMBB, EntryMBB->begin());
-
   B.buildConstant(LLT::scalar(32), 42);
   B.buildFConstant(LLT::scalar(32), 1.0);
 
   B.buildConstant(LLT::vector(2, 32), 99);
   B.buildFConstant(LLT::vector(2, 32), 2.0);
 
   auto CheckStr = R"(
   CHECK: [[CONST0:%[0-9]+]]:_(s32) = G_CONSTANT i32 42
   CHECK: [[FCONST0:%[0-9]+]]:_(s32) = G_FCONSTANT float 1.000000e+00
   CHECK: [[CONST1:%[0-9]+]]:_(s32) = G_CONSTANT i32 99
   CHECK: [[VEC0:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[CONST1]]:_(s32), [[CONST1]]:_(s32)
-  CHECK: [[FCONST1:%[0-9]+]]:_(s32) = G_FCONSTANT double 2.000000e+00
+  CHECK: [[FCONST1:%[0-9]+]]:_(s32) = G_FCONSTANT float 2.000000e+00
   CHECK: [[VEC1:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[FCONST1]]:_(s32), [[FCONST1]]:_(s32)
 
   )";
 
   ASSERT_TRUE(CheckMachineFunction(*MF, CheckStr));
 }
+
+
+#ifdef GTEST_HAS_DEATH_TEST
+#ifndef NDEBUG
+
+TEST_F(GISelMITest, TestBuildConstantFConstantDeath) {
+  if (!TM)
+    return;
+
+  LLVMContext &Ctx = MF->getFunction().getContext();
+  APInt APV32(32, 12345);
+
+  // Test APInt version breaks
+  EXPECT_DEATH(B.buildConstant(LLT::scalar(16), APV32),
+               "creating constant with the wrong size");
+  EXPECT_DEATH(B.buildConstant(LLT::vector(2, 16), APV32),
+               "creating constant with the wrong size");
+
+  // Test ConstantInt version breaks
+  ConstantInt *CI = ConstantInt::get(Ctx, APV32);
+  EXPECT_DEATH(B.buildConstant(LLT::scalar(16), *CI),
+               "creating constant with the wrong size");
+  EXPECT_DEATH(B.buildConstant(LLT::vector(2, 16), *CI),
+               "creating constant with the wrong size");
+
+  APFloat DoubleVal(APFloat::IEEEdouble());
+  ConstantFP *CF = ConstantFP::get(Ctx, DoubleVal);
+  EXPECT_DEATH(B.buildFConstant(LLT::scalar(16), *CF),
+               "creating fconstant with the wrong size");
+  EXPECT_DEATH(B.buildFConstant(LLT::vector(2, 16), *CF),
+               "creating fconstant with the wrong size");
+}
+
+#endif
+#endif