Lunderberg commented on code in PR #16584:
URL: https://github.com/apache/tvm/pull/16584#discussion_r1499995629
##########
tests/python/tir-base/test_tir_specialize.py:
##########
@@ -324,5 +326,41 @@ def expected(A_data: T.handle("float32")):
tvm.ir.assert_structural_equal(expected, after)
+def test_specialization_removes_struct_info():
+ """Reset struct info in specialization
+
+ While a PrimFunc usually doesn't have a `relax.StructInfo`, the
+ field can be populated in some edge cases. If that PrimFunc is
Review Comment:
There's a couple of MLC
([example](https://github.com/mlc-ai/mlc-llm/blob/fb4c52c749c4b32594f1e9382113fecd11847893/mlc_llm/transform/fuse_split_rotary_embedding.py#L100))
where the a dynamic PrimFunc is annotated with the `relax.FuncStructInfo`.
This allows the `R.call_tir` shape inference to work correctly, even after the
arguments are mutated.
IIRC, the primary reasons not to add `FuncStructInfo` to PrimFuncs was that
primfuncs didn't follow the relax requirements for functional purity. Since we
now have the purity expressed within `FuncStructInfo`, I think that means we
could add `FuncStructInfo` to every PrimFunc on construction, without making
any incorrect statements about the functions.
--
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]
