[AMDGPU] Introduce new ISel combine for trunc-slr patterns

[tsymalla]

Summary

In some cases, when selecting a (trunc (slr)) pattern, the slr gets translated
to a v_lshrrev_b3e2_e64 instruction whereas the truncation gets selected to
a sequence of v_and_b32_e64 and v_cmp_eq_u32_e64. In the final ISA, this appears
as selecting the nth-bit:

v_lshrrev_b32_e32 v0, 2, v1
v_and_b32_e32 v0, 1, v0
v_cmp_eq_u32_e32 vcc_lo, 1, v0

However, when the value used in the right shift is known at compilation time, the
whole sequence can be reduced to two VALUs when the constant operand in the v_and is adjusted to (1 << lshrrev_operand):

v_and_b32_e32 v0, (1 << 2), v1
v_cmp_ne_u32_e32 vcc_lo, 0, v0

In the example above, the following pseudo-code:

v0 = (v1 >> 2)
v0 = v0 & 1
vcc_lo = (v0 == 1)

would be translated to:

v0 = v1 & 0b100
vcc_lo = (v0 == 0b100)

which should yield an equivalent result.
This is a little bit hard to test as one needs to force the SelectionDAG to
contain the nodes before instruction selection, but the test sequence was
roughly derived from a production shader.

[foad]

It seems like the xor is getting in the way. Would something like D38161 help instead?

[tsymalla]

In D118461#3279351, @foad wrote:

It seems like the xor is getting in the way. Would something like D38161 help instead?

Currently, the xor gets combined to a setcc_ne which gets combined to the srl / trunc sequence.
Initially, there is the xor / setcc_eq sequence which could be simplified like in D38161, removing the need for the xor.
Probably that would clean up everything a bit.

[arsenm]

Can this be done as a combine instead? Plus if we handle this for VALU, should also for SALU

llvm/test/CodeGen/AMDGPU/dagcombine-lshr-and-cmp.ll

13

Don't need control flow in this test. Also should test pattern for scalar and vector inputs

[foad]

In D118461#3279539, @tsymalla wrote:

In D118461#3279351, @foad wrote:

It seems like the xor is getting in the way. Would something like D38161 help instead?

Currently, the xor gets combined to a setcc_ne which gets combined to the srl / trunc sequence.
Initially, there is the xor / setcc_eq sequence which could be simplified like in D38161, removing the need for the xor.
Probably that would clean up everything a bit.

Why do we even have the xor in the IR? Normally (if you run IR optimizations as well as backend passes) instcombine would combine it into the icmp. Why didn't this happen? Was it introduced late by StructurizeCFG? Does D118478 help?

[tsymalla]

In D118461#3279674, @foad wrote:

In D118461#3279539, @tsymalla wrote:

In D118461#3279351, @foad wrote:

It seems like the xor is getting in the way. Would something like D38161 help instead?

Currently, the xor gets combined to a setcc_ne which gets combined to the srl / trunc sequence.
Initially, there is the xor / setcc_eq sequence which could be simplified like in D38161, removing the need for the xor.
Probably that would clean up everything a bit.

Why do we even have the xor in the IR? Normally (if you run IR optimizations as well as backend passes) instcombine would combine it into the icmp. Why didn't this happen? Was it introduced late by StructurizeCFG? Does D118478 help?

You are correct, the xor is not needed in the IR. Simply inverting the predicate for the icmp should be sufficient. The xor is combined into a setcc, so this should be fine by just adjusting the test.

[tsymalla]

llvm/test/CodeGen/AMDGPU/dagcombine-lshr-and-cmp.ll

13

The control flow was used to prevent having the truncate in the SDag optimized away (which is used as part of the pattern match here). I am going to check if the adjustments to the test (check comment from @foad) are going to help here.
Going to test additional cases in the new revision.

[tsymalla]

Added handling for scalar cases, improved test case.

[foad]

I'm still not sure that we need this, if the xor can be cleaned up earlier. Does D118623 help?

[tsymalla]

In D118461#3286686, @foad wrote:

I'm still not sure that we need this, if the xor can be cleaned up earlier. Does D118623 help?

Jay, unfortunately it doesn't help. I tried your patch out, but for my test case, the matcher won't apply as there is no XOR in the LLVM IR, it only gets created as SDag node. By the way, your change relates to the function "buildConditions" in the comments which should be "insertConditions".
However, the idea here is not to remove the XOR, but to remove an additional VALU which gets created by translating the TRUNCATE and the AND in the MIR separately instead of handling them as one sequence. Replacing the "setcc ne" with its inverse and not introducing an additional XOR instead might remove the need for this change, but what about possible other cases where this pattern could get matched?

[foad]

In D118461#3289676, @tsymalla wrote:

