GlobalISel: Handle more cases for widenScalar of G_MERGE_VALUES

[arsenm]

Summary

Use an anyext to the requested type for the leftover operand to
produce a slightly wider type, and then truncate the final merge.

I have another implementation almost ready which handles arbitrary
widens, but I think it produces worse code in this example (which I
think is 90% due to not folding redundant copies or folding out
implicit_def users), so I wanted to add this as a baseline first.

[arsenm]

ping

[arsenm]

ping, this mostly doesn't matter as it is largely replaced by D64155

[aemerson]

I didn't bother reviewing this as the other patch seemed to supersede it. Can it be abandoned?

[arsenm]

In D64138#1588433, @aemerson wrote:

I didn't bother reviewing this as the other patch seemed to supersede it. Can it be abandoned?

I'll just commit at the same time. I wanted to preserve the test diff in history

[arsenm]

r366366

lib/CodeGen/GlobalISel/LegalizerHelper.cpp

if (!DstTy.isScalar())
     return UnableToLegalize;
 
   Register Src1 = MI.getOperand(1).getReg();
   LLT SrcTy = MRI.getType(Src1);
   int NumMerge = DstTy.getSizeInBits() / WideTy.getSizeInBits();
 
   // Try to turn this into a merge of merges if we can use the requested type as
   // the source.
