[Torch] add aten.bilinear op decomposing #3931
Conversation
|
Hi @vivekkhandelwal1. I don't have permission to request review yet. Do you mind taking a look at this change? Thanks! |
|
Hi @vivekkhandelwal1, a gentle reminder to review this pull request when you get a chance. Thank you! |
@dixinzhou Is this PR still active and needs a review? |
Hi @vivekkhandelwal1, thanks for the reply. Yes, this PR still needs to be reviewed and upstreamed. |
| int64_t size1 = inputSize1[i]; | ||
| int64_t size2 = inputSize2[i]; | ||
| if (size1 == kUnknownSize || size2 == kUnknownSize) { | ||
| mulShape.push_back(kUnknownSize); | ||
| } else { | ||
| mulShape.push_back(size1 == 1 ? size2 : size1); | ||
| } |
There was a problem hiding this comment.
Can you please explain this? I would suggest add a comment.
There was a problem hiding this comment.
Comment added in the code.
| if (isa<Torch::NoneType>(bias.getType())) { | ||
| rewriter.replaceOp(op, trilinear); | ||
| return success(); | ||
| } else { |
There was a problem hiding this comment.
Else is not required since the if ends with a return.
There was a problem hiding this comment.
The else is removed, and feedback resolved.
| // generate `aten._trilinear` op | ||
| unsigned n = inputType1.getSizes().size() - 1; | ||
| Type listOfInt = | ||
| rewriter.getType<Torch::ListType>(rewriter.getType<Torch::IntType>()); | ||
| Value zero = | ||
| rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(n)); | ||
| Value one = | ||
| rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(n + 1)); | ||
| Value two = | ||
| rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(n + 2)); | ||
| Value expand1 = rewriter.create<PrimListConstructOp>( | ||
| loc, listOfInt, SmallVector<Value>{zero, two}); | ||
| Value expand2 = rewriter.create<PrimListConstructOp>( | ||
| loc, listOfInt, SmallVector<Value>{zero, one}); | ||
| SmallVector<Value> expandWeightValue; | ||
| for (unsigned i = 0; i < n; i++) { | ||
| Value value = | ||
| rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(i)); | ||
| expandWeightValue.push_back(value); | ||
| } | ||
| Value expandw = | ||
| rewriter.create<PrimListConstructOp>(loc, listOfInt, expandWeightValue); | ||
| Value sumdim = rewriter.create<PrimListConstructOp>( | ||
| loc, listOfInt, SmallVector<Value>{one, two}); | ||
| Value constOne = | ||
| rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(1)); | ||
| Value trilinear = rewriter.create<Aten_TrilinearOp>( | ||
| loc, op.getType(), input1, weight, input2, expand1, expandw, expand2, | ||
| sumdim, constOne); |
There was a problem hiding this comment.
I did not understand this part. Can you please add a comment explaining this?
There was a problem hiding this comment.
Comment added in the code.
|
Hi @vivekkhandelwal1. I have updated the pull request per your review feedback. Do you mind taking another look at this change? Thanks! |
This PR adds support for
aten.bilinearop. Theaten.bilinearop is decomposed toaten._trilinearandaten.addaccording to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/Linear.cpp#L712.Additionally, this PR fixes a bug in
aten._trilinearop decomposition about tensor shape mismatch.