In D118461#3286686, @foad wrote:

I'm still not sure that we need this, if the xor can be cleaned up earlier. Does D118623 help?

Jay, unfortunately it doesn't help. I tried your patch out, but for my test case, the matcher won't apply as there is no XOR in the LLVM IR, it only gets created as SDag node. By the way, your change relates to the function "buildConditions" in the comments which should be "insertConditions".
However, the idea here is not to remove the XOR, but to remove an additional VALU which gets created by translating the TRUNCATE and the AND in the MIR separately instead of handling them as one sequence. Replacing the "setcc ne" with its inverse and not introducing an additional XOR instead might remove the need for this change, but what about possible other cases where this pattern could get matched?

Well there is an XOR in the LLVM IR in the test case in your patch! Can you share another test case?

As a rule of thumb, if there is a missed optimisation, I would like to try to fix it as early as possible in the pass pipeline. Otherwise SelectionDAG has to get more and more complicated, to try to optimise things that should really have been cleaned up in the IR before instruction selection.

llvm/lib/Target/AMDGPU/SIInstructions.td

2290

v_cmp_ne_u32_e64 $a, 0, $a is probably better because 0 is always an inline constant, but 1 << $b might not be.

llvm/test/CodeGen/AMDGPU/dagcombine-lshr-and-cmp.ll

1

Please auto-generate the checks with utils/update_llc_test_checks.py and pre-commit this test with the old codegen, so that this patch will clearly show how the codegen changes.

41

"No newline at end of file" :)

[tsymalla]

In D118461#3289697, @foad wrote:

In D118461#3289676, @tsymalla wrote:

In D118461#3286686, @foad wrote:

I'm still not sure that we need this, if the xor can be cleaned up earlier. Does D118623 help?

Jay, unfortunately it doesn't help. I tried your patch out, but for my test case, the matcher won't apply as there is no XOR in the LLVM IR, it only gets created as SDag node. By the way, your change relates to the function "buildConditions" in the comments which should be "insertConditions".
However, the idea here is not to remove the XOR, but to remove an additional VALU which gets created by translating the TRUNCATE and the AND in the MIR separately instead of handling them as one sequence. Replacing the "setcc ne" with its inverse and not introducing an additional XOR instead might remove the need for this change, but what about possible other cases where this pattern could get matched?

Well there is an XOR in the LLVM IR in the test case in your patch! Can you share another test case?

As a rule of thumb, if there is a missed optimisation, I would like to try to fix it as early as possible in the pass pipeline. Otherwise SelectionDAG has to get more and more complicated, to try to optimise things that should really have been cleaned up in the IR before instruction selection.

I think I misunderstood what you were saying. The pattern that gets matched here is inside the entry block of both functions. The XOR in the second test case (uniform) is there to prevent the truncate in the SDag from being optimized away and has no semantic relevance for the actual test. BTW, your patch still doesn't apply for me in this case. I agree with optimizing the XOR.
Thanks for your additional comments.

[foad]

The XOR in the second test case (uniform) is there to prevent the truncate in the SDag from being optimized away and has no semantic relevance for the actual test.

OK, sorry, I did not read the test carefully enough. I have looked at it more carefully now.

So the problem is that this DAG:

    t4: i32 = and t2, Constant:i32<2>
  t7: i1 = setcc t4, Constant:i32<0>, setne:ch
t9: i1,i64,ch = llvm.amdgcn.if t0, TargetConstant:i64<1324>, t7

gets optimized like this by a generic combine implemented in TargetLowering::SimplifySetCC:

    t24: i32 = srl t2, Constant:i64<1>
  t25: i1 = truncate t24
t9: i1,i64,ch = llvm.amdgcn.if t0, TargetConstant:i64<1324>, t25

I guess there is no way we can undo that with another combine, because they will end up fighting each other. So I think your patch is reasonable.

llvm/lib/Target/AMDGPU/SIInstructions.td

2290

Also, if you do this, there will only be one use of the constant not two, so I don't think you will have to "Restrict the range to prevent using an additional VGPR for the shifted value".

llvm/test/CodeGen/AMDGPU/dagcombine-lshr-and-cmp.ll

1

Actually it will have to use update_mir_test_checks.

[tsymalla]

Change requests from review addressed.

llvm/lib/Target/AMDGPU/SIInstructions.td

2290

The resulting value should still be checked to ensure no 32-bit overflow occurs, correct? For instance, if the shift value is something like 33, 1 << 33 would exceed Int32_Max.

[foad]

llvm/lib/Target/AMDGPU/SIInstructions.td

2290

I'm not sure there is any need to check. The result of a shift by 33 is undefined, so it doesn't really matter what code we generate in that case.

[tsymalla]

llvm/lib/Target/AMDGPU/SIInstructions.td

