[GlobalISel][AMDGPU] add legalization for G_FREEZE

[gargaroff]

Summary

Copy the legalization rules from SelectionDAG:
-widenScalar using anyext
-narrowScalar using intermediate merges
-scalarize/fewerElements using unmerge
-moreElements using G_IMPLICIT_DEF and insert

Add G_FREEZE legalization actions to AMDGPULegalizerInfo.
Use the same legalization actions as G_IMPLICIT_DEF.

Depends on D77795.

[gargaroff]

remove misleading comment

[gargaroff]

Our backend doesn't support vectors, so I'm really not too familiar with the legalization of them. I tried to implement vector legalization to the best of my understanding, but if there are any problems with those, I'd be very grateful if you could point me in the right direction.

[arsenm]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

843–846

I've been working on eliminating these somewhat awkward functions that end up producing G_INSERT/G_EXTRACT in odd breakdowns instead of creating intermediate merges to the requested type. fewerElementsVectorBasic should be pretty close to what this should do? (Although maybe pad with undef is the wrong choice here)

856–857

Should have a dedicated buildFreeze

1772

Is this what the DAG does? I would be slightly surprised that undef high bits would be OK

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

2970

If you want some tests that are less inconvenient than unit tests, for AMDGPU this should have the same legalization action as G_IMPLICIT_DEF (and probably true for all targets)

[gargaroff]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

1772

I'm not familiar with SelectionDAG, so I might have misunderstood the extension used.

This is the code in the DAG:

SDValue DAGTypeLegalizer::PromoteIntRes_FREEZE(SDNode *N) {
  SDValue V = GetPromotedInteger(N->getOperand(0));
  return DAG.getNode(ISD::FREEZE, SDLoc(N),
                     V.getValueType(), V);
}

In the documentation of GetPromotedInteger it says that the upper bits contain rubbish, which suggested to me that this corresponds to G_ANYEXT

[aqjune]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

1772

In case of SelDag, I think PromoteIntRes_FREEZE is correct: it is creating FREEZE instruction, so even if GetPromotedInteger returned undef bits, the result is frozen, so the high bits are frozen bits.

[gargaroff]

add legalization action for G_FREEZE to AMDGPULegalizerInfo

rewrite narrowScalar for G_FREEZE to use intermediate merges instead of extract/insert

extract G_FREEZE narrowing to narrowScalarFreeze

[gargaroff]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

843–846

It was pretty close to what was needed indeed. I just had to change some parameters to some function calls, but otherwise it's mostly a copy of fewerElementsVectorBasic.

I didn't bother with refactoring both to reuse most of the unchanged code yet to get a first feedback on the narrowing action. Also I need to check how well the code can be shared in the first place. Btw, if you have a suggestion on how such a shared function could be named, I'm open for suggestions because I'm struggling to come up with a meaningful name.

Also, I decided to use undef padding here, since @aqjune said that undef bits are fine, because they get frozen anyway.

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

2970

I went ahead and added the same legalization actions for G_FREEZE as G_IMPLICIT_DEF and have basically copied all tests from legalizeImplicitDef to legalizeFreeze and modified them to use G_FREEZE.

[arsenm]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

4151

Maybe reduceOperationWidth? We have a reduceLoadStoreWidth shared between fewerElementsVector and narrowScalar

[gargaroff]

refactor narrowScalarFreeze, fewerElementsVectorBasic to general reduceOperationWidth

[arsenm]

llvm/test/CodeGen/AMDGPU/GlobalISel/legalize-freeze.mir

92 ↗

(On Diff #258258)

_48->_s48

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

case TargetOpcode::G_TRUNC: {
     }
 
     auto Unmerge = MIRBuilder.buildUnmerge(NarrowTy, MI.getOperand(1));
     MIRBuilder.buildCopy(MI.getOperand(0), Unmerge.getReg(0));
     MI.eraseFromParent();
     return Legalized;
   }
 
