[llvm] [AMDGPU] Fix combineMasks dropping condition (PR #203180)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>

The problem is related to `in S_AND exec, (S_AND x, x)` case

When the nested mask op is the outer S_AND/S_OR first operand with two identical operands, combineMasks kept exec instead of a nested operand, folding to S_AND exec, exec and dropping the condition

---
Full diff: https://github.com/llvm/llvm-project/pull/203180.diff


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp (+15-6) 
- (modified) llvm/test/CodeGen/AMDGPU/lower-control-flow-other-terminators.mir (+36) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp b/llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp
index 9cc86e84407b1..51e370b3e5d3f 100644
--- a/llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp
+++ b/llvm/lib/Target/AMDGPU/SILowerControlFlow.cpp
@@ -607,15 +607,24 @@ void SILowerControlFlow::combineMasks(MachineInstr &MI) {
   findMaskOperands(MI, 2, Ops);
   if (Ops.size() != 3) return;
 
-  unsigned UniqueOpndIdx;
-  if (Ops[0].isIdenticalTo(Ops[1])) UniqueOpndIdx = 2;
-  else if (Ops[0].isIdenticalTo(Ops[2])) UniqueOpndIdx = 1;
-  else if (Ops[1].isIdenticalTo(Ops[2])) UniqueOpndIdx = 1;
-  else return;
+  // Ops holds the exec operand and the two operands of the nested op.
+  // Always keep a nested operand, never the exec operand.
+  unsigned ExecIdx = OpToReplace == 1 ? 2 : 0;
+  unsigned NestedLHS = OpToReplace == 1 ? 0 : 1;
+  unsigned NestedRHS = NestedLHS + 1;
+
+  unsigned KeepIdx;
+  if (Ops[ExecIdx].isIdenticalTo(Ops[NestedLHS]))
+    KeepIdx = NestedRHS;
+  else if (Ops[ExecIdx].isIdenticalTo(Ops[NestedRHS]) ||
+           Ops[NestedLHS].isIdenticalTo(Ops[NestedRHS]))
+    KeepIdx = NestedLHS;
+  else
+    return;
 
   Register Reg = MI.getOperand(OpToReplace).getReg();
   MI.removeOperand(OpToReplace);
-  MI.addOperand(Ops[UniqueOpndIdx]);
+  MI.addOperand(Ops[KeepIdx]);
   if (MRI->use_empty(Reg))
     MRI->getUniqueVRegDef(Reg)->eraseFromParent();
 }
diff --git a/llvm/test/CodeGen/AMDGPU/lower-control-flow-other-terminators.mir b/llvm/test/CodeGen/AMDGPU/lower-control-flow-other-terminators.mir
index eaf398fd51723..6feabac4bb2fd 100644
--- a/llvm/test/CodeGen/AMDGPU/lower-control-flow-other-terminators.mir
+++ b/llvm/test/CodeGen/AMDGPU/lower-control-flow-other-terminators.mir
@@ -270,3 +270,39 @@ body:             |
     S_BRANCH %bb.2
 
 ...
+
+# combineMasks must keep a nested operand, not the exec leaf, when the nested
+# S_AND is the first operand and has identical operands (S_AND %0, %0). The
+# result must be S_AND exec, %0, not S_AND exec, exec.
+
+---
+name: combine_masks_nested_first_operand
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: combine_masks_nested_first_operand
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.0(0x40000000), %bb.1(0x40000000)
+  ; CHECK-NEXT:   liveins: $sgpr0_sgpr1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:sreg_64 = COPY $sgpr0_sgpr1
+  ; CHECK-NEXT:   [[S_AND_B64_:%[0-9]+]]:sreg_64 = S_AND_B64 $exec, [[COPY]], implicit-def $scc
+  ; CHECK-NEXT:   $exec = S_ANDN2_B64_term $exec, [[S_AND_B64_]], implicit-def $scc
+  ; CHECK-NEXT:   S_CBRANCH_EXECNZ %bb.0, implicit $exec
+  ; CHECK-NEXT:   S_BRANCH %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   S_ENDPGM 0
+  bb.0:
+    successors: %bb.0, %bb.1
+    liveins: $sgpr0_sgpr1
+
+    %0:sreg_64 = COPY $sgpr0_sgpr1
+    %1:sreg_64 = S_AND_B64 %0, %0, implicit-def dead $scc
+    %2:sreg_64 = S_AND_B64 %1, $exec, implicit-def $scc
+    SI_LOOP %2, %bb.0, implicit-def $exec, implicit-def $scc, implicit $exec
+    S_BRANCH %bb.1
+
+  bb.1:
+    S_ENDPGM 0
+
+...

``````````

</details>


https://github.com/llvm/llvm-project/pull/203180