[DAG] getNode() - fold (zext (trunc x)) -> x iff the upper bits are known zero - add SRL support

[RKSimon]

Summary

This is part of the work to address the D155472 regressions, there's a number of issues with generalizing this fold which is why I'm just adding SRL support atm.

The fold encourages the creation of ISD::FSHR nodes which was breaking a lot of the AMDGPUISD::PERM tests, so I've added ISD::FSHR handling to the matchPERM helper

[goldstein.w.n]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5698 ↗

(On Diff #557123)

Why does the ISD::AssertZext need one use check?

But while you're at it think ISD::AND also work here right?

[jrbyrnes]

what is the best way to generalize the performOrCombine handling to support ISD::FSHR as well?

We should ideally produce AMDGPUISD::PERM whenever we have an i32 that is byte permutation of two sources (except some special cases) and it's best to do this via DAGCombining due to target-specific reasons. In practice, typically such byte permutations have ISD::OR as the root of the tree (which is why I have begun with that combine), however, other cases exist where the tree does not have ISD::OR as the root (e.g. AMDGPUISD::PERM, ISD::FSHR, etc). It seems to me, the best way to generalize is to extract the common code from ISD::OR, and use it to enable new roots / entry points -- this is an extension of the work we should do anyway. I hacked together an implementation of what I have just described and this patch no longer causes the v_perm miss regressions. @foad @arsenm objections?

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5698 ↗

(On Diff #557123)

Everything should work here - eventually the getOpcode() filters will go entirely - the problem we're having is a number of AArch64 mul folds create various zext(trunc()) patterns and then expect them to still exist later in the combine.

As usual with LLVM it comes down to yak shaving :)

[RKSimon]

@jrbyrnes I've pulled the AMDGPUISD::PERM matching out of performOrCombine, made it into a 'matchPERM' helper and setup ISD::FSHR to use it - is this what you had in mind?

[jrbyrnes]

Disregard previous version of comment -- I did not see the most recent changes.

[jrbyrnes]

Can you please bring in tests from https://github.com/llvm/llvm-project/pull/66965

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

10780

Can we also fail on vector types

10796

This seems to be broken if Index = 1; BytesProvided; = 2 and ByteShift = 3 -- we will attempt to get the 2nd byte of a 16 bit operand which is out of range (only 0,1 are valid bytes).

Perhaps it was meant to be: Index + (BytesProvided - ByteShift) ? But this still doesnt account for the case where ByteShift > n * BytesProvided.

How do you feel about

case ISD::FSHR: {
  auto ShiftOp = dyn_cast<ConstantSDNode>(Op->getOperand(2));
  if (!ShiftOp || Op.getValueType().isVector())
    return std::nullopt;

  uint64_t BitsProvided = Op.getValueSizeInBits();
  if (BitsProvided % 8 != 0)
    return std::nullopt;

  uint64_t BitShift = ShiftOp->getAPIntValue().urem(BitsProvided);
  if (BitShift % 8)
    return std::nullopt;

  uint64_t ConcatSizeInBytes = BitsProvided / 4;
  uint64_t ByteShift = BitShift / 8;

  auto NewIndex = (Index + ByteShift) % ConcatSizeInBytes;
  auto BytesProvided = BitsProvided / 8;
  auto NextOp = Op.getOperand(NewIndex >= BytesProvided ? 0 : 1);
  NewIndex %= BytesProvided;
  return calculateSrcByte(NextOp, StartingIndex, NewIndex);
}

llvm/test/CodeGen/AMDGPU/permute.ll

127 ↗

(On Diff #557146)

I can look into this after this patch lands.

[jrbyrnes]

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

10796

Oops -- should have been

return calculateByteProvider(NextOp, NewIndex, Depth + 1, StartingIndex

[RKSimon]

Update with @jrbyrnes suggestions and add early-out so we prefer an existing ISD::FSHR node to a AMDGPUISD::PERM equivalent.

[RKSimon]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

5698 ↗

(On Diff #557123)

@goldstein.w.n Hopefully D159537 will address the last regressions which will allow me to remove the OpOp.getOpcode() filter entirely.

[jrbyrnes]

Thanks -- the AMDGPUISD::PERM side LGTM

[RKSimon]

Thanks - any further comments?

[foad]

@jrbyrnes shouldn't we land the AMDGPU changes separately first? Or would it all be dead code without the rest of this patch?

[RKSimon]

In D159533#4649980, @foad wrote:

@jrbyrnes shouldn't we land the AMDGPU changes separately first? Or would it all be dead code without the rest of this patch?

I can pre-commit the matchPERM helper as a NFC-ish to reduce the diff if that helps?

[RKSimon]

rebase

[jrbyrnes]

In D159533#4649980, @foad wrote:

@jrbyrnes shouldn't we land the AMDGPU changes separately first? Or would it all be dead code without the rest of this patch?

The changes to calculateByteProvider may actually be NFC because sub 32bit ISD::FSHRs get legalized into ISD::OR which are already handled, but enabling ISD::FSHR combine and installing matchPERM there is not.

The changes look okay, but it seems it would be best to land them separately first.

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

setTargetDAGCombine({ISD::ADD,
                        ISD::SMIN,
                        ISD::SMAX,
                        ISD::UMIN,
                        ISD::UMAX,
                        ISD::SETCC,
                        ISD::AND,
                        ISD::OR,
                        ISD::XOR,
+                       ISD::FSHR,
                        ISD::SINT_TO_FP,
                        ISD::UINT_TO_FP,
                        ISD::FCANONICALIZE,
                        ISD::SCALAR_TO_VECTOR,
                        ISD::ZERO_EXTEND,
                        ISD::SIGN_EXTEND_INREG,
                        ISD::EXTRACT_VECTOR_ELT,
                        ISD::INSERT_VECTOR_ELT,
if ((IndexMask & BitMask) != IndexMask) {
       if (IndexMask & BitMask)
         return std::nullopt;
       return ByteProvider<SDValue>::getConstantZero();
     }
 
     return calculateSrcByte(Op->getOperand(0), StartingIndex, Index);
   }
 
+  case ISD::FSHR: {
+    // fshr(X,Y,Z): (X << (BW - (Z % BW))) | (Y >> (Z % BW))
+    auto ShiftOp = dyn_cast<ConstantSDNode>(Op->getOperand(2));
+    if (!ShiftOp || Op.getValueType().isVector())
+      return std::nullopt;
+
+    uint64_t BitsProvided = Op.getValueSizeInBits();
+    if (BitsProvided % 8 != 0)
+      return std::nullopt;
+
+    uint64_t BitShift = ShiftOp->getAPIntValue().urem(BitsProvided);
+    if (BitShift % 8)
+      return std::nullopt;
+
+    uint64_t ConcatSizeInBytes = BitsProvided / 4;
+    uint64_t ByteShift = BitShift / 8;
+
+    uint64_t NewIndex = (Index + ByteShift) % ConcatSizeInBytes;
+    uint64_t BytesProvided = BitsProvided / 8;
+    SDValue NextOp = Op.getOperand(NewIndex >= BytesProvided ? 0 : 1);
+    NewIndex %= BytesProvided;
+    return calculateByteProvider(NextOp, NewIndex, Depth + 1, StartingIndex);
+  }
+
   case ISD::SRA:
   case ISD::SRL: {
     auto ShiftOp = dyn_cast<ConstantSDNode>(Op->getOperand(1));
     if (!ShiftOp)
       return std::nullopt;
 
     uint64_t BitShift = ShiftOp->getZExtValue();
     if (BitShift % 8)
for (size_t i = 0; i < PermNodes.size(); i++) {
     assert(!DAG.getDataLayout().isBigEndian());
     PermMask |= (PermOp.SrcOffset + SrcByteAdjust) << (i * 8);
   }
 
   SDValue Op = *PermNodes[FirstSrc].Src;
   SDValue OtherOp = SecondSrc.has_value() ? *PermNodes[*SecondSrc].Src
                                           : *PermNodes[FirstSrc].Src;
 
+  // Check that we haven't just recreated the same FSHR node.
+  if (N->getOpcode() == ISD::FSHR &&
+      (N->getOperand(0) == Op || N->getOperand(0) == OtherOp) &&
+      (N->getOperand(1) == Op || N->getOperand(1) == OtherOp))
+    return SDValue();
+
   // Check that we are not just extracting the bytes in order from an op
   if (Op == OtherOp && Op.getValueSizeInBits() == 32) {
     int Low16 = PermMask & 0xffff;
     int Hi16 = (PermMask & 0xffff0000) >> 16;
 
     bool WellFormedLow = (Low16 == 0x0504) || (Low16 == 0x0100);
     bool WellFormedHi = (Hi16 == 0x0706) || (Hi16 == 0x0302);
 
SDValue SITargetLowering::PerformDAGCombine(SDNode *N,
   case AMDGPUISD::FMAX_LEGACY:
     return performMinMaxCombine(N, DCI);
   case ISD::FMA:
     return performFMACombine(N, DCI);
   case ISD::AND:
     return performAndCombine(N, DCI);
   case ISD::OR:
     return performOrCombine(N, DCI);
+  case ISD::FSHR: {
+    const SIInstrInfo *TII = getSubtarget()->getInstrInfo();
+    if (N->getValueType(0) == MVT::i32 && N->isDivergent() &&
+        TII->pseudoToMCOpcode(AMDGPU::V_PERM_B32_e64) != -1) {
+      return matchPERM(N, DCI);
+    }
+    break;
+  }
   case ISD::XOR:
     return performXorCombine(N, DCI);
   case ISD::ZERO_EXTEND:
     return performZeroExtendCombine(N, DCI);
   case ISD::SIGN_EXTEND_INREG:
     return performSignExtendInRegCombine(N , DCI);
   case AMDGPUISD::FP_CLASS:
     return performClassCombine(N, DCI);

llvm/test/CodeGen/AMDGPU/permute_i8.ll

 ; GFX10-NEXT:    global_load_dword v10, v[2:3], off
 ; GFX10-NEXT:    v_mov_b32_e32 v0, 16
 ; GFX10-NEXT:    v_mov_b32_e32 v1, 0xff
 ; GFX10-NEXT:    v_lshlrev_b16 v2, 8, v4
 ; GFX10-NEXT:    s_waitcnt vmcnt(1)
 ; GFX10-NEXT:    v_lshrrev_b32_sdwa v0, v0, v9 dst_sel:BYTE_1 dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:DWORD
 ; GFX10-NEXT:    s_waitcnt vmcnt(0)
 ; GFX10-NEXT:    v_and_b32_sdwa v1, v10, v1 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:DWORD
-; GFX10-NEXT:    v_or_b32_sdwa v0, v9, v0 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_0 src1_sel:DWORD
+; GFX10-NEXT:    v_or_b32_sdwa v0, v9, v0 dst_sel:WORD_1 dst_unused:UNUSED_PAD src0_sel:BYTE_0 src1_sel:DWORD
 ; GFX10-NEXT:    v_or_b32_e32 v1, v1, v2
-; GFX10-NEXT:    v_lshlrev_b32_e32 v2, 16, v0
-; GFX10-NEXT:    v_alignbit_b32 v0, v0, v10, 16
-; GFX10-NEXT:    v_or_b32_sdwa v1, v1, v2 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_0 src1_sel:DWORD
-; GFX10-NEXT:    global_store_dword v[5:6], v1, off
-; GFX10-NEXT:    global_store_dword v[7:8], v0, off
+; GFX10-NEXT:    v_or_b32_sdwa v0, v1, v0 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_0 src1_sel:DWORD
+; GFX10-NEXT:    v_perm_b32 v1, v10, v9, 0x2000706
+; GFX10-NEXT:    global_store_dword v[5:6], v0, off
+; GFX10-NEXT:    global_store_dword v[7:8], v1, off
 ; GFX10-NEXT:    s_setpc_b64 s[30:31]
 ;
 ; GFX9-LABEL: ive_store_div:
 ; GFX9:       ; %bb.0:
 ; GFX9-NEXT:    s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0)
 ; GFX9-NEXT:    v_and_b32_e32 v9, 0x3ff, v31
 ; GFX9-NEXT:    v_lshlrev_b32_e32 v9, 2, v9
 ; GFX9-NEXT:    v_add_co_u32_e32 v0, vcc, v0, v9
 ; GFX9-NEXT:    v_addc_co_u32_e32 v1, vcc, 0, v1, vcc
 ; GFX9-NEXT:    v_add_co_u32_e32 v2, vcc, v2, v9
 ; GFX9-NEXT:    v_addc_co_u32_e32 v3, vcc, 0, v3, vcc
 ; GFX9-NEXT:    global_load_dword v9, v[0:1], off
 ; GFX9-NEXT:    global_load_dword v10, v[2:3], off
 ; GFX9-NEXT:    s_movk_i32 s4, 0xff
 ; GFX9-NEXT:    v_lshlrev_b16_e32 v0, 8, v4
+; GFX9-NEXT:    s_mov_b32 s5, 0x2000706
 ; GFX9-NEXT:    s_waitcnt vmcnt(1)
 ; GFX9-NEXT:    v_lshrrev_b32_e32 v1, 16, v9
-; GFX9-NEXT:    v_lshlrev_b16_e32 v1, 8, v1
 ; GFX9-NEXT:    s_waitcnt vmcnt(0)
 ; GFX9-NEXT:    v_and_b32_sdwa v2, v10, s4 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:DWORD
-; GFX9-NEXT:    v_or_b32_sdwa v1, v9, v1 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:BYTE_0 src1_sel:DWORD
+; GFX9-NEXT:    v_lshlrev_b16_e32 v1, 8, v1
 ; GFX9-NEXT:    v_or_b32_e32 v0, v2, v0
-; GFX9-NEXT:    v_alignbit_b32 v2, v1, v10, 16
-; GFX9-NEXT:    v_lshlrev_b32_e32 v1, 16, v1
+; GFX9-NEXT:    v_or_b32_sdwa v1, v9, v1 dst_sel:WORD_1 dst_unused:UNUSED_PAD src0_sel:BYTE_0 src1_sel:DWORD
 ; GFX9-NEXT:    v_or_b32_sdwa v0, v0, v1 dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:WORD_0 src1_sel:DWORD
+; GFX9-NEXT:    v_perm_b32 v3, v10, v9, s5
 ; GFX9-NEXT:    global_store_dword v[5:6], v0, off
-; GFX9-NEXT:    global_store_dword v[7:8], v2, off
+; GFX9-NEXT:    global_store_dword v[7:8], v3, off
 ; GFX9-NEXT:    s_waitcnt vmcnt(0)
 ; GFX9-NEXT:    s_setpc_b64 s[30:31]
   %tid = call i32 @llvm.amdgcn.workitem.id.x()
   %gep0 = getelementptr <4 x i8>, ptr addrspace(1) %in0, i32 %tid
   %gep1 = getelementptr <4 x i8>, ptr addrspace(1) %in1, i32 %tid
   %vec0 = load <4 x i8>, ptr addrspace(1) %gep0, align 4
   %vec1 = load <4 x i8>, ptr addrspace(1) %gep1, align 4
   %shuffle0_0 = shufflevector <4 x i8> %vec0, <4 x i8> %vec1, <4 x i32> <i32 6, i32 7, i32 0, i32 2>