GlobalISel: Fix widenScalar for G_MERGE_VALUES to pointer

[arsenm]

Summary

AMDGPU testcase isn't broken now, but will be in a future patch
without this.

[dsanders]

LGTM with a nit

lib/CodeGen/GlobalISel/LegalizerHelper.cpp

995

I'd prefer it if this checked DstTy.isPointer() or at least asserted it. It's not very obvious that WideTy != DstTy => pointer

[arsenm]

r367591

include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

public:
   /// \return The newly created instruction.
   MachineInstrBuilder buildSExt(const DstOp &Res, const SrcOp &Op);
 
   /// Build and insert a G_PTRTOINT instruction.
   MachineInstrBuilder buildPtrToInt(const DstOp &Dst, const SrcOp &Src) {
     return buildInstr(TargetOpcode::G_PTRTOINT, {Dst}, {Src});
   }
 
+  /// Build and insert a G_INTTOPTR instruction.
+  MachineInstrBuilder buildIntToPtr(const DstOp &Dst, const SrcOp &Src) {
+    return buildInstr(TargetOpcode::G_INTTOPTR, {Dst}, {Src});
+  }
+
   /// Build and insert \p Dst = G_BITCAST \p Src
   MachineInstrBuilder buildBitcast(const DstOp &Dst, const SrcOp &Src) {
     return buildInstr(TargetOpcode::G_BITCAST, {Dst}, {Src});
   }
 
   /// \return The opcode of the extension the target wants to use for boolean
   /// values.
   unsigned getBoolExtOp(bool IsVec, bool IsFP) const;

lib/CodeGen/GlobalISel/LegalizerHelper.cpp

if (WideSize >= DstSize) {
     for (unsigned I = 2; I != NumOps; ++I) {
       const unsigned Offset = (I - 1) * PartSize;
 
       Register SrcReg = MI.getOperand(I).getReg();
       assert(MRI.getType(SrcReg) == LLT::scalar(PartSize));
 
       auto ZextInput = MIRBuilder.buildZExt(WideTy, SrcReg);
 
-      Register NextResult = I + 1 == NumOps && WideSize == DstSize ? DstReg :
+      Register NextResult = I + 1 == NumOps && WideTy == DstTy ? DstReg :
         MRI.createGenericVirtualRegister(WideTy);
 
       auto ShiftAmt = MIRBuilder.buildConstant(WideTy, Offset);
       auto Shl = MIRBuilder.buildShl(WideTy, ZextInput, ShiftAmt);
       MIRBuilder.buildOr(NextResult, ResultReg, Shl);
       ResultReg = NextResult;
     }
 
     if (WideSize > DstSize)
       MIRBuilder.buildTrunc(DstReg, ResultReg);
+    else if (WideTy != DstTy)
+      MIRBuilder.buildIntToPtr(DstReg, ResultReg);
 
     MI.eraseFromParent();
     return Legalized;
   }
 
   // Unmerge the original values to the GCD type, and recombine to the next
   // multiple greater than the original type.
   //

test/CodeGen/AMDGPU/GlobalISel/legalize-merge-values.mir

bb.0:
     ; CHECK: S_NOP 0, implicit [[TRUNC97]](s68)
     %0:_(s17) = G_CONSTANT i17 0
     %1:_(s17) = G_CONSTANT i17 1
     %2:_(s17) = G_CONSTANT i17 2
     %3:_(s17) = G_CONSTANT i17 3
     %4:_(s68) = G_MERGE_VALUES %0, %1, %2, %3
     S_NOP 0, implicit %4
 ...
+---
+name: test_merge_p3_s16_s16
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: test_merge_p3_s16_s16
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: [[TRUNC:%[0-9]+]]:_(s16) = G_TRUNC [[C]](s32)
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: [[TRUNC1:%[0-9]+]]:_(s16) = G_TRUNC [[C1]](s32)
+    ; CHECK: [[MV:%[0-9]+]]:_(p3) = G_MERGE_VALUES [[TRUNC]](s16), [[TRUNC1]](s16)
+    ; CHECK: $vgpr0 = COPY [[MV]](p3)
+    %0:_(s16) = G_CONSTANT i16 0
+    %1:_(s16) = G_CONSTANT i16 1
+    %2:_(p3) = G_MERGE_VALUES %0, %1
+    $vgpr0 = COPY %2
+...

unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

TEST_F(GISelMITest, LowerMergeValues) {
   [[K16:%[0-9]+]]:_(s32) = G_CONSTANT i32 8
   [[SHL:%[0-9]+]]:_(s32) = G_SHL [[ZEXT_K16]]:_, [[K16]]:_(s32)
   [[OR:%[0-9]+]]:_(s32) = G_OR [[ZEXT_K16]]:_, [[SHL]]:_
   (s16) = G_TRUNC [[OR]]:_(s32)
   )";
 
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
+
+TEST_F(GISelMITest, WidenScalarMergeValuesPointer) {
+  if (!TM)
+    return;
+
+  DefineLegalizerInfo(A, {});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+  B.setInsertPt(*EntryMBB, EntryMBB->end());
+
+  const LLT S32 = LLT::scalar(32);
+  const LLT S64 = LLT::scalar(64);
+  const LLT P0 = LLT::pointer(0, 64);
+
+  auto Lo = B.buildTrunc(S32, Copies[0]);
+  auto Hi = B.buildTrunc(S32, Copies[1]);
+
+  auto Merge = B.buildMerge(P0, {Lo.getReg(0), Hi.getReg(0)});
+
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.widenScalar(*Merge, 1, S64));
+
+  auto CheckStr = R"(
+   CHECK: [[TRUNC0:%[0-9]+]]:_(s32) = G_TRUNC
+   CHECK: [[TRUNC1:%[0-9]+]]:_(s32) = G_TRUNC
+   CHECK: [[ZEXT_TRUNC0:%[0-9]+]]:_(s64) = G_ZEXT [[TRUNC0]]
+   CHECK: [[ZEXT_TRUNC1:%[0-9]+]]:_(s64) = G_ZEXT [[TRUNC1]]
+   CHECK: [[SHIFT_AMT:%[0-9]+]]:_(s64) = G_CONSTANT i64 32
+   CHECK: [[SHL:%[0-9]+]]:_(s64) = G_SHL [[ZEXT_TRUNC1]]:_, [[SHIFT_AMT]]
+   CHECK: [[OR:%[0-9]+]]:_(s64) = G_OR [[ZEXT_TRUNC0]]:_, [[SHL]]
+   CHECK: [[INTTOPTR:%[0-9]+]]:_(p0) = G_INTTOPTR [[OR]]:_(s64)
+  )";
+
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}
+
 } // namespace