Missing dimensionality information in torch #3651
Comments
|
This is pretty interesting. Onnx shape inference can't seem to figure out that a slice op in dim 1 should preserve dim 0, so it leaves it as dynamic. When lowering in torch, we correct the slice shape to |
|
I believe that Here is something particularly strange: Let the shape of
I'll work on addressing the issue (1.) for now. If we can get the |
|
This is a small reproducer for the failure to infer shapes: module {
func.func @torch_jit(%arg0: !torch.vtensor<[?,768],f32>) -> !torch.vtensor<[?,?,768],f32> {
%int1 = torch.constant.int 1
%int0 = torch.constant.int 0
%0 = torch.aten.size.int %arg0, %int0 : !torch.vtensor<[?,768],f32>, !torch.int -> !torch.int
%1 = torch.prim.ListConstruct %int1, %0 : (!torch.int, !torch.int) -> !torch.list<int>
%2 = torch.aten.unflatten.int %arg0, %int0, %1 : !torch.vtensor<[?,768],f32>, !torch.int, !torch.list<int> -> !torch.vtensor<[?,?,768],f32>
return %2 : !torch.vtensor<[?,?,768],f32>
}
}Running does not successfully infer the output shape |
|
I'm guessing that since the |
|
I've added a canonicalization pattern for unflatten -> unsqueeze and tried it out on the original onnx ir in https://gist.github.com/nirvedhmeshram/f350fa447fdf5cdcbff45ced0dd77e6c. What is truly unfortunate is that the canonicalization patten can only take effect after shape inference determines that a certain slice op should have a 1 for the outermost dim. Therefore, shape inference would need to be run twice on this IR, since the unsqueeze op will only appear after the first shape inference is completed. I'm going to look into other options. If you want to check out the work I'm doing on this, see the branch https://github.com/zjgarvey/torch-mlir/tree/unflatten_fold |
|
Amazing. Using a static info cast op from the correct (static) unsqueeze type to the original dynamic type will resolve into using the correct (static) type after canonicalization. Adding this change to the branch and submitting a PR. |
Addresses an issue in <#3651> where some unflatten ops generated from onnx models weren't propagating static shape information. It may be necessary to add further optimizations for the more general case when some static information is present in the unflatten (or possibly reshape/view) op's `sizes` list, but not reflected in the output shape. These ops will only successfully infer shapes if the `sizes` list is gotten from a list of constant ints (with possibly one -1). A common example where this fails is when some of the `sizes` are determined from `aten.size.int` ops on dynamic tensors, and other `sizes` are known statically. This PR includes: - a canonicalizer for `aten.unflatten.int` which converts to `aten.unsqueeze` when it is expanding one dim to two, and one of the new dims is statically 1. - an improvement to the folder for `aten.__or__.bool` which does not rely on *both* operands being static.
|
I'm going to close this for now. If we find that there is still some ops with missing dimensionality, we can reopen. |
As a follow up from iree-org/iree#18229 it seems like there is some dimension information that is not being captured correctly in IR and recovering that in the program is pretty involved. This is the IR after
torch-finalizing-backend-type-conversionin particular this sequence of instructions
The second
tensor.reshapeoperation has information that the outer most dim of the result (%reshape_16) is of size 1. For consistency.%32should also have an outer dimension of size 1. Looking through the IR.. that basically means%0is of value 1. IIUC, torch should already know this information and it should have been materialized in the IR. This inconsistency between the same dimension being known statically in one-place and being dynamic in another causes issue downstream during compilation. The downstream compilation is being fixed, but this seems like something worth fixing in torch as well. If this is not possible, then we will have to build downstream of torch a way to constraint solve the dynamic dimensions for such cases (but I think Torch already has this).The text was updated successfully, but these errors were encountered: