Incorrect code for waterfall loop in VS shader

If I write a waterfall loop to implement my own version of NonUniformResourceIndex

// dxc /Tps_6_0 .\break.hlsl /DOUTPUT=SV_Target
// dxc /Tvs_6_0 .\break.hlsl /DOUTPUT=Z
StructuredBuffer<int> mainBuf[]: register(t2, space0);

[RootSignature("DescriptorTable(SRV(t2, numDescriptors=UNBOUNDED))")]
int main(int a : A, int b : B, int c : C) : OUTPUT
{
  int res = 0;

  for (;;) {
      int u = WaveReadLaneFirst(a);
      if (a == u) {
          res += mainBuf[u][b];
          break;
        }
    }
  return res;
}

It works fine if I compile the shader as a PS (dxc /Tps_6_0 .\break.hlsl /DOUTPUT=SV_Target). The buffer load is kept in the body of the loop using the dx.break mechanism.

  %1 = load i32, i32* getelementptr inbounds ([1 x i32], [1 x i32]* @dx.break.cond, i32 0, i32 0)
  %2 = icmp eq i32 %1, 0
  %3 = call i32 @dx.op.loadInput.i32(i32 4, i32 1, i32 0, i8 0, i32 undef)  ; LoadInput(inputSigId,rowIndex,colIndex,gsVertexAxis)
  %4 = call i32 @dx.op.loadInput.i32(i32 4, i32 0, i32 0, i8 0, i32 undef)  ; LoadInput(inputSigId,rowIndex,colIndex,gsVertexAxis)
  br label %5

; <label>:5                                       ; preds = %15, %0
  %6 = phi i32 [ 0, %0 ], [ %16, %15 ]
  %7 = call i32 @dx.op.waveReadLaneFirst.i32(i32 118, i32 %4)  ; WaveReadLaneFirst(value)
  %8 = icmp eq i32 %4, %7
  br i1 %8, label %9, label %15

; <label>:9                                       ; preds = %5
  %10 = add i32 %7, 2
  %11 = call %dx.types.Handle @dx.op.createHandle(i32 57, i8 0, i32 0, i32 %10, i1 false)  ; CreateHandle(resourceClass,rangeId,index,nonUniformIndex)
  %12 = call %dx.types.ResRet.i32 @dx.op.bufferLoad.i32(i32 68, %dx.types.Handle %11, i32 %3, i32 0)  ; BufferLoad(srv,index,wot)
  %13 = extractvalue %dx.types.ResRet.i32 %12, 0
  %14 = add nsw i32 %13, %6
  br i1 %2, label %17, label %15

; <label>:15                                      ; preds = %9, %5
  %16 = phi i32 [ %14, %9 ], [ %6, %5 ]
  br label %5

; <label>:17                                      ; preds = %9
  call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 0, i8 0, i32 %14)  ; StoreOutput(outputSigId,rowIndex,colIndex,value)
  ret void

However, if I compile the shader as a VS (dxc /Tvs_6_0 .\break.hlsl /DOUTPUT=Z) the buffer load is moved outside the loop body.

  %1 = call i32 @dx.op.loadInput.i32(i32 4, i32 0, i32 0, i8 0, i32 undef)  ; LoadInput(inputSigId,rowIndex,colIndex,gsVertexAxis)
  br label %2

; <label>:2                                       ; preds = %2, %0
  %3 = call i32 @dx.op.waveReadLaneFirst.i32(i32 118, i32 %1)  ; WaveReadLaneFirst(value)
  %4 = icmp eq i32 %1, %3
  br i1 %4, label %5, label %2

; <label>:5                                       ; preds = %2
  %6 = call i32 @dx.op.loadInput.i32(i32 4, i32 1, i32 0, i8 0, i32 undef)  ; LoadInput(inputSigId,rowIndex,colIndex,gsVertexAxis)
  %7 = add i32 %3, 2
  %8 = call %dx.types.Handle @dx.op.createHandle(i32 57, i8 0, i32 0, i32 %7, i1 false)  ; CreateHandle(resourceClass,rangeId,index,nonUniformIndex)
  %9 = call %dx.types.ResRet.i32 @dx.op.bufferLoad.i32(i32 68, %dx.types.Handle %8, i32 %6, i32 0)  ; BufferLoad(srv,index,wot)
  %10 = extractvalue %dx.types.ResRet.i32 %9, 0
  call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 0, i8 0, i32 %10)  ; StoreOutput(outputSigId,rowIndex,colIndex,value)
  ret void

I don't see any reason why the VS codegen should be different here.

The dx.break mechanism was added in #2795 to keep wave-sensitive values in a loop.