[GlobalISel][Legalizer] Implement lower action for G_FSHL/G_FSHR

[kschwarz]

Summary

This is a 1:1 port from the SelectionDAG code, modulo the recent optimizations for known non-zero shift amounts.
Additionally, no conversion between G_FSHL <-> G_FSHR are done yet.

[foad]

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

6246–6247

Doesn't this lowering work for vectors?

6252

Is DstTy guaranteed to be the same type as the shift amount operand?

6254–6255

Can't you just use auto ShAmt = MIRBuilder.buildWhatever(...) to avoid explicitly calling createGenericVirtualRegister here and below?

6255–6256

Do we have something that will combine this UREM into an AND if BW is a power of two?

[arsenm]

Duplicate D76500

[kschwarz]

In D89480#2332584, @arsenm wrote:

Duplicate D76500

I somehow forgot about this open review. I'll abandon this then.
Is there anything we can help with to get D76500 integrated? Since a recent upstream change we are seeing G_FSHL's being generated, which our GISel-only backend cannot handle.

[arsenm]

In D89480#2332914, @kschwarz wrote:

In D89480#2332584, @arsenm wrote:

Duplicate D76500

I somehow forgot about this open review. I'll abandon this then.
Is there anything we can help with to get D76500 integrated? Since a recent upstream change we are seeing G_FSHL's being generated, which our GISel-only backend cannot handle.

The main thing holding that up is I haven't had time to address the problem of avoiding expanding the division in the divide-by-constant case

llvm/include/llvm/CodeGen/GlobalISel/LegalizerHelper.h

public:
   LegalizeResult lowerAddSubSatToMinMax(MachineInstr &MI);
   LegalizeResult lowerAddSubSatToAddoSubo(MachineInstr &MI);
   LegalizeResult lowerShlSat(MachineInstr &MI);
   LegalizeResult lowerBswap(MachineInstr &MI);
   LegalizeResult lowerBitreverse(MachineInstr &MI);
   LegalizeResult lowerReadWriteRegister(MachineInstr &MI);
   LegalizeResult lowerSMULH_UMULH(MachineInstr &MI);
   LegalizeResult lowerSelect(MachineInstr &MI);
-
+  LegalizeResult lowerFunnelShift(MachineInstr &MI);
 };
 
 /// Helper function that creates a libcall to the given \p Name using the given
 /// calling convention \p CC.
 LegalizerHelper::LegalizeResult
 createLibcall(MachineIRBuilder &MIRBuilder, const char *Name,
               const CallLowering::ArgInfo &Result,
               ArrayRef<CallLowering::ArgInfo> Args, CallingConv::ID CC);

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

auto Shift =
         MIRBuilder.buildAShr(DstTy, OpReg, ShiftAmt);
     auto Add = MIRBuilder.buildAdd(DstTy, OpReg, Shift);
     MIRBuilder.buildXor(MI.getOperand(0).getReg(), Add, Shift);
     MI.eraseFromParent();
     return Legalized;
   }
   case G_SELECT:
     return lowerSelect(MI);
+
+  case G_FSHL:
+  case G_FSHR:
+    return lowerFunnelShift(MI);
   }
 }
 
 Align LegalizerHelper::getStackTemporaryAlignment(LLT Ty,
                                                   Align MinAlign) const {
   // FIXME: We're missing a way to go back from LLT to llvm::Type to query the
   // datalayout for the preferred alignment. Also there should be a target hook
   // for this to allow targets to reduce the alignment and ignore the
if (MaskTy.getSizeInBits() != Op1Ty.getSizeInBits())
     return UnableToLegalize;
 
   auto NotMask = MIRBuilder.buildNot(MaskTy, MaskReg);
   auto NewOp1 = MIRBuilder.buildAnd(MaskTy, Op1Reg, MaskReg);
   auto NewOp2 = MIRBuilder.buildAnd(MaskTy, Op2Reg, NotMask);
   MIRBuilder.buildOr(DstReg, NewOp1, NewOp2);
   MI.eraseFromParent();
   return Legalized;
-}
+}
+
+LegalizerHelper::LegalizeResult
+LegalizerHelper::lowerFunnelShift(MachineInstr &MI) {
+  Register DstReg = MI.getOperand(0).getReg();
+  Register XReg = MI.getOperand(1).getReg();
+  Register YReg = MI.getOperand(2).getReg();
+  Register ZReg = MI.getOperand(3).getReg();
+
+  bool IsFSHL = MI.getOpcode() == TargetOpcode::G_FSHL;
+
+  LLT DstTy = MRI.getType(DstReg);
+
+  if (DstTy.isVector())
+    return UnableToLegalize;
+
+  // fshl: X << (Z % BW) | Y >> 1 >> (BW - 1 - (Z % BW))
+  // fshr: X << 1 << (BW - 1 - (Z % BW)) | Y >> (Z % BW)
+
+  auto Mask = MIRBuilder.buildConstant(DstTy, DstTy.getScalarSizeInBits() - 1);
+  auto BitWidthC = MIRBuilder.buildConstant(DstTy, DstTy.getScalarSizeInBits());
+  Register ShAmtReg = MRI.createGenericVirtualRegister(DstTy);
+  MIRBuilder.buildInstr(TargetOpcode::G_UREM, {ShAmtReg},
+                        {ZReg, BitWidthC.getReg(0)});
+
+  Register InvShAmtReg = MRI.createGenericVirtualRegister(DstTy);
+  MIRBuilder.buildSub(InvShAmtReg, Mask.getReg(0), ShAmtReg);
+
+  auto One = MIRBuilder.buildConstant(DstTy, 1);
+
+  Register ShX = MRI.createGenericVirtualRegister(DstTy);
+  Register ShY = MRI.createGenericVirtualRegister(DstTy);
+  if (IsFSHL) {
+    MIRBuilder.buildShl(ShX, XReg, ShAmtReg);
+    Register ShY1 = MRI.createGenericVirtualRegister(DstTy);
+    MIRBuilder.buildLShr(ShY1, YReg, One);
+    MIRBuilder.buildLShr(ShY, ShY1, InvShAmtReg);
+  } else {
+    Register ShX1 = MRI.createGenericVirtualRegister(DstTy);
+    MIRBuilder.buildShl(ShX1, XReg, One);
+    MIRBuilder.buildShl(ShX, ShX1, InvShAmtReg);
+    MIRBuilder.buildLShr(ShY, YReg, ShAmtReg);
+  }
+
+  MIRBuilder.buildOr(DstReg, ShX, ShY);
+
+  MI.eraseFromParent();
+  return Legalized;
+}

