DAGCombine: (shl (or x, c1), c2) -> (or (shl x, c2), c1 << c2)

[RKSimon]

Summary

We already have a combine for this pattern when the input to shl
is add, so we just need to enable the transformation with the input
is or.

[arsenm]

Needs a test for the multiple use case. Tests separate from the ds_permute would also be good

[• tstellarAMD]

The pattern was matching correctly but always emitting add in the output dag,
even when int he input was or. This was fixed and I added a test case for this
behavior.

[• tstellarAMD]

Add some more tests.

[RKSimon]

Any movement on this?

[RKSimon]

Tom - do you mind if I commandeer this?

[RKSimon]

Rebaing - this is part of constant canonicalizations to try and combine MUL/SHL ops separated by ADD/OR ops.

If an AMDGPU expert can check the test changes, that'd be great.

[tstellar]

test/CodeGen/AMDGPU/fneg-fabs.f16.ll

91–95 ↗

(On Diff #102124)

Does this test generate extra shl instructions now, or did it generate those before the patch and there just weren't any check lines for it?

[RKSimon]

Refreshed test checks after rL312537 to make the diff clearer

[RKSimon]

Sorry for the late reply - this fell to the bottom of my todo pile for a while

test/CodeGen/AMDGPU/fneg-fabs.f16.ll

91–95 ↗

(On Diff #102124)

To confirm, the shifts were always there but the test didn't check for them

[RKSimon]

ping?

[efriedma]

For reference, instcombine does the same transform in InstCombiner::tryFactorization.

I can't really check the AMDGPU tests; otherwise LGTM.

[tstellar]

LGTM.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue HiBitsMask =
       DAG.getConstant(APInt::getHighBitsSet(BitSize,
                                             BitSize - N1C->getZExtValue()),
                       DL, VT);
     return DAG.getNode(ISD::AND, DL, VT, N0.getOperand(0),
                        HiBitsMask);
   }
 
   // fold (shl (add x, c1), c2) -> (add (shl x, c2), c1 << c2)
+  // fold (shl (or x, c1), c2) -> (or (shl x, c2), c1 << c2)
   // Variant of version done on multiply, except mul by a power of 2 is turned
   // into a shift.
   APInt Val;
-  if (N1C && N0.getOpcode() == ISD::ADD && N0.getNode()->hasOneUse() &&
+  if (N1C && (N0.getOpcode() == ISD::ADD || N0.getOpcode() == ISD::OR) &&
+      N0.getNode()->hasOneUse() &&
       (isa<ConstantSDNode>(N0.getOperand(1)) ||
        isConstantSplatVector(N0.getOperand(1).getNode(), Val))) {
     SDValue Shl0 = DAG.getNode(ISD::SHL, SDLoc(N0), VT, N0.getOperand(0), N1);
     SDValue Shl1 = DAG.getNode(ISD::SHL, SDLoc(N1), VT, N0.getOperand(1), N1);
-    return DAG.getNode(ISD::ADD, SDLoc(N), VT, Shl0, Shl1);
+    return DAG.getNode(N0.getOpcode(), SDLoc(N), VT, Shl0, Shl1);
   }
 
   // fold (shl (mul x, c1), c2) -> (mul x, c1 << c2)
   if (N1C && N0.getOpcode() == ISD::MUL && N0.getNode()->hasOneUse()) {
     if (ConstantSDNode *N0C1 = isConstOrConstSplat(N0.getOperand(1))) {
       if (SDValue Folded =
               DAG.FoldConstantArithmetic(ISD::SHL, SDLoc(N1), VT, N0C1, N1C))
         return DAG.getNode(ISD::MUL, SDLoc(N), VT, N0.getOperand(0), Folded);

test/CodeGen/AMDGPU/llvm.amdgcn.ds.bpermute.ll

 ; CHECK: s_waitcnt lgkmcnt
 define void @ds_bpermute_imm_offset(i32 addrspace(1)* %out, i32 %base_index, i32 %src) nounwind {
   %index = add i32 %base_index, 4
   %bpermute = call i32 @llvm.amdgcn.ds.bpermute(i32 %index, i32 %src) #0
   store i32 %bpermute, i32 addrspace(1)* %out, align 4
   ret void
 }
 
+; CHECK-LABEL: {{^}}ds_bpermute_add_shl:
+; CHECK: ds_bpermute_b32 v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}} offset:4
+; CHECK: s_waitcnt lgkmcnt
+define void @ds_bpermute_add_shl(i32 addrspace(1)* %out, i32 %base_index, i32 %src) nounwind {
+  %index = add i32 %base_index, 1
+  %byte_index = shl i32 %index, 2
+  %bpermute = call i32 @llvm.amdgcn.ds.bpermute(i32 %byte_index, i32 %src) #0
+  store i32 %bpermute, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
+; CHECK-LABEL: {{^}}ds_bpermute_or_shl:
+; CHECK: ds_bpermute_b32 v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}} offset:4
+; CHECK: s_waitcnt lgkmcnt
+define void @ds_bpermute_or_shl(i32 addrspace(1)* %out, i32 %base_index, i32 %src) nounwind {
+  %masked = and i32 %base_index, 62
+  %index = or i32 %masked, 1
+  %byte_index = shl i32 %index, 2
+  %bpermute = call i32 @llvm.amdgcn.ds.bpermute(i32 %byte_index, i32 %src) #0
+  store i32 %bpermute, i32 addrspace(1)* %out, align 4
+  ret void
+}
+
 attributes #0 = { nounwind readnone convergent }

test/CodeGen/AMDGPU/shl.ll

 ; SI-DAG: s_mov_b32 s[[K_LO:[0-9]+]], 0{{$}}
 ; SI: s_lshl_b64 s{{\[[0-9]+:[0-9]+\]}}, s{{\[}}[[K_LO]]:[[K_HI]]{{\]}}, s{{[0-9]+}}
 define void @s_shl_inline_high_imm_f32_neg_4.0_i64(i64 addrspace(1)* %out, i64 addrspace(1)* %aptr, i64 %a) {
   %shl = shl i64 13871086852301127680, %a
   store i64 %shl, i64 addrspace(1)* %out, align 8
   ret void
 }
 
+; FUNC-LABEL: {{^}}shl_or_k:
+; SI: s_lshl_b32 [[SHL:s[0-9]+]], s{{[0-9]+}}, 2
+; SI: s_or_b32 [[OR:s[0-9]+]], [[SHL]], 4
+; SI: v_mov_b32_e32 [[VOR:v[0-9]+]], [[OR]]
+; SI: buffer_store_dword [[VOR]]
+define void @shl_or_k(i32 addrspace(1)* %out, i32 %in) {
+  %tmp0 = or i32 %in, 1
+  %tmp2 = shl i32 %tmp0, 2
+  store i32 %tmp2, i32 addrspace(1)* %out
+  ret void
+}
+
+; FUNC-LABEL: {{^}}shl_or_k_two_uses:
+; SI: s_or_b32 [[OR:s[0-9]+]], s{{[0-9]+}}, 1
+; SI: s_lshl_b32 [[SHL:s[0-9]+]], [[OR]], 2
+; SI-DAG: v_mov_b32_e32 [[VOR:v[0-9]+]], [[OR]]
+; SI-DAG: v_mov_b32_e32 [[VSHL:v[0-9]+]], [[SHL]]
+; SI-DAG: buffer_store_dword [[VOR]]
+; SI-DAG: buffer_store_dword [[VSHL]]
+define void @shl_or_k_two_uses(i32 addrspace(1)* %out0, i32 addrspace(1)* %out1, i32 %in) {
+  %tmp0 = or i32 %in, 1
+  %tmp2 = shl i32 %tmp0, 2
+  store i32 %tmp2, i32 addrspace(1)* %out0
+  store i32 %tmp0, i32 addrspace(1)* %out1
+  ret void
+}
+
 attributes #0 = { nounwind readnone }