Use op.dtype to create aten.empty.memory_format during decomposition. #3941
Conversation
|
Hi @ramiro050, I see that I am reverting your change here. It's not clear to me why the |
|
Hello @ramiro050 and @vivekkhandelwal1 can you please review this PR when you have a chance? Thanks! |
| func.func @torch.aten.empty.memory_format$noneDtype() -> !torch.vtensor<[200,200,26],f64> attributes {torch.assume_strict_symbolic_shapes} { | ||
| %int200 = torch.constant.int 200 | ||
| %int26 = torch.constant.int 26 | ||
| %false = torch.constant.bool false | ||
| %none = torch.constant.none | ||
| %0 = torch.prim.ListConstruct %int200, %int200, %int26 : (!torch.int, !torch.int, !torch.int) -> !torch.list<int> | ||
| %1 = torch.aten.empty.memory_format %0, %none, %none, %none, %false, %none : !torch.list<int>, !torch.none, !torch.none, !torch.none, !torch.bool, !torch.none -> !torch.vtensor<[200,200,26],f64> | ||
| return %1 : !torch.vtensor<[200,200,26],f64> | ||
| } |
There was a problem hiding this comment.
Hi @sahas3, the above test case is an invalid test since the tensor element type is f32 when none is specifed for the aten.empty.memory_format op.
I just tried it. See:
>>> a = torch.ops.aten.empty.memory_format([2, 3])
>>> a
tensor([[-3.4054e+29, 4.2543e-41, -3.4085e+29],
[ 4.2543e-41, -3.4086e+29, 4.2543e-41]])
>>> a.dtype
torch.float32
There was a problem hiding this comment.
Also, the changes that you've done are not correct and hence this patch should not be merged.
There was a problem hiding this comment.
Great flag, @vivekkhandelwal1. I missed that torch.ops.aten.empty.memory_format produces f32 for dtype=none in PyTorch as I assumed the aten.empty.memory_format op being produced during conversion of ExportedProgram to torch dialect was correct.
Upon further investigation, I think the bug is instead in the DecomposeComplexOps pass where we decompose different torch ops to aten.empty.memory_format. For empty_like op if dtype is not specified then it defaults to input dtype as per https://pytorch.org/docs/stable/generated/torch.empty_like.html. This was not being captured when decomposing to empty.memory_format.
I have pushed new changes, reverting the original change that addresses this issue. Can you please take another look at these new changes? Thanks!
There was a problem hiding this comment.
Hi @vivekkhandelwal1 , can you please take a look at the changes after your initial feedback. Thanks!
sahas3
changed the title
|
|
||
| FailureOr<Value> dtype = getDtypeFromOp(rewriter, op); | ||
| if (failed(dtype)) { | ||
| return rewriter.notifyMatchFailure( | ||
| op, "could not determine dtype from the op."); | ||
| } | ||
|
|
There was a problem hiding this comment.
Hi @sahas3, as you mentioned that in this case, if the dtype is not present then it defaults to input dtype. But this function returns the result dtype if the dtype arg is none. So, I am not able to understand how does your comment and the changes relate.
There was a problem hiding this comment.
Hi @vivekkhandelwal1, my understanding is that dtype specifies the type of the returned tensor of any op. For ops like empty_like this dtype if not specified defaults to that of input tensor as per: https://pytorch.org/docs/stable/generated/torch.empty_like.html. So looking at the torch.empty_like op itself I think type of result tensor and input tensor will be the same if dtype is not specified.
On the other hand, for ops like empty_strided, dtype defaults to a global value as per https://pytorch.org/docs/stable/generated/torch.empty_strided.html. In this case, there's no input tensor but the result tensor should already have the type set correctly when it's in the torch dialect.
So for both cases, I think we can generalize the construction of AtenEmptyMemoryFormatOp based on the dtype of the result tensor when dtype for the op to be rewritten is not specified. We don't need to check for the input tensor dtype though I can add an assert that types of input and output tensors are same for ops like empty_like. The reason to generalize to using result tensor dtype is because ops like empty_strided don't have an input tensor and hence no op.getSelf() method on those ops.
I hope this clarifies the changes. I'll add a gist of this as comment as well for future reference.
Prior to the change in this PR
torch-mlir-opt --convert-torch-to-linalgwas running into the following error:This is because when
dtypeof theaten.empty.memory_formatisnone, by defaultf32was being selected as the element type of the resulting tensor which doesn't match with the actual element type of the result.