llvm/unittests/CodeGen/GlobalISel/LegalizerHelperTest.cpp

TEST_F(AArch64GISelMITest, widenScalarUnmerge) {
   CHECK: [[MERGE:%[0-9]+]]:_(s48) = G_MERGE_VALUES [[UNMERGE3]]:_(s16), [[UNMERGE4]]:_(s16), [[UNMERGE5]]:_(s16)
   CHECK: [[MERGE1:%[0-9]+]]:_(s48) = G_MERGE_VALUES [[UNMERGE6]]:_(s16), [[UNMERGE7]]:_(s16), [[UNMERGE8]]:_(s16)
   )";
 
   // Check
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
 
+// Test lower funnel shift left
+TEST_F(AArch64GISelMITest, lowerFunnelShiftLeft) {
+  setUp();
+  if (!TM)
+    return;
+
+  DefineLegalizerInfo(A, {});
+
+  LLT S64{LLT::scalar(64)};
+
+  auto FSHL = B.buildInstr(TargetOpcode::G_FSHL, {S64},
+                           {Copies[0], Copies[1], Copies[2]});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+
+  // Perform Legalization
+  B.setInsertPt(*EntryMBB, FSHL->getIterator());
+
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.lower(*FSHL, 0, LLT::scalar(64)));
+
+  const auto *CheckStr = R"(
+
+  CHECK: [[COPY0:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[COPY1:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[COPY2:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[SIXTY3:%[0-9]+]]:_(s64) = G_CONSTANT i64 63
+  CHECK: [[SIXTY4:%[0-9]+]]:_(s64) = G_CONSTANT i64 64
+  CHECK: [[UREM:%[0-9]+]]:_(s64) = G_UREM [[COPY2]]:_, [[SIXTY4]]:_
+  CHECK: [[SUB:%[0-9]+]]:_(s64) = G_SUB [[SIXTY3]]:_, [[UREM]]:_
+  CHECK: [[ONE:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
+  CHECK: [[SHL:%[0-9]+]]:_(s64) = G_SHL [[COPY0]]:_, [[UREM]]:_(s64)
+  CHECK: [[LSHR:%[0-9]+]]:_(s64) = G_LSHR [[COPY1]]:_, [[ONE]]:_(s64)
+  CHECK: [[LSHR2:%[0-9]+]]:_(s64) = G_LSHR [[LSHR]]:_, [[SUB]]:_(s64)
+  CHECK: [[OR:%[0-9]+]]:_(s64) = G_OR [[SHL]]:_, [[LSHR2]]:_
+  )";
+
+  // Check
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}
+
+// Test lower funnel shift rigth
+TEST_F(AArch64GISelMITest, lowerFunnelShiftRight) {
+  setUp();
+  if (!TM)
+    return;
+
+  DefineLegalizerInfo(A, {});
+
+  LLT S64{LLT::scalar(64)};
+
+  auto FSHR = B.buildInstr(TargetOpcode::G_FSHR, {S64},
+                           {Copies[0], Copies[1], Copies[2]});
+
+  AInfo Info(MF->getSubtarget());
+  DummyGISelObserver Observer;
+  LegalizerHelper Helper(*MF, Info, Observer, B);
+
+  // Perform Legalization
+  B.setInsertPt(*EntryMBB, FSHR->getIterator());
+
+  EXPECT_EQ(LegalizerHelper::LegalizeResult::Legalized,
+            Helper.lower(*FSHR, 0, LLT::scalar(64)));
+
+  const auto *CheckStr = R"(
+
+  CHECK: [[COPY0:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[COPY1:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[COPY2:%[0-9]+]]:_(s64) = COPY
+  CHECK: [[SIXTY3:%[0-9]+]]:_(s64) = G_CONSTANT i64 63
+  CHECK: [[SIXTY4:%[0-9]+]]:_(s64) = G_CONSTANT i64 64
+  CHECK: [[UREM:%[0-9]+]]:_(s64) = G_UREM [[COPY2]]:_, [[SIXTY4]]:_
+  CHECK: [[SUB:%[0-9]+]]:_(s64) = G_SUB [[SIXTY3]]:_, [[UREM]]:_
+  CHECK: [[ONE:%[0-9]+]]:_(s64) = G_CONSTANT i64 1
+  CHECK: [[SHL:%[0-9]+]]:_(s64) = G_SHL [[COPY0]]:_, [[ONE]]:_(s64)
+  CHECK: [[SHL2:%[0-9]+]]:_(s64) = G_SHL [[SHL]]:_, [[SUB]]:_(s64)
+  CHECK: [[LSHR:%[0-9]+]]:_(s64) = G_LSHR [[COPY1]]:_, [[UREM]]:_(s64)
+  CHECK: [[OR:%[0-9]+]]:_(s64) = G_OR [[SHL2]]:_, [[LSHR]]:_
+  )";
+
+  // Check
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}
+
 } // namespace