[Bug] PTODSL failed to compile when a variable carried by a while loop is conditionally reassigned before the loop #1558

MLIRError: Verification failed:
error: unknown: operand #0 does not dominate this use

from ptodsl import pto


@pto.jit(name="while_carry_merge", kernel_kind="vector", target="a5", mode="explicit")
def kernel(out: pto.ptr(pto.i32, "gm"), n: pto.i32):
    buf = pto.castptr(pto.const(0, dtype=pto.i64), pto.ptr(pto.ui8, "ub"))
    acc = pto.const(0.0, dtype=pto.f32)
    if n > 0:
        acc = pto.load(pto.castptr(buf, pto.ptr(pto.f32, "ub")), 100)
    i = pto.const(0, dtype=pto.i32)
    while i < n:
        acc = acc + pto.load(pto.castptr(buf, pto.ptr(pto.f32, "ub")), pto.cast(i, pto.i64))
        i = pto.cast(i + 1, pto.i32)
    pto.store(i, out, pto.const(0, dtype=pto.i64))


if __name__ == "__main__":
    kernel.compile()
    print("compiled OK")

Traceback (most recent call last):
  File "test_while_carry.py", line 18, in <module>
    while_carry_merge_repro.compile()
  File "ptodsl/ptodsl/_jit.py", line 343, in compile
    return self._compiler.compile(**constexpr_bindings)
  File "ptodsl/ptodsl/_kernel_compilation.py", line 142, in compile
    compiled.build()
  File "ptodsl/ptodsl/_tracing/artifacts.py", line 36, in build
    built = self._module_factory()
  File "ptodsl/ptodsl/_tracing/runtime.py", line 98, in build_module
    self.verify_module(module)
  File "ptodsl/ptodsl/_tracing/runtime.py", line 52, in verify_module
    module.operation.verify()
ptoas.mlir._mlir_libs._site_initialize.<locals>.MLIRError: Verification failed:
error: unknown: operand #0 does not dominate this use
 note: unknown: see current operation: 
  %6:2 = "scf.while"(%19, %5) ({
  ^bb0(%arg4: f32, %arg5: i32):
    %16 = "pto.cmpi"(%arg5, %arg1) <{predicate = #pto.scalar_cmp_predicate<lt>, signedness = #pto.signedness<signed>}> : (i32, i32) -> i1
    "scf.condition"(%16, %arg4, %arg5) : (i1, f32, i32) -> ()
  }, {
  ^bb0(%arg2: f32, %arg3: i32):
    %9 = "pto.castptr"(%1) : (!pto.ptr<ui8, ub>) -> !pto.ptr<f32, ub>
    %10 = "pto.exti"(%arg3) <{signedness = #pto.signedness<signed>}> : (i32) -> i64
    %11 = "pto.index_cast"(%10) <{signedness = #pto.signedness<signed>}> : (i64) -> index
    %12 = "pto.load"(%9, %11) : (!pto.ptr<f32, ub>, index) -> f32
    %13 = "pto.addf"(%arg2, %12) <{fastmath = #pto.fastmath<none>}> : (f32, f32) -> f32
    %14 = "pto.constant"() <{value = 1 : i32}> : () -> i32
    %15 = "pto.addi"(%arg3, %14) <{overflowFlags = #pto.overflow<none>}> : (i32, i32) -> i32
    "scf.yield"(%13, %15) : (f32, i32) -> ()
  }) : (f32, i32) -> (f32, i32)
 note: unknown: operand defined here (op in a child region)