[TORCH] Add support for aten.heaviside Op #4220
Conversation
|
@stellaraccident @vivekkhandelwal1 @penguin-wwy @zjgarvey @AmosLewis, I’d be grateful if any of you could take a look at this PR. Your feedback would be greatly appreciated! |
projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise_comparison.py
Outdated
Show resolved
Hide resolved
| SmallVector<int64_t> broadcastShape; | ||
| SmallVector<Value> broadcastShapeValue; | ||
| computeBroadcastShape(rewriter, loc, input, value, broadcastShape, | ||
| broadcastShapeValue); | ||
|
|
||
| auto broadcastType = ValueTensorType::get( | ||
| op.getContext(), llvm::ArrayRef(broadcastShape), resultTy.getDtype()); | ||
| auto boolBroadcastType = ValueTensorType::get( | ||
| op.getContext(), llvm::ArrayRef(broadcastShape), rewriter.getI1Type()); | ||
| Value indexBroadcastShapeTorchList = rewriter.create<PrimListConstructOp>( | ||
| loc, Torch::ListType::get(Torch::IntType::get(op.getContext())), | ||
| broadcastShapeValue); | ||
| auto inputBroadcasted = rewriter.create<AtenBroadcastToOp>( | ||
| loc, broadcastType, input, indexBroadcastShapeTorchList); | ||
| auto valueBroadcasted = rewriter.create<AtenBroadcastToOp>( | ||
| loc, broadcastType, value, indexBroadcastShapeTorchList); |
There was a problem hiding this comment.
I think this is not needed. Since you are decomposing this op into elementwise ops, the broadcasting part will be handled during Torch->Linalg lowering.
There was a problem hiding this comment.
@vivekkhandelwal1 You're right
But I ran into an issue in a specific case: when the input shape is [1, 2, 3] and the value shape is [1, 1, 1, 1], the broadcasted result shape becomes [1, 1, 2, 3].
Without explicitly broadcasting the inputs, some intermediate ops (like aten.eq.scalar or aten.isnan) end up producing tensors of shape [1, 2, 3], which causes this error:
'tensor.cast' op operand type 'tensor<1x2x3xi1>' and result type 'tensor<1x1x2x3xi1>' are cast incompatible
So to avoid this mismatch, I added explicit broadcasting to ensure all intermediate results match the final shape.
|
@vivekkhandelwal1 Thanks a lot for the feedback. I’ve updated the code. |
This implementation addresses and closes #4211