+  case TargetOpcode::G_FREEZE: {
+    const Register DstReg = MI.getOperand(0).getReg();
+    const Register SrcReg = MI.getOperand(1).getReg();
+    const LLT DstTy = MRI.getType(DstReg);
+    const LLT SrcTy = MRI.getType(SrcReg);
+
+    // FIXME: support narrowing to smaller vector element type
+    if (DstTy.isVector() && NarrowTy.isVector()) {
+      LLVM_DEBUG(
+          dbgs()
+          << "Cannot narrow G_FREEZE to smaller vector element type yet.");
+      return UnableToLegalize;
+    }
+
+    LLT LeftoverTy;
+    SmallVector<Register, 2> SrcRegs, LeftoverRegs, DstRegs, LeftoverDstRegs;
+    if (!extractParts(SrcReg, SrcTy, NarrowTy, LeftoverTy, SrcRegs,
+                      LeftoverRegs))
+      return UnableToLegalize;
+
+    for (Register PartReg : SrcRegs) {
+      const Register PartDstReg = MRI.createGenericVirtualRegister(NarrowTy);
+      MIRBuilder.buildInstr(TargetOpcode::G_FREEZE, {PartDstReg}, {PartReg});
+      DstRegs.emplace_back(PartDstReg);
+    }
+
+    for (Register LeftoverReg : LeftoverRegs) {
+      const Register PartDstReg = MRI.createGenericVirtualRegister(LeftoverTy);
+      MIRBuilder.buildInstr(TargetOpcode::G_FREEZE, {PartDstReg},
+                            {LeftoverReg});
+      LeftoverDstRegs.emplace_back(PartDstReg);
+    }
+
+    insertParts(DstReg, DstTy, NarrowTy, DstRegs, LeftoverTy, LeftoverDstRegs);
+
+    MI.eraseFromParent();
+    return Legalized;
+  }
+
   case TargetOpcode::G_ADD: {
     // FIXME: add support for when SizeOp0 isn't an exact multiple of
     // NarrowSize.
     if (SizeOp0 % NarrowSize != 0)
       return UnableToLegalize;
     // Expand in terms of carry-setting/consuming G_ADDE instructions.
     int NumParts = SizeOp0 / NarrowTy.getSizeInBits();
 
case TargetOpcode::G_BITREVERSE: {
     MIRBuilder.setInsertPt(MIRBuilder.getMBB(), ++MIRBuilder.getInsertPt());
 
     auto ShiftAmt = MIRBuilder.buildConstant(WideTy, DiffBits);
     auto Shift = MIRBuilder.buildLShr(WideTy, DstExt, ShiftAmt);
     MIRBuilder.buildTrunc(DstReg, Shift);
     Observer.changedInstr(MI);
     return Legalized;
   }
+  case TargetOpcode::G_FREEZE:
+    Observer.changingInstr(MI);
+    widenScalarSrc(MI, WideTy, 1, TargetOpcode::G_ANYEXT);
+    widenScalarDst(MI, WideTy);
+    Observer.changedInstr(MI);
+    return Legalized;
+
   case TargetOpcode::G_ADD:
   case TargetOpcode::G_AND:
   case TargetOpcode::G_MUL:
   case TargetOpcode::G_OR:
   case TargetOpcode::G_XOR:
   case TargetOpcode::G_SUB:
     // Perform operation at larger width (any extension is fines here, high bits
     // don't affect the result) and then truncate the result back to the
LegalizerHelper::fewerElementsVector(MachineInstr &MI, unsigned TypeIdx,
   case G_FMINNUM:
   case G_FMAXNUM:
   case G_FMINNUM_IEEE:
   case G_FMAXNUM_IEEE:
   case G_FMINIMUM:
   case G_FMAXIMUM:
   case G_FSHL:
   case G_FSHR:
+  case G_FREEZE:
     return fewerElementsVectorBasic(MI, TypeIdx, NarrowTy);
   case G_SHL:
   case G_LSHR:
   case G_ASHR:
   case G_CTLZ:
   case G_CTLZ_ZERO_UNDEF:
   case G_CTTZ:
   case G_CTTZ_ZERO_UNDEF:
LegalizerHelper::moreElementsVector(MachineInstr &MI, unsigned TypeIdx,
   case TargetOpcode::G_EXTRACT:
     if (TypeIdx != 1)
       return UnableToLegalize;
     Observer.changingInstr(MI);
     moreElementsVectorSrc(MI, MoreTy, 1);
     Observer.changedInstr(MI);
     return Legalized;
   case TargetOpcode::G_INSERT:
+  case TargetOpcode::G_FREEZE:
     if (TypeIdx != 0)
       return UnableToLegalize;
     Observer.changingInstr(MI);
     moreElementsVectorSrc(MI, MoreTy, 1);
     moreElementsVectorDst(MI, MoreTy, 0);
     Observer.changedInstr(MI);
     return Legalized;
   case TargetOpcode::G_SELECT:
if (SrcTy.isScalar() && SrcTy.getSizeInBits() == 2 * NarrowSize) {
     return Legalized;
   }
 
   return UnableToLegalize;
 }
 
 LegalizerHelper::LegalizeResult
 LegalizerHelper::lowerBitCount(MachineInstr &MI, unsigned TypeIdx, LLT Ty) {
   unsigned Opc = MI.getOpcode();
   auto &TII = *MI.getMF()->getSubtarget().getInstrInfo();
   auto isSupported = [this](const LegalityQuery &Q) {
     auto QAction = LI.getAction(Q).Action;
     return QAction == Legal || QAction == Libcall || QAction == Custom;
   };
   switch (Opc) {
   default:
     return UnableToLegalize;

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

TEST_F(AArch64GISelMITest, NarrowImplicitDef) {
   CHECK: [[DEF3:%[0-9]+]]:_(s32) = G_IMPLICIT_DEF
   CHECK: [[BV:%[0-9]+]]:_(<2 x s64>) = G_BUILD_VECTOR [[DEF]]:_(s32), [[DEF1]]:_(s32), [[DEF2]]:_(s32), [[DEF3]]:_(s32)
   )";
 
   // Check
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
 
+// Test widening of G_FREEZE
+TEST_F(AArch64GISelMITest, WidenFreeze) {
+  setUp();
+  if (!TM)
+    return;
+
+  DefineLegalizerInfo(A, {});
+
+  // Make sure that G_FREEZE is widened with anyext
+  LLT S64{LLT::scalar(64)};
+  LLT S128{LLT::scalar(128)};
+  LLT V2S32{LLT::vector(2, 32)};
+  LLT V2S64{LLT::vector(2, 64)};
+
+  auto Vector = B.buildBitcast(V2S32, Copies[0]);
+
+  auto FreezeScalar = B.buildInstr(TargetOpcode::G_FREEZE, {S64}, {Copies[0]});
+  auto FreezeVector = B.buildInstr(TargetOpcode::G_FREEZE, {V2S32}, {Vector});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+
+  // Perform Legalization
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.widenScalar(*FreezeScalar, 0, S128));
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.widenScalar(*FreezeVector, 0, V2S64));
+
+  const auto *CheckStr = R"(
+  CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[BITCAST:%[0-9]+]]:_(<2 x s32>) = G_BITCAST [[COPY]]
+
+  CHECK: [[ANYEXT:%[0-9]+]]:_(s128) = G_ANYEXT [[COPY]]
+  CHECK: [[FREEZE:%[0-9]+]]:_(s128) = G_FREEZE [[ANYEXT]]
+  CHECK: [[TRUNC:%[0-9]+]]:_(s64) = G_TRUNC [[FREEZE]]
+
+  CHECK: [[ANYEXT1:%[0-9]+]]:_(<2 x s64>) = G_ANYEXT [[BITCAST]]
+  CHECK: [[FREEZE1:%[0-9]+]]:_(<2 x s64>) = G_FREEZE [[ANYEXT1]]
+  CHECK: [[TRUNC1:%[0-9]+]]:_(<2 x s32>) = G_TRUNC [[FREEZE1]]
+  )";
+
+  // Check
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}
+
+// Test narrowing of G_FREEZE
+TEST_F(AArch64GISelMITest, NarrowFreeze) {
+  setUp();
+  if (!TM)
+    return;
+
+  DefineLegalizerInfo(A, {});
+
+  // Make sure that G_FREEZE is narrowed using unmerge/extract
+  LLT S16{LLT::scalar(16)};
+  LLT S32{LLT::scalar(32)};
+  LLT S33{LLT::scalar(33)};
+  LLT S64{LLT::scalar(64)};
+  LLT V2S16{LLT::vector(2, 16)};
+  LLT V2S32{LLT::vector(2, 32)};
+
+  auto Trunc = B.buildTrunc(S33, {Copies[0]});
+  auto Vector = B.buildBitcast(V2S32, Copies[0]);
+
+  auto FreezeScalar = B.buildInstr(TargetOpcode::G_FREEZE, {S64}, {Copies[0]});
+  auto FreezeOdd = B.buildInstr(TargetOpcode::G_FREEZE, {S33}, {Trunc});
+  auto FreezeVector = B.buildInstr(TargetOpcode::G_FREEZE, {V2S32}, {Vector});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+
+  // Perform Legalization
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.narrowScalar(*FreezeScalar, 0, S32));
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.narrowScalar(*FreezeOdd, 0, S32));
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::UnableToLegalize,
+            Helper.narrowScalar(*FreezeVector, 0, V2S16));
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.narrowScalar(*FreezeVector, 0, S16));
+
+  const auto *CheckStr = R"(
+  CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[TRUNC:%[0-9]+]]:_(s33) = G_TRUNC [[COPY]]
+  CHECK: [[BITCAST:%[0-9]+]]:_(<2 x s32>) = G_BITCAST [[COPY]]
+
+  CHECK: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[COPY]]
+  CHECK: [[FREEZE:%[0-9]+]]:_(s32) = G_FREEZE [[UV]]
+  CHECK: [[FREEZE1:%[0-9]+]]:_(s32) = G_FREEZE [[UV1]]
+  CHECK: [[MV:%[0-9]+]]:_(s64) = G_MERGE_VALUES [[FREEZE]]:_(s32), [[FREEZE1]]
+
+  CHECK: [[EXTR:%[0-9]+]]:_(s32) = G_EXTRACT [[TRUNC]]:_(s33), 0
+  CHECK: [[EXTR1:%[0-9]+]]:_(s1) = G_EXTRACT [[TRUNC]]:_(s33), 32
+  CHECK: [[FREEZE2:%[0-9]+]]:_(s32) = G_FREEZE [[EXTR]]
+  CHECK: [[FREEZE3:%[0-9]+]]:_(s1) = G_FREEZE [[EXTR1]]
+  CHECK: [[UNDEF:%[0-9]+]]:_(s33) = G_IMPLICIT_DEF
+  CHECK: [[INSERT:%[0-9]+]]:_(s33) = G_INSERT [[UNDEF]]:_, [[FREEZE2]]:_(s32), 0
+  CHECK: [[INSERT1:%[0-9]+]]:_(s33) = G_INSERT [[INSERT]]:_, [[FREEZE3]]:_(s1), 32
+
+  CHECK: [[UV:%[0-9]+]]:_(s16), [[UV1:%[0-9]+]]:_(s16), [[UV2:%[0-9]+]]:_(s16), [[UV3:%[0-9]+]]:_(s16) = G_UNMERGE_VALUES [[BITCAST]]
+  CHECK: [[FREEZE4:%[0-9]+]]:_(s16) = G_FREEZE [[UV]]
+  CHECK: [[FREEZE5:%[0-9]+]]:_(s16) = G_FREEZE [[UV1]]
+  CHECK: [[FREEZE6:%[0-9]+]]:_(s16) = G_FREEZE [[UV2]]
+  CHECK: [[FREEZE7:%[0-9]+]]:_(s16) = G_FREEZE [[UV3]]
+  CHECK: [[BV:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[FREEZE4]]:_(s16), [[FREEZE5]]:_(s16), [[FREEZE6]]:_(s16), [[FREEZE7]]
+  )";
+
+  // Check
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}
+
+// Test fewer elements of G_FREEZE
+TEST_F(AArch64GISelMITest, FewerElementsFreeze) {
+  setUp();
+  if (!TM)
+    return;
+
+  DefineLegalizerInfo(A, {});
+
+  // Make sure that G_FREEZE is narrowed using unmerge/extract
+  LLT S32{LLT::scalar(32)};
+  LLT V2S16{LLT::vector(2, 16)};
+  LLT V2S32{LLT::vector(2, 32)};
+  LLT V4S16{LLT::vector(4, 16)};
+
+  auto Vector1 = B.buildBitcast(V2S32, Copies[0]);
+  auto Vector2 = B.buildBitcast(V4S16, Copies[0]);
+
+  auto FreezeVector1 = B.buildInstr(TargetOpcode::G_FREEZE, {V2S32}, {Vector1});
+  auto FreezeVector2 = B.buildInstr(TargetOpcode::G_FREEZE, {V4S16}, {Vector2});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+
+  // Perform Legalization
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.fewerElementsVector(*FreezeVector1, 0, S32));
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.fewerElementsVector(*FreezeVector2, 0, V2S16));
+
+  const auto *CheckStr = R"(
+  CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[BITCAST:%[0-9]+]]:_(<2 x s32>) = G_BITCAST [[COPY]]
+  CHECK: [[BITCAST1:%[0-9]+]]:_(<4 x s16>) = G_BITCAST [[COPY]]
+
+  CHECK: [[UV:%[0-9]+]]:_(s32), [[UV1:%[0-9]+]]:_(s32) = G_UNMERGE_VALUES [[BITCAST]]
+  CHECK: [[FREEZE:%[0-9]+]]:_(s32) = G_FREEZE [[UV]]
+  CHECK: [[FREEZE1:%[0-9]+]]:_(s32) = G_FREEZE [[UV1]]
+  CHECK: [[MV:%[0-9]+]]:_(<2 x s32>) = G_BUILD_VECTOR [[FREEZE]]:_(s32), [[FREEZE1]]
+
+  CHECK: [[UV:%[0-9]+]]:_(<2 x s16>), [[UV1:%[0-9]+]]:_(<2 x s16>) = G_UNMERGE_VALUES [[BITCAST1]]
+  CHECK: [[FREEZE2:%[0-9]+]]:_(<2 x s16>) = G_FREEZE [[UV]]
+  CHECK: [[FREEZE3:%[0-9]+]]:_(<2 x s16>) = G_FREEZE [[UV1]]
+  CHECK: [[MV:%[0-9]+]]:_(<4 x s16>) = G_CONCAT_VECTORS [[FREEZE2]]:_(<2 x s16>), [[FREEZE3]]
+  )";
+
+  // Check
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}
+
+// Test more elements of G_FREEZE
+TEST_F(AArch64GISelMITest, MoreElementsFreeze) {
+  setUp();
+  if (!TM)
+    return;
+
+  DefineLegalizerInfo(A, {});
+
+  // Make sure that G_FREEZE is narrowed using unmerge/extract
+  LLT V2S32{LLT::vector(2, 32)};
+  LLT V4S32{LLT::vector(4, 32)};
+
+  auto Vector1 = B.buildBitcast(V2S32, Copies[0]);
+  auto FreezeVector1 = B.buildInstr(TargetOpcode::G_FREEZE, {V2S32}, {Vector1});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+
+  // Perform Legalization
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.moreElementsVector(*FreezeVector1, 0, V4S32));
+
+  const auto *CheckStr = R"(
+  CHECK: [[COPY:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[BITCAST:%[0-9]+]]:_(<2 x s32>) = G_BITCAST [[COPY]]
+
+  CHECK: [[UNDEF:%[0-9]+]]:_(<2 x s32>) = G_IMPLICIT_DEF
+  CHECK: [[CV:%[0-9]+]]:_(<4 x s32>) = G_CONCAT_VECTORS [[BITCAST]]:_(<2 x s32>), [[UNDEF]]
+  CHECK: [[FREEZE:%[0-9]+]]:_(<4 x s32>) = G_FREEZE [[CV]]
+  CHECK: [[EXTR:%[0-9]+]]:_(<2 x s32>) = G_EXTRACT [[FREEZE]]:_(<4 x s32>), 0
+  )";
+
+  // Check
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}
+
 } // namespace