quic-sanirudh commented on PR #16642:
URL: https://github.com/apache/tvm/pull/16642#issuecomment-2032679685
@Lunderberg After this pass introduced support for computing struct_info for
PrimFuncs, I'm seeing a case where a primfunc whose output buffer is modified
with the `sch.transform_layout` primitive and added back to the mod during
legalize like what happens in the legalization of `R.layout_transform` causes a
failure later in `call_tir_rewrite` pass.
A small test case that reproduces the issue is something like below:
```python
import tvm
from tvm.script import relax as R, tir as T, ir as I
from tvm import tir, relax
@I.ir_module
class Before:
@R.function
def main(x: R.Tensor((16,), dtype="float32"), y: R.Tensor((16,),
dtype="float32")) -> R.Tensor((16,), dtype="float32"):
with R.dataflow():
lv: R.Tensor((4, 4), dtype="float32") = R.layout_transform(x,
index_map=lambda i: (i // 4, i % 4), pad_value=None)
gv: R.Tensor((4, 4), dtype="float32") = lv
R.output(gv)
return gv
if __name__ == '__main__':
mod = Before
mod = relax.transform.LegalizeOps()(mod)
mod = relax.transform.CallTIRRewrite()(mod)
```
From what I've been able to understand, the `FuncStructInfo` for the
PrimFunc is not recomputed when it's output buffer got changed, which happens
through a schedule primitive in the case of the R.layout_transform
legalization. I can try to fix it, but wanted to check whether the right fix
would be to identify cases where a prim_func can change its struct_info through
any means (scheduling in this case) and call InferStructInfo in all those
cases, but I wanted to ask if there was any better solutions here.
--
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]
