Shape polymorphism - why not?


While looking at the tf-opt --tf-to-hlo-pipeline I discovered that it includes the application of --tf-guarantee-all-funcs-one-use. This pass replicates function symbols to ensure that shape analysis can be applied in a simple way. From afar this looks like a solution to a polymorphism-related problem that avoids making shape analysis more intelligent.

And in this sense it’s related to an issue that I have encountered in my work: Assume that I have function that adds 5.0 to each of the elements of a 1D memref. The definition of this function uses memref<?xf32> inputs and outputs, but when the caller function wants to use fixed shapes I have to always cast to/from the fixed size to memref<?xf32>.

My question: why having chosen the current solution, instead of allowing the specification of (linear) constraints between input and output shapes?

BTW: coming back to tf-opt --tf-guarantee-all-funcs-one-use: why instantiate a function multiple types, even when all instances have the same signature?



Correct this was a practical solution to defer more complete analysis/inference until later that just works well enough that we haven’t revisited it (given our input models it just hasn’t caused real issues either).

And indeed we have considered having a more relaxed “TF function” op here too given the calling allows for compatible shapes vs exact. Contrary to FuncOp. And could be convenient at “front end” level. Now memref side this may influence the calling convention (a memref of 100x100 you only need to pass in pointer to function as shape is in type, while ?x100 you also need to pass in one dim - now of course one could decide to just always pass the full shape across function boundaries and then it doesn’t).

Linear constraint is not sufficient here, we have some shape functions doing multi-dimensional contractions and data dependent behavior. But we could have linear too to avoid duplicating in those cases where supported. It’s not been a priority there yet unfortunately. But it is something we want to improve.

Re: instantiating multiple times even if same signature. I’m assuming you are referring to a static signature, and indeed if so that is not needed and I’ll take a look at that. One effect there is if a call sites pass in a constant value then that unlocks the data dependent shape functions to fire. The current version is also not as we want it as it mutates and we will change that which may move the duplication purely internal to the pass/analysis.



1 Like

I’m referring to both:

  • Fully static signatures, e.g. func @f(%x:i32)->(i32)
  • The fact that at most an instance should be needed per size of I/O data. For instance, if my function is func @f(%x:memref<?xf32>)->memref<?xf32>, and if all instances are of type (memref<100xf32>)->(memref<200xf32>), then only one instance is needed, not one per call. And this instance could be called f_100_200. This is something that is common in synchronous language design, e.g. in the Heptagon dialect of Lustre, which @albertcohen knows well.