[Mlir-commits] [mlir] [mlir][spirv] Split conditional basic blocks during deserialization (PR #127639)

[Igor Wodiany]

IgWod-IMG wrote:

I have replaced `std::distance()` and updated test, so hopefully it's more readable now.

Regarding the bigger issue. I have redesigned the code so now it only splits conditional blocks that are headers of selection regions. Loops are not touched. Looking at loops (https://mlir.llvm.org/docs/Dialects/SPIR-V/#loop) the entry block will always be unconditional, so splitting would not be applied here. Then conditional blocks that form the loop header and continue will be fully sunk anyway, regardless of splitting, so we can skip splitting them.

So, to sum up. This PR only targets selection regions (which so far were the only problem on my side when it comes to sinking), and if we need anything for loops, we can work it out in the future.

@mishaobu please let me know if the new patch causes any issues. Loops should be untouched now, so they should keep their structure (unless something else is going on). The `if` statement looked the same in both code snippets, so I hope splitting doesn't cause any issues there. Just to verify, I attempted to run your `raymarchwater` (GLSL -> SPIR-V -> MLIR -> SPIR-V -> GLSL) and I got the function as you would with your patch. Let me know your thoughts.

There is still an extra block in one of the tests:

```mlir
  spirv.func @selection_cf() "None" {
    %0 = spirv.Variable : !spirv.ptr<i1, Function>
    spirv.Branch ^bb1
  ^bb1:  // pred: ^bb0
    %true = spirv.Constant true
    spirv.Branch ^bb2
  ^bb2:  // pred: ^bb1
    spirv.mlir.selection {
      spirv.BranchConditional %true, ^bb1, ^bb2
    ^bb1:  // pred: ^bb0
      %true_1 = spirv.Constant true
      spirv.Store "Function" %0, %true_1 : i1
      spirv.Branch ^bb3
    ^bb2:  // pred: ^bb0
      %false = spirv.Constant false
      spirv.Store "Function" %0, %false : i1
      spirv.Branch ^bb3
    ^bb3:  // 2 preds: ^bb1, ^bb2
      spirv.mlir.merge
    }
```

Compared against:

```mlir
  spirv.func @selection_cf() "None" {
    %0 = spirv.Variable : !spirv.ptr<i1, Function>
    spirv.Branch ^bb1
  ^bb1:  // pred: ^bb0
    spirv.mlir.selection {
      %true = spirv.Constant true
      spirv.BranchConditional %true, ^bb1, ^bb2
    ^bb1:  // pred: ^bb0
      %true_1 = spirv.Constant true
      spirv.Store "Function" %0, %true_1 : i1
      spirv.Branch ^bb3
    ^bb2:  // pred: ^bb0
      %false = spirv.Constant false
      spirv.Store "Function" %0, %false : i1
      spirv.Branch ^bb3
    ^bb3:  // 2 preds: ^bb1, ^bb2
      spirv.mlir.merge
    }
```

It happens because the `%true` constant gets pushed outside the selection region. Arguably this works as intended; the header block gets split. From my perspective this is a rather minor issue (that will probably get optimised at some later stages anyway), and I think a price worth paying for a more robust deserialization. And improving it is something that can be addressed in the future. @kuhar do you have any thoughts on that?

Nevertheless, if we cannot reach a consensus, I'm happy to look for an alternative fix. Potentially taking @mishaobu implementation and getting it into a mergeable state.

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