apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.66k stars 3.45k forks source link

[Compute-inline] Prefer T.where for reverse compute-inlined block with predicate #17128

Closed wrongtest-intellif closed 3 months ago

wrongtest-intellif commented 3 months ago

Hi, we found a kind of corner bad cases for reverse inline primitive. Currently we encode predicate produced by inlining as block compute body, like

with T.block():
    if pred:
        body()

While other primitives (typically cache_write) currently may only respect the predicate of T.where for region estimation. The newly created blocks take the risk of region out-of-bound write accesses, lead to segfaults even in cpu target computations.

The change proposes to use equavalent T.where when possible to make it more consistent with s-tir system, like

with T.block():
    T.where(pred_on_loop_vars)
    body()
wrongtest-intellif commented 3 months ago

@tvm-bot re-run