2290

Sure, I’ll remove the check.

[tsymalla]

Removed range checks.

[foad]

LGTM, thanks! Just very minor comments inline.

llvm/lib/Target/AMDGPU/SIInstructions.td

2273

Don't need parens around the << expression.

llvm/test/CodeGen/AMDGPU/dagcombine-lshr-and-cmp.ll

1

"dagcombine" is not a very good name for the file, because this is isel not a combine, but I guess it's OK if you've already committed the file.

llvm/lib/Target/AMDGPU/SIInstructions.td

def : GCNPat <
   (V_CMP_EQ_U32_e64 (V_AND_B32_e64 (i32 1), $a), (i32 1))
 >;
 
 def : GCNPat <
   (i1 (DivergentUnaryFrag<trunc> i16:$a)),
   (V_CMP_EQ_U32_e64 (V_AND_B32_e64 (i32 1), $a), (i32 1))
 >;
 
+// Restrict the range to prevent using an additional VGPR
+// for the shifted value.
+def IMMBitSelRange : ImmLeaf <i32, [{
+  return Imm > 0 && Imm < 16;
+}]>;
+
+def IMMBitSelConst : SDNodeXForm<imm, [{
+  return CurDAG->getTargetConstant((1 << N->getZExtValue()), SDLoc(N),
+                                   MVT::i32);
+}]>;
+
+// Matching separate SRL and TRUNC instructions
+// with dependent operands (SRL dest is source of TRUNC)
+// generates three instructions. However, by using bit shifts,
+// the V_LSHRREV_B32_e64 result can be directly used in the
+// operand of the V_AND_B32_e64 instruction:
+// (trunc i32 (srl i32 $a, i32 $b)) ->
+// v_and_b32_e64 $a, (1 << $b), $a
+// v_cmp_eq_u32_e64 $a, (1 << $b), $a 
+
+// Handle the VALU case.
+def : GCNPat <
+  (i1 (DivergentUnaryFrag<trunc> (i32 (srl i32:$a, IMMBitSelRange:$b)))),
+  (V_CMP_EQ_U32_e64 (V_AND_B32_e64 (i32 (IMMBitSelConst $b)), $a),
+    (i32 (IMMBitSelConst $b)))
+>;
+
+// Handle the scalar case.
+def : GCNPat <
+  (i1 (UniformUnaryFrag<trunc> (i32 (srl i32:$a, IMMBitSelRange:$b)))),
+  (S_CMP_EQ_U32 (S_AND_B32 (i32 (IMMBitSelConst $b)), $a),
+    (i32 (IMMBitSelConst $b)))
+>;
+
 def : GCNPat <
   (i1 (DivergentUnaryFrag<trunc> i64:$a)),
   (V_CMP_EQ_U32_e64 (V_AND_B32_e64 (i32 1),
                     (i32 (EXTRACT_SUBREG $a, sub0))), (i32 1))
 >;
 
 def : GCNPat <
   (i32 (bswap i32:$a)),

llvm/test/CodeGen/AMDGPU/dagcombine-lshr-and-cmp.ll

+; RUN: llc -march=amdgcn -mtriple=amdgcn-- -stop-after=amdgpu-isel -verify-machineinstrs -O0 < %s | FileCheck -check-prefix=GCN %s
+
+; GCN-LABEL: bb.0.entry:
+; GCN-NOT:      V_LSHRREV_B32_e64
+; GCN:          V_AND_B32_e64 2
+; GCN:          V_CMP_EQ_U32_e64 killed {{.*}}, 2
+define i32 @divergent_lshr_and_cmp(i32 %x) {
+entry:
+  %0 = and i32 %x, 2
+  %1 = icmp ne i32 %0, 0
+  ; Prevent removal of truncate in SDag by inserting llvm.amdgcn.if
+  br i1 %1, label %out.true, label %out.else
+
+out.true:
+  %2 = shl i32 %x, 2
+  ret i32 %2
+
+out.else:
+  ret i32 %x
+}
+
+; GCN-LABEL: bb.0.entry:
+; GCN:          S_AND_B32 2
+; GCN:          S_CMP_EQ_U32 killed %{{.*}}, 2
+define amdgpu_kernel void @uniform_opt_lshr_and_cmp(i1 addrspace(1)* %out, i32 %x) {
+entry:
+  %0 = and i32 %x, 2
+  %1 = icmp ne i32 %0, 0
+  ; Don't optimize the truncate in the SDag away.
+  br i1 %1, label %out.true, label %out.else
+
+out.true:
+  %2 = xor i1 %1, -1
+  store i1 %2, i1 addrspace(1)* %out
+  ret void
+
+out.else:
+  store i1 %1, i1 addrspace(1)* %out
+  ret void
+}
 No newline at end of file