[Unity][TIR] Clear struct info when specializing PrimFunc #16584
Conversation
In rare cases, a `PrimFunc` may be annotated with `StructInfo`, to indicate that it is an impure function with specific shapes for the parameters. If struct info is present, it is invalidated when specializing a `PrimFunc`, and should be cleared.
slyubomirsky
left a comment
slyubomirsky
left a comment
There was a problem hiding this comment.
The change seems straightforward but I would be more curious to hear about why it's necessary (i.e., what problems were being caused) and when the StructInfo is being set for PrimFuncs.
| While a PrimFunc usually doesn't have a `relax.StructInfo`, the | ||
| field can be populated in some edge cases. If that PrimFunc is |
There was a problem hiding this comment.
Do you know of times when we're doing it? It should be something we do consistently or not at all.
There was a problem hiding this comment.
There's a couple of MLC (example) 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.
There was a problem hiding this comment.
Always specifying it sounds like a better approach, as long as it would not interfere with anything in TIR.
slyubomirsky
left a comment
slyubomirsky
left a comment
There was a problem hiding this comment.
This change is fine but we should ultimately be consistent about whether or not to include StructInfo on PrimFuncs. Per the above discussion, always including it might cause fewer headaches.
In rare cases, a `PrimFunc` may be annotated with `StructInfo`, to indicate that it is an impure function with specific shapes for the parameters. If struct info is present, it is invalidated when specializing a `PrimFunc`, and should be cleared.
In rare cases, a `PrimFunc` may be annotated with `StructInfo`, to indicate that it is an impure function with specific shapes for the parameters. If struct info is present, it is invalidated when specializing a `PrimFunc`, and should be cleared.
3 participants
In rare cases, a
PrimFuncmay be annotated withStructInfo, to indicate that it is an impure function with specific shapes for the parameters. If struct info is present, it is invalidated when specializing aPrimFunc, and should be cleared.