The interface is meant to allow for resolving the dimension of the result of an operation in terms of its operands. For example consider,
%0 = <some op that implements InferShapedTypeOpInterface>
%1 = memref/tensor.dim %0, ...
It is beneficial to resolve %1
in terms of operands of the operation that creates %0
. This has two benefits
- If the operation that creates
%0
has no other uses, the operation becomes dead and can be DCE-ed - If you can resolve value of
%1
in terms of operands of the first operation, and if those themselves are created by operations that implement theInferShapedTypeOpInterface
, you could further resolve those. Using a fixed point iteration, you can them resolve the shapes of the program (assuming every op implements this interface)
Currently the InferShapedTypeOpInterface
implements three methods
-
InferReturnTypeComponents
: This method is used to get information of the return values when they are shaped type. This is meant to get the return type during operation creation. -
reifyResultTypeShapes
: This method is used to resolvedim
of result in terms of dim of operands, and is also meant to target unranked types. In particular it returns (by reference) a list ofValue
s of typetensor<? x index>
(one for each result of the operation). The rank of the tensor is same as the rank of the result, and thei
-th element of the tensor is the shape of the result ati
-th dimenion. -
reifyResultTypeShapesPerResultDim
: This method is used to resolvedim
of result in terms of dim of operands but for ranked types. In particular it returns (by reference) a list ofSmallVector<Value>
(one for each result). Thei
-th element of theSmallVector
is the shape of the result of thei
-th dimension. This interface avoids creating (and later canonicalizing away) the operations to create atensor<? x index>
and then extracing thei
-th element by just directly forwarding the SSA value.
These seems like it could use some redesign/refactoring. Particularly
-
The
InferReturnTypeComponents
of this interface seems to be doing the same asInferTypeComponents
. I dont have visiblity into what this method is for, but can this be folded into theInferTypeComponents
interface? -
In any case the
reifyResultTypeShapes
andreifyResultTypeShapesPerResultDim
are similar but solving slightly different uses. It is currently expected that an operation implements only one of these methods. The first one takes precedence over the second when thedim
op is canonicalized using theResolveShapedTypeResultDimsPass
(here)
At this point I am more interested in the second one. Ideally these should be a separate interface since this is what is needed to resolve shapes using fixed point iterations (in the dynamic shape cases). In any case, there needs to be more synergy between the two methods. My proposal is
- The
reifyResultTypeShapesPerResultDim
has a default implementation that usesreifyResultTypeShapes
to get the dimensions for each dim by usingtensor.extract
. The issue here is that this would add a explicit dependence of the interface ontensor
dialect which has been a show stopper. Any suggestions here? - The dim canonicalization (linked above) can be flipped to use just the
reifyResultTypeShapesPerResultDim
interface. With the default implementation above, it would give the same behavior.
I am happy to implement these things when all the stake holders agree. @jpienaar @herhut @nicolasvasilache @stellaraccident (and maybe even @benvanik and @River707 )