-
-  // TODO: Pad with undef if DstTy > WideTy
-  if (NumMerge > 1 && WideTy.getSizeInBits() % SrcTy.getSizeInBits() == 0) {
+  if (NumMerge > 1) {
     int PartsPerMerge = WideTy.getSizeInBits() / SrcTy.getSizeInBits();
+    if (WideTy.getSizeInBits() % SrcTy.getSizeInBits() != 0)
+      return UnableToLegalize;
+
+    int RemainderBits = DstTy.getSizeInBits() % WideTy.getSizeInBits();
+    int RemainderParts = RemainderBits / SrcTy.getSizeInBits();
+
     SmallVector<Register, 4> Parts;
     SmallVector<Register, 4> SubMerges;
 
     for (int I = 0; I != NumMerge; ++I) {
       for (int J = 0; J != PartsPerMerge; ++J)
         Parts.push_back(MI.getOperand(I * PartsPerMerge + J + 1).getReg());
 
       auto SubMerge = MIRBuilder.buildMerge(WideTy, Parts);
       SubMerges.push_back(SubMerge.getReg(0));
       Parts.clear();
     }
 
+    if (RemainderParts == 0) {
       MIRBuilder.buildMerge(DstReg, SubMerges);
       MI.eraseFromParent();
       return Legalized;
     }
 
+    assert(RemainderParts == 1);
+
+    auto AnyExt = MIRBuilder.buildAnyExt(
+      WideTy, MI.getOperand(MI.getNumOperands() - 1).getReg());
+    SubMerges.push_back(AnyExt.getReg(0));
+
+    LLT WiderDstTy = LLT::scalar(SubMerges.size() * WideTy.getSizeInBits());
+    auto Merge = MIRBuilder.buildMerge(WiderDstTy, SubMerges);
+    MIRBuilder.buildTrunc(DstReg, Merge);
+
+    MI.eraseFromParent();
+    return Legalized;
+  }
+
   unsigned NumOps = MI.getNumOperands();
   unsigned NumSrc = MI.getNumOperands() - 1;
   unsigned PartSize = DstTy.getSizeInBits() / NumSrc;
 
   Register ResultReg = MIRBuilder.buildZExt(DstTy, Src1).getReg(0);
 
   for (unsigned I = 2; I != NumOps; ++I) {
     const unsigned Offset = (I - 1) * PartSize;

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

bb.0:
     %1:_(s16) = G_CONSTANT i16 1
     %2:_(s16) = G_CONSTANT i16 2
     %3:_(s16) = G_CONSTANT i16 3
     %4:_(s16) = G_CONSTANT i16 4
     %5:_(s16) = G_CONSTANT i16 5
     %8:_(s96) = G_MERGE_VALUES %0, %1, %2, %3, %4, %5
     $vgpr0_vgpr1_vgpr2 = COPY %8
 ...
+
+---
+name: test_merge_s56_s8_s8_s8_s8_s8_s8_s8
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: test_merge_s56_s8_s8_s8_s8_s8_s8_s8
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: [[C1:%[0-9]+]]:_(s32) = G_CONSTANT i32 1
+    ; CHECK: [[C2:%[0-9]+]]:_(s32) = G_CONSTANT i32 2
+    ; CHECK: [[C3:%[0-9]+]]:_(s32) = G_CONSTANT i32 3
+    ; CHECK: [[C4:%[0-9]+]]:_(s32) = G_CONSTANT i32 4
+    ; CHECK: [[C5:%[0-9]+]]:_(s32) = G_CONSTANT i32 5
+    ; CHECK: [[C6:%[0-9]+]]:_(s32) = G_CONSTANT i32 6
+    ; CHECK: [[C7:%[0-9]+]]:_(s32) = G_CONSTANT i32 8
+    ; CHECK: [[C8:%[0-9]+]]:_(s32) = G_CONSTANT i32 65535
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY [[C7]](s32)
+    ; CHECK: [[AND:%[0-9]+]]:_(s32) = G_AND [[COPY]], [[C8]]
+    ; CHECK: [[C9:%[0-9]+]]:_(s32) = G_CONSTANT i32 255
+    ; CHECK: [[COPY1:%[0-9]+]]:_(s32) = COPY [[C1]](s32)
+    ; CHECK: [[AND1:%[0-9]+]]:_(s32) = G_AND [[COPY1]], [[C9]]
+    ; CHECK: [[SHL:%[0-9]+]]:_(s32) = G_SHL [[AND1]], [[AND]](s32)
+    ; CHECK: [[COPY2:%[0-9]+]]:_(s32) = COPY [[C]](s32)
+    ; CHECK: [[AND2:%[0-9]+]]:_(s32) = G_AND [[COPY2]], [[C9]]
+    ; CHECK: [[COPY3:%[0-9]+]]:_(s32) = COPY [[SHL]](s32)
+    ; CHECK: [[OR:%[0-9]+]]:_(s32) = G_OR [[AND2]], [[COPY3]]
+    ; CHECK: [[TRUNC:%[0-9]+]]:_(s16) = G_TRUNC [[OR]](s32)
+    ; CHECK: [[C10:%[0-9]+]]:_(s32) = G_CONSTANT i32 8
+    ; CHECK: [[COPY4:%[0-9]+]]:_(s32) = COPY [[C10]](s32)
+    ; CHECK: [[AND3:%[0-9]+]]:_(s32) = G_AND [[COPY4]], [[C8]]
+    ; CHECK: [[COPY5:%[0-9]+]]:_(s32) = COPY [[C3]](s32)
+    ; CHECK: [[AND4:%[0-9]+]]:_(s32) = G_AND [[COPY5]], [[C9]]
+    ; CHECK: [[SHL1:%[0-9]+]]:_(s32) = G_SHL [[AND4]], [[AND3]](s32)
+    ; CHECK: [[COPY6:%[0-9]+]]:_(s32) = COPY [[C2]](s32)
+    ; CHECK: [[AND5:%[0-9]+]]:_(s32) = G_AND [[COPY6]], [[C9]]
+    ; CHECK: [[COPY7:%[0-9]+]]:_(s32) = COPY [[SHL1]](s32)
+    ; CHECK: [[OR1:%[0-9]+]]:_(s32) = G_OR [[AND5]], [[COPY7]]
+    ; CHECK: [[TRUNC1:%[0-9]+]]:_(s16) = G_TRUNC [[OR1]](s32)
+    ; CHECK: [[C11:%[0-9]+]]:_(s32) = G_CONSTANT i32 8
+    ; CHECK: [[COPY8:%[0-9]+]]:_(s32) = COPY [[C11]](s32)
+    ; CHECK: [[AND6:%[0-9]+]]:_(s32) = G_AND [[COPY8]], [[C8]]
+    ; CHECK: [[COPY9:%[0-9]+]]:_(s32) = COPY [[C5]](s32)
+    ; CHECK: [[AND7:%[0-9]+]]:_(s32) = G_AND [[COPY9]], [[C9]]
+    ; CHECK: [[SHL2:%[0-9]+]]:_(s32) = G_SHL [[AND7]], [[AND6]](s32)
+    ; CHECK: [[COPY10:%[0-9]+]]:_(s32) = COPY [[C4]](s32)
+    ; CHECK: [[AND8:%[0-9]+]]:_(s32) = G_AND [[COPY10]], [[C9]]
+    ; CHECK: [[COPY11:%[0-9]+]]:_(s32) = COPY [[SHL2]](s32)
+    ; CHECK: [[OR2:%[0-9]+]]:_(s32) = G_OR [[AND8]], [[COPY11]]
+    ; CHECK: [[TRUNC2:%[0-9]+]]:_(s16) = G_TRUNC [[OR2]](s32)
+    ; CHECK: [[TRUNC3:%[0-9]+]]:_(s16) = G_TRUNC [[C6]](s32)
+    ; CHECK: [[MV:%[0-9]+]]:_(s64) = G_MERGE_VALUES [[TRUNC]](s16), [[TRUNC1]](s16), [[TRUNC2]](s16), [[TRUNC3]](s16)
+    ; CHECK: [[TRUNC4:%[0-9]+]]:_(s56) = G_TRUNC [[MV]](s64)
+    ; CHECK: S_NOP 0, implicit [[TRUNC4]](s56)
+    %0:_(s8) = G_CONSTANT i8 0
+    %1:_(s8) = G_CONSTANT i8 1
+    %2:_(s8) = G_CONSTANT i8 2
+    %3:_(s8) = G_CONSTANT i8 3
+    %4:_(s8) = G_CONSTANT i8 4
+    %5:_(s8) = G_CONSTANT i8 5
+    %6:_(s8) = G_CONSTANT i8 6
+    %7:_(s56) = G_MERGE_VALUES %0, %1, %2, %3, %4, %5, %6
+    S_NOP 0, implicit %7
+...

unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

 //===- LegalizerHelperTest.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"
 
+using namespace LegalizeActions;
+using namespace LegalizeMutations;
+using namespace LegalityPredicates;
+
 namespace {
 
 class DummyGISelObserver : public GISelChangeObserver {
 public:
   void changingInstr(MachineInstr &MI) override {}
   void changedInstr(MachineInstr &MI) override {}
   void createdInstr(MachineInstr &MI) override {}
   void erasingInstr(MachineInstr &MI) override {}
TEST_F(GISelMITest, LowerMinMax) {
   CHECK: [[UMINV:%[0-9]+]]:_(<2 x s32>) = G_SELECT [[VCMP2]]:_(<2 x s1>), [[VEC0]]:_, [[VEC1]]:_
 
   CHECK: [[VCMP3:%[0-9]+]]:_(<2 x s1>) = G_ICMP intpred(ugt), [[VEC0]]:_(<2 x s32>), [[VEC1]]:_
   CHECK: [[UMAXV:%[0-9]+]]:_(<2 x s32>) = G_SELECT [[VCMP3]]:_(<2 x s1>), [[VEC0]]:_, [[VEC1]]:_
   )";
 
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
+
+TEST_F(GISelMITest, LowerMergeValues) {
+  if (!TM)
+    return;
+
+  const LLT S24 = LLT::scalar(24);
+  const LLT S9 = LLT::scalar(9);
+  const LLT S3 = LLT::scalar(3);
+
+  DefineLegalizerInfo(A, {
+    getActionDefinitionsBuilder(G_UNMERGE_VALUES)
+      .widenScalarIf(typeIs(1, LLT::scalar(3)), changeTo(1, LLT::scalar(9)));
+  });
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+  B.setInsertPt(*EntryMBB, EntryMBB->end());
+
+  // 24 = 3 3 3   3 3 3   3 3
+  //     => 9
+  //
+  // This can do 2 merges for the first parts, but has 2 leftover operands.
+
+
+  SmallVector<Register, 7> MergeOps;
+  for (int I = 0; I != 8; ++I)
+    MergeOps.push_back(B.buildConstant(S3, I).getReg(0));
+
+  auto Merge = B.buildMerge(S24, MergeOps);
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::UnableToLegalize,
+            Helper.lower(*Merge, 1, S9));
+}
+
 } // namespace