Re: [PR] [Relax] Allow R.Prim('bool') in relax::If and assert_op [tvm]

["quic-sanirudh (via GitHub)" <gi...@apache.org>]

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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org