MichaelJKlaiber opened a new issue, #16224:
URL: https://github.com/apache/tvm/issues/16224
T.if_then_else vanishes after blockize
```
from tvm.script import tir as T
from tvm import tir
from tvm.ir import IRModule
@T.prim_func
def func(A: T.Buffer((128, 129), "float32"), B: T.Buffer((128, 129),
"float32")):
for i, j in T.grid(128, 129):
with T.block("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = T.if_then_else(vj<129, A[vi, vj] * 2.0, 0,
dtype="int32")
sch = tir.Schedule(func)
i, j = sch.get_loops(sch.get_block("B"))
print(sch.mod)
sch.blockize(j)
print(sch.mod)
```
### Expected behavior
```
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 129), "float32"), B: T.Buffer((128, 129),
"float32")):
T.func_attr({"global_symbol": "func"})
# with T.block("root"):
for i in range(128):
with T.block("B_o"):
vi_o = T.axis.spatial(128, i)
vj_o = T.axis.spatial(1, 0)
T.reads(A[vi_o, 0:129])
T.writes(B[vi_o, 0:129])
for j in range(129):
with T.block("B"):
vj_i = T.axis.spatial(129, j)
T.reads(A[vi_o, vj_i])
T.writes(B[vi_o, vj_i])
B[vi, vj] = T.if_then_else(vj < 129, A[vi, vj] *
T.float32(2), T.float32(0))
```
### Actual behavior
```
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 129), "float32"), B: T.Buffer((128, 129),
"float32")):
T.func_attr({"global_symbol": "func"})
# with T.block("root"):
for i in range(128):
with T.block("B_o"):
vi_o = T.axis.spatial(128, i)
vj_o = T.axis.spatial(1, 0)
T.reads(A[vi_o, 0:129])
T.writes(B[vi_o, 0:129])
for j in range(129):
with T.block("B"):
vj_i = T.axis.spatial(129, j)
T.reads(A[vi_o, vj_i])
T.writes(B[vi_o, vj_i])
B[vi_o, vj_i] = A[vi_o, vj_i] * T.float32(2)
```
### Environment
main a9fcac1a47f4b57c1d9d828c7ed1d77452b02cc5
Dec10 23
### Steps to reproduce
```
from tvm.script import tir as T
from tvm import tir
from tvm.ir import IRModule
@T.prim_func
def func(A: T.Buffer((128, 129), "float32"), B: T.Buffer((128, 129),
"float32")):
for i, j in T.grid(128, 129):
with T.block("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = T.if_then_else(vj<129, A[vi, vj] * 2.0, 0,
dtype="int32")
sch = tir.Schedule(func)
i, j = sch.get_loops(sch.get_block("B"))
print(sch.mod)
sch.blockize(j)
print(sch.mod)
```
Output:
```
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 129), "float32"), B: T.Buffer((128, 129),
"float32")):
T.func_attr({"global_symbol": "func"})
# with T.block("root"):
for i, j in T.grid(128, 129):
with T.block("B"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(A[vi, vj])
T.writes(B[vi, vj])
B[vi, vj] = T.if_then_else(vj < 129, A[vi, vj] *
T.float32(2), T.float32(0))
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer((128, 129), "float32"), B: T.Buffer((128, 129),
"float32")):
T.func_attr({"global_symbol": "func"})
# with T.block("root"):
for i in range(128):
with T.block("B_o"):
vi_o = T.axis.spatial(128, i)
vj_o = T.axis.spatial(1, 0)
T.reads(A[vi_o, 0:129])
T.writes(B[vi_o, 0:129])
for j in range(129):
with T.block("B"):
vj_i = T.axis.spatial(129, j)
T.reads(A[vi_o, vj_i])
T.writes(B[vi_o, vj_i])
B[vi_o, vj_i] = A[vi_o, vj_i] * T.float32(2)
```
### Triage
* tir
* tir:schedule
* tir: transform
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
