[AMDGPU] Fixed incorrect break from loop

[tpr]

Summary

Lower control flow did not correctly handle the case that a loop break
in if/else was on a condition that was not guaranteed to be masked by
exec. The first test kernel shows an example of this going wrong; after
exiting the loop, exec is all ones, even if it was not before the loop.

The fix is for lowering of if-break and else-break to insert an
S_AND_B64 to mask the break condition with exec. This commit also
includes the optimization of not inserting that S_AND_B64 if it is
obviously not needed because the break condition is the result of a
V_CMP in the same basic block.

[arsenm]

lib/Target/AMDGPU/SILowerControlFlow.cpp

353 ↗

(On Diff #136857)

Braces

354 ↗

(On Diff #136857)

Can you add a testcase where this will fail if there is no defining instruction (e.g. the source value was a function argument)

356 ↗

(On Diff #136857)

This doesn't seem like the right condition to me. This should probably be a predicate that checks the specific operand for being a result that is masked with exec? Theoretically this could be something exotic like a write lane.

This also could be = parentmatch && isFoo() rather than an if

362–363 ↗

(On Diff #136857)

Identation of these is weird

[rampitec]

lib/Target/AMDGPU/SILowerControlFlow.cpp

354 ↗

(On Diff #136857)

Can this be a partial def?

[rampitec]

I am not sure why this new and instruction shall be needed. As far as I understand one of the arguments of SI_IF_BREAK shall be full copy of exec mask saved before the loop, isn't it?
If that is not so, shall we fix it instead?

[tpr]

The arguments of SI_IF_BREAK appear to be the "I want to break this time" mask and the running "already exited loop" mask. Therefore we need to ensure that the first one of those is masked with exec before it is ORed into the second one. See the test case for an example of where it is needed.

lib/Target/AMDGPU/SILowerControlFlow.cpp

354 ↗

(On Diff #136857)

I don't think it can be a partial def, as we know that it is something that was an i1 in the IR.

354 ↗

(On Diff #136857)

I'll give the suggested test a try, but I suspect it won't like me passing an i1 as an arg.

356 ↗

(On Diff #136857)

The operand is the break condition, which is something that was an i1 in the IR, but is now a pair of sgprs (including the case of vcc). If it's a valu, I don't think that can be anything other than an instruction with carry-out, i.e. v_cmp (or possibly the carry-out from v_add_u). I should probably add a comment to this effect.

As an optimization, it could maybe also spot other cases, principally an s_and_b64 with exec that's already there. But I don't think that is needed for this commit.

[tpr]

V2: Addressed some review comments.

[nhaehnle]

LGTM

[rampitec]

LGTM

[arsenm]

Is there anything holding up this patch? I've run into this problem in another testcase

[tpr]

Apologies for the delay. I found a higher than expected number of test cases where this adds an extra unnecessary "and" instruction. So I plan to revisit this.

Perhaps I should land it anyway, and address the unnecessary "and" instructions in a future patch.

llvm/trunk/lib/Target/AMDGPU/SILowerControlFlow.cpp

MachineInstr *Or = BuildMI(MBB, &MI, DL, TII->get(AMDGPU::S_OR_B64), Dst)
                          .add(MI.getOperand(1));
 
   if (LIS)
     LIS->ReplaceMachineInstrInMaps(MI, *Or);
   MI.eraseFromParent();
 }
 
 void SILowerControlFlow::emitIfBreak(MachineInstr &MI) {
-  MI.setDesc(TII->get(AMDGPU::S_OR_B64));
+  MachineBasicBlock &MBB = *MI.getParent();
+  const DebugLoc &DL = MI.getDebugLoc();
+  auto Dst = MI.getOperand(0).getReg();
+
+  // Skip ANDing with exec if the break condition is already masked by exec
+  // because it is a V_CMP in the same basic block. (We know the break
+  // condition operand was an i1 in IR, so if it is a VALU instruction it must
+  // be one with a carry-out.)
+  bool SkipAnding = false;
+  if (MI.getOperand(1).isReg()) {
+    if (MachineInstr *Def = MRI->getUniqueVRegDef(MI.getOperand(1).getReg())) {
+      SkipAnding = Def->getParent() == MI.getParent()
+          && SIInstrInfo::isVALU(*Def);
+    }
+  }
+
+  // AND the break condition operand with exec, then OR that into the "loop
+  // exit" mask.
+  MachineInstr *And = nullptr, *Or = nullptr;
+  if (!SkipAnding) {
+    And = BuildMI(MBB, &MI, DL, TII->get(AMDGPU::S_AND_B64), Dst)
+             .addReg(AMDGPU::EXEC)
+             .add(MI.getOperand(1));
+    Or = BuildMI(MBB, &MI, DL, TII->get(AMDGPU::S_OR_B64), Dst)
+             .addReg(Dst)
+             .add(MI.getOperand(2));
+  } else
+    Or = BuildMI(MBB, &MI, DL, TII->get(AMDGPU::S_OR_B64), Dst)
+             .add(MI.getOperand(1))
+             .add(MI.getOperand(2));
+
+  if (LIS) {
+    if (And)
+      LIS->InsertMachineInstrInMaps(*And);
+    LIS->ReplaceMachineInstrInMaps(MI, *Or);
+  }
+
+  MI.eraseFromParent();
 }
 
 void SILowerControlFlow::emitElseBreak(MachineInstr &MI) {
-  MI.setDesc(TII->get(AMDGPU::S_OR_B64));
+  // Lowered in the same way as emitIfBreak above.
+  emitIfBreak(MI);
 }
 
 void SILowerControlFlow::emitLoop(MachineInstr &MI) {
   MachineBasicBlock &MBB = *MI.getParent();
   const DebugLoc &DL = MI.getDebugLoc();
 
   MachineInstr *AndN2 =
       BuildMI(MBB, &MI, DL, TII->get(AMDGPU::S_ANDN2_B64_term), AMDGPU::EXEC)

llvm/trunk/test/CodeGen/AMDGPU/loop_exit_with_xor.ll

+; RUN: llc -mtriple=amdgcn--amdpal -mcpu=gfx803 -verify-machineinstrs < %s | FileCheck -check-prefix=GCN %s
+
+; Where the mask of lanes wanting to exit the loop on this iteration is not
+; obviously already masked by exec (in this case, the xor with -1 inserted by
+; control flow annotation), then lower control flow must insert an S_AND_B64
+; with exec.
+
+; GCN-LABEL: {{^}}needs_and:
+; GCN: s_xor_b64 [[REG1:[^ ,]*]], {{[^ ,]*, -1$}}
+; GCN: s_and_b64 [[REG2:[^ ,]*]], exec, [[REG1]]
+; GCN: s_or_b64 [[REG3:[^ ,]*]], [[REG2]],
+; GCN: s_andn2_b64 exec, exec, [[REG3]]
+
+define void @needs_and(i32 %arg) {
+entry:
+  br label %loop
+
+loop:
+  %tmp23phi = phi i32 [ %tmp23, %endif ], [ 0, %entry ]
+  %tmp23 = add nuw i32 %tmp23phi, 1
+  %tmp27 = icmp ult i32 %arg, %tmp23
+  br i1 %tmp27, label %then, label %endif
+
+then:                                             ; preds = %bb
+  call void @llvm.amdgcn.buffer.store.f32(float undef, <4 x i32> undef, i32 0, i32 undef, i1 false, i1 false) #1
+  br label %endif
+
+endif:                                             ; preds = %bb28, %bb
+  br i1 %tmp27, label %loop, label %loopexit
+
+loopexit:
+  ret void
+}
+
+; Where the mask of lanes wanting to exit the loop on this iteration is
+; obviously already masked by exec (a V_CMP), then lower control flow can omit
+; the S_AND_B64 to avoid an unnecessary instruction.
+
+; GCN-LABEL: {{^}}doesnt_need_and:
+; GCN: v_cmp{{[^ ]*}} [[REG1:[^ ,]*]]
+; GCN: s_or_b64 [[REG2:[^ ,]*]], [[REG1]],
+; GCN: s_andn2_b64 exec, exec, [[REG2]]
+
+define void @doesnt_need_and(i32 %arg) {
+entry:
+  br label %loop
+
+loop:
+  %tmp23phi = phi i32 [ %tmp23, %loop ], [ 0, %entry ]
+  %tmp23 = add nuw i32 %tmp23phi, 1
+  %tmp27 = icmp ult i32 %arg, %tmp23
+  call void @llvm.amdgcn.buffer.store.f32(float undef, <4 x i32> undef, i32 0, i32 undef, i1 false, i1 false) #1
+  br i1 %tmp27, label %loop, label %loopexit
+
+loopexit:
+  ret void
+}
+
+; Another case where the mask of lanes wanting to exit the loop is not masked
+; by exec, because it is a function parameter.
+
+; GCN-LABEL: {{^}}break_cond_is_arg:
+; GCN: s_xor_b64 [[REG1:[^ ,]*]], {{[^ ,]*, -1$}}
+; GCN: s_and_b64 [[REG2:[^ ,]*]], exec, [[REG1]]
+; GCN: s_or_b64 [[REG3:[^ ,]*]], [[REG2]],
+; GCN: s_andn2_b64 exec, exec, [[REG3]]
+
+define void @break_cond_is_arg(i32 %arg, i1 %breakcond) {
+entry:
+  br label %loop
+
+loop:
+  %tmp23phi = phi i32 [ %tmp23, %endif ], [ 0, %entry ]
+  %tmp23 = add nuw i32 %tmp23phi, 1
+  %tmp27 = icmp ult i32 %arg, %tmp23
+  br i1 %tmp27, label %then, label %endif
+
+then:                                             ; preds = %bb
+  call void @llvm.amdgcn.buffer.store.f32(float undef, <4 x i32> undef, i32 0, i32 undef, i1 false, i1 false) #1
+  br label %endif
+
+endif:                                             ; preds = %bb28, %bb
+  br i1 %breakcond, label %loop, label %loopexit
+
+loopexit:
+  ret void
+}
+
+
+declare void @llvm.amdgcn.buffer.store.f32(float, <4 x i32>, i32, i32, i1, i1) #3
+
+attributes #3 = { nounwind writeonly }
+

llvm/trunk/test/CodeGen/AMDGPU/multilevel-break.ll

 ; GCN: [[INNER_LOOP:BB[0-9]+_[0-9]+]]: ; %LOOP{{$}}
 ; GCN: s_and_saveexec_b64 [[SAVE_BREAK:s\[[0-9]+:[0-9]+\]]], vcc
 
 ; GCN: BB{{[0-9]+}}_{{[0-9]+}}: ; %Flow{{$}}
 ; GCN-NEXT: ; in Loop: Header=[[INNER_LOOP]] Depth=2
 
 ; Ensure extra or eliminated
 ; GCN-NEXT: s_or_b64 exec, exec, [[SAVE_BREAK]]
-; GCN-NEXT: s_or_b64 [[OR_BREAK:s\[[0-9]+:[0-9]+\]]], [[SAVE_BREAK]], s{{\[[0-9]+:[0-9]+\]}}
+; GCN-NEXT: s_mov_b64
+; GCN-NEXT: s_and_b64 [[MASKED_SAVE_BREAK:s\[[0-9]+:[0-9]+\]]], exec, [[SAVE_BREAK]]
+; GCN-NEXT: s_or_b64 [[OR_BREAK:s\[[0-9]+:[0-9]+\]]], [[MASKED_SAVE_BREAK]], s{{\[[0-9]+:[0-9]+\]}}
 ; GCN-NEXT: s_andn2_b64 exec, exec, [[OR_BREAK]]
 ; GCN-NEXT: s_cbranch_execnz [[INNER_LOOP]]
 
 ; GCN: ; %bb.{{[0-9]+}}: ; %Flow1{{$}}
 ; GCN-NEXT: ; in Loop: Header=[[OUTER_LOOP]] Depth=1
 
 ; Ensure copy is eliminated
 ; GCN-NEXT: s_or_b64 exec, exec, [[OR_BREAK]]
-; GCN-NEXT: s_or_b64 [[OUTER_OR_BREAK:s\[[0-9]+:[0-9]+\]]], [[SAVE_BREAK]], s{{\[[0-9]+:[0-9]+\]}}
+; GCN-NEXT: s_and_b64 [[MASKED2_SAVE_BREAK:s\[[0-9]+:[0-9]+\]]], exec, [[SAVE_BREAK]]
+; GCN-NEXT: s_or_b64 [[OUTER_OR_BREAK:s\[[0-9]+:[0-9]+\]]], [[MASKED2_SAVE_BREAK]], s{{\[[0-9]+:[0-9]+\]}}
 ; GCN-NEXT: s_andn2_b64 exec, exec, [[OUTER_OR_BREAK]]
 ; GCN-NEXT: s_cbranch_execnz [[OUTER_LOOP]]
 define amdgpu_vs void @multi_else_break(<4 x float> %vec, i32 %ub, i32 %cont) {
 main_body:
   br label %LOOP.outer
 
 LOOP.outer:                                       ; preds = %ENDIF, %main_body
   %tmp43 = phi i32 [ 0, %main_body ], [ %tmp47, %ENDIF ]

llvm/trunk/test/CodeGen/AMDGPU/valu-i1.ll

 ; SI: BB{{[0-9]+_[0-9]+}}: ; %bb20
 ; SI: buffer_store_dword
 ; SI: v_cmp_ge_i64_e{{32|64}} [[CMP:s\[[0-9]+:[0-9]+\]|vcc]]
 ; SI: s_or_b64 [[TMP:s\[[0-9]+:[0-9]+\]]], [[CMP]], [[COND_STATE]]
 
 ; SI: [[LABEL_FLOW]]:
 ; SI-NEXT: ; in Loop: Header=[[LABEL_LOOP]]
 ; SI-NEXT: s_or_b64 exec, exec, [[ORNEG3]]
-; SI-NEXT: s_or_b64 [[COND_STATE]], [[ORNEG3]], [[TMP]]
+; SI-NEXT: s_mov_b64 [[MOVED_TMP:s\[[0-9]+:[0-9]+\]]], [[TMP]]
+; SI-NEXT: s_and_b64 [[MASKED_ORNEG3:s\[[0-9]+:[0-9]+\]]], exec, [[ORNEG3]]
+; SI-NEXT: s_or_b64 [[COND_STATE]], [[MASKED_ORNEG3]], [[MOVED_TMP]]
 ; SI-NEXT: s_andn2_b64 exec, exec, [[COND_STATE]]
 ; SI-NEXT: s_cbranch_execnz [[LABEL_LOOP]]
 
 ; SI: [[LABEL_EXIT]]:
 ; SI-NOT: [[COND_STATE]]
 ; SI: s_endpgm
 
 define amdgpu_kernel void @multi_vcond_loop(i32 addrspace(1)* noalias nocapture %arg, i32 addrspace(1)* noalias nocapture readonly %arg1, i32 addrspace(1)* noalias nocapture readonly %arg2, i32 addrspace(1)* noalias nocapture readonly %arg3) #1 {